{-
(c) The AQUA Project, Glasgow University, 1994-1998

\section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
-}

{-# LANGUAGE CPP #-}
module LiberateCase ( liberateCase ) where

#include "HsVersions.h"

import GhcPrelude

import DynFlags
import CoreSyn
import CoreUnfold       ( couldBeSmallEnoughToInline )
import TysWiredIn       ( unitDataConId )
import Id
import VarEnv
import Util             ( notNull )

{-
The liberate-case transformation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This module walks over @Core@, and looks for @case@ on free variables.
The criterion is:
        if there is case on a free on the route to the recursive call,
        then the recursive call is replaced with an unfolding.

Example

   f = \ t -> case v of
                 V a b -> a : f t

=> the inner f is replaced.

   f = \ t -> case v of
                 V a b -> a : (letrec
                                f =  \ t -> case v of
                                               V a b -> a : f t
                               in f) t
(note the NEED for shadowing)

=> Simplify

  f = \ t -> case v of
                 V a b -> a : (letrec
                                f = \ t -> a : f t
                               in f t)

Better code, because 'a' is  free inside the inner letrec, rather
than needing projection from v.

Note that this deals with *free variables*.  SpecConstr deals with
*arguments* that are of known form.  E.g.

        last []     = error
        last (x:[]) = x
        last (x:xs) = last xs


Note [Scrutinee with cast]
~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
    f = \ t -> case (v `cast` co) of
                 V a b -> a : f t

Exactly the same optimisation (unrolling one call to f) will work here,
despite the cast.  See mk_alt_env in the Case branch of libCase.


To think about (Apr 94)
~~~~~~~~~~~~~~
Main worry: duplicating code excessively.  At the moment we duplicate
the entire binding group once at each recursive call.  But there may
be a group of recursive calls which share a common set of evaluated
free variables, in which case the duplication is a plain waste.

Another thing we could consider adding is some unfold-threshold thing,
so that we'll only duplicate if the size of the group rhss isn't too
big.

Data types
~~~~~~~~~~
The ``level'' of a binder tells how many
recursive defns lexically enclose the binding
A recursive defn "encloses" its RHS, not its
scope.  For example:
\begin{verbatim}
        letrec f = let g = ... in ...
        in
        let h = ...
        in ...
\end{verbatim}
Here, the level of @f@ is zero, the level of @g@ is one,
and the level of @h@ is zero (NB not one).


************************************************************************
*                                                                      *
         Top-level code
*                                                                      *
************************************************************************
-}

liberateCase :: DynFlags -> CoreProgram -> CoreProgram
liberateCase :: DynFlags -> CoreProgram -> CoreProgram
liberateCase dflags :: DynFlags
dflags binds :: CoreProgram
binds = LibCaseEnv -> CoreProgram -> CoreProgram
do_prog (DynFlags -> LibCaseEnv
initEnv DynFlags
dflags) CoreProgram
binds
  where
    do_prog :: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog _   [] = []
    do_prog env :: LibCaseEnv
env (bind :: CoreBind
bind:binds :: CoreProgram
binds) = CoreBind
bind' CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog LibCaseEnv
env' CoreProgram
binds
                             where
                               (env' :: LibCaseEnv
env', bind' :: CoreBind
bind') = LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env CoreBind
bind

{-
************************************************************************
*                                                                      *
         Main payload
*                                                                      *
************************************************************************

Bindings
~~~~~~~~
-}

libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)

libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind env :: LibCaseEnv
env (NonRec binder :: CoreBndr
binder rhs :: Expr CoreBndr
rhs)
  = (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr
binder], CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
binder (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
rhs))

libCaseBind env :: LibCaseEnv
env (Rec pairs :: [(CoreBndr, Expr CoreBndr)]
pairs)
  = (LibCaseEnv
env_body, [(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs')
  where
    binders :: [CoreBndr]
binders = ((CoreBndr, Expr CoreBndr) -> CoreBndr)
-> [(CoreBndr, Expr CoreBndr)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, Expr CoreBndr) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, Expr CoreBndr)]
pairs

    env_body :: LibCaseEnv
env_body = LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr]
binders

    pairs' :: [(CoreBndr, Expr CoreBndr)]
