{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

\section{SetLevels}

                ***************************
                        Overview
                ***************************

1. We attach binding levels to Core bindings, in preparation for floating
   outwards (@FloatOut@).

2. We also let-ify many expressions (notably case scrutinees), so they
   will have a fighting chance of being floated sensible.

3. Note [Need for cloning during float-out]
   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   We clone the binders of any floatable let-binding, so that when it is
   floated out it will be unique. Example
      (let x=2 in x) + (let x=3 in x)
   we must clone before floating so we get
      let x1=2 in
      let x2=3 in
      x1+x2

   NOTE: this can't be done using the uniqAway idea, because the variable
         must be unique in the whole program, not just its current scope,
         because two variables in different scopes may float out to the
         same top level place

   NOTE: Very tiresomely, we must apply this substitution to
         the rules stored inside a variable too.

   We do *not* clone top-level bindings, because some of them must not change,
   but we *do* clone bindings that are heading for the top level

4. Note [Binder-swap during float-out]
   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   In the expression
        case x of wild { p -> ...wild... }
   we substitute x for wild in the RHS of the case alternatives:
        case x of wild { p -> ...x... }
   This means that a sub-expression involving x is not "trapped" inside the RHS.
   And it's not inconvenient because we already have a substitution.

  Note that this is EXACTLY BACKWARDS from the what the simplifier does.
  The simplifier tries to get rid of occurrences of x, in favour of wild,
  in the hope that there will only be one remaining occurrence of x, namely
  the scrutinee of the case, and we can inline it.
-}

