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

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

{-# LANGUAGE CPP #-}
module GHC.Core.Opt.LiberateCase ( liberateCase ) where

#include "HsVersions.h"

import GHC.Prelude

import GHC.Driver.Session
import GHC.Core
import GHC.Core.Unfold  ( couldBeSmallEnoughToInline )
import GHC.Builtin.Types ( unitDataConId )
import GHC.Types.Id
import GHC.Types.Var.Env
import GHC.Utils.Misc    ( 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 DynFlags
dflags CoreProgram
binds = LibCaseEnv -> CoreProgram -> CoreProgram
do_prog (DynFlags -> LibCaseEnv
initEnv DynFlags
dflags) CoreProgram
binds
  where
    do_prog :: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog LibCaseEnv
_   [] = []
    do_prog LibCaseEnv
env (CoreBind
bind:CoreProgram
binds) = CoreBind
bind' CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
: LibCaseEnv -> CoreProgram -> CoreProgram
do_prog LibCaseEnv
env' CoreProgram
binds
                             where
                               (LibCaseEnv
env', 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 LibCaseEnv
env (NonRec CoreBndr
binder 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 LibCaseEnv
env (Rec [(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) | (CoreBndr
binder,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)
                | (CoreBndr
binder, 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 LibCaseLevel
bombOutSize LibCaseEnv
env of
                      Maybe LibCaseLevel
Nothing   -> Bool
True   -- Infinity
                      Just LibCaseLevel
size -> DynFlags -> LibCaseLevel -> Expr CoreBndr -> Bool
couldBeSmallEnoughToInline (LibCaseEnv -> DynFlags
lc_dflags LibCaseEnv
env) LibCaseLevel
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 (CoreBndr
id,b
_)
        =  CoreBndr -> LibCaseLevel
idArity CoreBndr
id LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
> LibCaseLevel
0       -- Note [Only functions!]
        Bool -> Bool -> Bool
&& Bool -> Bool
not (CoreBndr -> Bool
isDeadEndId CoreBndr
id) -- Note [Not bottoming ids]

{- Note [Not bottoming Ids]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
Do not specialise error-functions (this is unusual, but I once saw it,
(actually 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 LibCaseEnv
env (Var CoreBndr
v)             = LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v []
libCase LibCaseEnv
_   (Lit Literal
lit)           = Literal -> Expr CoreBndr
forall b. Literal -> Expr b
Lit Literal
lit
libCase LibCaseEnv
_   (Type Type
ty)           = Type -> Expr CoreBndr
forall b. Type -> Expr b
Type Type
ty
libCase LibCaseEnv
_   (Coercion Coercion
co)       = Coercion -> Expr CoreBndr
forall b. Coercion -> Expr b
Coercion Coercion
co
libCase LibCaseEnv
env e :: Expr CoreBndr
e@(App {})          | let (Expr CoreBndr
fun, [Expr CoreBndr]
args) = Expr CoreBndr -> (Expr CoreBndr, [Expr CoreBndr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs Expr CoreBndr
e
                                , Var CoreBndr
v <- Expr CoreBndr
fun
                                = LibCaseEnv -> CoreBndr -> [Expr CoreBndr] -> Expr CoreBndr
libCaseApp LibCaseEnv
env CoreBndr
v [Expr CoreBndr]
args
libCase LibCaseEnv
env (App Expr CoreBndr
fun 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 LibCaseEnv
env (Tick Tickish CoreBndr
tickish 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 LibCaseEnv
env (Cast Expr CoreBndr
e 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 LibCaseEnv
env (Lam CoreBndr
binder 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 LibCaseEnv
env (Let CoreBind
bind 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
    (LibCaseEnv
env_body, CoreBind
bind') = LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
libCaseBind LibCaseEnv
env CoreBind
bind

libCase LibCaseEnv
env (Case Expr CoreBndr
scrut CoreBndr
bndr Type
ty [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 CoreBndr
scrut_var) = LibCaseEnv -> CoreBndr -> LibCaseEnv
addScrutedVar LibCaseEnv
env CoreBndr
scrut_var
    mk_alt_env (Cast Expr b
scrut Coercion
_)  = Expr b -> LibCaseEnv
mk_alt_env Expr b
scrut       -- Note [Scrutinee with cast]
    mk_alt_env Expr b
_               = LibCaseEnv
env

libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
                         -> (AltCon, [CoreBndr], CoreExpr)
libCaseAlt :: LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
libCaseAlt LibCaseEnv
env (AltCon
con,[CoreBndr]
args,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 LibCaseEnv
env CoreBndr
v [Expr CoreBndr]
args
  | Just 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 :: LibCaseLevel
rec_id_level = LibCaseEnv -> CoreBndr -> LibCaseLevel
lookupLevel LibCaseEnv
env CoreBndr
v
    free_scruts :: [CoreBndr]
free_scruts  = LibCaseEnv -> LibCaseLevel -> [CoreBndr]
freeScruts LibCaseEnv
env LibCaseLevel
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 -> LibCaseLevel -> [CoreBndr]
freeScruts LibCaseEnv
env LibCaseLevel
rec_bind_lvl
  = [CoreBndr
v | (CoreBndr
v, LibCaseLevel
scrut_bind_lvl, LibCaseLevel
scrut_at_lvl) <- LibCaseEnv -> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts LibCaseEnv
env
       , LibCaseLevel
scrut_bind_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
<= LibCaseLevel
rec_bind_lvl
       , LibCaseLevel
scrut_at_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
> LibCaseLevel
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 -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env }) [CoreBndr]
binders
  = LibCaseEnv
env { lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env' }
  where
    lvl_env' :: IdEnv LibCaseLevel
lvl_env' = IdEnv LibCaseLevel
-> [(CoreBndr, LibCaseLevel)] -> IdEnv LibCaseLevel
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv LibCaseLevel
lvl_env ([CoreBndr]
binders [CoreBndr] -> [LibCaseLevel] -> [(CoreBndr, LibCaseLevel)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` LibCaseLevel -> [LibCaseLevel]
forall a. a -> [a]
repeat LibCaseLevel
lvl)

addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
addRecBinds :: LibCaseEnv -> [(CoreBndr, Expr CoreBndr)] -> LibCaseEnv
addRecBinds env :: LibCaseEnv
env@(LibCaseEnv {lc_lvl :: LibCaseEnv -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env,
                             lc_rec_env :: LibCaseEnv -> IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env}) [(CoreBndr, Expr CoreBndr)]
pairs
  = LibCaseEnv
env { lc_lvl :: LibCaseLevel
lc_lvl = LibCaseLevel
lvl', lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env', lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
rec_env' }
  where
    lvl' :: LibCaseLevel
lvl'     = LibCaseLevel
lvl LibCaseLevel -> LibCaseLevel -> LibCaseLevel
forall a. Num a => a -> a -> a
+ LibCaseLevel
1
    lvl_env' :: IdEnv LibCaseLevel
lvl_env' = IdEnv LibCaseLevel
-> [(CoreBndr, LibCaseLevel)] -> IdEnv LibCaseLevel
forall a. VarEnv a -> [(CoreBndr, a)] -> VarEnv a
extendVarEnvList IdEnv LibCaseLevel
lvl_env [(CoreBndr
binder,LibCaseLevel
lvl) | (CoreBndr
binder,Expr CoreBndr
_) <- [(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) | (CoreBndr
binder,Expr CoreBndr
_) <- [(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 -> LibCaseLevel
lc_lvl = LibCaseLevel
lvl, lc_lvl_env :: LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
lvl_env,
                                lc_scruts :: LibCaseEnv -> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts = [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts }) CoreBndr
scrut_var
  | LibCaseLevel
bind_lvl LibCaseLevel -> LibCaseLevel -> Bool
forall a. Ord a => a -> a -> Bool
< LibCaseLevel
lvl
  = LibCaseEnv
env { lc_scruts :: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
lc_scruts = [(CoreBndr, LibCaseLevel, LibCaseLevel)]
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, LibCaseLevel, LibCaseLevel)]
scruts'  = (CoreBndr
scrut_var, LibCaseLevel
bind_lvl, LibCaseLevel
lvl) (CoreBndr, LibCaseLevel, LibCaseLevel)
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
forall a. a -> [a] -> [a]
: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
scruts
    bind_lvl :: LibCaseLevel
bind_lvl = case IdEnv LibCaseLevel -> CoreBndr -> Maybe LibCaseLevel
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv IdEnv LibCaseLevel
lvl_env CoreBndr
scrut_var of
                 Just LibCaseLevel
lvl -> LibCaseLevel
lvl
                 Maybe LibCaseLevel
Nothing  -> LibCaseLevel
topLevel

lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
lookupRecId :: LibCaseEnv -> CoreBndr -> Maybe CoreBind
lookupRecId LibCaseEnv
env 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 -> LibCaseLevel
lookupLevel LibCaseEnv
env CoreBndr
id
  = case IdEnv LibCaseLevel -> CoreBndr -> Maybe LibCaseLevel
forall a. VarEnv a -> CoreBndr -> Maybe a
lookupVarEnv (LibCaseEnv -> IdEnv LibCaseLevel
lc_lvl_env LibCaseEnv
env) CoreBndr
id of
      Just LibCaseLevel
lvl -> LibCaseLevel
lvl
      Maybe LibCaseLevel
Nothing  -> LibCaseLevel
topLevel

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

type LibCaseLevel = Int

topLevel :: LibCaseLevel
topLevel :: LibCaseLevel
topLevel = LibCaseLevel
0

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

        LibCaseEnv -> LibCaseLevel
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 LibCaseLevel
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, LibCaseLevel, LibCaseLevel)]
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 DynFlags
dflags
  = LibCaseEnv :: DynFlags
-> LibCaseLevel
-> IdEnv LibCaseLevel
-> IdEnv CoreBind
-> [(CoreBndr, LibCaseLevel, LibCaseLevel)]
-> LibCaseEnv
LibCaseEnv { lc_dflags :: DynFlags
lc_dflags = DynFlags
dflags,
                 lc_lvl :: LibCaseLevel
lc_lvl = LibCaseLevel
0,
                 lc_lvl_env :: IdEnv LibCaseLevel
lc_lvl_env = IdEnv LibCaseLevel
forall a. VarEnv a
emptyVarEnv,
                 lc_rec_env :: IdEnv CoreBind
lc_rec_env = IdEnv CoreBind
forall a. VarEnv a
emptyVarEnv,
                 lc_scruts :: [(CoreBndr, LibCaseLevel, LibCaseLevel)]
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 LibCaseLevel
bombOutSize = DynFlags -> Maybe LibCaseLevel
liberateCaseThreshold (DynFlags -> Maybe LibCaseLevel)
-> (LibCaseEnv -> DynFlags) -> LibCaseEnv -> Maybe LibCaseLevel
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LibCaseEnv -> DynFlags
lc_dflags