pairs' = [(CoreBndr
binder, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_rhs Expr CoreBndr
rhs) | (binder :: CoreBndr
binder,rhs :: Expr CoreBndr
rhs) <- [(CoreBndr, Expr CoreBndr)]
pairs]

        -- We extend the rec-env by binding each Id to its rhs, first
        -- processing the rhs with an *un-extended* environment, so
        -- that the same process doesn't occur for ever!
    env_rhs :: LibCaseEnv
env_rhs | Bool
is_dupable_bind = LibCaseEnv -> [(CoreBndr, Expr CoreBndr)] -> LibCaseEnv
addRecBinds LibCaseEnv
env [(CoreBndr, Expr CoreBndr)]
dup_pairs
            | Bool
otherwise       = LibCaseEnv
env

    dup_pairs :: [(CoreBndr, Expr CoreBndr)]
dup_pairs = [ (CoreBndr -> CoreBndr
localiseId CoreBndr
binder, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_body Expr CoreBndr
rhs)
                | (binder :: CoreBndr
binder, rhs :: Expr CoreBndr
rhs) <- [(CoreBndr, Expr CoreBndr)]
pairs ]
        -- localiseID : see Note [Need to localiseId in libCaseBind]

    is_dupable_bind :: Bool
is_dupable_bind = Bool
small_enough Bool -> Bool -> Bool
&& ((CoreBndr, Expr CoreBndr) -> Bool)
-> [(CoreBndr, Expr CoreBndr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (CoreBndr, Expr CoreBndr) -> Bool
forall b. (CoreBndr, b) -> Bool
ok_pair [(CoreBndr, Expr CoreBndr)]
pairs

    -- Size: we are going to duplicate dup_pairs; to find their
    --       size, build a fake binding (let { dup_pairs } in (),
    --       and find the size of that
    -- See Note [Small enough]
    small_enough :: Bool
small_enough = case LibCaseEnv -> Maybe Int
bombOutSize LibCaseEnv
env of
                      Nothing   -> Bool
True   -- Infinity
                      Just size :: Int
size -> DynFlags -> Int -> Expr CoreBndr -> Bool
couldBeSmallEnoughToInline (LibCaseEnv -> DynFlags
lc_dflags LibCaseEnv
env) Int
size (Expr CoreBndr -> Bool) -> Expr CoreBndr -> Bool
forall a b. (a -> b) -> a -> b
$
                                   CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
dup_pairs) (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
unitDataConId)

    ok_pair :: (CoreBndr, b) -> Bool
ok_pair (id :: CoreBndr
id,_)
        =  CoreBndr -> Int
idArity CoreBndr
id Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0          -- Note [Only functions!]
        Bool -> Bool -> Bool
&& Bool -> Bool
not (CoreBndr -> Bool
isBottomingId CoreBndr
id)  -- Note [Not bottoming ids]

{- Note [Not bottoming Ids]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Do not specialise error-functions (this is unusual, but I once saw it,
(acually in Data.Typable.Internal)

Note [Only functions!]
~~~~~~~~~~~~~~~~~~~~~~
Consider the following code

       f = g (case v of V a b -> a : t f)

where g is expensive. If we aren't careful, liberate case will turn this into

       f = g (case v of
               V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
                                in f)
             )

Yikes! We evaluate g twice. This leads to a O(2^n) explosion
if g calls back to the same code recursively.

Solution: make sure that we only do the liberate-case thing on *functions*

Note [Small enough]
~~~~~~~~~~~~~~~~~~~
Consider
  \fv. letrec
         f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
         g = \y. SMALL...f...

Then we *can* in principle do liberate-case on 'g' (small RHS) but not
for 'f' (too big).  But doing so is not profitable, because duplicating
'g' at its call site in 'f' doesn't get rid of any cases.  So we just
ask for the whole group to be small enough.

Note [Need to localiseId in libCaseBind]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The call to localiseId is needed for two subtle reasons
(a)  Reset the export flags on the binders so
        that we don't get name clashes on exported things if the
        local binding floats out to top level.  This is most unlikely
        to happen, since the whole point concerns free variables.
        But resetting the export flag is right regardless.

(b)  Make the name an Internal one.  External Names should never be
        nested; if it were floated to the top level, we'd get a name
        clash at code generation time.

Expressions
~~~~~~~~~~~
-}

libCase :: LibCaseEnv
        -> CoreExpr
        -> CoreExpr

libCase :: LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase env :: LibCaseEnv
env (Var v :: CoreBndr
v)             = LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v []
libCase _   (Lit lit :: Literal
lit)           = Literal -> Expr CoreBndr
forall b. Literal -> Expr b
Lit Literal
lit
libCase _   (Type ty :: Type
ty)           = Type -> Expr CoreBndr
forall b. Type -> Expr b
Type Type
ty
libCase _   (Coercion co :: Coercion
co)       = Coercion -> Expr CoreBndr
forall b. Coercion -> Expr b
Coercion Coercion
co
libCase env :: LibCaseEnv
env e :: Expr CoreBndr
e@(App {})          | let (fun :: Expr CoreBndr
fun, args :: [Expr CoreBndr]
args) = Expr CoreBndr -> (Expr CoreBndr, [Expr CoreBndr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr CoreBndr
e
                                , Var v :: CoreBndr
v <- Expr CoreBndr
fun
                                = LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v [Expr CoreBndr]
args
libCase env :: LibCaseEnv
env (App fun :: Expr CoreBndr
fun arg :: Expr CoreBndr
arg)       = Expr CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Expr b -> Expr b -> Expr b
App (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
fun) (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
arg)
libCase env :: LibCaseEnv
env (Tick tickish :: Tickish CoreBndr
tickish body :: Expr CoreBndr
body) = Tickish CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick Tickish CoreBndr
tickish (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
body)
libCase env :: LibCaseEnv
env (Cast e :: Expr CoreBndr
e co :: Coercion
co)         = Expr CoreBndr -> Coercion -> Expr CoreBndr
forall b. Expr b -> Coercion -> Expr b
Cast (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
e) Coercion
co

libCase env :: LibCaseEnv
env (Lam binder :: CoreBndr
binder body :: Expr CoreBndr
body)
  = CoreBndr -> Expr CoreBndr -> Expr CoreBndr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
binder (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr
binder]) Expr CoreBndr
body)

libCase env :: LibCaseEnv
env (Let bind :: CoreBind
bind body :: Expr CoreBndr
body)
  = CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind' (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env_body Expr CoreBndr
body)
  where
    (env_body :: LibCaseEnv
env_body, bind' :: CoreBind
bind') = LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env CoreBind
bind

libCase env :: LibCaseEnv
env (Case scrut :: Expr CoreBndr
scrut bndr :: CoreBndr
bndr ty :: Type
ty alts :: [Alt CoreBndr]
alts)
  = Expr CoreBndr
-> CoreBndr -> Type -> [Alt CoreBndr] -> Expr CoreBndr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env Expr CoreBndr
scrut) CoreBndr
bndr Type
ty ((Alt CoreBndr -> Alt CoreBndr) -> [Alt CoreBndr] -> [Alt CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
libCaseAlt LibCaseEnv
env_alts) [Alt CoreBndr]
alts)
  where
    env_alts :: LibCaseEnv
env_alts = LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders (Expr CoreBndr -> LibCaseEnv
forall b. Expr b -> LibCaseEnv
mk_alt_env Expr CoreBndr
scrut) [CoreBndr
bndr]
    mk_alt_env :: Expr b -> LibCaseEnv
mk_alt_env (Var scrut_var :: CoreBndr
scrut_var) = LibCaseEnv -> CoreBndr -> LibCaseEnv
addScrutedVar LibCaseEnv
env CoreBndr
scrut_var
    mk_alt_env (Cast scrut :: Expr b
scrut _)  = Expr b -> LibCaseEnv
mk_alt_env Expr b
scrut       -- Note [Scrutinee with cast]
    mk_alt_env _               = LibCaseEnv
env

libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
                         -> (AltCon, [CoreBndr], CoreExpr)
libCaseAlt :: LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
libCaseAlt env :: LibCaseEnv
env (con :: AltCon
con,args :: [CoreBndr]
args,rhs :: Expr CoreBndr
rhs) = (AltCon
con, [CoreBndr]
args, LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase (LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders LibCaseEnv
env [CoreBndr]
args) Expr CoreBndr
rhs)

{-
Ids
~~~

To unfold, we can't just wrap the id itself in its binding if it's a join point:

  jump j a b c  =>  (joinrec j x y z = ... in jump j) a b c -- wrong!!!

Every jump must provide all arguments, so we have to be careful to wrap the
whole jump instead:

  jump j a b c  =>  joinrec j x y z = ... in jump j a b c -- right

-}

libCaseApp :: LibCaseEnv -> Id -> [CoreExpr] -> CoreExpr
libCaseApp :: LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp env :: LibCaseEnv
env v :: CoreBndr
v args :: [Expr CoreBndr]
args
  | Just the_bind :: CoreBind
the_bind <- LibCaseEnv -> CoreBndr -> Maybe CoreBind
lookupRecId LibCaseEnv
env CoreBndr
v  -- It's a use of a recursive thing
  , [CoreBndr] -> Bool
forall a. [a] -> Bool
notNull [CoreBndr]
free_scruts                 -- with free vars scrutinised in RHS
  = CoreBind -> Expr CoreBndr -> Expr CoreBndr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
the_bind Expr CoreBndr
expr'

  | Bool
otherwise
  = Expr CoreBndr
expr'

  where
    rec_id_level :: Int
rec_id_level = LibCaseEnv -> CoreBndr -> Int
lookupLevel LibCaseEnv
env CoreBndr
v
    free_scruts :: [CoreBndr]
free_scruts  = LibCaseEnv -> Int -> [CoreBndr]
freeScruts LibCaseEnv
env Int
rec_id_level
    expr' :: Expr CoreBndr
expr'        = Expr CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (CoreBndr -> Expr CoreBndr
forall b. CoreBndr -> Expr b
Var CoreBndr
v) ((Expr CoreBndr -> Expr CoreBndr)
-> [Expr CoreBndr] -> [Expr CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (LibCaseEnv -> Expr CoreBndr -> Expr CoreBndr
libCase LibCaseEnv
env) [Expr CoreBndr]
args)

freeScruts :: LibCaseEnv
           -> LibCaseLevel      -- Level of the recursive Id
           -> [Id]              -- Ids that are scrutinised between the binding
                                -- of the recursive Id and here
freeScruts :: LibCaseEnv -> Int -> [CoreBndr]
freeScruts env :: LibCaseEnv
env rec_bind_lvl :: Int
rec_bind_lvl
  = [CoreBndr
v | (v :: CoreBndr
v, scrut_bind_lvl :: Int
scrut_bind_lvl, scrut_at_lvl :: Int
scrut_at_lvl) <- LibCaseEnv -> [(CoreBndr, Int, Int)]
lc_scruts LibCaseEnv
env
       , Int
scrut_bind_lvl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
rec_bind_lvl
       , Int
scrut_at_lvl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
rec_bind_lvl]
        -- Note [When to specialise]
        -- Note [Avoiding fruitless liberate-case]

{-
Note [When to specialise]
~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  f = \x. letrec g = \y. case x of
                           True  -> ... (f a) ...
                           False -> ... (g b) ...

We get the following levels
          f  0
          x  1
          g  1
          y  2

Then 'x' is being scrutinised at a deeper level than its binding, so
it's added to lc_sruts:  [(x,1)]

We do *not* want to specialise the call to 'f', because 'x' is not free
in 'f'.  So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).

We *do* want to specialise the call to 'g', because 'x' is free in g.
Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).