{-# LANGUAGE CPP, MultiWayIf #-}
module SetLevels (
        setLevels,

        Level(..), LevelType(..), tOP_LEVEL, isJoinCeilLvl, asJoinCeilLvl,
        LevelledBind, LevelledExpr, LevelledBndr,
        FloatSpec(..), floatSpecLevel,

        incMinorLvl, ltMajLvl, ltLvl, isTopLvl
    ) where

#include "HsVersions.h"

import GhcPrelude

import CoreSyn
import CoreMonad        ( FloatOutSwitches(..) )
import CoreUtils        ( exprType, exprIsHNF
                        , exprOkForSpeculation
                        , exprIsTopLevelBindable
                        , isExprLevPoly
                        , collectMakeStaticArgs
                        )
import CoreArity        ( exprBotStrictness_maybe )
import CoreFVs          -- all of it
import CoreSubst
import MkCore           ( sortQuantVars )

import Id
import IdInfo
import Var
import VarSet
import UniqSet          ( nonDetFoldUniqSet )
import UniqDSet         ( getUniqDSet )
import VarEnv
import Literal          ( litIsTrivial )
import Demand           ( StrictSig, Demand, isStrictDmd, splitStrictSig, increaseStrictSigArity )
import Name             ( getOccName, mkSystemVarName )
import OccName          ( occNameString )
import Type             ( Type, mkLamTypes, splitTyConApp_maybe, tyCoVarsOfType
                        , isUnliftedType, closeOverKindsDSet )
import BasicTypes       ( Arity, RecFlag(..), isRec )
import DataCon          ( dataConOrigResTy )
import TysWiredIn
import UniqSupply
import Util
import Outputable
import FastString
import UniqDFM
import FV
import Data.Maybe
import MonadUtils       ( mapAccumLM )

{-
************************************************************************
*                                                                      *
\subsection{Level numbers}
*                                                                      *
************************************************************************
-}

type LevelledExpr = TaggedExpr FloatSpec
type LevelledBind = TaggedBind FloatSpec
type LevelledBndr = TaggedBndr FloatSpec

data Level = Level Int  -- Level number of enclosing lambdas
                   Int  -- Number of big-lambda and/or case expressions and/or
                        -- context boundaries between
                        -- here and the nearest enclosing lambda
                   LevelType -- Binder or join ceiling?
data LevelType = BndrLvl | JoinCeilLvl deriving (LevelType -> LevelType -> Bool
(LevelType -> LevelType -> Bool)
-> (LevelType -> LevelType -> Bool) -> Eq LevelType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: LevelType -> LevelType -> Bool
$c/= :: LevelType -> LevelType -> Bool
== :: LevelType -> LevelType -> Bool
$c== :: LevelType -> LevelType -> Bool
Eq)

data FloatSpec
  = FloatMe Level       -- Float to just inside the binding
                        --    tagged with this level
  | StayPut Level       -- Stay where it is; binding is
                        --     tagged with this level

floatSpecLevel :: FloatSpec -> Level
floatSpecLevel :: FloatSpec -> Level
floatSpecLevel (FloatMe l :: Level
l) = Level
l
floatSpecLevel (StayPut l :: Level
l) = Level
l

{-
The {\em level number} on a (type-)lambda-bound variable is the
nesting depth of the (type-)lambda which binds it.  The outermost lambda
has level 1, so (Level 0 0) means that the variable is bound outside any lambda.

On an expression, it's the maximum level number of its free
(type-)variables.  On a let(rec)-bound variable, it's the level of its
RHS.  On a case-bound variable, it's the number of enclosing lambdas.

Top-level variables: level~0.  Those bound on the RHS of a top-level
definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
as ``subscripts'')...
\begin{verbatim}
a_0 = let  b_? = ...  in
           x_1 = ... b ... in ...
\end{verbatim}

The main function @lvlExpr@ carries a ``context level'' (@le_ctxt_lvl@).
That's meant to be the level number of the enclosing binder in the
final (floated) program.  If the level number of a sub-expression is
less than that of the context, then it might be worth let-binding the
sub-expression so that it will indeed float.

If you can float to level @Level 0 0@ worth doing so because then your
allocation becomes static instead of dynamic.  We always start with
context @Level 0 0@.


Note [FloatOut inside INLINE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
to say "don't float anything out of here".  That's exactly what we
want for the body of an INLINE, where we don't want to float anything
out at all.  See notes with lvlMFE below.

But, check this out:

-- At one time I tried the effect of not floating anything out of an InlineMe,
-- but it sometimes works badly.  For example, consider PrelArr.done.  It
-- has the form         __inline (\d. e)
-- where e doesn't mention d.  If we float this to
--      __inline (let x = e in \d. x)
-- things are bad.  The inliner doesn't even inline it because it doesn't look
-- like a head-normal form.  So it seems a lesser evil to let things float.
-- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
-- which discourages floating out.

So the conclusion is: don't do any floating at all inside an InlineMe.
(In the above example, don't float the {x=e} out of the \d.)

One particular case is that of workers: we don't want to float the
call to the worker outside the wrapper, otherwise the worker might get
inlined into the floated expression, and an importing module won't see
the worker at all.

Note [Join ceiling]
~~~~~~~~~~~~~~~~~~~
Join points can't float very far; too far, and they can't remain join points
So, suppose we have:

  f x = (joinrec j y = ... x ... in jump j x) + 1

One may be tempted to float j out to the top of f's RHS, but then the jump
would not be a tail call. Thus we keep track of a level called the *join
ceiling* past which join points are not allowed to float.

The troublesome thing is that, unlike most levels to which something might
float, there is not necessarily an identifier to which the join ceiling is
attached. Fortunately, if something is to be floated to a join ceiling, it must
be dropped at the *nearest* join ceiling. Thus each level is marked as to
whether it is a join ceiling, so that FloatOut can tell which binders are being
floated to the nearest join ceiling and which to a particular binder (or set of
binders).
-}

instance Outputable FloatSpec where
  ppr :: FloatSpec -> SDoc
ppr (FloatMe l :: Level
l) = Char -> SDoc
char 'F' SDoc -> SDoc -> SDoc
<> Level -> SDoc
forall a. Outputable a => a -> SDoc
ppr Level
l
  ppr (StayPut l :: Level
l) = Level -> SDoc
forall a. Outputable a => a -> SDoc
ppr Level
l

tOP_LEVEL :: Level
tOP_LEVEL :: Level
tOP_LEVEL   = Int -> Int -> LevelType -> Level
Level 0 0 LevelType
BndrLvl

incMajorLvl :: Level -> Level
incMajorLvl :: Level -> Level
incMajorLvl (Level major :: Int
major _ _) = Int -> Int -> LevelType -> Level
Level (Int
major Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) 0 LevelType
BndrLvl

incMinorLvl :: Level -> Level
incMinorLvl :: Level -> Level
incMinorLvl (Level major :: Int
major minor :: Int
minor _) = Int -> Int -> LevelType -> Level
Level Int
major (Int
minorInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) LevelType
BndrLvl

asJoinCeilLvl :: Level -> Level
asJoinCeilLvl :: Level -> Level
asJoinCeilLvl (Level major :: Int
major minor :: Int
minor _) = Int -> Int -> LevelType -> Level
Level Int
major Int
minor LevelType
JoinCeilLvl

maxLvl :: Level -> Level -> Level
maxLvl :: Level -> Level -> Level
maxLvl l1 :: Level
l1@(Level maj1 :: Int
maj1 min1 :: Int
min1 _) l2 :: Level
l2@(Level maj2 :: Int
maj2 min2 :: Int
min2 _)
  | (Int
maj1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maj2) Bool -> Bool -> Bool
|| (Int
maj1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maj2 Bool -> Bool -> Bool
&& Int
min1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
min2) = Level
l1
  | Bool
otherwise                                      = Level
l2

ltLvl :: Level -> Level -> Bool
ltLvl :: Level -> Level -> Bool
ltLvl (Level maj1 :: Int
maj1 min1 :: Int
min1 _) (Level maj2 :: Int
maj2 min2 :: Int
min2 _)
  = (Int
maj1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maj2) Bool -> Bool -> Bool
|| (Int
maj1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maj2 Bool -> Bool -> Bool
&& Int
min1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
min2)

ltMajLvl :: Level -> Level -> Bool
    -- Tells if one level belongs to a difft *lambda* level to another
ltMajLvl :: Level -> Level -> Bool
ltMajLvl (Level maj1 :: Int
maj1 _ _) (Level maj2 :: Int
maj2 _ _) = Int
maj1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maj2

isTopLvl :: Level -> Bool
isTopLvl :: Level -> Bool
isTopLvl (Level 0 0 _) = Bool
True
isTopLvl _             = Bool
False

isJoinCeilLvl :: Level -> Bool
isJoinCeilLvl :: Level -> Bool
isJoinCeilLvl (Level _ _ t :: LevelType
t) = LevelType
t LevelType -> LevelType -> Bool
forall a. Eq a => a -> a -> Bool
== LevelType
JoinCeilLvl

instance Outputable Level where
  ppr :: Level -> SDoc
ppr (Level maj :: Int
maj min :: Int
min typ :: LevelType
typ)
    = [SDoc] -> SDoc
hcat [ Char -> SDoc
char '<', Int -> SDoc
int Int
maj, Char -> SDoc
char ',', Int -> SDoc
int Int
min, Char -> SDoc
char '>'
           , Bool -> SDoc -> SDoc
ppWhen (LevelType
typ LevelType -> LevelType -> Bool
forall a. Eq a => a -> a -> Bool
== LevelType
JoinCeilLvl) (Char -> SDoc
char 'C') ]

instance Eq Level where
  (Level maj1 :: Int
maj1 min1 :: Int
min1 _) == :: Level -> Level -> Bool
== (Level maj2 :: Int
maj2 min2 :: Int
min2 _) = Int
maj1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maj2 Bool -> Bool -> Bool
&& Int
min1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
min2

{-
************************************************************************
*                                                                      *
\subsection{Main level-setting code}
*                                                                      *
************************************************************************
-}

setLevels :: FloatOutSwitches
          -> CoreProgram
          -> UniqSupply
          -> [LevelledBind]

setLevels :: FloatOutSwitches -> CoreProgram -> UniqSupply -> [LevelledBind]
setLevels float_lams :: FloatOutSwitches
float_lams binds :: CoreProgram
binds us :: UniqSupply
us
  = UniqSupply -> UniqSM [LevelledBind] -> [LevelledBind]
forall a. UniqSupply -> UniqSM a -> a
initLvl UniqSupply
us (LevelEnv -> CoreProgram -> UniqSM [LevelledBind]
do_them LevelEnv
init_env CoreProgram
binds)
  where
    init_env :: LevelEnv
init_env = FloatOutSwitches -> LevelEnv
initialEnv FloatOutSwitches
float_lams

    do_them :: LevelEnv -> [CoreBind] -> LvlM [LevelledBind]
    do_them :: LevelEnv -> CoreProgram -> UniqSM [LevelledBind]
do_them _ [] = [LevelledBind] -> UniqSM [LevelledBind]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    do_them env :: LevelEnv
env (b :: CoreBind
b:bs :: CoreProgram
bs)
      = do { (lvld_bind :: LevelledBind
lvld_bind, env' :: LevelEnv
env') <- LevelEnv -> CoreBind -> LvlM (LevelledBind, LevelEnv)
lvlTopBind LevelEnv
env CoreBind
b
           ; [LevelledBind]
lvld_binds <- LevelEnv -> CoreProgram -> UniqSM [LevelledBind]
do_them LevelEnv
env' CoreProgram
bs
           ; [LevelledBind] -> UniqSM [LevelledBind]
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBind
lvld_bind LevelledBind -> [LevelledBind] -> [LevelledBind]
forall a. a -> [a] -> [a]
: [LevelledBind]
lvld_binds) }

lvlTopBind :: LevelEnv -> Bind Id -> LvlM (LevelledBind, LevelEnv)
lvlTopBind :: LevelEnv -> CoreBind -> LvlM (LevelledBind, LevelEnv)
lvlTopBind env :: LevelEnv
env (NonRec bndr :: Id
bndr rhs :: Expr Id
rhs)
  = do { LevelledExpr
rhs' <- LevelEnv -> RecFlag -> Id -> Expr Id -> LvlM LevelledExpr
lvl_top LevelEnv
env RecFlag
NonRecursive Id
bndr Expr Id
rhs
       ; let (env' :: LevelEnv
env', [bndr' :: LevelledBndr
bndr']) = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
NonRecursive LevelEnv
env Level
tOP_LEVEL [Id
bndr]
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec LevelledBndr
bndr' LevelledExpr
rhs', LevelEnv
env') }

lvlTopBind env :: LevelEnv
env (Rec pairs :: [(Id, Expr Id)]
pairs)
  = do { let (env' :: LevelEnv
env', bndrs' :: [LevelledBndr]
bndrs') = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
Recursive LevelEnv
env Level
tOP_LEVEL
                                               (((Id, Expr Id) -> Id) -> [(Id, Expr Id)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, Expr Id) -> Id
forall a b. (a, b) -> a
fst [(Id, Expr Id)]
pairs)
       ; [LevelledExpr]
rhss' <- ((Id, Expr Id) -> LvlM LevelledExpr)
-> [(Id, Expr Id)] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(b :: Id
b,r :: Expr Id
r) -> LevelEnv -> RecFlag -> Id -> Expr Id -> LvlM LevelledExpr
lvl_top LevelEnv
env' RecFlag
Recursive Id
b Expr Id
r) [(Id, Expr Id)]
pairs
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec ([LevelledBndr]
bndrs' [LevelledBndr] -> [LevelledExpr] -> [(LevelledBndr, LevelledExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [LevelledExpr]
rhss'), LevelEnv
env') }

lvl_top :: LevelEnv -> RecFlag -> Id -> CoreExpr -> LvlM LevelledExpr
lvl_top :: LevelEnv -> RecFlag -> Id -> Expr Id -> LvlM LevelledExpr
lvl_top env :: LevelEnv
env is_rec :: RecFlag
is_rec bndr :: Id
bndr rhs :: Expr Id
rhs
  = LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlRhs LevelEnv
env RecFlag
is_rec
           (Id -> Bool
isBottomingId Id
bndr)
           Maybe Int
forall a. Maybe a
Nothing  -- Not a join point
           (Expr Id -> CoreExprWithFVs
freeVars Expr Id
rhs)

{-
************************************************************************
*                                                                      *
\subsection{Setting expression levels}
*                                                                      *
************************************************************************

Note [Floating over-saturated applications]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If we see (f x y), and (f x) is a redex (ie f's arity is 1),
we call (f x) an "over-saturated application"

Should we float out an over-sat app, if can escape a value lambda?
It is sometimes very beneficial (-7% runtime -4% alloc over nofib -O2).
But we don't want to do it for class selectors, because the work saved
is minimal, and the extra local thunks allocated cost money.

Arguably we could float even class-op applications if they were going to
top level -- but then they must be applied to a constant dictionary and
will almost certainly be optimised away anyway.
-}

lvlExpr :: LevelEnv             -- Context
        -> CoreExprWithFVs      -- Input expression
        -> LvlM LevelledExpr    -- Result expression

{-
The @le_ctxt_lvl@ is, roughly, the level of the innermost enclosing
binder.  Here's an example

        v = \x -> ...\y -> let r = case (..x..) of
                                        ..x..
                           in ..

When looking at the rhs of @r@, @le_ctxt_lvl@ will be 1 because that's
the level of @r@, even though it's inside a level-2 @\y@.  It's
important that @le_ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
--- because it isn't a *maximal* free expression.

If there were another lambda in @r@'s rhs, it would get level-2 as well.
-}

lvlExpr :: LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr env :: LevelEnv
env (_, AnnType ty :: Type
ty)     = LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> LevelledExpr
forall b. Type -> Expr b
Type (Subst -> Type -> Type
CoreSubst.substTy (LevelEnv -> Subst
le_subst LevelEnv
env) Type
ty))
lvlExpr env :: LevelEnv
env (_, AnnCoercion co :: Coercion
co) = LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Coercion -> LevelledExpr
forall b. Coercion -> Expr b
Coercion (Subst -> Coercion -> Coercion
substCo (LevelEnv -> Subst
le_subst LevelEnv
env) Coercion
co))
lvlExpr env :: LevelEnv
env (_, AnnVar v :: Id
v)       = LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelEnv -> Id -> LevelledExpr
lookupVar LevelEnv
env Id
v)
lvlExpr _   (_, AnnLit lit :: Literal
lit)     = LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Literal -> LevelledExpr
forall b. Literal -> Expr b
Lit Literal
lit)

lvlExpr env :: LevelEnv
env (_, AnnCast expr :: CoreExprWithFVs
expr (_, co :: Coercion
co)) = do
    LevelledExpr
expr' <- LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailExpr LevelEnv
env CoreExprWithFVs
expr
    LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledExpr -> Coercion -> LevelledExpr
forall b. Expr b -> Coercion -> Expr b
Cast LevelledExpr
expr' (Subst -> Coercion -> Coercion
substCo (LevelEnv -> Subst
le_subst LevelEnv
env) Coercion
co))

lvlExpr env :: LevelEnv
env (_, AnnTick tickish :: Tickish Id
tickish expr :: CoreExprWithFVs
expr) = do
    LevelledExpr
expr' <- LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailExpr LevelEnv
env CoreExprWithFVs
expr
    let tickish' :: Tickish Id
tickish' = Subst -> Tickish Id -> Tickish Id
substTickish (LevelEnv -> Subst
le_subst LevelEnv
env) Tickish Id
tickish
    LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Tickish Id -> LevelledExpr -> LevelledExpr
forall b. Tickish Id -> Expr b -> Expr b
Tick Tickish Id
tickish' LevelledExpr
expr')

lvlExpr env :: LevelEnv
env expr :: CoreExprWithFVs
expr@(_, AnnApp _ _) = LevelEnv
-> CoreExprWithFVs
-> (CoreExprWithFVs, [CoreExprWithFVs])
-> LvlM LevelledExpr
lvlApp LevelEnv
env CoreExprWithFVs
expr (CoreExprWithFVs -> (CoreExprWithFVs, [CoreExprWithFVs])
forall b a. AnnExpr b a -> (AnnExpr b a, [AnnExpr b a])
collectAnnArgs CoreExprWithFVs
expr)

-- We don't split adjacent lambdas.  That is, given
--      \x y -> (x+1,y)
-- we don't float to give
--      \x -> let v = x+1 in \y -> (v,y)
-- Why not?  Because partial applications are fairly rare, and splitting
-- lambdas makes them more expensive.

lvlExpr env :: LevelEnv
env expr :: CoreExprWithFVs
expr@(_, AnnLam {})
  = do { LevelledExpr
new_body <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE LevelEnv
new_env Bool
True CoreExprWithFVs
body
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ([LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
new_bndrs LevelledExpr
new_body) }
  where
    (bndrs :: [Id]
bndrs, body :: CoreExprWithFVs
body)        = CoreExprWithFVs -> ([Id], CoreExprWithFVs)
forall bndr annot.
AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectAnnBndrs CoreExprWithFVs
expr
    (env1 :: LevelEnv
env1, bndrs1 :: [Id]
bndrs1)       = RecFlag -> LevelEnv -> [Id] -> (LevelEnv, [Id])
substBndrsSL RecFlag
NonRecursive LevelEnv
env [Id]
bndrs
    (new_env :: LevelEnv
new_env, new_bndrs :: [LevelledBndr]
new_bndrs) = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs LevelEnv
env1 (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env) [Id]
bndrs1
        -- At one time we called a special version of collectBinders,
        -- which ignored coercions, because we don't want to split
        -- a lambda like this (\x -> coerce t (\s -> ...))
        -- This used to happen quite a bit in state-transformer programs,
        -- but not nearly so much now non-recursive newtypes are transparent.
        -- [See SetLevels rev 1.50 for a version with this approach.]

lvlExpr env :: LevelEnv
env (_, AnnLet bind :: AnnBind Id FVAnn
bind body :: CoreExprWithFVs
body)
  = do { (bind' :: LevelledBind
bind', new_env :: LevelEnv
new_env) <- LevelEnv -> AnnBind Id FVAnn -> LvlM (LevelledBind, LevelEnv)
lvlBind LevelEnv
env AnnBind Id FVAnn
bind
       ; LevelledExpr
body' <- LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
new_env CoreExprWithFVs
body
           -- No point in going via lvlMFE here.  If the binding is alive
           -- (mentioned in body), and the whole let-expression doesn't
           -- float, then neither will the body
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBind -> LevelledExpr -> LevelledExpr
forall b. Bind b -> Expr b -> Expr b
Let LevelledBind
bind' LevelledExpr
body') }

lvlExpr env :: LevelEnv
env (_, AnnCase scrut :: CoreExprWithFVs
scrut case_bndr :: Id
case_bndr ty :: Type
ty alts :: [AnnAlt Id FVAnn]
alts)
  = do { LevelledExpr
scrut' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE LevelEnv
env Bool
True CoreExprWithFVs
scrut
       ; LevelEnv
-> FVAnn
-> LevelledExpr
-> Id
-> Type
-> [AnnAlt Id FVAnn]
-> LvlM LevelledExpr
lvlCase LevelEnv
env (CoreExprWithFVs -> FVAnn
freeVarsOf CoreExprWithFVs
scrut) LevelledExpr
scrut' Id
case_bndr Type
ty [AnnAlt Id FVAnn]
alts }

lvlNonTailExpr :: LevelEnv             -- Context
               -> CoreExprWithFVs      -- Input expression
               -> LvlM LevelledExpr    -- Result expression
lvlNonTailExpr :: LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailExpr env :: LevelEnv
env expr :: CoreExprWithFVs
expr
  = LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr (LevelEnv -> LevelEnv
placeJoinCeiling LevelEnv
env) CoreExprWithFVs
expr

-------------------------------------------
lvlApp :: LevelEnv
       -> CoreExprWithFVs
       -> (CoreExprWithFVs, [CoreExprWithFVs]) -- Input application
        -> LvlM LevelledExpr                   -- Result expression
lvlApp :: LevelEnv
-> CoreExprWithFVs
-> (CoreExprWithFVs, [CoreExprWithFVs])
-> LvlM LevelledExpr
lvlApp env :: LevelEnv
env orig_expr :: CoreExprWithFVs
orig_expr ((_,AnnVar fn :: Id
fn), args :: [CoreExprWithFVs]
args)
  | LevelEnv -> Bool
floatOverSat LevelEnv
env   -- See Note [Floating over-saturated applications]
  , Int
arity Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0
  , Int
arity Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n_val_args
  , Maybe Class
Nothing <- Id -> Maybe Class
isClassOpId_maybe Id
fn
  =  do { [LevelledExpr]
rargs' <- (CoreExprWithFVs -> LvlM LevelledExpr)
-> [CoreExprWithFVs] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE LevelEnv
env Bool
False) [CoreExprWithFVs]
rargs
        ; LevelledExpr
lapp'  <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE LevelEnv
env Bool
False CoreExprWithFVs
lapp
        ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ((LevelledExpr -> LevelledExpr -> LevelledExpr)
-> LevelledExpr -> [LevelledExpr] -> LevelledExpr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' LevelledExpr -> LevelledExpr -> LevelledExpr
forall b. Expr b -> Expr b -> Expr b
App LevelledExpr
lapp' [LevelledExpr]
rargs') }

  | Bool
otherwise
  = do { (_, args' :: [LevelledExpr]
args') <- ([Demand] -> CoreExprWithFVs -> UniqSM ([Demand], LevelledExpr))
-> [Demand]
-> [CoreExprWithFVs]
-> UniqSM ([Demand], [LevelledExpr])
forall (m :: * -> *) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM [Demand] -> CoreExprWithFVs -> UniqSM ([Demand], LevelledExpr)
lvl_arg [Demand]
stricts [CoreExprWithFVs]
args
            -- Take account of argument strictness; see
            -- Note [Floating to the top]
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ((LevelledExpr -> LevelledExpr -> LevelledExpr)
-> LevelledExpr -> [LevelledExpr] -> LevelledExpr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' LevelledExpr -> LevelledExpr -> LevelledExpr
forall b. Expr b -> Expr b -> Expr b
App (LevelEnv -> Id -> LevelledExpr
lookupVar LevelEnv
env Id
fn) [LevelledExpr]
args') }
  where
    n_val_args :: Int
n_val_args = (CoreExprWithFVs -> Bool) -> [CoreExprWithFVs] -> Int
forall a. (a -> Bool) -> [a] -> Int
count (Expr Id -> Bool
forall b. Expr b -> Bool
isValArg (Expr Id -> Bool)
-> (CoreExprWithFVs -> Expr Id) -> CoreExprWithFVs -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreExprWithFVs -> Expr Id
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate) [CoreExprWithFVs]
args
    arity :: Int
arity      = Id -> Int
idArity Id
fn

    stricts :: [Demand]   -- True for strict /value/ arguments
    stricts :: [Demand]
stricts = case StrictSig -> ([Demand], DmdResult)
splitStrictSig (Id -> StrictSig
idStrictness Id
fn) of
                (arg_ds :: [Demand]
arg_ds, _) | [Demand]
arg_ds [Demand] -> Int -> Bool
forall a. [a] -> Int -> Bool
`lengthExceeds` Int
n_val_args
                            -> []
                            | Bool
otherwise
                            -> [Demand]
arg_ds

    -- Separate out the PAP that we are floating from the extra
    -- arguments, by traversing the spine until we have collected
    -- (n_val_args - arity) value arguments.
    (lapp :: CoreExprWithFVs
lapp, rargs :: [CoreExprWithFVs]
rargs) = Int
-> CoreExprWithFVs
-> [CoreExprWithFVs]
-> (CoreExprWithFVs, [CoreExprWithFVs])
forall t b annot.
(Eq t, Num t) =>
t
-> AnnExpr b annot
-> [AnnExpr b annot]
-> (AnnExpr b annot, [AnnExpr b annot])
left (Int
n_val_args Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
arity) CoreExprWithFVs
orig_expr []

    left :: t
-> AnnExpr b annot
-> [AnnExpr b annot]
-> (AnnExpr b annot, [AnnExpr b annot])
left 0 e :: AnnExpr b annot
e               rargs :: [AnnExpr b annot]
rargs = (AnnExpr b annot
e, [AnnExpr b annot]
rargs)
    left n :: t
n (_, AnnApp f :: AnnExpr b annot
f a :: AnnExpr b annot
a) rargs :: [AnnExpr b annot]
rargs
       | Expr b -> Bool
forall b. Expr b -> Bool
isValArg (AnnExpr b annot -> Expr b
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate AnnExpr b annot
a) = t
-> AnnExpr b annot
-> [AnnExpr b annot]
-> (AnnExpr b annot, [AnnExpr b annot])
left (t
nt -> t -> t
forall a. Num a => a -> a -> a
-1) AnnExpr b annot
f (AnnExpr b annot
aAnnExpr b annot -> [AnnExpr b annot] -> [AnnExpr b annot]
forall a. a -> [a] -> [a]
:[AnnExpr b annot]
rargs)
       | Bool
otherwise               = t
-> AnnExpr b annot
-> [AnnExpr b annot]
-> (AnnExpr b annot, [AnnExpr b annot])
left t
n     AnnExpr b annot
f (AnnExpr b annot
aAnnExpr b annot -> [AnnExpr b annot] -> [AnnExpr b annot]
forall a. a -> [a] -> [a]
:[AnnExpr b annot]
rargs)
    left _ _ _                   = String -> (AnnExpr b annot, [AnnExpr b annot])
forall a. String -> a
panic "SetLevels.lvlExpr.left"

    is_val_arg :: CoreExprWithFVs -> Bool
    is_val_arg :: CoreExprWithFVs -> Bool
is_val_arg (_, AnnType {}) = Bool
False
    is_val_arg _               = Bool
True

    lvl_arg :: [Demand] -> CoreExprWithFVs -> LvlM ([Demand], LevelledExpr)
    lvl_arg :: [Demand] -> CoreExprWithFVs -> UniqSM ([Demand], LevelledExpr)
lvl_arg strs :: [Demand]
strs arg :: CoreExprWithFVs
arg | (str1 :: Demand
str1 : strs' :: [Demand]
strs') <- [Demand]
strs
                     , CoreExprWithFVs -> Bool
is_val_arg CoreExprWithFVs
arg
                     = do { LevelledExpr
arg' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
env (Demand -> Bool
forall s u. JointDmd (Str s) (Use u) -> Bool
isStrictDmd Demand
str1) CoreExprWithFVs
arg
                          ; ([Demand], LevelledExpr) -> UniqSM ([Demand], LevelledExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Demand]
strs', LevelledExpr
arg') }
                     | Bool
otherwise
                     = do { LevelledExpr
arg' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
env Bool
False CoreExprWithFVs
arg
                          ; ([Demand], LevelledExpr) -> UniqSM ([Demand], LevelledExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Demand]
strs, LevelledExpr
arg') }

lvlApp env :: LevelEnv
env _ (fun :: CoreExprWithFVs
fun, args :: [CoreExprWithFVs]
args)
  =  -- No PAPs that we can float: just carry on with the
     -- arguments and the function.
     do { [LevelledExpr]
args' <- (CoreExprWithFVs -> LvlM LevelledExpr)
-> [CoreExprWithFVs] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE LevelEnv
env Bool
False) [CoreExprWithFVs]
args
        ; LevelledExpr
fun'  <- LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailExpr LevelEnv
env CoreExprWithFVs
fun
        ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ((LevelledExpr -> LevelledExpr -> LevelledExpr)
-> LevelledExpr -> [LevelledExpr] -> LevelledExpr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' LevelledExpr -> LevelledExpr -> LevelledExpr
forall b. Expr b -> Expr b -> Expr b
App LevelledExpr
fun' [LevelledExpr]
args') }

-------------------------------------------
lvlCase :: LevelEnv             -- Level of in-scope names/tyvars
        -> DVarSet              -- Free vars of input scrutinee
        -> LevelledExpr         -- Processed scrutinee
        -> Id -> Type           -- Case binder and result type
        -> [CoreAltWithFVs]     -- Input alternatives
        -> LvlM LevelledExpr    -- Result expression
lvlCase :: LevelEnv
-> FVAnn
-> LevelledExpr
-> Id
-> Type
-> [AnnAlt Id FVAnn]
-> LvlM LevelledExpr
lvlCase env :: LevelEnv
env scrut_fvs :: FVAnn
scrut_fvs scrut' :: LevelledExpr
scrut' case_bndr :: Id
case_bndr ty :: Type
ty alts :: [AnnAlt Id FVAnn]
alts
  -- See Note [Floating single-alternative cases]
  | [(con :: AltCon
con@(DataAlt {}), bs :: [Id]
bs, body :: CoreExprWithFVs
body)] <- [AnnAlt Id FVAnn]
alts
  , Expr Id -> Bool
exprIsHNF (LevelledExpr -> Expr Id
forall t. TaggedExpr t -> Expr Id
deTagExpr LevelledExpr
scrut')  -- See Note [Check the output scrutinee for exprIsHNF]
  , Bool -> Bool
not (Level -> Bool
isTopLvl Level
dest_lvl)       -- Can't have top-level cases
  , Bool -> Bool
not (LevelEnv -> Bool
floatTopLvlOnly LevelEnv
env)     -- Can float anywhere
  =     -- Always float the case if possible
        -- Unlike lets we don't insist that it escapes a value lambda
    do { (env1 :: LevelEnv
env1, (case_bndr' :: Id
case_bndr' : bs' :: [Id]
bs')) <- LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneCaseBndrs LevelEnv
env Level
dest_lvl (Id
case_bndr Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
: [Id]
bs)
       ; let rhs_env :: LevelEnv
rhs_env = LevelEnv -> Id -> LevelledExpr -> LevelEnv
extendCaseBndrEnv LevelEnv
env1 Id
case_bndr LevelledExpr
scrut'
       ; LevelledExpr
body' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
rhs_env Bool
True CoreExprWithFVs
body
       ; let alt' :: (AltCon, [LevelledBndr], LevelledExpr)
alt' = (AltCon
con, (Id -> LevelledBndr) -> [Id] -> [LevelledBndr]
forall a b. (a -> b) -> [a] -> [b]
map (Level -> Id -> LevelledBndr
stayPut Level
dest_lvl) [Id]
bs', LevelledExpr
body')
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledExpr
-> LevelledBndr
-> Type
-> [(AltCon, [LevelledBndr], LevelledExpr)]
-> LevelledExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case LevelledExpr
scrut' (Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
case_bndr' (Level -> FloatSpec
FloatMe Level
dest_lvl)) Type
ty' [(AltCon, [LevelledBndr], LevelledExpr)
alt']) }

  | Bool
otherwise     -- Stays put
  = do { let (alts_env1 :: LevelEnv
alts_env1, [case_bndr' :: LevelledBndr
case_bndr']) = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
NonRecursive LevelEnv
env Level
incd_lvl [Id
case_bndr]
             alts_env :: LevelEnv
alts_env = LevelEnv -> Id -> LevelledExpr -> LevelEnv
extendCaseBndrEnv LevelEnv
alts_env1 Id
case_bndr LevelledExpr
scrut'
       ; [(AltCon, [LevelledBndr], LevelledExpr)]
alts' <- (AnnAlt Id FVAnn -> UniqSM (AltCon, [LevelledBndr], LevelledExpr))
-> [AnnAlt Id FVAnn]
-> UniqSM [(AltCon, [LevelledBndr], LevelledExpr)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LevelEnv
-> AnnAlt Id FVAnn -> UniqSM (AltCon, [LevelledBndr], LevelledExpr)
forall a.
LevelEnv
-> (a, [Id], CoreExprWithFVs)
-> UniqSM (a, [LevelledBndr], LevelledExpr)
lvl_alt LevelEnv
alts_env) [AnnAlt Id FVAnn]
alts
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledExpr
-> LevelledBndr
-> Type
-> [(AltCon, [LevelledBndr], LevelledExpr)]
-> LevelledExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case LevelledExpr
scrut' LevelledBndr
case_bndr' Type
ty' [(AltCon, [LevelledBndr], LevelledExpr)]
alts') }
  where
    ty' :: Type
ty' = Subst -> Type -> Type
substTy (LevelEnv -> Subst
le_subst LevelEnv
env) Type
ty

    incd_lvl :: Level
incd_lvl = Level -> Level
incMinorLvl (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)
    dest_lvl :: Level
dest_lvl = (Id -> Bool) -> LevelEnv -> FVAnn -> Level
maxFvLevel (Bool -> Id -> Bool
forall a b. a -> b -> a
const Bool
True) LevelEnv
env FVAnn
scrut_fvs
            -- Don't abstract over type variables, hence const True

    lvl_alt :: LevelEnv
-> (a, [Id], CoreExprWithFVs)
-> UniqSM (a, [LevelledBndr], LevelledExpr)
lvl_alt alts_env :: LevelEnv
alts_env (con :: a
con, bs :: [Id]
bs, rhs :: CoreExprWithFVs
rhs)
      = do { LevelledExpr
rhs' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
new_env Bool
True CoreExprWithFVs
rhs
           ; (a, [LevelledBndr], LevelledExpr)
-> UniqSM (a, [LevelledBndr], LevelledExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
con, [LevelledBndr]
bs', LevelledExpr
rhs') }
      where
        (new_env :: LevelEnv
new_env, bs' :: [LevelledBndr]
bs') = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
NonRecursive LevelEnv
alts_env Level
incd_lvl [Id]
bs

{- Note [Floating single-alternative cases]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
  data T a = MkT !a
  f :: T Int -> blah
  f x vs = case x of { MkT y ->
             let f vs = ...(case y of I# w -> e)...f..
             in f vs

Here we can float the (case y ...) out, because y is sure
to be evaluated, to give
  f x vs = case x of { MkT y ->
           case y of I# w ->
             let f vs = ...(e)...f..
             in f vs

That saves unboxing it every time round the loop.  It's important in
some DPH stuff where we really want to avoid that repeated unboxing in
the inner loop.

Things to note:

 * The test we perform is exprIsHNF, and /not/ exprOkForSpeculation.

     - exrpIsHNF catches the key case of an evaluated variable

     - exprOkForSpeculation is /false/ of an evaluated variable;
       See Note [exprOkForSpeculation and evaluated variables] in CoreUtils
       So we'd actually miss the key case!

     - Nothing is gained from the extra generality of exprOkForSpeculation
       since we only consider floating a case whose single alternative
       is a DataAlt   K a b -> rhs

 * We can't float a case to top level

 * It's worth doing this float even if we don't float
   the case outside a value lambda.  Example
     case x of {
       MkT y -> (case y of I# w2 -> ..., case y of I# w2 -> ...)
   If we floated the cases out we could eliminate one of them.

 * We only do this with a single-alternative case

Note [Check the output scrutinee for exprIsHNF]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
  case x of y {
    A -> ....(case y of alts)....
  }

Because of the binder-swap, the inner case will get substituted to
(case x of ..).  So when testing whether the scrutinee is in HNF we
must be careful to test the *result* scrutinee ('x' in this case), not
the *input* one 'y'.  The latter *is* in HNF here (because y is
evaluated), but the former is not -- and indeed we can't float the
inner case out, at least not unless x is also evaluated at its binding
site.  See Trac #5453.

That's why we apply exprIsHNF to scrut' and not to scrut.

See Note [Floating single-alternative cases] for why
we use exprIsHNF in the first place.
-}

lvlNonTailMFE :: LevelEnv             -- Level of in-scope names/tyvars
              -> Bool                 -- True <=> strict context [body of case
                                      --   or let]
              -> CoreExprWithFVs      -- input expression
              -> LvlM LevelledExpr    -- Result expression
lvlNonTailMFE :: LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlNonTailMFE env :: LevelEnv
env strict_ctxt :: Bool
strict_ctxt ann_expr :: CoreExprWithFVs
ann_expr
  = LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE (LevelEnv -> LevelEnv
placeJoinCeiling LevelEnv
env) Bool
strict_ctxt CoreExprWithFVs
ann_expr

lvlMFE ::  LevelEnv             -- Level of in-scope names/tyvars
        -> Bool                 -- True <=> strict context [body of case or let]
        -> CoreExprWithFVs      -- input expression
        -> LvlM LevelledExpr    -- Result expression
-- lvlMFE is just like lvlExpr, except that it might let-bind
-- the expression, so that it can itself be floated.

lvlMFE :: LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE env :: LevelEnv
env _ (_, AnnType ty :: Type
ty)
  = LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> LevelledExpr
forall b. Type -> Expr b
Type (Subst -> Type -> Type
CoreSubst.substTy (LevelEnv -> Subst
le_subst LevelEnv
env) Type
ty))

-- No point in floating out an expression wrapped in a coercion or note
-- If we do we'll transform  lvl = e |> co
--                       to  lvl' = e; lvl = lvl' |> co
-- and then inline lvl.  Better just to float out the payload.
lvlMFE env :: LevelEnv
env strict_ctxt :: Bool
strict_ctxt (_, AnnTick t :: Tickish Id
t e :: CoreExprWithFVs
e)
  = do { LevelledExpr
e' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
env Bool
strict_ctxt CoreExprWithFVs
e
       ; let t' :: Tickish Id
t' = Subst -> Tickish Id -> Tickish Id
substTickish (LevelEnv -> Subst
le_subst LevelEnv
env) Tickish Id
t
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Tickish Id -> LevelledExpr -> LevelledExpr
forall b. Tickish Id -> Expr b -> Expr b
Tick Tickish Id
t' LevelledExpr
e') }

lvlMFE env :: LevelEnv
env strict_ctxt :: Bool
strict_ctxt (_, AnnCast e :: CoreExprWithFVs
e (_, co :: Coercion
co))
  = do  { LevelledExpr
e' <- LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE LevelEnv
env Bool
strict_ctxt CoreExprWithFVs
e
        ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledExpr -> Coercion -> LevelledExpr
forall b. Expr b -> Coercion -> Expr b
Cast LevelledExpr
e' (Subst -> Coercion -> Coercion
substCo (LevelEnv -> Subst
le_subst LevelEnv
env) Coercion
co)) }

lvlMFE env :: LevelEnv
env strict_ctxt :: Bool
strict_ctxt e :: CoreExprWithFVs
e@(_, AnnCase {})
  | Bool
strict_ctxt       -- Don't share cases in a strict context
  = LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
env CoreExprWithFVs
e     -- See Note [Case MFEs]

lvlMFE env :: LevelEnv
env strict_ctxt :: Bool
strict_ctxt ann_expr :: CoreExprWithFVs
ann_expr
  |  LevelEnv -> Bool
floatTopLvlOnly LevelEnv
env Bool -> Bool -> Bool
&& Bool -> Bool
not (Level -> Bool
isTopLvl Level
dest_lvl)
         -- Only floating to the top level is allowed.
  Bool -> Bool -> Bool
|| (Id -> Bool) -> FVAnn -> Bool
anyDVarSet Id -> Bool
isJoinId FVAnn
fvs   -- If there is a free join, don't float
                               -- See Note [Free join points]
  Bool -> Bool -> Bool
|| Expr Id -> Bool
isExprLevPoly Expr Id
expr
         -- We can't let-bind levity polymorphic expressions
         -- See Note [Levity polymorphism invariants] in CoreSyn
  Bool -> Bool -> Bool
|| Expr Id -> [Id] -> Bool
notWorthFloating Expr Id
expr [Id]
abs_vars
  Bool -> Bool -> Bool
|| Bool -> Bool
not Bool
float_me
  =     -- Don't float it out
    LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
env CoreExprWithFVs
ann_expr

  |  Bool
float_is_new_lam Bool -> Bool -> Bool
|| Expr Id -> Type -> Bool
exprIsTopLevelBindable Expr Id
expr Type
expr_ty
         -- No wrapping needed if the type is lifted, or is a literal string
         -- or if we are wrapping it in one or more value lambdas
  = do { LevelledExpr
expr1 <- [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs [Id]
abs_vars Level
dest_lvl LevelEnv
rhs_env RecFlag
NonRecursive
                              (Maybe (Int, StrictSig) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Int, StrictSig)
mb_bot_str)
                              Maybe Int
forall a. Maybe a
join_arity_maybe
                              CoreExprWithFVs
ann_expr
                  -- Treat the expr just like a right-hand side
       ; Id
var <- LevelledExpr -> Maybe Int -> Bool -> LvlM Id
newLvlVar LevelledExpr
expr1 Maybe Int
forall a. Maybe a
join_arity_maybe Bool
is_mk_static
       ; let var2 :: Id
var2 = Id -> Int -> Maybe (Int, StrictSig) -> Id
annotateBotStr Id
var Int
float_n_lams Maybe (Int, StrictSig)
mb_bot_str
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBind -> LevelledExpr -> LevelledExpr
forall b. Bind b -> Expr b -> Expr b
Let (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
var2 (Level -> FloatSpec
FloatMe Level
dest_lvl)) LevelledExpr
expr1)
                     (LevelledExpr -> [Id] -> LevelledExpr
forall b. Expr b -> [Id] -> Expr b
mkVarApps (Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
var2) [Id]
abs_vars)) }

  -- OK, so the float has an unlifted type (not top-level bindable)
  --     and no new value lambdas (float_is_new_lam is False)
  -- Try for the boxing strategy
  -- See Note [Floating MFEs of unlifted type]
  | Bool
escapes_value_lam
  , Bool -> Bool
not Bool
expr_ok_for_spec -- Boxing/unboxing isn't worth it for cheap expressions
                         -- See Note [Test cheapness with exprOkForSpeculation]
  , Just (tc :: TyCon
tc, _) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
expr_ty
  , Just dc :: DataCon
dc <- TyCon -> Maybe DataCon
boxingDataCon_maybe TyCon
tc
  , let dc_res_ty :: Type
dc_res_ty = DataCon -> Type
dataConOrigResTy DataCon
dc  -- No free type variables
        [bx_bndr :: Id
bx_bndr, ubx_bndr :: Id
ubx_bndr] = [Type] -> [Id]
mkTemplateLocals [Type
dc_res_ty, Type
expr_ty]
  = do { LevelledExpr
expr1 <- LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
rhs_env CoreExprWithFVs
ann_expr
       ; let l1r :: Level
l1r       = LevelEnv -> Level
incMinorLvlFrom LevelEnv
rhs_env
             float_rhs :: LevelledExpr
float_rhs = [LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
abs_vars_w_lvls (LevelledExpr -> LevelledExpr) -> LevelledExpr -> LevelledExpr
forall a b. (a -> b) -> a -> b
$
                         LevelledExpr
-> LevelledBndr
-> Type
-> [(AltCon, [LevelledBndr], LevelledExpr)]
-> LevelledExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case LevelledExpr
expr1 (Level -> Id -> LevelledBndr
stayPut Level
l1r Id
ubx_bndr) Type
dc_res_ty
                             [(AltCon
DEFAULT, [], DataCon -> [LevelledExpr] -> LevelledExpr
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
dc [Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
ubx_bndr])]

       ; Id
var <- LevelledExpr -> Maybe Int -> Bool -> LvlM Id
newLvlVar LevelledExpr
float_rhs Maybe Int
forall a. Maybe a
Nothing Bool
is_mk_static
       ; let l1u :: Level
l1u      = LevelEnv -> Level
incMinorLvlFrom LevelEnv
env
             use_expr :: LevelledExpr
use_expr = LevelledExpr
-> LevelledBndr
-> Type
-> [(AltCon, [LevelledBndr], LevelledExpr)]
-> LevelledExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (LevelledExpr -> [Id] -> LevelledExpr
forall b. Expr b -> [Id] -> Expr b
mkVarApps (Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
var) [Id]
abs_vars)
                             (Level -> Id -> LevelledBndr
stayPut Level
l1u Id
bx_bndr) Type
expr_ty
                             [(DataCon -> AltCon
DataAlt DataCon
dc, [Level -> Id -> LevelledBndr
stayPut Level
l1u Id
ubx_bndr], Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
ubx_bndr)]
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBind -> LevelledExpr -> LevelledExpr
forall b. Bind b -> Expr b -> Expr b
Let (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
var (Level -> FloatSpec
FloatMe Level
dest_lvl)) LevelledExpr
float_rhs)
                     LevelledExpr
use_expr) }

  | Bool
otherwise          -- e.g. do not float unboxed tuples
  = LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
env CoreExprWithFVs
ann_expr

  where
    expr :: Expr Id
expr         = CoreExprWithFVs -> Expr Id
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate CoreExprWithFVs
ann_expr
    expr_ty :: Type
expr_ty      = Expr Id -> Type
exprType Expr Id
expr
    fvs :: FVAnn
fvs          = CoreExprWithFVs -> FVAnn
freeVarsOf CoreExprWithFVs
ann_expr
    fvs_ty :: TyCoVarSet
fvs_ty       = Type -> TyCoVarSet
tyCoVarsOfType Type
expr_ty
    is_bot :: Bool
is_bot       = Maybe (Int, StrictSig) -> Bool
forall s. Maybe (Int, s) -> Bool
isBottomThunk Maybe (Int, StrictSig)
mb_bot_str
    is_function :: Bool
is_function  = CoreExprWithFVs -> Bool
isFunction CoreExprWithFVs
ann_expr
    mb_bot_str :: Maybe (Int, StrictSig)
mb_bot_str   = Expr Id -> Maybe (Int, StrictSig)
exprBotStrictness_maybe Expr Id
expr
                           -- See Note [Bottoming floats]
                           -- esp Bottoming floats (2)
    expr_ok_for_spec :: Bool
expr_ok_for_spec = Expr Id -> Bool
exprOkForSpeculation Expr Id
expr
    dest_lvl :: Level
dest_lvl     = LevelEnv -> FVAnn -> TyCoVarSet -> Bool -> Bool -> Bool -> Level
destLevel LevelEnv
env FVAnn
fvs TyCoVarSet
fvs_ty Bool
is_function Bool
is_bot Bool
False
    abs_vars :: [Id]
abs_vars     = Level -> LevelEnv -> FVAnn -> [Id]
abstractVars Level
dest_lvl LevelEnv
env FVAnn
fvs

    -- float_is_new_lam: the floated thing will be a new value lambda
    -- replacing, say (g (x+4)) by (lvl x).  No work is saved, nor is
    -- allocation saved.  The benefit is to get it to the top level
    -- and hence out of the body of this function altogether, making
    -- it smaller and more inlinable
    float_is_new_lam :: Bool
float_is_new_lam = Int
float_n_lams Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0
    float_n_lams :: Int
float_n_lams     = (Id -> Bool) -> [Id] -> Int
forall a. (a -> Bool) -> [a] -> Int
count Id -> Bool
isId [Id]
abs_vars

    (rhs_env :: LevelEnv
rhs_env, abs_vars_w_lvls :: [LevelledBndr]
abs_vars_w_lvls) = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs LevelEnv
env Level
dest_lvl [Id]
abs_vars

    join_arity_maybe :: Maybe a
join_arity_maybe = Maybe a
forall a. Maybe a
Nothing

    is_mk_static :: Bool
is_mk_static = Maybe (Expr Id, Type, Expr Id, Expr Id) -> Bool
forall a. Maybe a -> Bool
isJust (Expr Id -> Maybe (Expr Id, Type, Expr Id, Expr Id)
collectMakeStaticArgs Expr Id
expr)
        -- Yuk: See Note [Grand plan for static forms] in main/StaticPtrTable

        -- A decision to float entails let-binding this thing, and we only do
        -- that if we'll escape a value lambda, or will go to the top level.
    float_me :: Bool
float_me = Bool
saves_work Bool -> Bool -> Bool
|| Bool
saves_alloc Bool -> Bool -> Bool
|| Bool
is_mk_static

    -- We can save work if we can move a redex outside a value lambda
    -- But if float_is_new_lam is True, then the redex is wrapped in a
    -- a new lambda, so no work is saved
    saves_work :: Bool
saves_work = Bool
escapes_value_lam Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
float_is_new_lam

    escapes_value_lam :: Bool
escapes_value_lam = Level
dest_lvl Level -> Level -> Bool
`ltMajLvl` (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)
                  -- See Note [Escaping a value lambda]

    -- See Note [Floating to the top]
    saves_alloc :: Bool
saves_alloc =  Level -> Bool
isTopLvl Level
dest_lvl
                Bool -> Bool -> Bool
&& LevelEnv -> Bool
floatConsts LevelEnv
env
                Bool -> Bool -> Bool
&& (Bool -> Bool
not Bool
strict_ctxt Bool -> Bool -> Bool
|| Bool
is_bot Bool -> Bool -> Bool
|| Expr Id -> Bool
exprIsHNF Expr Id
expr)

isBottomThunk :: Maybe (Arity, s) -> Bool
-- See Note [Bottoming floats] (2)
isBottomThunk :: Maybe (Int, s) -> Bool
isBottomThunk (Just (0, _)) = Bool
True   -- Zero arity
isBottomThunk _             = Bool
False

{- Note [Floating to the top]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We are keen to float something to the top level, even if it does not
escape a value lambda (and hence save work), for two reasons:

  * Doing so makes the function smaller, by floating out
    bottoming expressions, or integer or string literals.  That in
    turn makes it easier to inline, with less duplication.

  * (Minor) Doing so may turn a dynamic allocation (done by machine
    instructions) into a static one. Minor because we are assuming
    we are not escaping a value lambda.

But do not so if:
     - the context is a strict, and
     - the expression is not a HNF, and
     - the expression is not bottoming

Exammples:

* Bottoming
      f x = case x of
              0 -> error <big thing>
              _ -> x+1
  Here we want to float (error <big thing>) to top level, abstracting
  over 'x', so as to make f's RHS smaller.

* HNF
      f = case y of
            True  -> p:q
            False -> blah
  We may as well float the (p:q) so it becomes a static data structure.

* Case scrutinee
      f = case g True of ....
  Don't float (g True) to top level; then we have the admin of a
  top-level thunk to worry about, with zero gain.

* Case alternative
      h = case y of
             True  -> g True
             False -> False
  Don't float (g True) to the top level

* Arguments
     t = f (g True)
  If f is lazy, we /do/ float (g True) because then we can allocate
  the thunk statically rather than dynamically.  But if f is strict
  we don't (see the use of idStrictness in lvlApp).  It's not clear
  if this test is worth the bother: it's only about CAFs!

It's controlled by a flag (floatConsts), because doing this too
early loses opportunities for RULES which (needless to say) are
important in some nofib programs (gcd is an example).  [SPJ note:
I think this is obselete; the flag seems always on.]

Note [Floating join point bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Mostly we only float a join point if it can /stay/ a join point.  But
there is one exception: if it can go to the top level (Trac #13286).
Consider
  f x = joinrec j y n = <...j y' n'...>
        in jump j x 0

Here we may just as well produce
  j y n = <....j y' n'...>
  f x = j x 0

and now there is a chance that 'f' will be inlined at its call sites.
It shouldn't make a lot of difference, but thes tests
  perf/should_run/MethSharing
  simplCore/should_compile/spec-inline
and one nofib program, all improve if you do float to top, because
of the resulting inlining of f.  So ok, let's do it.

Note [Free join points]
~~~~~~~~~~~~~~~~~~~~~~~
We never float a MFE that has a free join-point variable.  You mght think
this can never occur.  After all, consider
     join j x = ...
     in ....(jump j x)....
How might we ever want to float that (jump j x)?
  * If it would escape a value lambda, thus
        join j x = ... in (\y. ...(jump j x)... )
    then 'j' isn't a valid join point in the first place.

But consider
     join j x = .... in
     joinrec j2 y =  ...(jump j x)...(a+b)....

Since j2 is recursive, it /is/ worth floating (a+b) out of the joinrec.
But it is emphatically /not/ good to float the (jump j x) out:
 (a) 'j' will stop being a join point
 (b) In any case, jumping to 'j' must be an exit of the j2 loop, so no
     work would be saved by floating it out of the \y.

Even if we floated 'j' to top level, (b) would still hold.

Bottom line: never float a MFE that has a free JoinId.

Note [Floating MFEs of unlifted type]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have
   case f x of (r::Int#) -> blah
we'd like to float (f x). But it's not trivial because it has type
Int#, and we don't want to evaluate it too early.  But we can instead
float a boxed version
   y = case f x of r -> I# r
and replace the original (f x) with
   case (case y of I# r -> r) of r -> blah

Being able to float unboxed expressions is sometimes important; see
Trac #12603.  I'm not sure how /often/ it is important, but it's
not hard to achieve.

We only do it for a fixed collection of types for which we have a
convenient boxing constructor (see boxingDataCon_maybe).  In
particular we /don't/ do it for unboxed tuples; it's better to float
the components of the tuple individually.

I did experiment with a form of boxing that works for any type, namely
wrapping in a function.  In our example

   let y = case f x of r -> \v. f x
   in case y void of r -> blah

It works fine, but it's 50% slower (based on some crude benchmarking).
I suppose we could do it for types not covered by boxingDataCon_maybe,
but it's more code and I'll wait to see if anyone wants it.

Note [Test cheapness with exprOkForSpeculation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We don't want to float very cheap expressions by boxing and unboxing.
But we use exprOkForSpeculation for the test, not exprIsCheap.
Why?  Because it's important /not/ to transform
     f (a /# 3)
to
     f (case bx of I# a -> a /# 3)
and float bx = I# (a /# 3), because the application of f no
longer obeys the let/app invariant.  But (a /# 3) is ok-for-spec
due to a special hack that says division operators can't fail
when the denominator is definitely non-zero.  And yet that
same expression says False to exprIsCheap.  Simplest way to
guarantee the let/app invariant is to use the same function!

If an expression is okay for speculation, we could also float it out
*without* boxing and unboxing, since evaluating it early is okay.
However, it turned out to usually be better not to float such expressions,
since they tend to be extremely cheap things like (x +# 1#). Even the
cost of spilling the let-bound variable to the stack across a call may
exceed the cost of recomputing such an expression. (And we can't float
unlifted bindings to top-level.)

We could try to do something smarter here, and float out expensive yet
okay-for-speculation things, such as division by non-zero constants.
But I suspect it's a narrow target.

Note [Bottoming floats]
~~~~~~~~~~~~~~~~~~~~~~~
If we see
        f = \x. g (error "urk")
we'd like to float the call to error, to get
        lvl = error "urk"
        f = \x. g lvl

But, as ever, we need to be careful:

(1) We want to float a bottoming
    expression even if it has free variables:
        f = \x. g (let v = h x in error ("urk" ++ v))
    Then we'd like to abstract over 'x' can float the whole arg of g:
        lvl = \x. let v = h x in error ("urk" ++ v)
        f = \x. g (lvl x)
    To achieve this we pass is_bot to destLevel

(2) We do not do this for lambdas that return
    bottom.  Instead we treat the /body/ of such a function specially,
    via point (1).  For example:
        f = \x. ....(\y z. if x then error y else error z)....
    ===>
        lvl = \x z y. if b then error y else error z
        f = \x. ...(\y z. lvl x z y)...
    (There is no guarantee that we'll choose the perfect argument order.)

(3) If we have a /binding/ that returns bottom, we want to float it to top
    level, even if it has free vars (point (1)), and even it has lambdas.
    Example:
       ... let { v = \y. error (show x ++ show y) } in ...
    We want to abstract over x and float the whole thing to top:
       lvl = \xy. errror (show x ++ show y)
       ...let {v = lvl x} in ...

    Then of course we don't want to separately float the body (error ...)
    as /another/ MFE, so we tell lvlFloatRhs not to do that, via the is_bot
    argument.

See Maessen's paper 1999 "Bottom extraction: factoring error handling out
of functional programs" (unpublished I think).

When we do this, we set the strictness and arity of the new bottoming
Id, *immediately*, for three reasons:

  * To prevent the abstracted thing being immediately inlined back in again
    via preInlineUnconditionally.  The latter has a test for bottoming Ids
    to stop inlining them, so we'd better make sure it *is* a bottoming Id!

  * So that it's properly exposed as such in the interface file, even if
    this is all happening after strictness analysis.

  * In case we do CSE with the same expression that *is* marked bottom
        lvl          = error "urk"
          x{str=bot) = error "urk"
    Here we don't want to replace 'x' with 'lvl', else we may get Lint
    errors, e.g. via a case with empty alternatives:  (case x of {})
    Lint complains unless the scrutinee of such a case is clearly bottom.

    This was reported in Trac #11290.   But since the whole bottoming-float
    thing is based on the cheap-and-cheerful exprIsBottom, I'm not sure
    that it'll nail all such cases.

Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Tiresomely, though, the simplifier has an invariant that the manifest
arity of the RHS should be the same as the arity; but we can't call
etaExpand during SetLevels because it works over a decorated form of
CoreExpr.  So we do the eta expansion later, in FloatOut.

Note [Case MFEs]
~~~~~~~~~~~~~~~~
We don't float a case expression as an MFE from a strict context.  Why not?
Because in doing so we share a tiny bit of computation (the switch) but
in exchange we build a thunk, which is bad.  This case reduces allocation
by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
Doesn't change any other allocation at all.

We will make a separate decision for the scrutinee and alternatives.

However this can have a knock-on effect for fusion: consider
    \v -> foldr k z (case x of I# y -> build ..y..)
Perhaps we can float the entire (case x of ...) out of the \v.  Then
fusion will not happen, but we will get more sharing.  But if we don't
float the case (as advocated here) we won't float the (build ...y..)
either, so fusion will happen.  It can be a big effect, esp in some
artificial benchmarks (e.g. integer, queens), but there is no perfect
answer.

-}

annotateBotStr :: Id -> Arity -> Maybe (Arity, StrictSig) -> Id
-- See Note [Bottoming floats] for why we want to add
-- bottoming information right now
--
-- n_extra are the number of extra value arguments added during floating
annotateBotStr :: Id -> Int -> Maybe (Int, StrictSig) -> Id
annotateBotStr id :: Id
id n_extra :: Int
n_extra mb_str :: Maybe (Int, StrictSig)
mb_str
  = case Maybe (Int, StrictSig)
mb_str of
      Nothing           -> Id
id
      Just (arity :: Int
arity, sig :: StrictSig
sig) -> Id
id Id -> Int -> Id
`setIdArity`      (Int
arity Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n_extra)
                              Id -> StrictSig -> Id
`setIdStrictness` (Int -> StrictSig -> StrictSig
increaseStrictSigArity Int
n_extra StrictSig
sig)

notWorthFloating :: CoreExpr -> [Var] -> Bool
-- Returns True if the expression would be replaced by
-- something bigger than it is now.  For example:
--   abs_vars = tvars only:  return True if e is trivial,
--                           but False for anything bigger
--   abs_vars = [x] (an Id): return True for trivial, or an application (f x)
--                           but False for (f x x)
--
-- One big goal is that floating should be idempotent.  Eg if
-- we replace e with (lvl79 x y) and then run FloatOut again, don't want
-- to replace (lvl79 x y) with (lvl83 x y)!

notWorthFloating :: Expr Id -> [Id] -> Bool
notWorthFloating e :: Expr Id
e abs_vars :: [Id]
abs_vars
  = Expr Id -> Int -> Bool
forall a b. (Ord a, Num a) => Expr b -> a -> Bool
go Expr Id
e ((Id -> Bool) -> [Id] -> Int
forall a. (a -> Bool) -> [a] -> Int
count Id -> Bool
isId [Id]
abs_vars)
  where
    go :: Expr b -> a -> Bool
go (Var {}) n :: a
n    = a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= 0
    go (Lit lit :: Literal
lit) n :: a
n   = ASSERT( n==0 )
                       Literal -> Bool
litIsTrivial Literal
lit   -- Note [Floating literals]
    go (Tick t :: Tickish Id
t e :: Expr b
e) n :: a
n  = Bool -> Bool
not (Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish Id
t) Bool -> Bool -> Bool
&& Expr b -> a -> Bool
go Expr b
e a
n
    go (Cast e :: Expr b
e _)  n :: a
n = Expr b -> a -> Bool
go Expr b
e a
n
    go (App e :: Expr b
e arg :: Expr b
arg) n :: a
n
       | Type {}     <- Expr b
arg = Expr b -> a -> Bool
go Expr b
e a
n
       | Coercion {} <- Expr b
arg = Expr b -> a -> Bool
go Expr b
e a
n
       | a
na -> a -> Bool
forall a. Eq a => a -> a -> Bool
==0               = Bool
False
       | Expr b -> Bool
forall b. Expr b -> Bool
is_triv Expr b
arg        = Expr b -> a -> Bool
go Expr b
e (a
na -> a -> a
forall a. Num a => a -> a -> a
-1)
       | Bool
otherwise          = Bool
False
    go _ _                  = Bool
False

    is_triv :: Expr b -> Bool
is_triv (Lit {})              = Bool
True        -- Treat all literals as trivial
    is_triv (Var {})              = Bool
True        -- (ie not worth floating)
    is_triv (Cast e :: Expr b
e _)            = Expr b -> Bool
is_triv Expr b
e
    is_triv (App e :: Expr b
e (Type {}))     = Expr b -> Bool
is_triv Expr b
e
    is_triv (App e :: Expr b
e (Coercion {})) = Expr b -> Bool
is_triv Expr b
e
    is_triv (Tick t :: Tickish Id
t e :: Expr b
e)            = Bool -> Bool
not (Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish Id
t) Bool -> Bool -> Bool
&& Expr b -> Bool
is_triv Expr b
e
    is_triv _                     = Bool
False

{-
Note [Floating literals]
~~~~~~~~~~~~~~~~~~~~~~~~
It's important to float Integer literals, so that they get shared,
rather than being allocated every time round the loop.
Hence the litIsTrivial.

Ditto literal strings (LitString), which we'd like to float to top
level, which is now possible.


Note [Escaping a value lambda]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to float even cheap expressions out of value lambdas,
because that saves allocation.  Consider
        f = \x.  .. (\y.e) ...
Then we'd like to avoid allocating the (\y.e) every time we call f,
(assuming e does not mention x). An example where this really makes a
difference is simplrun009.

Another reason it's good is because it makes SpecContr fire on functions.
Consider
        f = \x. ....(f (\y.e))....
After floating we get
        lvl = \y.e
        f = \x. ....(f lvl)...
and that is much easier for SpecConstr to generate a robust
specialisation for.

However, if we are wrapping the thing in extra value lambdas (in
abs_vars), then nothing is saved.  E.g.
        f = \xyz. ...(e1[y],e2)....
If we float
        lvl = \y. (e1[y],e2)
        f = \xyz. ...(lvl y)...
we have saved nothing: one pair will still be allocated for each
call of 'f'.  Hence the (not float_is_lam) in float_me.


************************************************************************
*                                                                      *
\subsection{Bindings}
*                                                                      *
************************************************************************

The binding stuff works for top level too.
-}

lvlBind :: LevelEnv
        -> CoreBindWithFVs
        -> LvlM (LevelledBind, LevelEnv)

lvlBind :: LevelEnv -> AnnBind Id FVAnn -> LvlM (LevelledBind, LevelEnv)
lvlBind env :: LevelEnv
env (AnnNonRec bndr :: Id
bndr rhs :: CoreExprWithFVs
rhs)
  | Id -> Bool
isTyVar Id
bndr    -- Don't do anything for TyVar binders
                    --   (simplifier gets rid of them pronto)
  Bool -> Bool -> Bool
|| Id -> Bool
isCoVar Id
bndr   -- Difficult to fix up CoVar occurrences (see extendPolyLvlEnv)
                    -- so we will ignore this case for now
  Bool -> Bool -> Bool
|| Bool -> Bool
not (LevelEnv -> Level -> Bool
profitableFloat LevelEnv
env Level
dest_lvl)
  Bool -> Bool -> Bool
|| (Level -> Bool
isTopLvl Level
dest_lvl Bool -> Bool -> Bool
&& Bool -> Bool
not (Expr Id -> Type -> Bool
exprIsTopLevelBindable Expr Id
deann_rhs Type
bndr_ty))
          -- We can't float an unlifted binding to top level (except
          -- literal strings), so we don't float it at all.  It's a
          -- bit brutal, but unlifted bindings aren't expensive either

  = -- No float
    do { LevelledExpr
rhs' <- LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlRhs LevelEnv
env RecFlag
NonRecursive Bool
is_bot Maybe Int
mb_join_arity CoreExprWithFVs
rhs
       ; let  bind_lvl :: Level
bind_lvl        = Level -> Level
incMinorLvl (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)
              (env' :: LevelEnv
env', [bndr' :: LevelledBndr
bndr']) = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
NonRecursive LevelEnv
env Level
bind_lvl [Id
bndr]
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec LevelledBndr
bndr' LevelledExpr
rhs', LevelEnv
env') }

  -- Otherwise we are going to float
  | [Id] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Id]
abs_vars
  = do {  -- No type abstraction; clone existing binder
         LevelledExpr
rhs' <- [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs [] Level
dest_lvl LevelEnv
env RecFlag
NonRecursive
                             Bool
is_bot Maybe Int
mb_join_arity CoreExprWithFVs
rhs
       ; (env' :: LevelEnv
env', [bndr' :: Id
bndr']) <- RecFlag -> LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneLetVars RecFlag
NonRecursive LevelEnv
env Level
dest_lvl [Id
bndr]
       ; let bndr2 :: Id
bndr2 = Id -> Int -> Maybe (Int, StrictSig) -> Id
annotateBotStr Id
bndr' 0 Maybe (Int, StrictSig)
mb_bot_str
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
bndr2 (Level -> FloatSpec
FloatMe Level
dest_lvl)) LevelledExpr
rhs', LevelEnv
env') }

  | Bool
otherwise
  = do {  -- Yes, type abstraction; create a new binder, extend substitution, etc
         LevelledExpr
rhs' <- [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs [Id]
abs_vars Level
dest_lvl LevelEnv
env RecFlag
NonRecursive
                             Bool
is_bot Maybe Int
mb_join_arity CoreExprWithFVs
rhs
       ; (env' :: LevelEnv
env', [bndr' :: Id
bndr']) <- Level -> LevelEnv -> [Id] -> [Id] -> LvlM (LevelEnv, [Id])
newPolyBndrs Level
dest_lvl LevelEnv
env [Id]
abs_vars [Id
bndr]
       ; let bndr2 :: Id
bndr2 = Id -> Int -> Maybe (Int, StrictSig) -> Id
annotateBotStr Id
bndr' Int
n_extra Maybe (Int, StrictSig)
mb_bot_str
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelledBndr -> LevelledExpr -> LevelledBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
bndr2 (Level -> FloatSpec
FloatMe Level
dest_lvl)) LevelledExpr
rhs', LevelEnv
env') }

  where
    bndr_ty :: Type
bndr_ty    = Id -> Type
idType Id
bndr
    ty_fvs :: TyCoVarSet
ty_fvs     = Type -> TyCoVarSet
tyCoVarsOfType Type
bndr_ty
    rhs_fvs :: FVAnn
rhs_fvs    = CoreExprWithFVs -> FVAnn
freeVarsOf CoreExprWithFVs
rhs
    bind_fvs :: FVAnn
bind_fvs   = FVAnn
rhs_fvs FVAnn -> FVAnn -> FVAnn
`unionDVarSet` Id -> FVAnn
dIdFreeVars Id
bndr
    abs_vars :: [Id]
abs_vars   = Level -> LevelEnv -> FVAnn -> [Id]
abstractVars Level
dest_lvl LevelEnv
env FVAnn
bind_fvs
    dest_lvl :: Level
dest_lvl   = LevelEnv -> FVAnn -> TyCoVarSet -> Bool -> Bool -> Bool -> Level
destLevel LevelEnv
env FVAnn
bind_fvs TyCoVarSet
ty_fvs (CoreExprWithFVs -> Bool
isFunction CoreExprWithFVs
rhs) Bool
is_bot Bool
is_join

    deann_rhs :: Expr Id
deann_rhs  = CoreExprWithFVs -> Expr Id
forall bndr annot. AnnExpr bndr annot -> Expr bndr
deAnnotate CoreExprWithFVs
rhs
    mb_bot_str :: Maybe (Int, StrictSig)
mb_bot_str = Expr Id -> Maybe (Int, StrictSig)
exprBotStrictness_maybe Expr Id
deann_rhs
    is_bot :: Bool
is_bot     = Maybe (Int, StrictSig) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (Int, StrictSig)
mb_bot_str
        -- NB: not isBottomThunk!  See Note [Bottoming floats] point (3)

    n_extra :: Int
n_extra    = (Id -> Bool) -> [Id] -> Int
forall a. (a -> Bool) -> [a] -> Int
count Id -> Bool
isId [Id]
abs_vars
    mb_join_arity :: Maybe Int
mb_join_arity = Id -> Maybe Int
isJoinId_maybe Id
bndr
    is_join :: Bool
is_join       = Maybe Int -> Bool
forall a. Maybe a -> Bool
isJust Maybe Int
mb_join_arity

lvlBind env :: LevelEnv
env (AnnRec pairs :: [(Id, CoreExprWithFVs)]
pairs)
  |  LevelEnv -> Bool
floatTopLvlOnly LevelEnv
env Bool -> Bool -> Bool
&& Bool -> Bool
not (Level -> Bool
isTopLvl Level
dest_lvl)
         -- Only floating to the top level is allowed.
  Bool -> Bool -> Bool
|| Bool -> Bool
not (LevelEnv -> Level -> Bool
profitableFloat LevelEnv
env Level
dest_lvl)
  Bool -> Bool -> Bool
|| (Level -> Bool
isTopLvl Level
dest_lvl Bool -> Bool -> Bool
&& (Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType (Type -> Bool) -> (Id -> Type) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Type
idType) [Id]
bndrs)
       -- This isUnliftedType stuff is the same test as in the non-rec case
       -- You might wonder whether we can have a recursive binding for
       -- an unlifted value -- but we can if it's a /join binding/ (#16978)
       -- (Ultimately I think we should not use SetLevels to
       -- float join bindings at all, but that's another story.)
  =    -- No float
    do { let bind_lvl :: Level
bind_lvl       = Level -> Level
incMinorLvl (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)
             (env' :: LevelEnv
env', bndrs' :: [LevelledBndr]
bndrs') = RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs RecFlag
Recursive LevelEnv
env Level
bind_lvl [Id]
bndrs
             lvl_rhs :: (Id, CoreExprWithFVs) -> LvlM LevelledExpr
lvl_rhs (b :: Id
b,r :: CoreExprWithFVs
r)  = LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlRhs LevelEnv
env' RecFlag
Recursive Bool
is_bot (Id -> Maybe Int
isJoinId_maybe Id
b) CoreExprWithFVs
r
       ; [LevelledExpr]
rhss' <- ((Id, CoreExprWithFVs) -> LvlM LevelledExpr)
-> [(Id, CoreExprWithFVs)] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Id, CoreExprWithFVs) -> LvlM LevelledExpr
lvl_rhs [(Id, CoreExprWithFVs)]
pairs
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec ([LevelledBndr]
bndrs' [LevelledBndr] -> [LevelledExpr] -> [(LevelledBndr, LevelledExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [LevelledExpr]
rhss'), LevelEnv
env') }

  -- Otherwise we are going to float
  | [Id] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Id]
abs_vars
  = do { (new_env :: LevelEnv
new_env, new_bndrs :: [Id]
new_bndrs) <- RecFlag -> LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneLetVars RecFlag
Recursive LevelEnv
env Level
dest_lvl [Id]
bndrs
       ; [LevelledExpr]
new_rhss <- ((Id, CoreExprWithFVs) -> LvlM LevelledExpr)
-> [(Id, CoreExprWithFVs)] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LevelEnv -> (Id, CoreExprWithFVs) -> LvlM LevelledExpr
do_rhs LevelEnv
new_env) [(Id, CoreExprWithFVs)]
pairs
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return ( [(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec ([Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
b (Level -> FloatSpec
FloatMe Level
dest_lvl) | Id
b <- [Id]
new_bndrs] [LevelledBndr] -> [LevelledExpr] -> [(LevelledBndr, LevelledExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [LevelledExpr]
new_rhss)
                , LevelEnv
new_env) }

-- ToDo: when enabling the floatLambda stuff,
--       I think we want to stop doing this
  | [(bndr :: Id
bndr,rhs :: CoreExprWithFVs
rhs)] <- [(Id, CoreExprWithFVs)]
pairs
  , (Id -> Bool) -> [Id] -> Int
forall a. (a -> Bool) -> [a] -> Int
count Id -> Bool
isId [Id]
abs_vars Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 1
  = do  -- Special case for self recursion where there are
        -- several variables carried around: build a local loop:
        --      poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
        -- This just makes the closures a bit smaller.  If we don't do
        -- this, allocation rises significantly on some programs
        --
        -- We could elaborate it for the case where there are several
        -- mutually recursive functions, but it's quite a bit more complicated
        --
        -- This all seems a bit ad hoc -- sigh
    let (rhs_env :: LevelEnv
rhs_env, abs_vars_w_lvls :: [LevelledBndr]
abs_vars_w_lvls) = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs LevelEnv
env Level
dest_lvl [Id]
abs_vars
        rhs_lvl :: Level
rhs_lvl = LevelEnv -> Level
le_ctxt_lvl LevelEnv
rhs_env

    (rhs_env' :: LevelEnv
rhs_env', [new_bndr :: Id
new_bndr]) <- RecFlag -> LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneLetVars RecFlag
Recursive LevelEnv
rhs_env Level
rhs_lvl [Id
bndr]
    let
        (lam_bndrs :: [Id]
lam_bndrs, rhs_body :: CoreExprWithFVs
rhs_body)   = CoreExprWithFVs -> ([Id], CoreExprWithFVs)
forall bndr annot.
AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectAnnBndrs CoreExprWithFVs
rhs
        (body_env1 :: LevelEnv
body_env1, lam_bndrs1 :: [Id]
lam_bndrs1) = RecFlag -> LevelEnv -> [Id] -> (LevelEnv, [Id])
substBndrsSL RecFlag
NonRecursive LevelEnv
rhs_env' [Id]
lam_bndrs
        (body_env2 :: LevelEnv
body_env2, lam_bndrs2 :: [LevelledBndr]
lam_bndrs2) = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs LevelEnv
body_env1 Level
rhs_lvl [Id]
lam_bndrs1
    LevelledExpr
new_rhs_body <- LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlRhs LevelEnv
body_env2 RecFlag
Recursive Bool
is_bot (Id -> Maybe Int
get_join Id
bndr) CoreExprWithFVs
rhs_body
    (poly_env :: LevelEnv
poly_env, [poly_bndr :: Id
poly_bndr]) <- Level -> LevelEnv -> [Id] -> [Id] -> LvlM (LevelEnv, [Id])
newPolyBndrs Level
dest_lvl LevelEnv
env [Id]
abs_vars [Id
bndr]
    (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
poly_bndr (Level -> FloatSpec
FloatMe Level
dest_lvl)
                 , [LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
abs_vars_w_lvls (LevelledExpr -> LevelledExpr) -> LevelledExpr -> LevelledExpr
forall a b. (a -> b) -> a -> b
$
                   [LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
lam_bndrs2 (LevelledExpr -> LevelledExpr) -> LevelledExpr -> LevelledExpr
forall a b. (a -> b) -> a -> b
$
                   LevelledBind -> LevelledExpr -> LevelledExpr
forall b. Bind b -> Expr b -> Expr b
Let ([(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec [( Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
new_bndr (Level -> FloatSpec
StayPut Level
rhs_lvl)
                             , [LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
lam_bndrs2 LevelledExpr
new_rhs_body)])
                       (LevelledExpr -> [Id] -> LevelledExpr
forall b. Expr b -> [Id] -> Expr b
mkVarApps (Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
new_bndr) [Id]
lam_bndrs1))]
           , LevelEnv
poly_env)

  | Bool
otherwise  -- Non-null abs_vars
  = do { (new_env :: LevelEnv
new_env, new_bndrs :: [Id]
new_bndrs) <- Level -> LevelEnv -> [Id] -> [Id] -> LvlM (LevelEnv, [Id])
newPolyBndrs Level
dest_lvl LevelEnv
env [Id]
abs_vars [Id]
bndrs
       ; [LevelledExpr]
new_rhss <- ((Id, CoreExprWithFVs) -> LvlM LevelledExpr)
-> [(Id, CoreExprWithFVs)] -> UniqSM [LevelledExpr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (LevelEnv -> (Id, CoreExprWithFVs) -> LvlM LevelledExpr
do_rhs LevelEnv
new_env) [(Id, CoreExprWithFVs)]
pairs
       ; (LevelledBind, LevelEnv) -> LvlM (LevelledBind, LevelEnv)
forall (m :: * -> *) a. Monad m => a -> m a
return ( [(LevelledBndr, LevelledExpr)] -> LevelledBind
forall b. [(b, Expr b)] -> Bind b
Rec ([Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
b (Level -> FloatSpec
FloatMe Level
dest_lvl) | Id
b <- [Id]
new_bndrs] [LevelledBndr] -> [LevelledExpr] -> [(LevelledBndr, LevelledExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [LevelledExpr]
new_rhss)
                , LevelEnv
new_env) }

  where
    (bndrs :: [Id]
bndrs,rhss :: [CoreExprWithFVs]
rhss) = [(Id, CoreExprWithFVs)] -> ([Id], [CoreExprWithFVs])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, CoreExprWithFVs)]
pairs
    is_join :: Bool
is_join  = Id -> Bool
isJoinId ([Id] -> Id
forall a. [a] -> a
head [Id]
bndrs)
                -- bndrs is always non-empty and if one is a join they all are
                -- Both are checked by Lint
    is_fun :: Bool
is_fun   = (CoreExprWithFVs -> Bool) -> [CoreExprWithFVs] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all CoreExprWithFVs -> Bool
isFunction [CoreExprWithFVs]
rhss
    is_bot :: Bool
is_bot   = Bool
False  -- It's odd to have an unconditionally divergent
                      -- function in a Rec, and we don't much care what
                      -- happens to it.  False is simple!

    do_rhs :: LevelEnv -> (Id, CoreExprWithFVs) -> LvlM LevelledExpr
do_rhs env :: LevelEnv
env (bndr :: Id
bndr,rhs :: CoreExprWithFVs
rhs) = [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs [Id]
abs_vars Level
dest_lvl LevelEnv
env RecFlag
Recursive
                                        Bool
is_bot (Id -> Maybe Int
get_join Id
bndr)
                                        CoreExprWithFVs
rhs

    get_join :: Id -> Maybe Int
get_join bndr :: Id
bndr | Bool
need_zap  = Maybe Int
forall a. Maybe a
Nothing
                  | Bool
otherwise = Id -> Maybe Int
isJoinId_maybe Id
bndr
    need_zap :: Bool
need_zap = Level
dest_lvl Level -> Level -> Bool
`ltLvl` LevelEnv -> Level
joinCeilingLevel LevelEnv
env

        -- Finding the free vars of the binding group is annoying
    bind_fvs :: FVAnn
bind_fvs = (([FVAnn] -> FVAnn
unionDVarSets [ CoreExprWithFVs -> FVAnn
freeVarsOf CoreExprWithFVs
rhs | (_, rhs :: CoreExprWithFVs
rhs) <- [(Id, CoreExprWithFVs)]
pairs])
                FVAnn -> FVAnn -> FVAnn
`unionDVarSet`
                (FV -> FVAnn
fvDVarSet (FV -> FVAnn) -> FV -> FVAnn
forall a b. (a -> b) -> a -> b
$ [FV] -> FV
unionsFV [ Id -> FV
idFVs Id
bndr
                                      | (bndr :: Id
bndr, (_,_)) <- [(Id, CoreExprWithFVs)]
pairs]))
               FVAnn -> [Id] -> FVAnn
`delDVarSetList`
                [Id]
bndrs

    ty_fvs :: TyCoVarSet
ty_fvs   = (Id -> TyCoVarSet -> TyCoVarSet)
-> TyCoVarSet -> [Id] -> TyCoVarSet
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (TyCoVarSet -> TyCoVarSet -> TyCoVarSet
unionVarSet (TyCoVarSet -> TyCoVarSet -> TyCoVarSet)
-> (Id -> TyCoVarSet) -> Id -> TyCoVarSet -> TyCoVarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> TyCoVarSet
tyCoVarsOfType (Type -> TyCoVarSet) -> (Id -> Type) -> Id -> TyCoVarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Type
idType) TyCoVarSet
emptyVarSet [Id]
bndrs
    dest_lvl :: Level
dest_lvl = LevelEnv -> FVAnn -> TyCoVarSet -> Bool -> Bool -> Bool -> Level
destLevel LevelEnv
env FVAnn
bind_fvs TyCoVarSet
ty_fvs Bool
is_fun Bool
is_bot Bool
is_join
    abs_vars :: [Id]
abs_vars = Level -> LevelEnv -> FVAnn -> [Id]
abstractVars Level
dest_lvl LevelEnv
env FVAnn
bind_fvs

profitableFloat :: LevelEnv -> Level -> Bool
profitableFloat :: LevelEnv -> Level -> Bool
profitableFloat env :: LevelEnv
env dest_lvl :: Level
dest_lvl
  =  (Level
dest_lvl Level -> Level -> Bool
`ltMajLvl` LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)  -- Escapes a value lambda
  Bool -> Bool -> Bool
|| Level -> Bool
isTopLvl Level
dest_lvl                      -- Going all the way to top level


----------------------------------------------------
-- Three help functions for the type-abstraction case

lvlRhs :: LevelEnv
       -> RecFlag
       -> Bool               -- Is this a bottoming function
       -> Maybe JoinArity
       -> CoreExprWithFVs
       -> LvlM LevelledExpr
lvlRhs :: LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlRhs env :: LevelEnv
env rec_flag :: RecFlag
rec_flag is_bot :: Bool
is_bot mb_join_arity :: Maybe Int
mb_join_arity expr :: CoreExprWithFVs
expr
  = [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs [] (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env) LevelEnv
env
                RecFlag
rec_flag Bool
is_bot Maybe Int
mb_join_arity CoreExprWithFVs
expr

lvlFloatRhs :: [OutVar] -> Level -> LevelEnv -> RecFlag
            -> Bool   -- Binding is for a bottoming function
            -> Maybe JoinArity
            -> CoreExprWithFVs
            -> LvlM (Expr LevelledBndr)
-- Ignores the le_ctxt_lvl in env; treats dest_lvl as the baseline
lvlFloatRhs :: [Id]
-> Level
-> LevelEnv
-> RecFlag
-> Bool
-> Maybe Int
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlFloatRhs abs_vars :: [Id]
abs_vars dest_lvl :: Level
dest_lvl env :: LevelEnv
env rec :: RecFlag
rec is_bot :: Bool
is_bot mb_join_arity :: Maybe Int
mb_join_arity rhs :: CoreExprWithFVs
rhs
  = do { LevelledExpr
body' <- if Bool -> Bool
not Bool
is_bot  -- See Note [Floating from a RHS]
                     Bool -> Bool -> Bool
&& (Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Id -> Bool
isId [Id]
bndrs
                  then LevelEnv -> Bool -> CoreExprWithFVs -> LvlM LevelledExpr
lvlMFE  LevelEnv
body_env Bool
True CoreExprWithFVs
body
                  else LevelEnv -> CoreExprWithFVs -> LvlM LevelledExpr
lvlExpr LevelEnv
body_env      CoreExprWithFVs
body
       ; LevelledExpr -> LvlM LevelledExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ([LevelledBndr] -> LevelledExpr -> LevelledExpr
forall b. [b] -> Expr b -> Expr b
mkLams [LevelledBndr]
bndrs' LevelledExpr
body') }
  where
    (bndrs :: [Id]
bndrs, body :: CoreExprWithFVs
body)     | Just join_arity :: Int
join_arity <- Maybe Int
mb_join_arity
                      = Int -> CoreExprWithFVs -> ([Id], CoreExprWithFVs)
forall bndr annot.
Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectNAnnBndrs Int
join_arity CoreExprWithFVs
rhs
                      | Bool
otherwise
                      = CoreExprWithFVs -> ([Id], CoreExprWithFVs)
forall bndr annot.
AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
collectAnnBndrs CoreExprWithFVs
rhs
    (env1 :: LevelEnv
env1, bndrs1 :: [Id]
bndrs1)    = RecFlag -> LevelEnv -> [Id] -> (LevelEnv, [Id])
substBndrsSL RecFlag
NonRecursive LevelEnv
env [Id]
bndrs
    all_bndrs :: [Id]
all_bndrs         = [Id]
abs_vars [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
bndrs1
    (body_env :: LevelEnv
body_env, bndrs' :: [LevelledBndr]
bndrs') | Just _ <- Maybe Int
mb_join_arity
                      = LevelEnv -> Level -> RecFlag -> [Id] -> (LevelEnv, [LevelledBndr])
lvlJoinBndrs LevelEnv
env1 Level
dest_lvl RecFlag
rec [Id]
all_bndrs
                      | Bool
otherwise
                      = case LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs LevelEnv
env1 Level
dest_lvl [Id]
all_bndrs of
                          (env2 :: LevelEnv
env2, bndrs' :: [LevelledBndr]
bndrs') -> (LevelEnv -> LevelEnv
placeJoinCeiling LevelEnv
env2, [LevelledBndr]
bndrs')
        -- The important thing here is that we call lvlLamBndrs on
        -- all these binders at once (abs_vars and bndrs), so they
        -- all get the same major level.  Otherwise we create stupid
        -- let-bindings inside, joyfully thinking they can float; but
        -- in the end they don't because we never float bindings in
        -- between lambdas

{- Note [Floating from a RHS]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When floating the RHS of a let-binding, we don't always want to apply
lvlMFE to the body of a lambda, as we usually do, because the entire
binding body is already going to the right place (dest_lvl).

A particular example is the top level.  Consider
   concat = /\ a -> foldr ..a.. (++) []
We don't want to float the body of the lambda to get
   lvl    = /\ a -> foldr ..a.. (++) []
   concat = /\ a -> lvl a
That would be stupid.

Previously this was avoided in a much nastier way, by testing strict_ctxt
in float_me in lvlMFE.  But that wasn't even right because it would fail
to float out the error sub-expression in
    f = \x. case x of
              True  -> error ("blah" ++ show x)
              False -> ...

But we must be careful:

* If we had
    f = \x -> factorial 20
  we /would/ want to float that (factorial 20) out!  Functions are treated
  differently: see the use of isFunction in the calls to destLevel. If
  there are only type lambdas, then destLevel will say "go to top, and
  abstract over the free tyvars" and we don't want that here.

* But if we had
    f = \x -> error (...x....)
  we would NOT want to float the bottoming expression out to give
    lvl = \x -> error (...x...)
    f = \x -> lvl x

Conclusion: use lvlMFE if there are
  * any value lambdas in the original function, and
  * this is not a bottoming function (the is_bot argument)
Use lvlExpr otherwise.  A little subtle, and I got it wrong at least twice
(e.g. Trac #13369).
-}

{-
************************************************************************
*                                                                      *
\subsection{Deciding floatability}
*                                                                      *
************************************************************************
-}

substAndLvlBndrs :: RecFlag -> LevelEnv -> Level -> [InVar] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs :: RecFlag -> LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
substAndLvlBndrs is_rec :: RecFlag
is_rec env :: LevelEnv
env lvl :: Level
lvl bndrs :: [Id]
bndrs
  = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlBndrs LevelEnv
subst_env Level
lvl [Id]
subst_bndrs
  where
    (subst_env :: LevelEnv
subst_env, subst_bndrs :: [Id]
subst_bndrs) = RecFlag -> LevelEnv -> [Id] -> (LevelEnv, [Id])
substBndrsSL RecFlag
is_rec LevelEnv
env [Id]
bndrs

substBndrsSL :: RecFlag -> LevelEnv -> [InVar] -> (LevelEnv, [OutVar])
-- So named only to avoid the name clash with CoreSubst.substBndrs
substBndrsSL :: RecFlag -> LevelEnv -> [Id] -> (LevelEnv, [Id])
substBndrsSL is_rec :: RecFlag
is_rec env :: LevelEnv
env@(LE { le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env }) bndrs :: [Id]
bndrs
  = ( LevelEnv
env { le_subst :: Subst
le_subst    = Subst
subst'
          , le_env :: IdEnv ([Id], LevelledExpr)
le_env      = (IdEnv ([Id], LevelledExpr)
 -> (Id, Id) -> IdEnv ([Id], LevelledExpr))
-> IdEnv ([Id], LevelledExpr)
-> [(Id, Id)]
-> IdEnv ([Id], LevelledExpr)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
add_id  IdEnv ([Id], LevelledExpr)
id_env ([Id]
bndrs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
bndrs') }
    , [Id]
bndrs')
  where
    (subst' :: Subst
subst', bndrs' :: [Id]
bndrs') = case RecFlag
is_rec of
                         NonRecursive -> Subst -> [Id] -> (Subst, [Id])
substBndrs    Subst
subst [Id]
bndrs
                         Recursive    -> Subst -> [Id] -> (Subst, [Id])
substRecBndrs Subst
subst [Id]
bndrs

lvlLamBndrs :: LevelEnv -> Level -> [OutVar] -> (LevelEnv, [LevelledBndr])
-- Compute the levels for the binders of a lambda group
lvlLamBndrs :: LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlLamBndrs env :: LevelEnv
env lvl :: Level
lvl bndrs :: [Id]
bndrs
  = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlBndrs LevelEnv
env Level
new_lvl [Id]
bndrs
  where
    new_lvl :: Level
new_lvl | (Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Id -> Bool
is_major [Id]
bndrs = Level -> Level
incMajorLvl Level
lvl
            | Bool
otherwise          = Level -> Level
incMinorLvl Level
lvl

    is_major :: Id -> Bool
is_major bndr :: Id
bndr = Id -> Bool
isId Id
bndr Bool -> Bool -> Bool
&& Bool -> Bool
not (Id -> Bool
isProbablyOneShotLambda Id
bndr)
       -- The "probably" part says "don't float things out of a
       -- probable one-shot lambda"
       -- See Note [Computing one-shot info] in Demand.hs

lvlJoinBndrs :: LevelEnv -> Level -> RecFlag -> [OutVar]
             -> (LevelEnv, [LevelledBndr])
lvlJoinBndrs :: LevelEnv -> Level -> RecFlag -> [Id] -> (LevelEnv, [LevelledBndr])
lvlJoinBndrs env :: LevelEnv
env lvl :: Level
lvl rec :: RecFlag
rec bndrs :: [Id]
bndrs
  = LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlBndrs LevelEnv
env Level
new_lvl [Id]
bndrs
  where
    new_lvl :: Level
new_lvl | RecFlag -> Bool
isRec RecFlag
rec = Level -> Level
incMajorLvl Level
lvl
            | Bool
otherwise = Level -> Level
incMinorLvl Level
lvl
      -- Non-recursive join points are one-shot; recursive ones are not

lvlBndrs :: LevelEnv -> Level -> [CoreBndr] -> (LevelEnv, [LevelledBndr])
-- The binders returned are exactly the same as the ones passed,
-- apart from applying the substitution, but they are now paired
-- with a (StayPut level)
--
-- The returned envt has le_ctxt_lvl updated to the new_lvl
--
-- All the new binders get the same level, because
-- any floating binding is either going to float past
-- all or none.  We never separate binders.
lvlBndrs :: LevelEnv -> Level -> [Id] -> (LevelEnv, [LevelledBndr])
lvlBndrs env :: LevelEnv
env@(LE { le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env }) new_lvl :: Level
new_lvl bndrs :: [Id]
bndrs
  = ( LevelEnv
env { le_ctxt_lvl :: Level
le_ctxt_lvl = Level
new_lvl
          , le_join_ceil :: Level
le_join_ceil = Level
new_lvl
          , le_lvl_env :: VarEnv Level
le_lvl_env  = Level -> VarEnv Level -> [Id] -> VarEnv Level
addLvls Level
new_lvl VarEnv Level
lvl_env [Id]
bndrs }
    , (Id -> LevelledBndr) -> [Id] -> [LevelledBndr]
forall a b. (a -> b) -> [a] -> [b]
map (Level -> Id -> LevelledBndr
stayPut Level
new_lvl) [Id]
bndrs)

stayPut :: Level -> OutVar -> LevelledBndr
stayPut :: Level -> Id -> LevelledBndr
stayPut new_lvl :: Level
new_lvl bndr :: Id
bndr = Id -> FloatSpec -> LevelledBndr
forall t. Id -> t -> TaggedBndr t
TB Id
bndr (Level -> FloatSpec
StayPut Level
new_lvl)

  -- Destination level is the max Id level of the expression
  -- (We'll abstract the type variables, if any.)
destLevel :: LevelEnv
          -> DVarSet    -- Free vars of the term
          -> TyCoVarSet -- Free in the /type/ of the term
                        -- (a subset of the previous argument)
          -> Bool   -- True <=> is function
          -> Bool   -- True <=> is bottom
          -> Bool   -- True <=> is a join point
          -> Level
-- INVARIANT: if is_join=True then result >= join_ceiling
destLevel :: LevelEnv -> FVAnn -> TyCoVarSet -> Bool -> Bool -> Bool -> Level
destLevel env :: LevelEnv
env fvs :: FVAnn
fvs fvs_ty :: TyCoVarSet
fvs_ty is_function :: Bool
is_function is_bot :: Bool
is_bot is_join :: Bool
is_join
  | Level -> Bool
isTopLvl Level
max_fv_id_level  -- Float even joins if they get to top level
                              -- See Note [Floating join point bindings]
  = Level
tOP_LEVEL

  | Bool
is_join  -- Never float a join point past the join ceiling
             -- See Note [Join points] in FloatOut
  = if Level
max_fv_id_level Level -> Level -> Bool
`ltLvl` Level
join_ceiling
    then Level
join_ceiling
    else Level
max_fv_id_level

  | Bool
is_bot              -- Send bottoming bindings to the top
  = Level
as_far_as_poss      -- regardless; see Note [Bottoming floats]
                        -- Esp Bottoming floats (1)

  | Just n_args :: Int
n_args <- LevelEnv -> Maybe Int
floatLams LevelEnv
env
  , Int
n_args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0  -- n=0 case handled uniformly by the 'otherwise' case
  , Bool
is_function
  , FVAnn -> Int
countFreeIds FVAnn
fvs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n_args
  = Level
as_far_as_poss  -- Send functions to top level; see
                    -- the comments with isFunction

  | Bool
otherwise = Level
max_fv_id_level
  where
    join_ceiling :: Level
join_ceiling    = LevelEnv -> Level
joinCeilingLevel LevelEnv
env
    max_fv_id_level :: Level
max_fv_id_level = (Id -> Bool) -> LevelEnv -> FVAnn -> Level
maxFvLevel Id -> Bool
isId LevelEnv
env FVAnn
fvs -- Max over Ids only; the
                                              -- tyvars will be abstracted

    as_far_as_poss :: Level
as_far_as_poss = (Id -> Bool) -> LevelEnv -> TyCoVarSet -> Level
maxFvLevel' Id -> Bool
isId LevelEnv
env TyCoVarSet
fvs_ty
                     -- See Note [Floating and kind casts]

{- Note [Floating and kind casts]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this
   case x of
     K (co :: * ~# k) -> let v :: Int |> co
                             v = e
                         in blah

Then, even if we are abstracting over Ids, or if e is bottom, we can't
float v outside the 'co' binding.  Reason: if we did we'd get
    v' :: forall k. (Int ~# Age) => Int |> co
and now 'co' isn't in scope in that type. The underlying reason is
that 'co' is a value-level thing and we can't abstract over that in a
type (else we'd get a dependent type).  So if v's /type/ mentions 'co'
we can't float it out beyond the binding site of 'co'.

That's why we have this as_far_as_poss stuff.  Usually as_far_as_poss
is just tOP_LEVEL; but occasionally a coercion variable (which is an
Id) mentioned in type prevents this.

Example Trac #14270 comment:15.
-}


isFunction :: CoreExprWithFVs -> Bool
-- The idea here is that we want to float *functions* to
-- the top level.  This saves no work, but
--      (a) it can make the host function body a lot smaller,
--              and hence inlinable.
--      (b) it can also save allocation when the function is recursive:
--          h = \x -> letrec f = \y -> ...f...y...x...
--                    in f x
--     becomes
--          f = \x y -> ...(f x)...y...x...
--          h = \x -> f x x
--     No allocation for f now.
-- We may only want to do this if there are sufficiently few free
-- variables.  We certainly only want to do it for values, and not for
-- constructors.  So the simple thing is just to look for lambdas
isFunction :: CoreExprWithFVs -> Bool
isFunction (_, AnnLam b :: Id
b e :: CoreExprWithFVs
e) | Id -> Bool
isId Id
b    = Bool
True
                           | Bool
otherwise = CoreExprWithFVs -> Bool
isFunction CoreExprWithFVs
e
-- isFunction (_, AnnTick _ e)         = isFunction e  -- dubious
isFunction _                           = Bool
False

countFreeIds :: DVarSet -> Int
countFreeIds :: FVAnn -> Int
countFreeIds = (Id -> Int -> Int) -> Int -> UniqDFM Id -> Int
forall elt a. (elt -> a -> a) -> a -> UniqDFM elt -> a
nonDetFoldUDFM Id -> Int -> Int
add 0 (UniqDFM Id -> Int) -> (FVAnn -> UniqDFM Id) -> FVAnn -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FVAnn -> UniqDFM Id
forall a. UniqDSet a -> UniqDFM a
getUniqDSet
  -- It's OK to use nonDetFoldUDFM here because we're just counting things.
  where
    add :: Var -> Int -> Int
    add :: Id -> Int -> Int
add v :: Id
v n :: Int
n | Id -> Bool
isId Id
v    = Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+1
            | Bool
otherwise = Int
n

{-
************************************************************************
*                                                                      *
\subsection{Free-To-Level Monad}
*                                                                      *
************************************************************************
-}

data LevelEnv
  = LE { LevelEnv -> FloatOutSwitches
le_switches :: FloatOutSwitches
       , LevelEnv -> Level
le_ctxt_lvl :: Level           -- The current level
       , LevelEnv -> VarEnv Level
le_lvl_env  :: VarEnv Level    -- Domain is *post-cloned* TyVars and Ids
       , LevelEnv -> Level
le_join_ceil:: Level           -- Highest level to which joins float
                                        -- Invariant: always >= le_ctxt_lvl

       -- See Note [le_subst and le_env]
       , LevelEnv -> Subst
le_subst    :: Subst           -- Domain is pre-cloned TyVars and Ids
                                        -- The Id -> CoreExpr in the Subst is ignored
                                        -- (since we want to substitute a LevelledExpr for
                                        -- an Id via le_env) but we do use the Co/TyVar substs
       , LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env      :: IdEnv ([OutVar], LevelledExpr)  -- Domain is pre-cloned Ids
    }

{- Note [le_subst and le_env]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We clone let- and case-bound variables so that they are still distinct
when floated out; hence the le_subst/le_env.  (see point 3 of the
module overview comment).  We also use these envs when making a
variable polymorphic because we want to float it out past a big
lambda.

The le_subst and le_env always implement the same mapping,
     in_x :->  out_x a b
where out_x is an OutVar, and a,b are its arguments (when
we perform abstraction at the same time as floating).

  le_subst maps to CoreExpr
  le_env   maps to LevelledExpr

Since the range is always a variable or application, there is never
any difference between the two, but sadly the types differ.  The
le_subst is used when substituting in a variable's IdInfo; the le_env
when we find a Var.

In addition the le_env records a [OutVar] of variables free in the
OutExpr/LevelledExpr, just so we don't have to call freeVars
repeatedly.  This list is always non-empty, and the first element is
out_x

The domain of the both envs is *pre-cloned* Ids, though

The domain of the le_lvl_env is the *post-cloned* Ids
-}

initialEnv :: FloatOutSwitches -> LevelEnv
initialEnv :: FloatOutSwitches -> LevelEnv
initialEnv float_lams :: FloatOutSwitches
float_lams
  = LE :: FloatOutSwitches
-> Level
-> VarEnv Level
-> Level
-> Subst
-> IdEnv ([Id], LevelledExpr)
-> LevelEnv
LE { le_switches :: FloatOutSwitches
le_switches = FloatOutSwitches
float_lams
       , le_ctxt_lvl :: Level
le_ctxt_lvl = Level
tOP_LEVEL
       , le_join_ceil :: Level
le_join_ceil = String -> Level
forall a. String -> a
panic "initialEnv"
       , le_lvl_env :: VarEnv Level
le_lvl_env = VarEnv Level
forall a. VarEnv a
emptyVarEnv
       , le_subst :: Subst
le_subst = Subst
emptySubst
       , le_env :: IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
forall a. VarEnv a
emptyVarEnv }

addLvl :: Level -> VarEnv Level -> OutVar -> VarEnv Level
addLvl :: Level -> VarEnv Level -> Id -> VarEnv Level
addLvl dest_lvl :: Level
dest_lvl env :: VarEnv Level
env v' :: Id
v' = VarEnv Level -> Id -> Level -> VarEnv Level
forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv VarEnv Level
env Id
v' Level
dest_lvl

addLvls :: Level -> VarEnv Level -> [OutVar] -> VarEnv Level
addLvls :: Level -> VarEnv Level -> [Id] -> VarEnv Level
addLvls dest_lvl :: Level
dest_lvl env :: VarEnv Level
env vs :: [Id]
vs = (VarEnv Level -> Id -> VarEnv Level)
-> VarEnv Level -> [Id] -> VarEnv Level
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Level -> VarEnv Level -> Id -> VarEnv Level
addLvl Level
dest_lvl) VarEnv Level
env [Id]
vs

floatLams :: LevelEnv -> Maybe Int
floatLams :: LevelEnv -> Maybe Int
floatLams le :: LevelEnv
le = FloatOutSwitches -> Maybe Int
floatOutLambdas (LevelEnv -> FloatOutSwitches
le_switches LevelEnv
le)

floatConsts :: LevelEnv -> Bool
floatConsts :: LevelEnv -> Bool
floatConsts le :: LevelEnv
le = FloatOutSwitches -> Bool
floatOutConstants (LevelEnv -> FloatOutSwitches
le_switches LevelEnv
le)

floatOverSat :: LevelEnv -> Bool
floatOverSat :: LevelEnv -> Bool
floatOverSat le :: LevelEnv
le = FloatOutSwitches -> Bool
floatOutOverSatApps (LevelEnv -> FloatOutSwitches
le_switches LevelEnv
le)

floatTopLvlOnly :: LevelEnv -> Bool
floatTopLvlOnly :: LevelEnv -> Bool
floatTopLvlOnly le :: LevelEnv
le = FloatOutSwitches -> Bool
floatToTopLevelOnly (LevelEnv -> FloatOutSwitches
le_switches LevelEnv
le)

incMinorLvlFrom :: LevelEnv -> Level
incMinorLvlFrom :: LevelEnv -> Level
incMinorLvlFrom env :: LevelEnv
env = Level -> Level
incMinorLvl (LevelEnv -> Level
le_ctxt_lvl LevelEnv
env)

-- extendCaseBndrEnv adds the mapping case-bndr->scrut-var if it can
-- See Note [Binder-swap during float-out]
extendCaseBndrEnv :: LevelEnv
                  -> Id                 -- Pre-cloned case binder
                  -> Expr LevelledBndr  -- Post-cloned scrutinee
                  -> LevelEnv
extendCaseBndrEnv :: LevelEnv -> Id -> LevelledExpr -> LevelEnv
extendCaseBndrEnv le :: LevelEnv
le@(LE { le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env })
                  case_bndr :: Id
case_bndr (Var scrut_var :: Id
scrut_var)
  = LevelEnv
le { le_subst :: Subst
le_subst   = Subst -> Id -> Id -> Subst
extendSubstWithVar Subst
subst Id
case_bndr Id
scrut_var
       , le_env :: IdEnv ([Id], LevelledExpr)
le_env     = IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
add_id IdEnv ([Id], LevelledExpr)
id_env (Id
case_bndr, Id
scrut_var) }
extendCaseBndrEnv env :: LevelEnv
env _ _ = LevelEnv
env

-- See Note [Join ceiling]
placeJoinCeiling :: LevelEnv -> LevelEnv
placeJoinCeiling :: LevelEnv -> LevelEnv
placeJoinCeiling le :: LevelEnv
le@(LE { le_ctxt_lvl :: LevelEnv -> Level
le_ctxt_lvl = Level
lvl })
  = LevelEnv
le { le_ctxt_lvl :: Level
le_ctxt_lvl = Level
lvl', le_join_ceil :: Level
le_join_ceil = Level
lvl' }
  where
    lvl' :: Level
lvl' = Level -> Level
asJoinCeilLvl (Level -> Level
incMinorLvl Level
lvl)

maxFvLevel :: (Var -> Bool) -> LevelEnv -> DVarSet -> Level
maxFvLevel :: (Id -> Bool) -> LevelEnv -> FVAnn -> Level
maxFvLevel max_me :: Id -> Bool
max_me env :: LevelEnv
env var_set :: FVAnn
var_set
  = (Id -> Level -> Level) -> Level -> FVAnn -> Level
forall a. (Id -> a -> a) -> a -> FVAnn -> a
foldDVarSet ((Id -> Bool) -> LevelEnv -> Id -> Level -> Level
maxIn Id -> Bool
max_me LevelEnv
env) Level
tOP_LEVEL FVAnn
var_set

maxFvLevel' :: (Var -> Bool) -> LevelEnv -> TyCoVarSet -> Level
-- Same but for TyCoVarSet
maxFvLevel' :: (Id -> Bool) -> LevelEnv -> TyCoVarSet -> Level
maxFvLevel' max_me :: Id -> Bool
max_me env :: LevelEnv
env var_set :: TyCoVarSet
var_set
  = (Id -> Level -> Level) -> Level -> TyCoVarSet -> Level
forall elt a. (elt -> a -> a) -> a -> UniqSet elt -> a
nonDetFoldUniqSet ((Id -> Bool) -> LevelEnv -> Id -> Level -> Level
maxIn Id -> Bool
max_me LevelEnv
env) Level
tOP_LEVEL TyCoVarSet
var_set

maxIn :: (Var -> Bool) -> LevelEnv -> InVar -> Level -> Level
maxIn :: (Id -> Bool) -> LevelEnv -> Id -> Level -> Level
maxIn max_me :: Id -> Bool
max_me (LE { le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env }) in_var :: Id
in_var lvl :: Level
lvl
  = case IdEnv ([Id], LevelledExpr) -> Id -> Maybe ([Id], LevelledExpr)
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv IdEnv ([Id], LevelledExpr)
id_env Id
in_var of
      Just (abs_vars :: [Id]
abs_vars, _) -> (Id -> Level -> Level) -> Level -> [Id] -> Level
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Id -> Level -> Level
max_out Level
lvl [Id]
abs_vars
      Nothing            -> Id -> Level -> Level
max_out Id
in_var Level
lvl
  where
    max_out :: Id -> Level -> Level
max_out out_var :: Id
out_var lvl :: Level
lvl
        | Id -> Bool
max_me Id
out_var = case VarEnv Level -> Id -> Maybe Level
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv VarEnv Level
lvl_env Id
out_var of
                                Just lvl' :: Level
lvl' -> Level -> Level -> Level
maxLvl Level
lvl' Level
lvl
                                Nothing   -> Level
lvl
        | Bool
otherwise = Level
lvl       -- Ignore some vars depending on max_me

lookupVar :: LevelEnv -> Id -> LevelledExpr
lookupVar :: LevelEnv -> Id -> LevelledExpr
lookupVar le :: LevelEnv
le v :: Id
v = case IdEnv ([Id], LevelledExpr) -> Id -> Maybe ([Id], LevelledExpr)
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env LevelEnv
le) Id
v of
                    Just (_, expr :: LevelledExpr
expr) -> LevelledExpr
expr
                    _              -> Id -> LevelledExpr
forall b. Id -> Expr b
Var Id
v

-- Level to which join points are allowed to float (boundary of current tail
-- context). See Note [Join ceiling]
joinCeilingLevel :: LevelEnv -> Level
joinCeilingLevel :: LevelEnv -> Level
joinCeilingLevel = LevelEnv -> Level
le_join_ceil

abstractVars :: Level -> LevelEnv -> DVarSet -> [OutVar]
        -- Find the variables in fvs, free vars of the target expression,
        -- whose level is greater than the destination level
        -- These are the ones we are going to abstract out
        --
        -- Note that to get reproducible builds, the variables need to be
        -- abstracted in deterministic order, not dependent on the values of
        -- Uniques. This is achieved by using DVarSets, deterministic free
        -- variable computation and deterministic sort.
        -- See Note [Unique Determinism] in Unique for explanation of why
        -- Uniques are not deterministic.
abstractVars :: Level -> LevelEnv -> FVAnn -> [Id]
abstractVars dest_lvl :: Level
dest_lvl (LE { le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env }) in_fvs :: FVAnn
in_fvs
  =  -- NB: sortQuantVars might not put duplicates next to each other
    (Id -> Id) -> [Id] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Id
zap ([Id] -> [Id]) -> [Id] -> [Id]
forall a b. (a -> b) -> a -> b
$ [Id] -> [Id]
sortQuantVars ([Id] -> [Id]) -> [Id] -> [Id]
forall a b. (a -> b) -> a -> b
$
    (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter Id -> Bool
abstract_me      ([Id] -> [Id]) -> [Id] -> [Id]
forall a b. (a -> b) -> a -> b
$
    FVAnn -> [Id]
dVarSetElems            (FVAnn -> [Id]) -> FVAnn -> [Id]
forall a b. (a -> b) -> a -> b
$
    FVAnn -> FVAnn
closeOverKindsDSet      (FVAnn -> FVAnn) -> FVAnn -> FVAnn
forall a b. (a -> b) -> a -> b
$
    Subst -> FVAnn -> FVAnn
substDVarSet Subst
subst FVAnn
in_fvs
        -- NB: it's important to call abstract_me only on the OutIds the
        -- come from substDVarSet (not on fv, which is an InId)
  where
    abstract_me :: Id -> Bool
abstract_me v :: Id
v = case VarEnv Level -> Id -> Maybe Level
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv VarEnv Level
lvl_env Id
v of
                        Just lvl :: Level
lvl -> Level
dest_lvl Level -> Level -> Bool
`ltLvl` Level
lvl
                        Nothing  -> Bool
False

        -- We are going to lambda-abstract, so nuke any IdInfo,
        -- and add the tyvars of the Id (if necessary)
    zap :: Id -> Id
zap v :: Id
v | Id -> Bool
isId Id
v = WARN( isStableUnfolding (idUnfolding v) ||
                           not (isEmptyRuleInfo (idSpecialisation v)),
                           text "absVarsOf: discarding info on" <+> ppr v )
                     Id -> IdInfo -> Id
setIdInfo Id
v IdInfo
vanillaIdInfo
          | Bool
otherwise = Id
v

type LvlM result = UniqSM result

initLvl :: UniqSupply -> UniqSM a -> a
initLvl :: UniqSupply -> UniqSM a -> a
initLvl = UniqSupply -> UniqSM a -> a
forall a. UniqSupply -> UniqSM a -> a
initUs_

newPolyBndrs :: Level -> LevelEnv -> [OutVar] -> [InId]
             -> LvlM (LevelEnv, [OutId])
-- The envt is extended to bind the new bndrs to dest_lvl, but
-- the le_ctxt_lvl is unaffected
newPolyBndrs :: Level -> LevelEnv -> [Id] -> [Id] -> LvlM (LevelEnv, [Id])
newPolyBndrs dest_lvl :: Level
dest_lvl
             env :: LevelEnv
env@(LE { le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env, le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env })
             abs_vars :: [Id]
abs_vars bndrs :: [Id]
bndrs
 = ASSERT( all (not . isCoVar) bndrs )   -- What would we add to the CoSubst in this case. No easy answer.
   do { [Unique]
uniqs <- UniqSM [Unique]
forall (m :: * -> *). MonadUnique m => m [Unique]
getUniquesM
      ; let new_bndrs :: [Id]
new_bndrs = (Id -> Unique -> Id) -> [Id] -> [Unique] -> [Id]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Id -> Unique -> Id
mk_poly_bndr [Id]
bndrs [Unique]
uniqs
            bndr_prs :: [(Id, Id)]
bndr_prs  = [Id]
bndrs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
new_bndrs
            env' :: LevelEnv
env' = LevelEnv
env { le_lvl_env :: VarEnv Level
le_lvl_env = Level -> VarEnv Level -> [Id] -> VarEnv Level
addLvls Level
dest_lvl VarEnv Level
lvl_env [Id]
new_bndrs
                       , le_subst :: Subst
le_subst   = (Subst -> (Id, Id) -> Subst) -> Subst -> [(Id, Id)] -> Subst
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Subst -> (Id, Id) -> Subst
add_subst Subst
subst   [(Id, Id)]
bndr_prs
                       , le_env :: IdEnv ([Id], LevelledExpr)
le_env     = (IdEnv ([Id], LevelledExpr)
 -> (Id, Id) -> IdEnv ([Id], LevelledExpr))
-> IdEnv ([Id], LevelledExpr)
-> [(Id, Id)]
-> IdEnv ([Id], LevelledExpr)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
forall b.
VarEnv ([Id], Expr b) -> (Id, Id) -> VarEnv ([Id], Expr b)
add_id    IdEnv ([Id], LevelledExpr)
id_env  [(Id, Id)]
bndr_prs }
      ; (LevelEnv, [Id]) -> LvlM (LevelEnv, [Id])
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelEnv
env', [Id]
new_bndrs) }
  where
    add_subst :: Subst -> (Id, Id) -> Subst
add_subst env :: Subst
env (v :: Id
v, v' :: Id
v') = Subst -> Id -> Expr Id -> Subst
extendIdSubst Subst
env Id
v (Expr Id -> [Id] -> Expr Id
forall b. Expr b -> [Id] -> Expr b
mkVarApps (Id -> Expr Id
forall b. Id -> Expr b
Var Id
v') [Id]
abs_vars)
    add_id :: VarEnv ([Id], Expr b) -> (Id, Id) -> VarEnv ([Id], Expr b)
add_id    env :: VarEnv ([Id], Expr b)
env (v :: Id
v, v' :: Id
v') = VarEnv ([Id], Expr b)
-> Id -> ([Id], Expr b) -> VarEnv ([Id], Expr b)
forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv VarEnv ([Id], Expr b)
env Id
v ((Id
v'Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
:[Id]
abs_vars), Expr b -> [Id] -> Expr b
forall b. Expr b -> [Id] -> Expr b
mkVarApps (Id -> Expr b
forall b. Id -> Expr b
Var Id
v') [Id]
abs_vars)

    mk_poly_bndr :: Id -> Unique -> Id
mk_poly_bndr bndr :: Id
bndr uniq :: Unique
uniq = Id -> [Id] -> Id -> Id
transferPolyIdInfo Id
bndr [Id]
abs_vars (Id -> Id) -> Id -> Id
forall a b. (a -> b) -> a -> b
$         -- Note [transferPolyIdInfo] in Id.hs
                             Id -> Id -> Id
transfer_join_info Id
bndr (Id -> Id) -> Id -> Id
forall a b. (a -> b) -> a -> b
$
                             FastString -> Unique -> Type -> Id
mkSysLocalOrCoVar (String -> FastString
mkFastString String
str) Unique
uniq Type
poly_ty
                           where
                             str :: String
str     = "poly_" String -> String -> String
forall a. [a] -> [a] -> [a]
++ OccName -> String
occNameString (Id -> OccName
forall a. NamedThing a => a -> OccName
getOccName Id
bndr)
                             poly_ty :: Type
poly_ty = [Id] -> Type -> Type
mkLamTypes [Id]
abs_vars (Subst -> Type -> Type
CoreSubst.substTy Subst
subst (Id -> Type
idType Id
bndr))

    -- If we are floating a join point to top level, it stops being
    -- a join point.  Otherwise it continues to be a join point,
    -- but we may need to adjust its arity
    dest_is_top :: Bool
dest_is_top = Level -> Bool
isTopLvl Level
dest_lvl
    transfer_join_info :: Id -> Id -> Id
transfer_join_info bndr :: Id
bndr new_bndr :: Id
new_bndr
      | Just join_arity :: Int
join_arity <- Id -> Maybe Int
isJoinId_maybe Id
bndr
      , Bool -> Bool
not Bool
dest_is_top
      = Id
new_bndr Id -> Int -> Id
`asJoinId` Int
join_arity Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Id] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
abs_vars
      | Bool
otherwise
      = Id
new_bndr

newLvlVar :: LevelledExpr        -- The RHS of the new binding
          -> Maybe JoinArity     -- Its join arity, if it is a join point
          -> Bool                -- True <=> the RHS looks like (makeStatic ...)
          -> LvlM Id
newLvlVar :: LevelledExpr -> Maybe Int -> Bool -> LvlM Id
newLvlVar lvld_rhs :: LevelledExpr
lvld_rhs join_arity_maybe :: Maybe Int
join_arity_maybe is_mk_static :: Bool
is_mk_static
  = do { Unique
uniq <- UniqSM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
       ; Id -> LvlM Id
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> Id
add_join_info (Unique -> Type -> Id
mk_id Unique
uniq Type
rhs_ty))
       }
  where
    add_join_info :: Id -> Id
add_join_info var :: Id
var = Id
var Id -> Maybe Int -> Id
`asJoinId_maybe` Maybe Int
join_arity_maybe
    de_tagged_rhs :: Expr Id
de_tagged_rhs = LevelledExpr -> Expr Id
forall t. TaggedExpr t -> Expr Id
deTagExpr LevelledExpr
lvld_rhs
    rhs_ty :: Type
rhs_ty        = Expr Id -> Type
exprType Expr Id
de_tagged_rhs

    mk_id :: Unique -> Type -> Id
mk_id uniq :: Unique
uniq rhs_ty :: Type
rhs_ty
      -- See Note [Grand plan for static forms] in StaticPtrTable.
      | Bool
is_mk_static
      = Name -> Type -> Id
mkExportedVanillaId (Unique -> FastString -> Name
mkSystemVarName Unique
uniq (String -> FastString
mkFastString "static_ptr"))
                            Type
rhs_ty
      | Bool
otherwise
      = FastString -> Unique -> Type -> Id
mkSysLocalOrCoVar (String -> FastString
mkFastString "lvl") Unique
uniq Type
rhs_ty

cloneCaseBndrs :: LevelEnv -> Level -> [Var] -> LvlM (LevelEnv, [Var])
cloneCaseBndrs :: LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneCaseBndrs env :: LevelEnv
env@(LE { le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env })
               new_lvl :: Level
new_lvl vs :: [Id]
vs
  = do { UniqSupply
us <- UniqSM UniqSupply
forall (m :: * -> *). MonadUnique m => m UniqSupply
getUniqueSupplyM
       ; let (subst' :: Subst
subst', vs' :: [Id]
vs') = Subst -> UniqSupply -> [Id] -> (Subst, [Id])
cloneBndrs Subst
subst UniqSupply
us [Id]
vs
             env' :: LevelEnv
env' = LevelEnv
env { le_lvl_env :: VarEnv Level
le_lvl_env   = Level -> VarEnv Level -> [Id] -> VarEnv Level
addLvls Level
new_lvl VarEnv Level
lvl_env [Id]
vs'
                        , le_subst :: Subst
le_subst     = Subst
subst'
                        , le_env :: IdEnv ([Id], LevelledExpr)
le_env       = (IdEnv ([Id], LevelledExpr)
 -> (Id, Id) -> IdEnv ([Id], LevelledExpr))
-> IdEnv ([Id], LevelledExpr)
-> [(Id, Id)]
-> IdEnv ([Id], LevelledExpr)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
add_id IdEnv ([Id], LevelledExpr)
id_env ([Id]
vs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
vs') }

       ; (LevelEnv, [Id]) -> LvlM (LevelEnv, [Id])
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelEnv
env', [Id]
vs') }

cloneLetVars :: RecFlag -> LevelEnv -> Level -> [InVar]
             -> LvlM (LevelEnv, [OutVar])
-- See Note [Need for cloning during float-out]
-- Works for Ids bound by let(rec)
-- The dest_lvl is attributed to the binders in the new env,
-- but cloneVars doesn't affect the le_ctxt_lvl of the incoming env
cloneLetVars :: RecFlag -> LevelEnv -> Level -> [Id] -> LvlM (LevelEnv, [Id])
cloneLetVars is_rec :: RecFlag
is_rec
          env :: LevelEnv
env@(LE { le_subst :: LevelEnv -> Subst
le_subst = Subst
subst, le_lvl_env :: LevelEnv -> VarEnv Level
le_lvl_env = VarEnv Level
lvl_env, le_env :: LevelEnv -> IdEnv ([Id], LevelledExpr)
le_env = IdEnv ([Id], LevelledExpr)
id_env })
          dest_lvl :: Level
dest_lvl vs :: [Id]
vs
  = do { UniqSupply
us <- UniqSM UniqSupply
forall (m :: * -> *). MonadUnique m => m UniqSupply
getUniqueSupplyM
       ; let vs1 :: [Id]
vs1  = (Id -> Id) -> [Id] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Id
zap [Id]
vs
                      -- See Note [Zapping the demand info]
             (subst' :: Subst
subst', vs2 :: [Id]
vs2) = case RecFlag
is_rec of
                               NonRecursive -> Subst -> UniqSupply -> [Id] -> (Subst, [Id])
cloneBndrs      Subst
subst UniqSupply
us [Id]
vs1
                               Recursive    -> Subst -> UniqSupply -> [Id] -> (Subst, [Id])
cloneRecIdBndrs Subst
subst UniqSupply
us [Id]
vs1
             prs :: [(Id, Id)]
prs  = [Id]
vs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
vs2
             env' :: LevelEnv
env' = LevelEnv
env { le_lvl_env :: VarEnv Level
le_lvl_env = Level -> VarEnv Level -> [Id] -> VarEnv Level
addLvls Level
dest_lvl VarEnv Level
lvl_env [Id]
vs2
                        , le_subst :: Subst
le_subst   = Subst
subst'
                        , le_env :: IdEnv ([Id], LevelledExpr)
le_env     = (IdEnv ([Id], LevelledExpr)
 -> (Id, Id) -> IdEnv ([Id], LevelledExpr))
-> IdEnv ([Id], LevelledExpr)
-> [(Id, Id)]
-> IdEnv ([Id], LevelledExpr)
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
add_id IdEnv ([Id], LevelledExpr)
id_env [(Id, Id)]
prs }

       ; (LevelEnv, [Id]) -> LvlM (LevelEnv, [Id])
forall (m :: * -> *) a. Monad m => a -> m a
return (LevelEnv
env', [Id]
vs2) }
  where
    zap :: Var -> Var
    zap :: Id -> Id
zap v :: Id
v | Id -> Bool
isId Id
v    = Id -> Id
zap_join (Id -> Id
zapIdDemandInfo Id
v)
          | Bool
otherwise = Id
v

    zap_join :: Id -> Id
zap_join | Level -> Bool
isTopLvl Level
dest_lvl = Id -> Id
zapJoinId
             | Bool
otherwise         = Id -> Id
forall a. a -> a
id

add_id :: IdEnv ([Var], LevelledExpr) -> (Var, Var) -> IdEnv ([Var], LevelledExpr)
add_id :: IdEnv ([Id], LevelledExpr)
-> (Id, Id) -> IdEnv ([Id], LevelledExpr)
add_id id_env :: IdEnv ([Id], LevelledExpr)
id_env (v :: Id
v, v1 :: Id
v1)
  | Id -> Bool
isTyVar Id
v = IdEnv ([Id], LevelledExpr) -> Id -> IdEnv ([Id], LevelledExpr)
forall a. VarEnv a -> Id -> VarEnv a
delVarEnv    IdEnv ([Id], LevelledExpr)
id_env Id
v
  | Bool
otherwise = IdEnv ([Id], LevelledExpr)
-> Id -> ([Id], LevelledExpr) -> IdEnv ([Id], LevelledExpr)
forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv IdEnv ([Id], LevelledExpr)
id_env Id
v ([Id
v1], ASSERT(not (isCoVar v1)) Var v1)

{-
Note [Zapping the demand info]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
VERY IMPORTANT: we must zap the demand info if the thing is going to
float out, because it may be less demanded than at its original
binding site.  Eg
   f :: Int -> Int
   f x = let v = 3*4 in v+x
Here v is strict; but if we float v to top level, it isn't any more.

Similarly, if we're floating a join point, it won't be one anymore, so we zap
join point information as well.
-}