Note [Avoiding fruitless liberate-case]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider also:
  f = \x. case top_lvl_thing of
                I# _ -> let g = \y. ... g ...
                        in ...

Here, top_lvl_thing is scrutinised at a level (1) deeper than its
binding site (0).  Nevertheless, we do NOT want to specialise the call
to 'g' because all the structure in its free variables is already
visible at the definition site for g.  Hence, when considering specialising
an occurrence of 'g', we want to check that there's a scruted-var v st

   a) v's binding site is *outside* g
   b) v's scrutinisation site is *inside* g


************************************************************************
*                                                                      *
        Utility functions
*                                                                      *
************************************************************************
-}

addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders env :: LibCaseEnv
env@(LibCaseEnv { lc_lvl :: LibCaseEnv -> Int
lc_lvl = Int
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv Int
lc_lvl_env = IdEnv Int
lvl_env }) binders :: [CoreBndr]
binders
  = LibCaseEnv
env { lc_lvl_env :: IdEnv Int
lc_lvl_env = IdEnv Int
lvl_env' }
  where
    lvl_env' :: IdEnv Int
lvl_env' = IdEnv Int -> [(CoreBndr, Int)] -> IdEnv Int
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv Int
lvl_env ([CoreBndr]
binders [CoreBndr] -> [Int] -> [(CoreBndr, Int)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` Int -> [Int]
forall a. a -> [a]
repeat Int
lvl)

addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
addRecBinds :: LibCaseEnv -> [(CoreBndr, Expr CoreBndr)] -> LibCaseEnv
addRecBinds env :: LibCaseEnv
env@(LibCaseEnv {lc_lvl :: LibCaseEnv -> Int
lc_lvl = Int
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv Int
lc_lvl_env = IdEnv Int
lvl_env,
                             lc_rec_env :: LibCaseEnv -> IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env}) pairs :: [(CoreBndr, Expr CoreBndr)]
pairs
  = LibCaseEnv
env { lc_lvl :: Int
lc_lvl = Int
lvl', lc_lvl_env :: IdEnv Int
lc_lvl_env = IdEnv Int
lvl_env', lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env' }
  where
    lvl' :: Int
lvl'     = Int
lvl Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1
    lvl_env' :: IdEnv Int
lvl_env' = IdEnv Int -> [(CoreBndr, Int)] -> IdEnv Int
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv Int
lvl_env [(CoreBndr
binder,Int
lvl) | (binder :: CoreBndr
binder,_) <- [(CoreBndr, Expr CoreBndr)]
pairs]
    rec_env' :: IdEnv CoreBind
rec_env' = IdEnv CoreBind -> [(CoreBndr, CoreBind)] -> IdEnv CoreBind
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv CoreBind
rec_env [(CoreBndr
binder, [(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(CoreBndr, Expr CoreBndr)]
pairs) | (binder :: CoreBndr
binder,_) <- [(CoreBndr, Expr CoreBndr)]
pairs]

addScrutedVar :: LibCaseEnv
              -> Id             -- This Id is being scrutinised by a case expression
              -> LibCaseEnv

addScrutedVar :: LibCaseEnv -> CoreBndr -> LibCaseEnv
addScrutedVar env :: LibCaseEnv
env@(LibCaseEnv { lc_lvl :: LibCaseEnv -> Int
lc_lvl = Int
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv Int
lc_lvl_env = IdEnv Int
lvl_env,
                                lc_scruts :: LibCaseEnv -> [(CoreBndr, Int, Int)]
lc_scruts = [(CoreBndr, Int, Int)]
scruts }) scrut_var :: CoreBndr
scrut_var
  | Int
bind_lvl Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
lvl
  = LibCaseEnv
env { lc_scruts :: [(CoreBndr, Int, Int)]
lc_scruts = [(CoreBndr, Int, Int)]
scruts' }
        -- Add to scruts iff the scrut_var is being scrutinised at
        -- a deeper level than its defn

  | Bool
otherwise = LibCaseEnv
env
  where
    scruts' :: [(CoreBndr, Int, Int)]
scruts'  = (CoreBndr
scrut_var, Int
bind_lvl, Int
lvl) (CoreBndr, Int, Int)
-> [(CoreBndr, Int, Int)] -> [(CoreBndr, Int, Int)]
forall a. a -> [a] -> [a]
: [(CoreBndr, Int, Int)]
scruts
    bind_lvl :: Int
bind_lvl = case IdEnv Int -> CoreBndr -> Maybe Int
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv IdEnv Int
lvl_env CoreBndr
scrut_var of
                 Just lvl :: Int
lvl -> Int
lvl
                 Nothing  -> Int
topLevel

lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
lookupRecId :: LibCaseEnv -> CoreBndr -> Maybe CoreBind
lookupRecId env :: LibCaseEnv
env id :: CoreBndr
id = IdEnv CoreBind -> CoreBndr -> Maybe CoreBind
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv (LibCaseEnv -> IdEnv CoreBind
lc_rec_env LibCaseEnv
env) CoreBndr
id

lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
lookupLevel :: LibCaseEnv -> CoreBndr -> Int
lookupLevel env :: LibCaseEnv
env id :: CoreBndr
id
  = case IdEnv Int -> CoreBndr -> Maybe Int
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv (LibCaseEnv -> IdEnv Int
lc_lvl_env LibCaseEnv
env) CoreBndr
id of
      Just lvl :: Int
lvl -> Int
lvl
      Nothing  -> Int
topLevel

{-
************************************************************************
*                                                                      *
         The environment
*                                                                      *
************************************************************************
-}

type LibCaseLevel = Int

topLevel :: LibCaseLevel
topLevel :: Int
topLevel = 0

data LibCaseEnv
  = LibCaseEnv {
        LibCaseEnv -> DynFlags
lc_dflags :: DynFlags,

        LibCaseEnv -> Int
lc_lvl :: LibCaseLevel, -- Current level
                -- The level is incremented when (and only when) going
                -- inside the RHS of a (sufficiently small) recursive
                -- function.

        LibCaseEnv -> IdEnv Int
lc_lvl_env :: IdEnv LibCaseLevel,
                -- Binds all non-top-level in-scope Ids (top-level and
                -- imported things have a level of zero)

        LibCaseEnv -> IdEnv CoreBind
lc_rec_env :: IdEnv CoreBind,
                -- Binds *only* recursively defined ids, to their own
                -- binding group, and *only* in their own RHSs

        LibCaseEnv -> [(CoreBndr, Int, Int)]
lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
                -- Each of these Ids was scrutinised by an enclosing
                -- case expression, at a level deeper than its binding
                -- level.
                --
                -- The first LibCaseLevel is the *binding level* of
                --   the scrutinised Id,
                -- The second is the level *at which it was scrutinised*.
                --   (see Note [Avoiding fruitless liberate-case])
                -- The former is a bit redundant, since you could always
                -- look it up in lc_lvl_env, but it's just cached here
                --
                -- The order is insignificant; it's a bag really
                --
                -- There's one element per scrutinisation;
                --    in principle the same Id may appear multiple times,
                --    although that'd be unusual:
                --       case x of { (a,b) -> ....(case x of ...) .. }
        }

initEnv :: DynFlags -> LibCaseEnv
initEnv :: DynFlags -> LibCaseEnv
initEnv dflags :: DynFlags
dflags
  = LibCaseEnv :: DynFlags
-> Int
-> IdEnv Int
-> IdEnv CoreBind
-> [(CoreBndr, Int, Int)]
-> LibCaseEnv
LibCaseEnv { lc_dflags :: DynFlags
lc_dflags = DynFlags
dflags,
                 lc_lvl :: Int
lc_lvl = 0,
                 lc_lvl_env :: IdEnv Int
lc_lvl_env = IdEnv Int
forall a. VarEnv a
emptyVarEnv,
                 lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
forall a. VarEnv a
emptyVarEnv,
                 lc_scruts :: [(CoreBndr, Int, Int)]
lc_scruts = [] }

-- Bomb-out size for deciding if
-- potential liberatees are too big.
-- (passed in from cmd-line args)
bombOutSize :: LibCaseEnv -> Maybe Int
bombOutSize :: LibCaseEnv -> Maybe Int
bombOutSize = DynFlags -> Maybe Int
liberateCaseThreshold (DynFlags -> Maybe Int)
-> (LibCaseEnv -> DynFlags) -> LibCaseEnv -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LibCaseEnv -> DynFlags
lc_dflags