{-
(c) The University of Glasgow, 1994-2006


Core pass to saturate constructors and PrimOps
-}

{-# LANGUAGE BangPatterns, CPP, MultiWayIf #-}

module CorePrep (
      corePrepPgm, corePrepExpr, cvtLitInteger, cvtLitNatural,
      lookupMkIntegerName, lookupIntegerSDataConName,
      lookupMkNaturalName, lookupNaturalSDataConName
  ) where

#include "HsVersions.h"

import GhcPrelude

import OccurAnal

import HscTypes
import PrelNames
import MkId             ( realWorldPrimId )
import CoreUtils
import CoreArity
import CoreFVs
import CoreMonad        ( CoreToDo(..) )
import CoreLint         ( endPassIO )
import CoreSyn
import CoreSubst
import MkCore hiding( FloatBind(..) )   -- We use our own FloatBind here
import Type
import Literal
import Coercion
import TcEnv
import TyCon
import Demand
import Var
import VarSet
import VarEnv
import Id
import IdInfo
import TysWiredIn
import DataCon
import BasicTypes
import Module
import UniqSupply
import Maybes
import OrdList
import ErrUtils
import DynFlags
import Util
import Pair
import Outputable
import Platform
import FastString
import Config
import Name             ( NamedThing(..), nameSrcSpan )
import SrcLoc           ( SrcSpan(..), realSrcLocSpan, mkRealSrcLoc )
import Data.Bits
import MonadUtils       ( mapAccumLM )
import Data.List        ( mapAccumL )
import Control.Monad
import CostCentre       ( CostCentre, ccFromThisModule )
import qualified Data.Set as S

{-
-- ---------------------------------------------------------------------------
-- Note [CorePrep Overview]
-- ---------------------------------------------------------------------------

The goal of this pass is to prepare for code generation.

1.  Saturate constructor applications.

2.  Convert to A-normal form; that is, function arguments
    are always variables.

    * Use case for strict arguments:
        f E ==> case E of x -> f x
        (where f is strict)

    * Use let for non-trivial lazy arguments
        f E ==> let x = E in f x
        (were f is lazy and x is non-trivial)

3.  Similarly, convert any unboxed lets into cases.
    [I'm experimenting with leaving 'ok-for-speculation'
     rhss in let-form right up to this point.]

4.  Ensure that *value* lambdas only occur as the RHS of a binding
    (The code generator can't deal with anything else.)
    Type lambdas are ok, however, because the code gen discards them.

5.  [Not any more; nuked Jun 2002] Do the seq/par munging.

6.  Clone all local Ids.
    This means that all such Ids are unique, rather than the
    weaker guarantee of no clashes which the simplifier provides.
    And that is what the code generator needs.

    We don't clone TyVars or CoVars. The code gen doesn't need that,
    and doing so would be tiresome because then we'd need
    to substitute in types and coercions.

7.  Give each dynamic CCall occurrence a fresh unique; this is
    rather like the cloning step above.

8.  Inject bindings for the "implicit" Ids:
        * Constructor wrappers
        * Constructor workers
    We want curried definitions for all of these in case they
    aren't inlined by some caller.

9.  Replace (lazy e) by e.  See Note [lazyId magic] in MkId.hs
    Also replace (noinline e) by e.

10. Convert (LitInteger i t) into the core representation
    for the Integer i. Normally this uses mkInteger, but if
    we are using the integer-gmp implementation then there is a
    special case where we use the S# constructor for Integers that
    are in the range of Int.

11. Same for LitNatural.

12. Uphold tick consistency while doing this: We move ticks out of
    (non-type) applications where we can, and make sure that we
    annotate according to scoping rules when floating.

13. Collect cost centres (including cost centres in unfoldings) if we're in
    profiling mode. We have to do this here beucase we won't have unfoldings
    after this pass (see `zapUnfolding` and Note [Drop unfoldings and rules].

This is all done modulo type applications and abstractions, so that
when type erasure is done for conversion to STG, we don't end up with
any trivial or useless bindings.


Note [CorePrep invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~
Here is the syntax of the Core produced by CorePrep:

    Trivial expressions
       arg ::= lit |  var
              | arg ty  |  /\a. arg
              | truv co  |  /\c. arg  |  arg |> co

    Applications
       app ::= lit  |  var  |  app arg  |  app ty  | app co | app |> co

    Expressions
       body ::= app
              | let(rec) x = rhs in body     -- Boxed only
              | case body of pat -> body
              | /\a. body | /\c. body
              | body |> co

    Right hand sides (only place where value lambdas can occur)
       rhs ::= /\a.rhs  |  \x.rhs  |  body

We define a synonym for each of these non-terminals.  Functions
with the corresponding name produce a result in that syntax.
-}

type CpeArg  = CoreExpr    -- Non-terminal 'arg'
type CpeApp  = CoreExpr    -- Non-terminal 'app'
type CpeBody = CoreExpr    -- Non-terminal 'body'
type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'

{-
************************************************************************
*                                                                      *
                Top level stuff
*                                                                      *
************************************************************************
-}

corePrepPgm :: HscEnv -> Module -> ModLocation -> CoreProgram -> [TyCon]
            -> IO (CoreProgram, S.Set CostCentre)
corePrepPgm :: HscEnv
-> Module
-> ModLocation
-> CoreProgram
-> [TyCon]
-> IO (CoreProgram, Set CostCentre)
corePrepPgm hsc_env :: HscEnv
hsc_env this_mod :: Module
this_mod mod_loc :: ModLocation
mod_loc binds :: CoreProgram
binds data_tycons :: [TyCon]
data_tycons =
    IO DynFlags
-> SDoc
-> ((CoreProgram, Set CostCentre) -> ())
-> IO (CoreProgram, Set CostCentre)
-> IO (CoreProgram, Set CostCentre)
forall (m :: * -> *) a.
MonadIO m =>
m DynFlags -> SDoc -> (a -> ()) -> m a -> m a
withTiming (DynFlags -> IO DynFlags
forall (f :: * -> *) a. Applicative f => a -> f a
pure DynFlags
dflags)
               (String -> SDoc
text "CorePrep"SDoc -> SDoc -> SDoc
<+>SDoc -> SDoc
brackets (Module -> SDoc
forall a. Outputable a => a -> SDoc
ppr Module
this_mod))
               (() -> (CoreProgram, Set CostCentre) -> ()
forall a b. a -> b -> a
const ()) (IO (CoreProgram, Set CostCentre)
 -> IO (CoreProgram, Set CostCentre))
-> IO (CoreProgram, Set CostCentre)
-> IO (CoreProgram, Set CostCentre)
forall a b. (a -> b) -> a -> b
$ do
    UniqSupply
us <- Char -> IO UniqSupply
mkSplitUniqSupply 's'
    CorePrepEnv
initialCorePrepEnv <- DynFlags -> HscEnv -> IO CorePrepEnv
mkInitialCorePrepEnv DynFlags
dflags HscEnv
hsc_env

    let cost_centres :: Set CostCentre
cost_centres
          | Way
WayProf Way -> [Way] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` DynFlags -> [Way]
ways DynFlags
dflags
          = Module -> CoreProgram -> Set CostCentre
collectCostCentres Module
this_mod CoreProgram
binds
          | Bool
otherwise
          = Set CostCentre
forall a. Set a
S.empty

        implicit_binds :: CoreProgram
implicit_binds = DynFlags -> ModLocation -> [TyCon] -> CoreProgram
mkDataConWorkers DynFlags
dflags ModLocation
mod_loc [TyCon]
data_tycons
            -- NB: we must feed mkImplicitBinds through corePrep too
            -- so that they are suitably cloned and eta-expanded

        binds_out :: CoreProgram
binds_out = UniqSupply -> UniqSM CoreProgram -> CoreProgram
forall a. UniqSupply -> UniqSM a -> a
initUs_ UniqSupply
us (UniqSM CoreProgram -> CoreProgram)
-> UniqSM CoreProgram -> CoreProgram
forall a b. (a -> b) -> a -> b
$ do
                      Floats
floats1 <- CorePrepEnv -> CoreProgram -> UniqSM Floats
corePrepTopBinds CorePrepEnv
initialCorePrepEnv CoreProgram
binds
                      Floats
floats2 <- CorePrepEnv -> CoreProgram -> UniqSM Floats
corePrepTopBinds CorePrepEnv
initialCorePrepEnv CoreProgram
implicit_binds
                      CoreProgram -> UniqSM CoreProgram
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats -> CoreProgram
deFloatTop (Floats
floats1 Floats -> Floats -> Floats
`appendFloats` Floats
floats2))

    HscEnv
-> PrintUnqualified
-> CoreToDo
-> CoreProgram
-> [CoreRule]
-> IO ()
endPassIO HscEnv
hsc_env PrintUnqualified
alwaysQualify CoreToDo
CorePrep CoreProgram
binds_out []
    (CoreProgram, Set CostCentre) -> IO (CoreProgram, Set CostCentre)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreProgram
binds_out, Set CostCentre
cost_centres)
  where
    dflags :: DynFlags
dflags = HscEnv -> DynFlags
hsc_dflags HscEnv
hsc_env

corePrepExpr :: DynFlags -> HscEnv -> CoreExpr -> IO CoreExpr
corePrepExpr :: DynFlags -> HscEnv -> CoreExpr -> IO CoreExpr
corePrepExpr dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env expr :: CoreExpr
expr =
    IO DynFlags
-> SDoc -> (CoreExpr -> ()) -> IO CoreExpr -> IO CoreExpr
forall (m :: * -> *) a.
MonadIO m =>
m DynFlags -> SDoc -> (a -> ()) -> m a -> m a
withTiming (DynFlags -> IO DynFlags
forall (f :: * -> *) a. Applicative f => a -> f a
pure DynFlags
dflags) (String -> SDoc
text "CorePrep [expr]") (() -> CoreExpr -> ()
forall a b. a -> b -> a
const ()) (IO CoreExpr -> IO CoreExpr) -> IO CoreExpr -> IO CoreExpr
forall a b. (a -> b) -> a -> b
$ do
    UniqSupply
us <- Char -> IO UniqSupply
mkSplitUniqSupply 's'
    CorePrepEnv
initialCorePrepEnv <- DynFlags -> HscEnv -> IO CorePrepEnv
mkInitialCorePrepEnv DynFlags
dflags HscEnv
hsc_env
    let new_expr :: CoreExpr
new_expr = UniqSupply -> UniqSM CoreExpr -> CoreExpr
forall a. UniqSupply -> UniqSM a -> a
initUs_ UniqSupply
us (CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF CorePrepEnv
initialCorePrepEnv CoreExpr
expr)
    DynFlags -> DumpFlag -> String -> SDoc -> IO ()
dumpIfSet_dyn DynFlags
dflags DumpFlag
Opt_D_dump_prep "CorePrep" (CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
new_expr)
    CoreExpr -> IO CoreExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
new_expr

corePrepTopBinds :: CorePrepEnv -> [CoreBind] -> UniqSM Floats
-- Note [Floating out of top level bindings]
corePrepTopBinds :: CorePrepEnv -> CoreProgram -> UniqSM Floats
corePrepTopBinds initialCorePrepEnv :: CorePrepEnv
initialCorePrepEnv binds :: CoreProgram
binds
  = CorePrepEnv -> CoreProgram -> UniqSM Floats
go CorePrepEnv
initialCorePrepEnv CoreProgram
binds
  where
    go :: CorePrepEnv -> CoreProgram -> UniqSM Floats
go _   []             = Floats -> UniqSM Floats
forall (m :: * -> *) a. Monad m => a -> m a
return Floats
emptyFloats
    go env :: CorePrepEnv
env (bind :: CoreBind
bind : binds :: CoreProgram
binds) = do (env' :: CorePrepEnv
env', floats :: Floats
floats, maybe_new_bind :: Maybe CoreBind
maybe_new_bind)
                                 <- TopLevelFlag
-> CorePrepEnv
-> CoreBind
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
cpeBind TopLevelFlag
TopLevel CorePrepEnv
env CoreBind
bind
                               MASSERT(isNothing maybe_new_bind)
                                 -- Only join points get returned this way by
                                 -- cpeBind, and no join point may float to top
                               Floats
floatss <- CorePrepEnv -> CoreProgram -> UniqSM Floats
go CorePrepEnv
env' CoreProgram
binds
                               Floats -> UniqSM Floats
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats Floats -> Floats -> Floats
`appendFloats` Floats
floatss)

mkDataConWorkers :: DynFlags -> ModLocation -> [TyCon] -> [CoreBind]
-- See Note [Data constructor workers]
-- c.f. Note [Injecting implicit bindings] in TidyPgm
mkDataConWorkers :: DynFlags -> ModLocation -> [TyCon] -> CoreProgram
mkDataConWorkers dflags :: DynFlags
dflags mod_loc :: ModLocation
mod_loc data_tycons :: [TyCon]
data_tycons
  = [ Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
id (Name -> CoreExpr -> CoreExpr
forall b. Name -> Expr b -> Expr b
tick_it (DataCon -> Name
forall a. NamedThing a => a -> Name
getName DataCon
data_con) (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
id))
                                -- The ice is thin here, but it works
    | TyCon
tycon <- [TyCon]
data_tycons,     -- CorePrep will eta-expand it
      DataCon
data_con <- TyCon -> [DataCon]
tyConDataCons TyCon
tycon,
      let id :: Id
id = DataCon -> Id
dataConWorkId DataCon
data_con
    ]
 where
   -- If we want to generate debug info, we put a source note on the
   -- worker. This is useful, especially for heap profiling.
   tick_it :: Name -> Expr b -> Expr b
tick_it name :: Name
name
     | DynFlags -> Int
debugLevel DynFlags
dflags Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0                = Expr b -> Expr b
forall a. a -> a
id
     | RealSrcSpan span :: RealSrcSpan
span <- Name -> SrcSpan
nameSrcSpan Name
name  = RealSrcSpan -> Expr b -> Expr b
forall b. RealSrcSpan -> Expr b -> Expr b
tick RealSrcSpan
span
     | Just file :: String
file <- ModLocation -> Maybe String
ml_hs_file ModLocation
mod_loc       = RealSrcSpan -> Expr b -> Expr b
forall b. RealSrcSpan -> Expr b -> Expr b
tick (String -> RealSrcSpan
span1 String
file)
     | Bool
otherwise                             = RealSrcSpan -> Expr b -> Expr b
forall b. RealSrcSpan -> Expr b -> Expr b
tick (String -> RealSrcSpan
span1 "???")
     where tick :: RealSrcSpan -> Expr b -> Expr b
tick span :: RealSrcSpan
span  = Tickish Id -> Expr b -> Expr b
forall b. Tickish Id -> Expr b -> Expr b
Tick (RealSrcSpan -> String -> Tickish Id
forall id. RealSrcSpan -> String -> Tickish id
SourceNote RealSrcSpan
span (String -> Tickish Id) -> String -> Tickish Id
forall a b. (a -> b) -> a -> b
$ DynFlags -> SDoc -> String
showSDoc DynFlags
dflags (Name -> SDoc
forall a. Outputable a => a -> SDoc
ppr Name
name))
           span1 :: String -> RealSrcSpan
span1 file :: String
file = RealSrcLoc -> RealSrcSpan
realSrcLocSpan (RealSrcLoc -> RealSrcSpan) -> RealSrcLoc -> RealSrcSpan
forall a b. (a -> b) -> a -> b
$ FastString -> Int -> Int -> RealSrcLoc
mkRealSrcLoc (String -> FastString
mkFastString String
file) 1 1

{-
Note [Floating out of top level bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
NB: we do need to float out of top-level bindings
Consider        x = length [True,False]
We want to get
                s1 = False : []
                s2 = True  : s1
                x  = length s2

We return a *list* of bindings, because we may start with
        x* = f (g y)
where x is demanded, in which case we want to finish with
        a = g y
        x* = f a
And then x will actually end up case-bound

Note [CafInfo and floating]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
What happens when we try to float bindings to the top level?  At this
point all the CafInfo is supposed to be correct, and we must make certain
that is true of the new top-level bindings.  There are two cases
to consider

a) The top-level binding is marked asCafRefs.  In that case we are
   basically fine.  The floated bindings had better all be lazy lets,
   so they can float to top level, but they'll all have HasCafRefs
   (the default) which is safe.

b) The top-level binding is marked NoCafRefs.  This really happens
   Example.  CoreTidy produces
      $fApplicativeSTM [NoCafRefs] = D:Alternative retry# ...blah...
   Now CorePrep has to eta-expand to
      $fApplicativeSTM = let sat = \xy. retry x y
                         in D:Alternative sat ...blah...
   So what we *want* is
      sat [NoCafRefs] = \xy. retry x y
      $fApplicativeSTM [NoCafRefs] = D:Alternative sat ...blah...

   So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
   *and* substitute the modified 'sat' into the old RHS.

   It should be the case that 'sat' is itself [NoCafRefs] (a value, no
   cafs) else the original top-level binding would not itself have been
   marked [NoCafRefs].  The DEBUG check in CoreToStg for
   consistentCafInfo will find this.

This is all very gruesome and horrible. It would be better to figure
out CafInfo later, after CorePrep.  We'll do that in due course.
Meanwhile this horrible hack works.

Note [Join points and floating]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Join points can float out of other join points but not out of value bindings:

  let z =
    let  w = ... in -- can float
    join k = ... in -- can't float
    ... jump k ...
  join j x1 ... xn =
    let  y = ... in -- can float (but don't want to)
    join h = ... in -- can float (but not much point)
    ... jump h ...
  in ...

Here, the jump to h remains valid if h is floated outward, but the jump to k
does not.

We don't float *out* of join points. It would only be safe to float out of
nullary join points (or ones where the arguments are all either type arguments
or dead binders). Nullary join points aren't ever recursive, so they're always
effectively one-shot functions, which we don't float out of. We *could* float
join points from nullary join points, but there's no clear benefit at this
stage.

Note [Data constructor workers]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Create any necessary "implicit" bindings for data con workers.  We
create the rather strange (non-recursive!) binding

        $wC = \x y -> $wC x y

i.e. a curried constructor that allocates.  This means that we can
treat the worker for a constructor like any other function in the rest
of the compiler.  The point here is that CoreToStg will generate a
StgConApp for the RHS, rather than a call to the worker (which would
give a loop).  As Lennart says: the ice is thin here, but it works.

Hmm.  Should we create bindings for dictionary constructors?  They are
always fully applied, and the bindings are just there to support
partial applications. But it's easier to let them through.


Note [Dead code in CorePrep]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Imagine that we got an input program like this (see Trac #4962):

  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
  f x = (g True (Just x) + g () (Just x), g)
    where
      g :: Show a => a -> Maybe Int -> Int
      g _ Nothing = x
      g y (Just z) = if z > 100 then g y (Just (z + length (show y))) else g y unknown

After specialisation and SpecConstr, we would get something like this:

  f :: Show b => Int -> (Int, b -> Maybe Int -> Int)
  f x = (g$Bool_True_Just x + g$Unit_Unit_Just x, g)
    where
      {-# RULES g $dBool = g$Bool
                g $dUnit = g$Unit #-}
      g = ...
      {-# RULES forall x. g$Bool True (Just x) = g$Bool_True_Just x #-}
      g$Bool = ...
      {-# RULES forall x. g$Unit () (Just x) = g$Unit_Unit_Just x #-}
      g$Unit = ...
      g$Bool_True_Just = ...
      g$Unit_Unit_Just = ...

Note that the g$Bool and g$Unit functions are actually dead code: they
are only kept alive by the occurrence analyser because they are
referred to by the rules of g, which is being kept alive by the fact
that it is used (unspecialised) in the returned pair.

However, at the CorePrep stage there is no way that the rules for g
will ever fire, and it really seems like a shame to produce an output
program that goes to the trouble of allocating a closure for the
unreachable g$Bool and g$Unit functions.

The way we fix this is to:
 * In cloneBndr, drop all unfoldings/rules

 * In deFloatTop, run a simple dead code analyser on each top-level
   RHS to drop the dead local bindings. For that call to OccAnal, we
   disable the binder swap, else the occurrence analyser sometimes
   introduces new let bindings for cased binders, which lead to the bug
   in #5433.

The reason we don't just OccAnal the whole output of CorePrep is that
the tidier ensures that all top-level binders are GlobalIds, so they
don't show up in the free variables any longer. So if you run the
occurrence analyser on the output of CoreTidy (or later) you e.g. turn
this program:

  Rec {
  f = ... f ...
  }

Into this one:

  f = ... f ...

(Since f is not considered to be free in its own RHS.)


************************************************************************
*                                                                      *
                The main code
*                                                                      *
************************************************************************
-}

cpeBind :: TopLevelFlag -> CorePrepEnv -> CoreBind
        -> UniqSM (CorePrepEnv,
                   Floats,         -- Floating value bindings
                   Maybe CoreBind) -- Just bind' <=> returned new bind; no float
                                   -- Nothing <=> added bind' to floats instead
cpeBind :: TopLevelFlag
-> CorePrepEnv
-> CoreBind
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
cpeBind top_lvl :: TopLevelFlag
top_lvl env :: CorePrepEnv
env (NonRec bndr :: Id
bndr rhs :: CoreExpr
rhs)
  | Bool -> Bool
not (Id -> Bool
isJoinId Id
bndr)
  = do { (_, bndr1 :: Id
bndr1) <- CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id)
cpCloneBndr CorePrepEnv
env Id
bndr
       ; let dmd :: Demand
dmd         = Id -> Demand
idDemandInfo Id
bndr
             is_unlifted :: Bool
is_unlifted = HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType (Id -> Type
idType Id
bndr)
       ; (floats :: Floats
floats, rhs1 :: CoreExpr
rhs1) <- TopLevelFlag
-> RecFlag
-> Demand
-> Bool
-> CorePrepEnv
-> Id
-> CoreExpr
-> UniqSM (Floats, CoreExpr)
cpePair TopLevelFlag
top_lvl RecFlag
NonRecursive
                                   Demand
dmd Bool
is_unlifted
                                   CorePrepEnv
env Id
bndr1 CoreExpr
rhs
       -- See Note [Inlining in CorePrep]
       ; if CoreExpr -> Bool
exprIsTrivial CoreExpr
rhs1 Bool -> Bool -> Bool
&& TopLevelFlag -> Bool
isNotTopLevel TopLevelFlag
top_lvl
            then (CorePrepEnv, Floats, Maybe CoreBind)
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> Id -> CoreExpr -> CorePrepEnv
extendCorePrepEnvExpr CorePrepEnv
env Id
bndr CoreExpr
rhs1, Floats
floats, Maybe CoreBind
forall a. Maybe a
Nothing)
            else do {

       ; let new_float :: FloatingBind
new_float = Demand -> Bool -> Id -> CoreExpr -> FloatingBind
mkFloat Demand
dmd Bool
is_unlifted Id
bndr1 CoreExpr
rhs1

       ; (CorePrepEnv, Floats, Maybe CoreBind)
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv CorePrepEnv
env Id
bndr Id
bndr1,
                 Floats -> FloatingBind -> Floats
addFloat Floats
floats FloatingBind
new_float,
                 Maybe CoreBind
forall a. Maybe a
Nothing) }}

  | Bool
otherwise -- A join point; see Note [Join points and floating]
  = ASSERT(not (isTopLevel top_lvl)) -- can't have top-level join point
    do { (_, bndr1 :: Id
bndr1) <- CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id)
cpCloneBndr CorePrepEnv
env Id
bndr
       ; (bndr2 :: Id
bndr2, rhs1 :: CoreExpr
rhs1) <- CorePrepEnv -> Id -> CoreExpr -> UniqSM (Id, CoreExpr)
cpeJoinPair CorePrepEnv
env Id
bndr1 CoreExpr
rhs
       ; (CorePrepEnv, Floats, Maybe CoreBind)
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv CorePrepEnv
env Id
bndr Id
bndr2,
                 Floats
emptyFloats,
                 CoreBind -> Maybe CoreBind
forall a. a -> Maybe a
Just (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
bndr2 CoreExpr
rhs1)) }

cpeBind top_lvl :: TopLevelFlag
top_lvl env :: CorePrepEnv
env (Rec pairs :: [(Id, CoreExpr)]
pairs)
  | Bool -> Bool
not (Id -> Bool
isJoinId ([Id] -> Id
forall a. [a] -> a
head [Id]
bndrs))
  = do { (env' :: CorePrepEnv
env', bndrs1 :: [Id]
bndrs1) <- CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs CorePrepEnv
env [Id]
bndrs
       ; [(Floats, CoreExpr)]
stuff <- (Id -> CoreExpr -> UniqSM (Floats, CoreExpr))
-> [Id] -> [CoreExpr] -> UniqSM [(Floats, CoreExpr)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (TopLevelFlag
-> RecFlag
-> Demand
-> Bool
-> CorePrepEnv
-> Id
-> CoreExpr
-> UniqSM (Floats, CoreExpr)
cpePair TopLevelFlag
top_lvl RecFlag
Recursive Demand
topDmd Bool
False CorePrepEnv
env')
                           [Id]
bndrs1 [CoreExpr]
rhss

       ; let (floats_s :: [Floats]
floats_s, rhss1 :: [CoreExpr]
rhss1) = [(Floats, CoreExpr)] -> ([Floats], [CoreExpr])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Floats, CoreExpr)]
stuff
             all_pairs :: [(Id, CoreExpr)]
all_pairs = (FloatingBind -> [(Id, CoreExpr)] -> [(Id, CoreExpr)])
-> [(Id, CoreExpr)] -> OrdList FloatingBind -> [(Id, CoreExpr)]
forall a b. (a -> b -> b) -> b -> OrdList a -> b
foldrOL FloatingBind -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
add_float ([Id]
bndrs1 [Id] -> [CoreExpr] -> [(Id, CoreExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [CoreExpr]
rhss1)
                                           ([Floats] -> OrdList FloatingBind
concatFloats [Floats]
floats_s)

       ; (CorePrepEnv, Floats, Maybe CoreBind)
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> [(Id, Id)] -> CorePrepEnv
extendCorePrepEnvList CorePrepEnv
env ([Id]
bndrs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
bndrs1),
                 FloatingBind -> Floats
unitFloat (CoreBind -> FloatingBind
FloatLet ([(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, CoreExpr)]
all_pairs)),
                 Maybe CoreBind
forall a. Maybe a
Nothing) }

  | Bool
otherwise -- See Note [Join points and floating]
  = do { (env' :: CorePrepEnv
env', bndrs1 :: [Id]
bndrs1) <- CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs CorePrepEnv
env [Id]
bndrs
       ; [(Id, CoreExpr)]
pairs1 <- (Id -> CoreExpr -> UniqSM (Id, CoreExpr))
-> [Id] -> [CoreExpr] -> UniqSM [(Id, CoreExpr)]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM (CorePrepEnv -> Id -> CoreExpr -> UniqSM (Id, CoreExpr)
cpeJoinPair CorePrepEnv
env') [Id]
bndrs1 [CoreExpr]
rhss

       ; let bndrs2 :: [Id]
bndrs2 = ((Id, CoreExpr) -> Id) -> [(Id, CoreExpr)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CoreExpr) -> Id
forall a b. (a, b) -> a
fst [(Id, CoreExpr)]
pairs1
       ; (CorePrepEnv, Floats, Maybe CoreBind)
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> [(Id, Id)] -> CorePrepEnv
extendCorePrepEnvList CorePrepEnv
env' ([Id]
bndrs [Id] -> [Id] -> [(Id, Id)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [Id]
bndrs2),
                 Floats
emptyFloats,
                 CoreBind -> Maybe CoreBind
forall a. a -> Maybe a
Just ([(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, CoreExpr)]
pairs1)) }
  where
    (bndrs :: [Id]
bndrs, rhss :: [CoreExpr]
rhss) = [(Id, CoreExpr)] -> ([Id], [CoreExpr])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, CoreExpr)]
pairs

        -- Flatten all the floats, and the current
        -- group into a single giant Rec
    add_float :: FloatingBind -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
add_float (FloatLet (NonRec b :: Id
b r :: CoreExpr
r)) prs2 :: [(Id, CoreExpr)]
prs2 = (Id
b,CoreExpr
r) (Id, CoreExpr) -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
forall a. a -> [a] -> [a]
: [(Id, CoreExpr)]
prs2
    add_float (FloatLet (Rec prs1 :: [(Id, CoreExpr)]
prs1))   prs2 :: [(Id, CoreExpr)]
prs2 = [(Id, CoreExpr)]
prs1 [(Id, CoreExpr)] -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
forall a. [a] -> [a] -> [a]
++ [(Id, CoreExpr)]
prs2
    add_float b :: FloatingBind
b                       _    = String -> SDoc -> [(Id, CoreExpr)]
forall a. HasCallStack => String -> SDoc -> a
pprPanic "cpeBind" (FloatingBind -> SDoc
forall a. Outputable a => a -> SDoc
ppr FloatingBind
b)

---------------
cpePair :: TopLevelFlag -> RecFlag -> Demand -> Bool
        -> CorePrepEnv -> OutId -> CoreExpr
        -> UniqSM (Floats, CpeRhs)
-- Used for all bindings
-- The binder is already cloned, hence an OutId
cpePair :: TopLevelFlag
-> RecFlag
-> Demand
-> Bool
-> CorePrepEnv
-> Id
-> CoreExpr
-> UniqSM (Floats, CoreExpr)
cpePair top_lvl :: TopLevelFlag
top_lvl is_rec :: RecFlag
is_rec dmd :: Demand
dmd is_unlifted :: Bool
is_unlifted env :: CorePrepEnv
env bndr :: Id
bndr rhs :: CoreExpr
rhs
  = ASSERT(not (isJoinId bndr)) -- those should use cpeJoinPair
    do { (floats1 :: Floats
floats1, rhs1 :: CoreExpr
rhs1) <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
rhs

       -- See if we are allowed to float this stuff out of the RHS
       ; (floats2 :: Floats
floats2, rhs2 :: CoreExpr
rhs2) <- Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_from_rhs Floats
floats1 CoreExpr
rhs1

       -- Make the arity match up
       ; (floats3 :: Floats
floats3, rhs3 :: CoreExpr
rhs3)
            <- if CoreExpr -> Int
manifestArity CoreExpr
rhs1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
arity
               then (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats2, Int -> CoreExpr -> CoreExpr
cpeEtaExpand Int
arity CoreExpr
rhs2)
               else WARN(True, text "CorePrep: silly extra arguments:" <+> ppr bndr)
                               -- Note [Silly extra arguments]
                    (do { Id
v <- Type -> UniqSM Id
newVar (Id -> Type
idType Id
bndr)
                        ; let float :: FloatingBind
float = Demand -> Bool -> Id -> CoreExpr -> FloatingBind
mkFloat Demand
topDmd Bool
False Id
v CoreExpr
rhs2
                        ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return ( Floats -> FloatingBind -> Floats
addFloat Floats
floats2 FloatingBind
float
                                 , Int -> CoreExpr -> CoreExpr
cpeEtaExpand Int
arity (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
v)) })

        -- Wrap floating ticks
       ; let (floats4 :: Floats
floats4, rhs4 :: CoreExpr
rhs4) = Floats -> CoreExpr -> (Floats, CoreExpr)
wrapTicks Floats
floats3 CoreExpr
rhs3

       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats4, CoreExpr
rhs4) }
  where
    platform :: Platform
platform = DynFlags -> Platform
targetPlatform (CorePrepEnv -> DynFlags
cpe_dynFlags CorePrepEnv
env)

    arity :: Int
arity = Id -> Int
idArity Id
bndr        -- We must match this arity

    ---------------------
    float_from_rhs :: Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_from_rhs floats :: Floats
floats rhs :: CoreExpr
rhs
      | Floats -> Bool
isEmptyFloats Floats
floats = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
rhs)
      | TopLevelFlag -> Bool
isTopLevel TopLevelFlag
top_lvl   = Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_top    Floats
floats CoreExpr
rhs
      | Bool
otherwise            = Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_nested Floats
floats CoreExpr
rhs

    ---------------------
    float_nested :: Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_nested floats :: Floats
floats rhs :: CoreExpr
rhs
      | RecFlag -> Demand -> Bool -> Floats -> CoreExpr -> Bool
wantFloatNested RecFlag
is_rec Demand
dmd Bool
is_unlifted Floats
floats CoreExpr
rhs
                  = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, CoreExpr
rhs)
      | Bool
otherwise = Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
dontFloat Floats
floats CoreExpr
rhs

    ---------------------
    float_top :: Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
float_top floats :: Floats
floats rhs :: CoreExpr
rhs        -- Urhgh!  See Note [CafInfo and floating]
      | CafInfo -> Bool
mayHaveCafRefs (Id -> CafInfo
idCafInfo Id
bndr)
      , Floats -> Bool
allLazyTop Floats
floats
      = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, CoreExpr
rhs)

      -- So the top-level binding is marked NoCafRefs
      | Just (floats' :: Floats
floats', rhs' :: CoreExpr
rhs') <- Platform -> Floats -> CoreExpr -> Maybe (Floats, CoreExpr)
canFloatFromNoCaf Platform
platform Floats
floats CoreExpr
rhs
      = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats', CoreExpr
rhs')

      | Bool
otherwise
      = Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
dontFloat Floats
floats CoreExpr
rhs

dontFloat :: Floats -> CpeRhs -> UniqSM (Floats, CpeBody)
-- Non-empty floats, but do not want to float from rhs
-- So wrap the rhs in the floats
-- But: rhs1 might have lambdas, and we can't
--      put them inside a wrapBinds
dontFloat :: Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
dontFloat floats1 :: Floats
floats1 rhs :: CoreExpr
rhs
  = do { (floats2 :: Floats
floats2, body :: CoreExpr
body) <- CoreExpr -> UniqSM (Floats, CoreExpr)
rhsToBody CoreExpr
rhs
        ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, Floats -> CoreExpr -> CoreExpr
wrapBinds Floats
floats1 (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$
                               Floats -> CoreExpr -> CoreExpr
wrapBinds Floats
floats2 CoreExpr
body) }

{- Note [Silly extra arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we had this
        f{arity=1} = \x\y. e
We *must* match the arity on the Id, so we have to generate
        f' = \x\y. e
        f  = \x. f' x

It's a bizarre case: why is the arity on the Id wrong?  Reason
(in the days of __inline_me__):
        f{arity=0} = __inline_me__ (let v = expensive in \xy. e)
When InlineMe notes go away this won't happen any more.  But
it seems good for CorePrep to be robust.
-}

---------------
cpeJoinPair :: CorePrepEnv -> JoinId -> CoreExpr
            -> UniqSM (JoinId, CpeRhs)
-- Used for all join bindings
cpeJoinPair :: CorePrepEnv -> Id -> CoreExpr -> UniqSM (Id, CoreExpr)
cpeJoinPair env :: CorePrepEnv
env bndr :: Id
bndr rhs :: CoreExpr
rhs
  = ASSERT(isJoinId bndr)
    do { let Just join_arity :: Int
join_arity = Id -> Maybe Int
isJoinId_maybe Id
bndr
             (bndrs :: [Id]
bndrs, body :: CoreExpr
body)   = Int -> CoreExpr -> ([Id], CoreExpr)
forall b. Int -> Expr b -> ([b], Expr b)
collectNBinders Int
join_arity CoreExpr
rhs

       ; (env' :: CorePrepEnv
env', bndrs' :: [Id]
bndrs') <- CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs CorePrepEnv
env [Id]
bndrs

       ; CoreExpr
body' <- CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF CorePrepEnv
env' CoreExpr
body -- Will let-bind the body if it starts
                                      -- with a lambda

       ; let rhs' :: CoreExpr
rhs'  = [Id] -> CoreExpr -> CoreExpr
mkCoreLams [Id]
bndrs' CoreExpr
body'
             bndr' :: Id
bndr' = Id
bndr Id -> Unfolding -> Id
`setIdUnfolding` Unfolding
evaldUnfolding
                          Id -> Int -> Id
`setIdArity` (Id -> Bool) -> [Id] -> Int
forall a. (a -> Bool) -> [a] -> Int
count Id -> Bool
isId [Id]
bndrs
                            -- See Note [Arity and join points]

       ; (Id, CoreExpr) -> UniqSM (Id, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Id
bndr', CoreExpr
rhs') }

{-
Note [Arity and join points]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Up to now, we've allowed a join point to have an arity greater than its join
arity (minus type arguments), since this is what's useful for eta expansion.
However, for code gen purposes, its arity must be exactly the number of value
arguments it will be called with, and it must have exactly that many value
lambdas. Hence if there are extra lambdas we must let-bind the body of the RHS:

  join j x y z = \w -> ... in ...
    =>
  join j x y z = (let f = \w -> ... in f) in ...

This is also what happens with Note [Silly extra arguments]. Note that it's okay
for us to mess with the arity because a join point is never exported.
-}

-- ---------------------------------------------------------------------------
--              CpeRhs: produces a result satisfying CpeRhs
-- ---------------------------------------------------------------------------

cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
-- If
--      e  ===>  (bs, e')
-- then
--      e = let bs in e'        (semantically, that is!)
--
-- For example
--      f (g x)   ===>   ([v = g x], f v)

cpeRhsE :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE _env :: CorePrepEnv
_env expr :: CoreExpr
expr@(Type {})      = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
expr)
cpeRhsE _env :: CorePrepEnv
_env expr :: CoreExpr
expr@(Coercion {})  = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
expr)
cpeRhsE env :: CorePrepEnv
env (Lit (LitNumber LitNumInteger i :: Integer
i _))
    = CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env (DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
cvtLitInteger (CorePrepEnv -> DynFlags
cpe_dynFlags CorePrepEnv
env) (CorePrepEnv -> Id
getMkIntegerId CorePrepEnv
env)
                   (CorePrepEnv -> Maybe DataCon
cpe_integerSDataCon CorePrepEnv
env) Integer
i)
cpeRhsE env :: CorePrepEnv
env (Lit (LitNumber LitNumNatural i :: Integer
i _))
    = CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env (DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
cvtLitNatural (CorePrepEnv -> DynFlags
cpe_dynFlags CorePrepEnv
env) (CorePrepEnv -> Id
getMkNaturalId CorePrepEnv
env)
                   (CorePrepEnv -> Maybe DataCon
cpe_naturalSDataCon CorePrepEnv
env) Integer
i)
cpeRhsE _env :: CorePrepEnv
_env expr :: CoreExpr
expr@(Lit {}) = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
expr)
cpeRhsE env :: CorePrepEnv
env expr :: CoreExpr
expr@(Var {})  = CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeApp CorePrepEnv
env CoreExpr
expr
cpeRhsE env :: CorePrepEnv
env expr :: CoreExpr
expr@(App {}) = CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeApp CorePrepEnv
env CoreExpr
expr

cpeRhsE env :: CorePrepEnv
env (Let bind :: CoreBind
bind body :: CoreExpr
body)
  = do { (env' :: CorePrepEnv
env', bind_floats :: Floats
bind_floats, maybe_bind' :: Maybe CoreBind
maybe_bind') <- TopLevelFlag
-> CorePrepEnv
-> CoreBind
-> UniqSM (CorePrepEnv, Floats, Maybe CoreBind)
cpeBind TopLevelFlag
NotTopLevel CorePrepEnv
env CoreBind
bind
       ; (body_floats :: Floats
body_floats, body' :: CoreExpr
body') <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env' CoreExpr
body
       ; let expr' :: CoreExpr
expr' = case Maybe CoreBind
maybe_bind' of Just bind' :: CoreBind
bind' -> CoreBind -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind' CoreExpr
body'
                                         Nothing    -> CoreExpr
body'
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
bind_floats Floats -> Floats -> Floats
`appendFloats` Floats
body_floats, CoreExpr
expr') }

cpeRhsE env :: CorePrepEnv
env (Tick tickish :: Tickish Id
tickish expr :: CoreExpr
expr)
  | Tickish Id -> TickishPlacement
forall id. Tickish id -> TickishPlacement
tickishPlace Tickish Id
tickish TickishPlacement -> TickishPlacement -> Bool
forall a. Eq a => a -> a -> Bool
== TickishPlacement
PlaceNonLam Bool -> Bool -> Bool
&& Tickish Id
tickish Tickish Id -> TickishScoping -> Bool
forall id. Tickish id -> TickishScoping -> Bool
`tickishScopesLike` TickishScoping
SoftScope
  = do { (floats :: Floats
floats, body :: CoreExpr
body) <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
expr
         -- See [Floating Ticks in CorePrep]
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (FloatingBind -> Floats
unitFloat (Tickish Id -> FloatingBind
FloatTick Tickish Id
tickish) Floats -> Floats -> Floats
`appendFloats` Floats
floats, CoreExpr
body) }
  | Bool
otherwise
  = do { CoreExpr
body <- CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF CorePrepEnv
env CoreExpr
expr
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
tickish' CoreExpr
body) }
  where
    tickish' :: Tickish Id
tickish' | Breakpoint n :: Int
n fvs :: [Id]
fvs <- Tickish Id
tickish
             -- See also 'substTickish'
             = Int -> [Id] -> Tickish Id
forall id. Int -> [id] -> Tickish id
Breakpoint Int
n ((Id -> Id) -> [Id] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (HasDebugCallStack => CoreExpr -> Id
CoreExpr -> Id
getIdFromTrivialExpr (CoreExpr -> Id) -> (Id -> CoreExpr) -> Id -> Id
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CorePrepEnv -> Id -> CoreExpr
lookupCorePrepEnv CorePrepEnv
env) [Id]
fvs)
             | Bool
otherwise
             = Tickish Id
tickish

cpeRhsE env :: CorePrepEnv
env (Cast expr :: CoreExpr
expr co :: Coercion
co)
   = do { (floats :: Floats
floats, expr' :: CoreExpr
expr') <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
expr
        ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
expr' Coercion
co) }

cpeRhsE env :: CorePrepEnv
env expr :: CoreExpr
expr@(Lam {})
   = do { let (bndrs :: [Id]
bndrs,body :: CoreExpr
body) = CoreExpr -> ([Id], CoreExpr)
forall b. Expr b -> ([b], Expr b)
collectBinders CoreExpr
expr
        ; (env' :: CorePrepEnv
env', bndrs' :: [Id]
bndrs') <- CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs CorePrepEnv
env [Id]
bndrs
        ; CoreExpr
body' <- CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF CorePrepEnv
env' CoreExpr
body
        ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, [Id] -> CoreExpr -> CoreExpr
forall b. [b] -> Expr b -> Expr b
mkLams [Id]
bndrs' CoreExpr
body') }

cpeRhsE env :: CorePrepEnv
env (Case scrut :: CoreExpr
scrut bndr :: Id
bndr ty :: Type
ty alts :: [Alt Id]
alts)
  = do { (floats :: Floats
floats, scrut' :: CoreExpr
scrut') <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeBody CorePrepEnv
env CoreExpr
scrut
       ; (env' :: CorePrepEnv
env', bndr2 :: Id
bndr2) <- CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id)
cpCloneBndr CorePrepEnv
env Id
bndr
       ; let alts' :: [Alt Id]
alts'
                 -- This flag is intended to aid in debugging strictness
                 -- analysis bugs. These are particularly nasty to chase down as
                 -- they may manifest as segmentation faults. When this flag is
                 -- enabled we instead produce an 'error' expression to catch
                 -- the case where a function we think should bottom
                 -- unexpectedly returns.
               | GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_CatchBottoms (CorePrepEnv -> DynFlags
cpe_dynFlags CorePrepEnv
env)
               , Bool -> Bool
not ([Alt Id] -> Bool
forall b. [Alt b] -> Bool
altsAreExhaustive [Alt Id]
alts)
               = [Alt Id] -> Maybe CoreExpr -> [Alt Id]
forall a b. [(AltCon, [a], b)] -> Maybe b -> [(AltCon, [a], b)]
addDefault [Alt Id]
alts (CoreExpr -> Maybe CoreExpr
forall a. a -> Maybe a
Just CoreExpr
err)
               | Bool
otherwise = [Alt Id]
alts
               where err :: CoreExpr
err = Id -> Type -> String -> CoreExpr
mkRuntimeErrorApp Id
rUNTIME_ERROR_ID Type
ty
                                             "Bottoming expression returned"
       ; [Alt Id]
alts'' <- (Alt Id -> UniqSM (Alt Id)) -> [Alt Id] -> UniqSM [Alt Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (CorePrepEnv -> Alt Id -> UniqSM (Alt Id)
forall a.
CorePrepEnv -> (a, [Id], CoreExpr) -> UniqSM (a, [Id], CoreExpr)
sat_alt CorePrepEnv
env') [Alt Id]
alts'
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
scrut' Id
bndr2 Type
ty [Alt Id]
alts'') }
  where
    sat_alt :: CorePrepEnv -> (a, [Id], CoreExpr) -> UniqSM (a, [Id], CoreExpr)
sat_alt env :: CorePrepEnv
env (con :: a
con, bs :: [Id]
bs, rhs :: CoreExpr
rhs)
       = do { (env2 :: CorePrepEnv
env2, bs' :: [Id]
bs') <- CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs CorePrepEnv
env [Id]
bs
            ; CoreExpr
rhs' <- CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF CorePrepEnv
env2 CoreExpr
rhs
            ; (a, [Id], CoreExpr) -> UniqSM (a, [Id], CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
con, [Id]
bs', CoreExpr
rhs') }

cvtLitInteger :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
-- Here we convert a literal Integer to the low-level
-- representation. Exactly how we do this depends on the
-- library that implements Integer.  If it's GMP we
-- use the S# data constructor for small literals.
-- See Note [Integer literals] in Literal
cvtLitInteger :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
cvtLitInteger dflags :: DynFlags
dflags _ (Just sdatacon :: DataCon
sdatacon) i :: Integer
i
  | DynFlags -> Integer -> Bool
inIntRange DynFlags
dflags Integer
i -- Special case for small integers
    = DataCon -> [CoreExpr] -> CoreExpr
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
sdatacon [Literal -> CoreExpr
forall b. Literal -> Expr b
Lit (DynFlags -> Integer -> Literal
mkLitInt DynFlags
dflags Integer
i)]

cvtLitInteger dflags :: DynFlags
dflags mk_integer :: Id
mk_integer _ i :: Integer
i
    = CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
mk_integer) [CoreExpr
forall b. Expr b
isNonNegative, CoreExpr
ints]
  where isNonNegative :: Expr b
isNonNegative = if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< 0 then DataCon -> [Expr b] -> Expr b
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
falseDataCon []
                                 else DataCon -> [Expr b] -> Expr b
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
trueDataCon  []
        ints :: CoreExpr
ints = Type -> [CoreExpr] -> CoreExpr
mkListExpr Type
intTy (Integer -> [CoreExpr]
forall b. Integer -> [Expr b]
f (Integer -> Integer
forall a. Num a => a -> a
abs Integer
i))
        f :: Integer -> [Expr b]
f 0 = []
        f x :: Integer
x = let low :: Integer
low  = Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask
                  high :: Integer
high = Integer
x Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftR` Int
bits
              in DataCon -> [Expr b] -> Expr b
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
intDataCon [Literal -> Expr b
forall b. Literal -> Expr b
Lit (DynFlags -> Integer -> Literal
mkLitInt DynFlags
dflags Integer
low)] Expr b -> [Expr b] -> [Expr b]
forall a. a -> [a] -> [a]
: Integer -> [Expr b]
f Integer
high
        bits :: Int
bits = 31
        mask :: Integer
mask = 2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
bits Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- 1

cvtLitNatural :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
-- Here we convert a literal Natural to the low-level
-- representation.
-- See Note [Natural literals] in Literal
cvtLitNatural :: DynFlags -> Id -> Maybe DataCon -> Integer -> CoreExpr
cvtLitNatural dflags :: DynFlags
dflags _ (Just sdatacon :: DataCon
sdatacon) i :: Integer
i
  | DynFlags -> Integer -> Bool
inWordRange DynFlags
dflags Integer
i -- Special case for small naturals
    = DataCon -> [CoreExpr] -> CoreExpr
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
sdatacon [Literal -> CoreExpr
forall b. Literal -> Expr b
Lit (DynFlags -> Integer -> Literal
mkLitWord DynFlags
dflags Integer
i)]

cvtLitNatural dflags :: DynFlags
dflags mk_natural :: Id
mk_natural _ i :: Integer
i
    = CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
mk_natural) [CoreExpr
words]
  where words :: CoreExpr
words = Type -> [CoreExpr] -> CoreExpr
mkListExpr Type
wordTy (Integer -> [CoreExpr]
forall b. Integer -> [Expr b]
f Integer
i)
        f :: Integer -> [Expr b]
f 0 = []
        f x :: Integer
x = let low :: Integer
low  = Integer
x Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
mask
                  high :: Integer
high = Integer
x Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftR` Int
bits
              in DataCon -> [Expr b] -> Expr b
forall b. DataCon -> [Arg b] -> Arg b
mkConApp DataCon
wordDataCon [Literal -> Expr b
forall b. Literal -> Expr b
Lit (DynFlags -> Integer -> Literal
mkLitWord DynFlags
dflags Integer
low)] Expr b -> [Expr b] -> [Expr b]
forall a. a -> [a] -> [a]
: Integer -> [Expr b]
f Integer
high
        bits :: Int
bits = 32
        mask :: Integer
mask = 2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
bits Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- 1

-- ---------------------------------------------------------------------------
--              CpeBody: produces a result satisfying CpeBody
-- ---------------------------------------------------------------------------

-- | Convert a 'CoreExpr' so it satisfies 'CpeBody', without
-- producing any floats (any generated floats are immediately
-- let-bound using 'wrapBinds').  Generally you want this, esp.
-- when you've reached a binding form (e.g., a lambda) and
-- floating any further would be incorrect.
cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CpeBody
cpeBodyNF :: CorePrepEnv -> CoreExpr -> UniqSM CoreExpr
cpeBodyNF env :: CorePrepEnv
env expr :: CoreExpr
expr
  = do { (floats :: Floats
floats, body :: CoreExpr
body) <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeBody CorePrepEnv
env CoreExpr
expr
       ; CoreExpr -> UniqSM CoreExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats -> CoreExpr -> CoreExpr
wrapBinds Floats
floats CoreExpr
body) }

-- | Convert a 'CoreExpr' so it satisfies 'CpeBody'; also produce
-- a list of 'Floats' which are being propagated upwards.  In
-- fact, this function is used in only two cases: to
-- implement 'cpeBodyNF' (which is what you usually want),
-- and in the case when a let-binding is in a case scrutinee--here,
-- we can always float out:
--
--      case (let x = y in z) of ...
--      ==> let x = y in case z of ...
--
cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeBody)
cpeBody :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeBody env :: CorePrepEnv
env expr :: CoreExpr
expr
  = do { (floats1 :: Floats
floats1, rhs :: CoreExpr
rhs) <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
expr
       ; (floats2 :: Floats
floats2, body :: CoreExpr
body) <- CoreExpr -> UniqSM (Floats, CoreExpr)
rhsToBody CoreExpr
rhs
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats1 Floats -> Floats -> Floats
`appendFloats` Floats
floats2, CoreExpr
body) }

--------
rhsToBody :: CpeRhs -> UniqSM (Floats, CpeBody)
-- Remove top level lambdas by let-binding

rhsToBody :: CoreExpr -> UniqSM (Floats, CoreExpr)
rhsToBody (Tick t :: Tickish Id
t expr :: CoreExpr
expr)
  | Tickish Id -> TickishScoping
forall id. Tickish id -> TickishScoping
tickishScoped Tickish Id
t TickishScoping -> TickishScoping -> Bool
forall a. Eq a => a -> a -> Bool
== TickishScoping
NoScope  -- only float out of non-scoped annotations
  = do { (floats :: Floats
floats, expr' :: CoreExpr
expr') <- CoreExpr -> UniqSM (Floats, CoreExpr)
rhsToBody CoreExpr
expr
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
t CoreExpr
expr') }

rhsToBody (Cast e :: CoreExpr
e co :: Coercion
co)
        -- You can get things like
        --      case e of { p -> coerce t (\s -> ...) }
  = do { (floats :: Floats
floats, e' :: CoreExpr
e') <- CoreExpr -> UniqSM (Floats, CoreExpr)
rhsToBody CoreExpr
e
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats, CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
e' Coercion
co) }

rhsToBody expr :: CoreExpr
expr@(Lam {})
  | Just no_lam_result :: CoreExpr
no_lam_result <- [Id] -> CoreExpr -> Maybe CoreExpr
tryEtaReducePrep [Id]
bndrs CoreExpr
body
  = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
no_lam_result)
  | (Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Id -> Bool
isTyVar [Id]
bndrs           -- Type lambdas are ok
  = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
expr)
  | Bool
otherwise                   -- Some value lambdas
  = do { Id
fn <- Type -> UniqSM Id
newVar (CoreExpr -> Type
exprType CoreExpr
expr)
       ; let rhs :: CoreExpr
rhs   = Int -> CoreExpr -> CoreExpr
cpeEtaExpand (CoreExpr -> Int
exprArity CoreExpr
expr) CoreExpr
expr
             float :: FloatingBind
float = CoreBind -> FloatingBind
FloatLet (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
fn CoreExpr
rhs)
       ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (FloatingBind -> Floats
unitFloat FloatingBind
float, Id -> CoreExpr
forall b. Id -> Expr b
Var Id
fn) }
  where
    (bndrs :: [Id]
bndrs,body :: CoreExpr
body) = CoreExpr -> ([Id], CoreExpr)
forall b. Expr b -> ([b], Expr b)
collectBinders CoreExpr
expr

rhsToBody expr :: CoreExpr
expr = (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
emptyFloats, CoreExpr
expr)



-- ---------------------------------------------------------------------------
--              CpeApp: produces a result satisfying CpeApp
-- ---------------------------------------------------------------------------

data ArgInfo = CpeApp  CoreArg
             | CpeCast Coercion
             | CpeTick (Tickish Id)

{- Note [runRW arg]
~~~~~~~~~~~~~~~~~~~
If we got, say
   runRW# (case bot of {})
which happened in Trac #11291, we do /not/ want to turn it into
   (case bot of {}) realWorldPrimId#
because that gives a panic in CoreToStg.myCollectArgs, which expects
only variables in function position.  But if we are sure to make
runRW# strict (which we do in MkId), this can't happen
-}

cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
-- May return a CpeRhs because of saturating primops
cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeApp top_env :: CorePrepEnv
top_env expr :: CoreExpr
expr
  = do { let (terminal :: CoreExpr
terminal, args :: [ArgInfo]
args, depth :: Int
depth) = CoreExpr -> (CoreExpr, [ArgInfo], Int)
collect_args CoreExpr
expr
       ; CorePrepEnv
-> CoreExpr -> [ArgInfo] -> Int -> UniqSM (Floats, CoreExpr)
cpe_app CorePrepEnv
top_env CoreExpr
terminal [ArgInfo]
args Int
depth
       }

  where
    -- We have a nested data structure of the form
    -- e `App` a1 `App` a2 ... `App` an, convert it into
    -- (e, [CpeApp a1, CpeApp a2, ..., CpeApp an], depth)
    -- We use 'ArgInfo' because we may also need to
    -- record casts and ticks.  Depth counts the number
    -- of arguments that would consume strictness information
    -- (so, no type or coercion arguments.)
    collect_args :: CoreExpr -> (CoreExpr, [ArgInfo], Int)
    collect_args :: CoreExpr -> (CoreExpr, [ArgInfo], Int)
collect_args e :: CoreExpr
e = CoreExpr -> [ArgInfo] -> Int -> (CoreExpr, [ArgInfo], Int)
forall c.
Num c =>
CoreExpr -> [ArgInfo] -> c -> (CoreExpr, [ArgInfo], c)
go CoreExpr
e [] 0
      where
        go :: CoreExpr -> [ArgInfo] -> c -> (CoreExpr, [ArgInfo], c)
go (App fun :: CoreExpr
fun arg :: CoreExpr
arg)      as :: [ArgInfo]
as !c
depth
            = CoreExpr -> [ArgInfo] -> c -> (CoreExpr, [ArgInfo], c)
go CoreExpr
fun (CoreExpr -> ArgInfo
CpeApp CoreExpr
arg ArgInfo -> [ArgInfo] -> [ArgInfo]
forall a. a -> [a] -> [a]
: [ArgInfo]
as)
                (if CoreExpr -> Bool
forall b. Expr b -> Bool
isTyCoArg CoreExpr
arg then c
depth else c
depth c -> c -> c
forall a. Num a => a -> a -> a
+ 1)
        go (Cast fun :: CoreExpr
fun co :: Coercion
co)      as :: [ArgInfo]
as depth :: c
depth
            = CoreExpr -> [ArgInfo] -> c -> (CoreExpr, [ArgInfo], c)
go CoreExpr
fun (Coercion -> ArgInfo
CpeCast Coercion
co ArgInfo -> [ArgInfo] -> [ArgInfo]
forall a. a -> [a] -> [a]
: [ArgInfo]
as) c
depth
        go (Tick tickish :: Tickish Id
tickish fun :: CoreExpr
fun) as :: [ArgInfo]
as depth :: c
depth
            | Tickish Id -> TickishPlacement
forall id. Tickish id -> TickishPlacement
tickishPlace Tickish Id
tickish TickishPlacement -> TickishPlacement -> Bool
forall a. Eq a => a -> a -> Bool
== TickishPlacement
PlaceNonLam
            Bool -> Bool -> Bool
&& Tickish Id
tickish Tickish Id -> TickishScoping -> Bool
forall id. Tickish id -> TickishScoping -> Bool
`tickishScopesLike` TickishScoping
SoftScope
            = CoreExpr -> [ArgInfo] -> c -> (CoreExpr, [ArgInfo], c)
go CoreExpr
fun (Tickish Id -> ArgInfo
CpeTick Tickish Id
tickish ArgInfo -> [ArgInfo] -> [ArgInfo]
forall a. a -> [a] -> [a]
: [ArgInfo]
as) c
depth
        go terminal :: CoreExpr
terminal as :: [ArgInfo]
as depth :: c
depth = (CoreExpr
terminal, [ArgInfo]
as, c
depth)

    cpe_app :: CorePrepEnv
            -> CoreExpr
            -> [ArgInfo]
            -> Int
            -> UniqSM (Floats, CpeRhs)
    cpe_app :: CorePrepEnv
-> CoreExpr -> [ArgInfo] -> Int -> UniqSM (Floats, CoreExpr)
cpe_app env :: CorePrepEnv
env (Var f :: Id
f) (CpeApp Type{} : CpeApp arg :: CoreExpr
arg : args :: [ArgInfo]
args) depth :: Int
depth
        | Id
f Id -> Unique -> Bool
forall a. Uniquable a => a -> Unique -> Bool
`hasKey` Unique
lazyIdKey          -- Replace (lazy a) with a, and
       Bool -> Bool -> Bool
|| Id
f Id -> Unique -> Bool
forall a. Uniquable a => a -> Unique -> Bool
`hasKey` Unique
noinlineIdKey      -- Replace (noinline a) with a
        -- Consider the code:
        --
        --      lazy (f x) y
        --
        -- We need to make sure that we need to recursively collect arguments on
        -- "f x", otherwise we'll float "f x" out (it's not a variable) and
        -- end up with this awful -ddump-prep:
        --
        --      case f x of f_x {
        --        __DEFAULT -> f_x y
        --      }
        --
        -- rather than the far superior "f x y".  Test case is par01.
        = let (terminal :: CoreExpr
terminal, args' :: [ArgInfo]
args', depth' :: Int
depth') = CoreExpr -> (CoreExpr, [ArgInfo], Int)
collect_args CoreExpr
arg
          in CorePrepEnv
-> CoreExpr -> [ArgInfo] -> Int -> UniqSM (Floats, CoreExpr)
cpe_app CorePrepEnv
env CoreExpr
terminal ([ArgInfo]
args' [ArgInfo] -> [ArgInfo] -> [ArgInfo]
forall a. [a] -> [a] -> [a]
++ [ArgInfo]
args) (Int
depth Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
depth' Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)
    cpe_app env :: CorePrepEnv
env (Var f :: Id
f) [CpeApp _runtimeRep :: CoreExpr
_runtimeRep@Type{}, CpeApp _type :: CoreExpr
_type@Type{}, CpeApp arg :: CoreExpr
arg] 1
        | Id
f Id -> Unique -> Bool
forall a. Uniquable a => a -> Unique -> Bool
`hasKey` Unique
runRWKey
        -- See Note [runRW magic]
        -- Replace (runRW# f) by (f realWorld#), beta reducing if possible (this
        -- is why we return a CorePrepEnv as well)
        = case CoreExpr
arg of
            Lam s :: Id
s body :: CoreExpr
body -> CorePrepEnv
-> CoreExpr -> [ArgInfo] -> Int -> UniqSM (Floats, CoreExpr)
cpe_app (CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv CorePrepEnv
env Id
s Id
realWorldPrimId) CoreExpr
body [] 0
            _          -> CorePrepEnv
-> CoreExpr -> [ArgInfo] -> Int -> UniqSM (Floats, CoreExpr)
cpe_app CorePrepEnv
env CoreExpr
arg [CoreExpr -> ArgInfo
CpeApp (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
realWorldPrimId)] 1
    cpe_app env :: CorePrepEnv
env (Var v :: Id
v) args :: [ArgInfo]
args depth :: Int
depth
      = do { Id
v1 <- Id -> UniqSM Id
fiddleCCall Id
v
           ; let e2 :: CoreExpr
e2 = CorePrepEnv -> Id -> CoreExpr
lookupCorePrepEnv CorePrepEnv
env Id
v1
                 hd :: Maybe Id
hd = CoreExpr -> Maybe Id
getIdFromTrivialExpr_maybe CoreExpr
e2
           -- NB: depth from collect_args is right, because e2 is a trivial expression
           -- and thus its embedded Id *must* be at the same depth as any
           -- Apps it is under are type applications only (c.f.
           -- exprIsTrivial).  But note that we need the type of the
           -- expression, not the id.
           ; (app :: CoreExpr
app, floats :: Floats
floats) <- [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
args CoreExpr
e2 (CoreExpr -> Type
exprType CoreExpr
e2) Floats
emptyFloats [Demand]
stricts
           ; Maybe Id -> CoreExpr -> Floats -> Int -> UniqSM (Floats, CoreExpr)
forall a. Maybe Id -> CoreExpr -> a -> Int -> UniqSM (a, CoreExpr)
mb_saturate Maybe Id
hd CoreExpr
app Floats
floats Int
depth }
        where
          stricts :: [Demand]
stricts = case Id -> StrictSig
idStrictness Id
v of
                            StrictSig (DmdType _ demands :: [Demand]
demands _)
                              | [Demand] -> Int -> Ordering
forall a. [a] -> Int -> Ordering
listLengthCmp [Demand]
demands Int
depth Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
/= Ordering
GT -> [Demand]
demands
                                    -- length demands <= depth
                              | Bool
otherwise                         -> []
                -- If depth < length demands, then we have too few args to
                -- satisfy strictness  info so we have to  ignore all the
                -- strictness info, e.g. + (error "urk")
                -- Here, we can't evaluate the arg strictly, because this
                -- partial application might be seq'd

        -- We inlined into something that's not a var and has no args.
        -- Bounce it back up to cpeRhsE.
    cpe_app env :: CorePrepEnv
env fun :: CoreExpr
fun [] _ = CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
fun

        -- N-variable fun, better let-bind it
    cpe_app env :: CorePrepEnv
env fun :: CoreExpr
fun args :: [ArgInfo]
args depth :: Int
depth
      = do { (fun_floats :: Floats
fun_floats, fun' :: CoreExpr
fun') <- CorePrepEnv
-> Demand -> CoreExpr -> Type -> UniqSM (Floats, CoreExpr)
cpeArg CorePrepEnv
env Demand
evalDmd CoreExpr
fun Type
ty
                          -- The evalDmd says that it's sure to be evaluated,
                          -- so we'll end up case-binding it
           ; (app :: CoreExpr
app, floats :: Floats
floats) <- [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
args CoreExpr
fun' Type
ty Floats
fun_floats []
           ; Maybe Id -> CoreExpr -> Floats -> Int -> UniqSM (Floats, CoreExpr)
forall a. Maybe Id -> CoreExpr -> a -> Int -> UniqSM (a, CoreExpr)
mb_saturate Maybe Id
forall a. Maybe a
Nothing CoreExpr
app Floats
floats Int
depth }
        where
          ty :: Type
ty = CoreExpr -> Type
exprType CoreExpr
fun

    -- Saturate if necessary
    mb_saturate :: Maybe Id -> CoreExpr -> a -> Int -> UniqSM (a, CoreExpr)
mb_saturate head :: Maybe Id
head app :: CoreExpr
app floats :: a
floats depth :: Int
depth =
       case Maybe Id
head of
         Just fn_id :: Id
fn_id -> do { CoreExpr
sat_app <- Id -> CoreExpr -> Int -> UniqSM CoreExpr
maybeSaturate Id
fn_id CoreExpr
app Int
depth
                          ; (a, CoreExpr) -> UniqSM (a, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
floats, CoreExpr
sat_app) }
         _other :: Maybe Id
_other              -> (a, CoreExpr) -> UniqSM (a, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
floats, CoreExpr
app)

    -- Deconstruct and rebuild the application, floating any non-atomic
    -- arguments to the outside.  We collect the type of the expression,
    -- the head of the application, and the number of actual value arguments,
    -- all of which are used to possibly saturate this application if it
    -- has a constructor or primop at the head.
    rebuild_app
        :: [ArgInfo]                  -- The arguments (inner to outer)
        -> CpeApp
        -> Type
        -> Floats
        -> [Demand]
        -> UniqSM (CpeApp, Floats)
    rebuild_app :: [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [] app :: CoreExpr
app _ floats :: Floats
floats ss :: [Demand]
ss = do
      MASSERT(null ss) -- make sure we used all the strictness info
      (CoreExpr, Floats) -> UniqSM (CoreExpr, Floats)
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
app, Floats
floats)
    rebuild_app (a :: ArgInfo
a : as :: [ArgInfo]
as) fun' :: CoreExpr
fun' fun_ty :: Type
fun_ty floats :: Floats
floats ss :: [Demand]
ss = case ArgInfo
a of
      CpeApp arg :: CoreExpr
arg@(Type arg_ty :: Type
arg_ty) ->
        [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
as (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
fun' CoreExpr
arg) (HasDebugCallStack => Type -> Type -> Type
Type -> Type -> Type
piResultTy Type
fun_ty Type
arg_ty) Floats
floats [Demand]
ss
      CpeApp arg :: CoreExpr
arg@(Coercion {}) ->
        [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
as (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
fun' CoreExpr
arg) (Type -> Type
funResultTy Type
fun_ty) Floats
floats [Demand]
ss
      CpeApp arg :: CoreExpr
arg -> do
        let (ss1 :: Demand
ss1, ss_rest :: [Demand]
ss_rest)  -- See Note [lazyId magic] in MkId
               = case ([Demand]
ss, CoreExpr -> Bool
isLazyExpr CoreExpr
arg) of
                   (_   : ss_rest :: [Demand]
ss_rest, True)  -> (Demand
topDmd, [Demand]
ss_rest)
                   (ss1 :: Demand
ss1 : ss_rest :: [Demand]
ss_rest, False) -> (Demand
ss1,    [Demand]
ss_rest)
                   ([],            _)     -> (Demand
topDmd, [])
            (arg_ty :: Type
arg_ty, res_ty :: Type
res_ty) = String -> Maybe (Type, Type) -> (Type, Type)
forall a. HasCallStack => String -> Maybe a -> a
expectJust "cpeBody:collect_args" (Maybe (Type, Type) -> (Type, Type))
-> Maybe (Type, Type) -> (Type, Type)
forall a b. (a -> b) -> a -> b
$
                               Type -> Maybe (Type, Type)
splitFunTy_maybe Type
fun_ty
        (fs :: Floats
fs, arg' :: CoreExpr
arg') <- CorePrepEnv
-> Demand -> CoreExpr -> Type -> UniqSM (Floats, CoreExpr)
cpeArg CorePrepEnv
top_env Demand
ss1 CoreExpr
arg Type
arg_ty
        [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
as (CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
fun' CoreExpr
arg') Type
res_ty (Floats
fs Floats -> Floats -> Floats
`appendFloats` Floats
floats) [Demand]
ss_rest
      CpeCast co :: Coercion
co ->
        let Pair _ty1 :: Type
_ty1 ty2 :: Type
ty2 = Coercion -> Pair Type
coercionKind Coercion
co
        in [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
as (CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
fun' Coercion
co) Type
ty2 Floats
floats [Demand]
ss
      CpeTick tickish :: Tickish Id
tickish ->
        -- See [Floating Ticks in CorePrep]
        [ArgInfo]
-> CoreExpr
-> Type
-> Floats
-> [Demand]
-> UniqSM (CoreExpr, Floats)
rebuild_app [ArgInfo]
as CoreExpr
fun' Type
fun_ty (Floats -> FloatingBind -> Floats
addFloat Floats
floats (Tickish Id -> FloatingBind
FloatTick Tickish Id
tickish)) [Demand]
ss

isLazyExpr :: CoreExpr -> Bool
-- See Note [lazyId magic] in MkId
isLazyExpr :: CoreExpr -> Bool
isLazyExpr (Cast e :: CoreExpr
e _)              = CoreExpr -> Bool
isLazyExpr CoreExpr
e
isLazyExpr (Tick _ e :: CoreExpr
e)              = CoreExpr -> Bool
isLazyExpr CoreExpr
e
isLazyExpr (Var f :: Id
f `App` _ `App` _) = Id
f Id -> Unique -> Bool
forall a. Uniquable a => a -> Unique -> Bool
`hasKey` Unique
lazyIdKey
isLazyExpr _                       = Bool
False

{- Note [runRW magic]
~~~~~~~~~~~~~~~~~~~~~
Some definitions, for instance @runST@, must have careful control over float out
of the bindings in their body. Consider this use of @runST@,

    f x = runST ( \ s -> let (a, s')  = newArray# 100 [] s
                             (_, s'') = fill_in_array_or_something a x s'
                         in freezeArray# a s'' )

If we inline @runST@, we'll get:

    f x = let (a, s')  = newArray# 100 [] realWorld#{-NB-}
              (_, s'') = fill_in_array_or_something a x s'
          in freezeArray# a s''

And now if we allow the @newArray#@ binding to float out to become a CAF,
we end up with a result that is totally and utterly wrong:

    f = let (a, s')  = newArray# 100 [] realWorld#{-NB-} -- YIKES!!!
        in \ x ->
            let (_, s'') = fill_in_array_or_something a x s'
            in freezeArray# a s''

All calls to @f@ will share a {\em single} array! Clearly this is nonsense and
must be prevented.

This is what @runRW#@ gives us: by being inlined extremely late in the
optimization (right before lowering to STG, in CorePrep), we can ensure that
no further floating will occur. This allows us to safely inline things like
@runST@, which are otherwise needlessly expensive (see #10678 and #5916).

'runRW' is defined (for historical reasons) in GHC.Magic, with a NOINLINE
pragma.  It is levity-polymorphic.

    runRW# :: forall (r1 :: RuntimeRep). (o :: TYPE r)
           => (State# RealWorld -> (# State# RealWorld, o #))
                              -> (# State# RealWorld, o #)

It needs no special treatment in GHC except this special inlining here
in CorePrep (and in ByteCodeGen).

-- ---------------------------------------------------------------------------
--      CpeArg: produces a result satisfying CpeArg
-- ---------------------------------------------------------------------------

Note [ANF-ising literal string arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Consider a program like,

    data Foo = Foo Addr#

    foo = Foo "turtle"#

When we go to ANFise this we might think that we want to float the string
literal like we do any other non-trivial argument. This would look like,

    foo = u\ [] case "turtle"# of s { __DEFAULT__ -> Foo s }

However, this 1) isn't necessary since strings are in a sense "trivial"; and 2)
wreaks havoc on the CAF annotations that we produce here since we the result
above is caffy since it is updateable. Ideally at some point in the future we
would like to just float the literal to the top level as suggested in #11312,

    s = "turtle"#
    foo = Foo s

However, until then we simply add a special case excluding literals from the
floating done by cpeArg.
-}

-- | Is an argument okay to CPE?
okCpeArg :: CoreExpr -> Bool
-- Don't float literals. See Note [ANF-ising literal string arguments].
okCpeArg :: CoreExpr -> Bool
okCpeArg (Lit _) = Bool
False
-- Do not eta expand a trivial argument
okCpeArg expr :: CoreExpr
expr    = Bool -> Bool
not (CoreExpr -> Bool
exprIsTrivial CoreExpr
expr)

-- This is where we arrange that a non-trivial argument is let-bound
cpeArg :: CorePrepEnv -> Demand
       -> CoreArg -> Type -> UniqSM (Floats, CpeArg)
cpeArg :: CorePrepEnv
-> Demand -> CoreExpr -> Type -> UniqSM (Floats, CoreExpr)
cpeArg env :: CorePrepEnv
env dmd :: Demand
dmd arg :: CoreExpr
arg arg_ty :: Type
arg_ty
  = do { (floats1 :: Floats
floats1, arg1 :: CoreExpr
arg1) <- CorePrepEnv -> CoreExpr -> UniqSM (Floats, CoreExpr)
cpeRhsE CorePrepEnv
env CoreExpr
arg     -- arg1 can be a lambda
       ; (floats2 :: Floats
floats2, arg2 :: CoreExpr
arg2) <- if Floats -> CoreExpr -> Bool
want_float Floats
floats1 CoreExpr
arg1
                            then (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats1, CoreExpr
arg1)
                            else Floats -> CoreExpr -> UniqSM (Floats, CoreExpr)
dontFloat Floats
floats1 CoreExpr
arg1
                -- Else case: arg1 might have lambdas, and we can't
                --            put them inside a wrapBinds

       ; if CoreExpr -> Bool
okCpeArg CoreExpr
arg2
         then do { Id
v <- Type -> UniqSM Id
newVar Type
arg_ty
                 ; let arg3 :: CoreExpr
arg3      = Int -> CoreExpr -> CoreExpr
cpeEtaExpand (CoreExpr -> Int
exprArity CoreExpr
arg2) CoreExpr
arg2
                       arg_float :: FloatingBind
arg_float = Demand -> Bool -> Id -> CoreExpr -> FloatingBind
mkFloat Demand
dmd Bool
is_unlifted Id
v CoreExpr
arg3
                 ; (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats -> FloatingBind -> Floats
addFloat Floats
floats2 FloatingBind
arg_float, Id -> CoreExpr
forall b. Id -> Expr b
varToCoreExpr Id
v) }
         else (Floats, CoreExpr) -> UniqSM (Floats, CoreExpr)
forall (m :: * -> *) a. Monad m => a -> m a
return (Floats
floats2, CoreExpr
arg2)
       }
  where
    is_unlifted :: Bool
is_unlifted = HasDebugCallStack => Type -> Bool
Type -> Bool
isUnliftedType Type
arg_ty
    want_float :: Floats -> CoreExpr -> Bool
want_float  = RecFlag -> Demand -> Bool -> Floats -> CoreExpr -> Bool
wantFloatNested RecFlag
NonRecursive Demand
dmd Bool
is_unlifted

{-
Note [Floating unlifted arguments]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider    C (let v* = expensive in v)

where the "*" indicates "will be demanded".  Usually v will have been
inlined by now, but let's suppose it hasn't (see Trac #2756).  Then we
do *not* want to get

     let v* = expensive in C v

because that has different strictness.  Hence the use of 'allLazy'.
(NB: the let v* turns into a FloatCase, in mkLocalNonRec.)


------------------------------------------------------------------------------
-- Building the saturated syntax
-- ---------------------------------------------------------------------------

Note [Eta expansion of hasNoBinding things in CorePrep]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
maybeSaturate deals with eta expanding to saturate things that can't deal with
unsaturated applications (identified by 'hasNoBinding', currently just
foreign calls and unboxed tuple/sum constructors).

Note that eta expansion in CorePrep is very fragile due to the "prediction" of
CAFfyness made by TidyPgm (see Note [CAFfyness inconsistencies due to eta
expansion in CorePrep] in TidyPgm for details.  We previously saturated primop
applications here as well but due to this fragility (see #16846) we now deal
with this another way, as described in Note [Primop wrappers] in PrimOp.

It's quite likely that eta expansion of constructor applications will
eventually break in a similar way to how primops did. We really should
eliminate this case as well.
-}

maybeSaturate :: Id -> CpeApp -> Int -> UniqSM CpeRhs
maybeSaturate :: Id -> CoreExpr -> Int -> UniqSM CoreExpr
maybeSaturate fn :: Id
fn expr :: CoreExpr
expr n_args :: Int
n_args
  | Id -> Bool
hasNoBinding Id
fn        -- There's no binding
  = CoreExpr -> UniqSM CoreExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
sat_expr

  | Bool
otherwise
  = CoreExpr -> UniqSM CoreExpr
forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
expr
  where
    fn_arity :: Int
fn_arity     = Id -> Int
idArity Id
fn
    excess_arity :: Int
excess_arity = Int
fn_arity Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n_args
    sat_expr :: CoreExpr
sat_expr     = Int -> CoreExpr -> CoreExpr
cpeEtaExpand Int
excess_arity CoreExpr
expr

{-
************************************************************************
*                                                                      *
                Simple CoreSyn operations
*                                                                      *
************************************************************************
-}

{-
-- -----------------------------------------------------------------------------
--      Eta reduction
-- -----------------------------------------------------------------------------

Note [Eta expansion]
~~~~~~~~~~~~~~~~~~~~~
Eta expand to match the arity claimed by the binder Remember,
CorePrep must not change arity

Eta expansion might not have happened already, because it is done by
the simplifier only when there at least one lambda already.

NB1:we could refrain when the RHS is trivial (which can happen
    for exported things).  This would reduce the amount of code
    generated (a little) and make things a little words for
    code compiled without -O.  The case in point is data constructor
    wrappers.

NB2: we have to be careful that the result of etaExpand doesn't
   invalidate any of the assumptions that CorePrep is attempting
   to establish.  One possible cause is eta expanding inside of
   an SCC note - we're now careful in etaExpand to make sure the
   SCC is pushed inside any new lambdas that are generated.

Note [Eta expansion and the CorePrep invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It turns out to be much much easier to do eta expansion
*after* the main CorePrep stuff.  But that places constraints
on the eta expander: given a CpeRhs, it must return a CpeRhs.

For example here is what we do not want:
                f = /\a -> g (h 3)      -- h has arity 2
After ANFing we get
                f = /\a -> let s = h 3 in g s
and now we do NOT want eta expansion to give
                f = /\a -> \ y -> (let s = h 3 in g s) y

Instead CoreArity.etaExpand gives
                f = /\a -> \y -> let s = h 3 in g s y
-}

cpeEtaExpand :: Arity -> CpeRhs -> CpeRhs
cpeEtaExpand :: Int -> CoreExpr -> CoreExpr
cpeEtaExpand arity :: Int
arity expr :: CoreExpr
expr
  | Int
arity Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = CoreExpr
expr
  | Bool
otherwise  = Int -> CoreExpr -> CoreExpr
etaExpand Int
arity CoreExpr
expr

{-
-- -----------------------------------------------------------------------------
--      Eta reduction
-- -----------------------------------------------------------------------------

Why try eta reduction?  Hasn't the simplifier already done eta?
But the simplifier only eta reduces if that leaves something
trivial (like f, or f Int).  But for deLam it would be enough to
get to a partial application:
        case x of { p -> \xs. map f xs }
    ==> case x of { p -> map f }
-}

tryEtaReducePrep :: [CoreBndr] -> CoreExpr -> Maybe CoreExpr
tryEtaReducePrep :: [Id] -> CoreExpr -> Maybe CoreExpr
tryEtaReducePrep bndrs :: [Id]
bndrs expr :: CoreExpr
expr@(App _ _)
  | CoreExpr -> Bool
forall b. Expr b -> Bool
ok_to_eta_reduce CoreExpr
f
  , Int
n_remaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= 0
  , [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Id -> CoreExpr -> Bool) -> [Id] -> [CoreExpr] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Id -> CoreExpr -> Bool
forall b. Id -> Expr b -> Bool
ok [Id]
bndrs [CoreExpr]
last_args)
  , Bool -> Bool
not ((Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Id -> VarSet -> Bool
`elemVarSet` VarSet
fvs_remaining) [Id]
bndrs)
  , CoreExpr -> Bool
exprIsHNF CoreExpr
remaining_expr   -- Don't turn value into a non-value
                               -- else the behaviour with 'seq' changes
  = CoreExpr -> Maybe CoreExpr
forall a. a -> Maybe a
Just CoreExpr
remaining_expr
  where
    (f :: CoreExpr
f, args :: [CoreExpr]
args) = CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
    remaining_expr :: CoreExpr
remaining_expr = CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps CoreExpr
f [CoreExpr]
remaining_args
    fvs_remaining :: VarSet
fvs_remaining = CoreExpr -> VarSet
exprFreeVars CoreExpr
remaining_expr
    (remaining_args :: [CoreExpr]
remaining_args, last_args :: [CoreExpr]
last_args) = Int -> [CoreExpr] -> ([CoreExpr], [CoreExpr])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n_remaining [CoreExpr]
args
    n_remaining :: Int
n_remaining = [CoreExpr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [CoreExpr]
args Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Id] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Id]
bndrs

    ok :: Id -> Expr b -> Bool
ok bndr :: Id
bndr (Var arg :: Id
arg) = Id
bndr Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
arg
    ok _    _         = Bool
False

          -- We can't eta reduce something which must be saturated.
    ok_to_eta_reduce :: Expr b -> Bool
ok_to_eta_reduce (Var f :: Id
f) = Bool -> Bool
not (Id -> Bool
hasNoBinding Id
f)
    ok_to_eta_reduce _       = Bool
False -- Safe. ToDo: generalise

tryEtaReducePrep bndrs :: [Id]
bndrs (Let bind :: CoreBind
bind@(NonRec _ r :: CoreExpr
r) body :: CoreExpr
body)
  | Bool -> Bool
not ((Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Id -> VarSet -> Bool
`elemVarSet` VarSet
fvs) [Id]
bndrs)
  = case [Id] -> CoreExpr -> Maybe CoreExpr
tryEtaReducePrep [Id]
bndrs CoreExpr
body of
        Just e :: CoreExpr
e -> CoreExpr -> Maybe CoreExpr
forall a. a -> Maybe a
Just (CoreBind -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind CoreExpr
e)
        Nothing -> Maybe CoreExpr
forall a. Maybe a
Nothing
  where
    fvs :: VarSet
fvs = CoreExpr -> VarSet
exprFreeVars CoreExpr
r

-- NB: do not attempt to eta-reduce across ticks
-- Otherwise we risk reducing
--       \x. (Tick (Breakpoint {x}) f x)
--   ==> Tick (breakpoint {x}) f
-- which is bogus (Trac #17228)
-- tryEtaReducePrep bndrs (Tick tickish e)
--   = fmap (mkTick tickish) $ tryEtaReducePrep bndrs e

tryEtaReducePrep _ _ = Maybe CoreExpr
forall a. Maybe a
Nothing

{-
************************************************************************
*                                                                      *
                Floats
*                                                                      *
************************************************************************

Note [Pin demand info on floats]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We pin demand info on floated lets, so that we can see the one-shot thunks.
-}

data FloatingBind
  = FloatLet CoreBind    -- Rhs of bindings are CpeRhss
                         -- They are always of lifted type;
                         -- unlifted ones are done with FloatCase

 | FloatCase
      Id CpeBody
      Bool              -- The bool indicates "ok-for-speculation"

 -- | See Note [Floating Ticks in CorePrep]
 | FloatTick (Tickish Id)

data Floats = Floats OkToSpec (OrdList FloatingBind)

instance Outputable FloatingBind where
  ppr :: FloatingBind -> SDoc
ppr (FloatLet b :: CoreBind
b) = CoreBind -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreBind
b
  ppr (FloatCase b :: Id
b r :: CoreExpr
r ok :: Bool
ok) = SDoc -> SDoc
brackets (Bool -> SDoc
forall a. Outputable a => a -> SDoc
ppr Bool
ok) SDoc -> SDoc -> SDoc
<+> Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr Id
b SDoc -> SDoc -> SDoc
<+> SDoc
equals SDoc -> SDoc -> SDoc
<+> CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
r
  ppr (FloatTick t :: Tickish Id
t) = Tickish Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr Tickish Id
t

instance Outputable Floats where
  ppr :: Floats -> SDoc
ppr (Floats flag :: OkToSpec
flag fs :: OrdList FloatingBind
fs) = String -> SDoc
text "Floats" SDoc -> SDoc -> SDoc
<> SDoc -> SDoc
brackets (OkToSpec -> SDoc
forall a. Outputable a => a -> SDoc
ppr OkToSpec
flag) SDoc -> SDoc -> SDoc
<+>
                         SDoc -> SDoc
braces ([SDoc] -> SDoc
vcat ((FloatingBind -> SDoc) -> [FloatingBind] -> [SDoc]
forall a b. (a -> b) -> [a] -> [b]
map FloatingBind -> SDoc
forall a. Outputable a => a -> SDoc
ppr (OrdList FloatingBind -> [FloatingBind]
forall a. OrdList a -> [a]
fromOL OrdList FloatingBind
fs)))

instance Outputable OkToSpec where
  ppr :: OkToSpec -> SDoc
ppr OkToSpec    = String -> SDoc
text "OkToSpec"
  ppr IfUnboxedOk = String -> SDoc
text "IfUnboxedOk"
  ppr NotOkToSpec = String -> SDoc
text "NotOkToSpec"

-- Can we float these binds out of the rhs of a let?  We cache this decision
-- to avoid having to recompute it in a non-linear way when there are
-- deeply nested lets.
data OkToSpec
   = OkToSpec           -- Lazy bindings of lifted type
   | IfUnboxedOk        -- A mixture of lazy lifted bindings and n
                        -- ok-to-speculate unlifted bindings
   | NotOkToSpec        -- Some not-ok-to-speculate unlifted bindings

mkFloat :: Demand -> Bool -> Id -> CpeRhs -> FloatingBind
mkFloat :: Demand -> Bool -> Id -> CoreExpr -> FloatingBind
mkFloat dmd :: Demand
dmd is_unlifted :: Bool
is_unlifted bndr :: Id
bndr rhs :: CoreExpr
rhs
  | Bool
use_case  = Id -> CoreExpr -> Bool -> FloatingBind
FloatCase Id
bndr CoreExpr
rhs (CoreExpr -> Bool
exprOkForSpeculation CoreExpr
rhs)
  | Bool
is_hnf    = CoreBind -> FloatingBind
FloatLet (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
bndr                       CoreExpr
rhs)
  | Bool
otherwise = CoreBind -> FloatingBind
FloatLet (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec (Id -> Demand -> Id
setIdDemandInfo Id
bndr Demand
dmd) CoreExpr
rhs)
                   -- See Note [Pin demand info on floats]
  where
    is_hnf :: Bool
is_hnf    = CoreExpr -> Bool
exprIsHNF CoreExpr
rhs
    is_strict :: Bool
is_strict = Demand -> Bool
forall s u. JointDmd (Str s) (Use u) -> Bool
isStrictDmd Demand
dmd
    use_case :: Bool
use_case  = Bool
is_unlifted Bool -> Bool -> Bool
|| Bool
is_strict Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
is_hnf
                -- Don't make a case for a value binding,
                -- even if it's strict.  Otherwise we get
                --      case (\x -> e) of ...!

emptyFloats :: Floats
emptyFloats :: Floats
emptyFloats = OkToSpec -> OrdList FloatingBind -> Floats
Floats OkToSpec
OkToSpec OrdList FloatingBind
forall a. OrdList a
nilOL

isEmptyFloats :: Floats -> Bool
isEmptyFloats :: Floats -> Bool
isEmptyFloats (Floats _ bs :: OrdList FloatingBind
bs) = OrdList FloatingBind -> Bool
forall a. OrdList a -> Bool
isNilOL OrdList FloatingBind
bs

wrapBinds :: Floats -> CpeBody -> CpeBody
wrapBinds :: Floats -> CoreExpr -> CoreExpr
wrapBinds (Floats _ binds :: OrdList FloatingBind
binds) body :: CoreExpr
body
  = (FloatingBind -> CoreExpr -> CoreExpr)
-> CoreExpr -> OrdList FloatingBind -> CoreExpr
forall a b. (a -> b -> b) -> b -> OrdList a -> b
foldrOL FloatingBind -> CoreExpr -> CoreExpr
mk_bind CoreExpr
body OrdList FloatingBind
binds
  where
    mk_bind :: FloatingBind -> CoreExpr -> CoreExpr
mk_bind (FloatCase bndr :: Id
bndr rhs :: CoreExpr
rhs _) body :: CoreExpr
body = CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
rhs Id
bndr (CoreExpr -> Type
exprType CoreExpr
body) [(AltCon
DEFAULT, [], CoreExpr
body)]
    mk_bind (FloatLet bind :: CoreBind
bind)        body :: CoreExpr
body = CoreBind -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind CoreExpr
body
    mk_bind (FloatTick tickish :: Tickish Id
tickish)    body :: CoreExpr
body = Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
tickish CoreExpr
body

addFloat :: Floats -> FloatingBind -> Floats
addFloat :: Floats -> FloatingBind -> Floats
addFloat (Floats ok_to_spec :: OkToSpec
ok_to_spec floats :: OrdList FloatingBind
floats) new_float :: FloatingBind
new_float
  = OkToSpec -> OrdList FloatingBind -> Floats
Floats (OkToSpec -> OkToSpec -> OkToSpec
combine OkToSpec
ok_to_spec (FloatingBind -> OkToSpec
check FloatingBind
new_float)) (OrdList FloatingBind
floats OrdList FloatingBind -> FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> a -> OrdList a
`snocOL` FloatingBind
new_float)
  where
    check :: FloatingBind -> OkToSpec
check (FloatLet _) = OkToSpec
OkToSpec
    check (FloatCase _ _ ok_for_spec :: Bool
ok_for_spec)
        | Bool
ok_for_spec  =  OkToSpec
IfUnboxedOk
        | Bool
otherwise    =  OkToSpec
NotOkToSpec
    check FloatTick{}  = OkToSpec
OkToSpec
        -- The ok-for-speculation flag says that it's safe to
        -- float this Case out of a let, and thereby do it more eagerly
        -- We need the top-level flag because it's never ok to float
        -- an unboxed binding to the top level

unitFloat :: FloatingBind -> Floats
unitFloat :: FloatingBind -> Floats
unitFloat = Floats -> FloatingBind -> Floats
addFloat Floats
emptyFloats

appendFloats :: Floats -> Floats -> Floats
appendFloats :: Floats -> Floats -> Floats
appendFloats (Floats spec1 :: OkToSpec
spec1 floats1 :: OrdList FloatingBind
floats1) (Floats spec2 :: OkToSpec
spec2 floats2 :: OrdList FloatingBind
floats2)
  = OkToSpec -> OrdList FloatingBind -> Floats
Floats (OkToSpec -> OkToSpec -> OkToSpec
combine OkToSpec
spec1 OkToSpec
spec2) (OrdList FloatingBind
floats1 OrdList FloatingBind
-> OrdList FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> OrdList a -> OrdList a
`appOL` OrdList FloatingBind
floats2)

concatFloats :: [Floats] -> OrdList FloatingBind
concatFloats :: [Floats] -> OrdList FloatingBind
concatFloats = (Floats -> OrdList FloatingBind -> OrdList FloatingBind)
-> OrdList FloatingBind -> [Floats] -> OrdList FloatingBind
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Floats _ bs1 :: OrdList FloatingBind
bs1) bs2 :: OrdList FloatingBind
bs2 -> OrdList FloatingBind
-> OrdList FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> OrdList a -> OrdList a
appOL OrdList FloatingBind
bs1 OrdList FloatingBind
bs2) OrdList FloatingBind
forall a. OrdList a
nilOL

combine :: OkToSpec -> OkToSpec -> OkToSpec
combine :: OkToSpec -> OkToSpec -> OkToSpec
combine NotOkToSpec _ = OkToSpec
NotOkToSpec
combine _ NotOkToSpec = OkToSpec
NotOkToSpec
combine IfUnboxedOk _ = OkToSpec
IfUnboxedOk
combine _ IfUnboxedOk = OkToSpec
IfUnboxedOk
combine _ _           = OkToSpec
OkToSpec

deFloatTop :: Floats -> [CoreBind]
-- For top level only; we don't expect any FloatCases
deFloatTop :: Floats -> CoreProgram
deFloatTop (Floats _ floats :: OrdList FloatingBind
floats)
  = (FloatingBind -> CoreProgram -> CoreProgram)
-> CoreProgram -> OrdList FloatingBind -> CoreProgram
forall a b. (a -> b -> b) -> b -> OrdList a -> b
foldrOL FloatingBind -> CoreProgram -> CoreProgram
get [] OrdList FloatingBind
floats
  where
    get :: FloatingBind -> CoreProgram -> CoreProgram
get (FloatLet b :: CoreBind
b) bs :: CoreProgram
bs = CoreBind -> CoreBind
occurAnalyseRHSs CoreBind
b CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
: CoreProgram
bs
    get (FloatCase var :: Id
var body :: CoreExpr
body _) bs :: CoreProgram
bs  =
      CoreBind -> CoreBind
occurAnalyseRHSs (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
var CoreExpr
body) CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
: CoreProgram
bs
    get b :: FloatingBind
b _ = String -> SDoc -> CoreProgram
forall a. HasCallStack => String -> SDoc -> a
pprPanic "corePrepPgm" (FloatingBind -> SDoc
forall a. Outputable a => a -> SDoc
ppr FloatingBind
b)

    -- See Note [Dead code in CorePrep]
    occurAnalyseRHSs :: CoreBind -> CoreBind
occurAnalyseRHSs (NonRec x :: Id
x e :: CoreExpr
e) = Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
x (CoreExpr -> CoreExpr
occurAnalyseExpr_NoBinderSwap CoreExpr
e)
    occurAnalyseRHSs (Rec xes :: [(Id, CoreExpr)]
xes)    = [(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id
x, CoreExpr -> CoreExpr
occurAnalyseExpr_NoBinderSwap CoreExpr
e) | (x :: Id
x, e :: CoreExpr
e) <- [(Id, CoreExpr)]
xes]

---------------------------------------------------------------------------

canFloatFromNoCaf :: Platform -> Floats -> CpeRhs -> Maybe (Floats, CpeRhs)
       -- Note [CafInfo and floating]
canFloatFromNoCaf :: Platform -> Floats -> CoreExpr -> Maybe (Floats, CoreExpr)
canFloatFromNoCaf platform :: Platform
platform (Floats ok_to_spec :: OkToSpec
ok_to_spec fs :: OrdList FloatingBind
fs) rhs :: CoreExpr
rhs
  | OkToSpec
OkToSpec <- OkToSpec
ok_to_spec           -- Worth trying
  , Just (subst :: Subst
subst, fs' :: OrdList FloatingBind
fs') <- (Subst, OrdList FloatingBind)
-> [FloatingBind] -> Maybe (Subst, OrdList FloatingBind)
go (Subst
emptySubst, OrdList FloatingBind
forall a. OrdList a
nilOL) (OrdList FloatingBind -> [FloatingBind]
forall a. OrdList a -> [a]
fromOL OrdList FloatingBind
fs)
  = (Floats, CoreExpr) -> Maybe (Floats, CoreExpr)
forall a. a -> Maybe a
Just (OkToSpec -> OrdList FloatingBind -> Floats
Floats OkToSpec
OkToSpec OrdList FloatingBind
fs', Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
rhs)
  | Bool
otherwise
  = Maybe (Floats, CoreExpr)
forall a. Maybe a
Nothing
  where
    subst_expr :: Subst -> CoreExpr -> CoreExpr
subst_expr = SDoc -> Subst -> CoreExpr -> CoreExpr
substExpr (String -> SDoc
text "CorePrep")

    go :: (Subst, OrdList FloatingBind) -> [FloatingBind]
       -> Maybe (Subst, OrdList FloatingBind)

    go :: (Subst, OrdList FloatingBind)
-> [FloatingBind] -> Maybe (Subst, OrdList FloatingBind)
go (subst :: Subst
subst, fbs_out :: OrdList FloatingBind
fbs_out) [] = (Subst, OrdList FloatingBind)
-> Maybe (Subst, OrdList FloatingBind)
forall a. a -> Maybe a
Just (Subst
subst, OrdList FloatingBind
fbs_out)

    go (subst :: Subst
subst, fbs_out :: OrdList FloatingBind
fbs_out) (FloatLet (NonRec b :: Id
b r :: CoreExpr
r) : fbs_in :: [FloatingBind]
fbs_in)
      | CoreExpr -> Bool
rhs_ok CoreExpr
r
      = (Subst, OrdList FloatingBind)
-> [FloatingBind] -> Maybe (Subst, OrdList FloatingBind)
go (Subst
subst', OrdList FloatingBind
fbs_out OrdList FloatingBind -> FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> a -> OrdList a
`snocOL` FloatingBind
new_fb) [FloatingBind]
fbs_in
      where
        (subst' :: Subst
subst', b' :: Id
b') = Subst -> Id -> (Subst, Id)
set_nocaf_bndr Subst
subst Id
b
        new_fb :: FloatingBind
new_fb = CoreBind -> FloatingBind
FloatLet (Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
b' (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
r))

    go (subst :: Subst
subst, fbs_out :: OrdList FloatingBind
fbs_out) (FloatLet (Rec prs :: [(Id, CoreExpr)]
prs) : fbs_in :: [FloatingBind]
fbs_in)
      | (CoreExpr -> Bool) -> [CoreExpr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all CoreExpr -> Bool
rhs_ok [CoreExpr]
rs
      = (Subst, OrdList FloatingBind)
-> [FloatingBind] -> Maybe (Subst, OrdList FloatingBind)
go (Subst
subst', OrdList FloatingBind
fbs_out OrdList FloatingBind -> FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> a -> OrdList a
`snocOL` FloatingBind
new_fb) [FloatingBind]
fbs_in
      where
        (bs :: [Id]
bs,rs :: [CoreExpr]
rs) = [(Id, CoreExpr)] -> ([Id], [CoreExpr])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Id, CoreExpr)]
prs
        (subst' :: Subst
subst', bs' :: [Id]
bs') = (Subst -> Id -> (Subst, Id)) -> Subst -> [Id] -> (Subst, [Id])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL Subst -> Id -> (Subst, Id)
set_nocaf_bndr Subst
subst [Id]
bs
        rs' :: [CoreExpr]
rs' = (CoreExpr -> CoreExpr) -> [CoreExpr] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst') [CoreExpr]
rs
        new_fb :: FloatingBind
new_fb = CoreBind -> FloatingBind
FloatLet ([(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([Id]
bs' [Id] -> [CoreExpr] -> [(Id, CoreExpr)]
forall a b. [a] -> [b] -> [(a, b)]
`zip` [CoreExpr]
rs'))

    go (subst :: Subst
subst, fbs_out :: OrdList FloatingBind
fbs_out) (ft :: FloatingBind
ft@FloatTick{} : fbs_in :: [FloatingBind]
fbs_in)
      = (Subst, OrdList FloatingBind)
-> [FloatingBind] -> Maybe (Subst, OrdList FloatingBind)
go (Subst
subst, OrdList FloatingBind
fbs_out OrdList FloatingBind -> FloatingBind -> OrdList FloatingBind
forall a. OrdList a -> a -> OrdList a
`snocOL` FloatingBind
ft) [FloatingBind]
fbs_in

    go _ _ = Maybe (Subst, OrdList FloatingBind)
forall a. Maybe a
Nothing      -- Encountered a caffy binding

    ------------
    set_nocaf_bndr :: Subst -> Id -> (Subst, Id)
set_nocaf_bndr subst :: Subst
subst bndr :: Id
bndr
      = (Subst -> Id -> CoreExpr -> Subst
extendIdSubst Subst
subst Id
bndr (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
bndr'), Id
bndr')
      where
        bndr' :: Id
bndr' = Id
bndr Id -> CafInfo -> Id
`setIdCafInfo` CafInfo
NoCafRefs

    ------------
    rhs_ok :: CoreExpr -> Bool
    -- We can only float to top level from a NoCaf thing if
    -- the new binding is static. However it can't mention
    -- any non-static things or it would *already* be Caffy
    rhs_ok :: CoreExpr -> Bool
rhs_ok = Platform
-> (Name -> Bool)
-> (LitNumType -> Integer -> Maybe CoreExpr)
-> CoreExpr
-> Bool
rhsIsStatic Platform
platform (\_ -> Bool
False)
                         (\_nt :: LitNumType
_nt i :: Integer
i -> String -> SDoc -> Maybe CoreExpr
forall a. HasCallStack => String -> SDoc -> a
pprPanic "rhsIsStatic" (Integer -> SDoc
integer Integer
i))
                         -- Integer or Natural literals should not show up

wantFloatNested :: RecFlag -> Demand -> Bool -> Floats -> CpeRhs -> Bool
wantFloatNested :: RecFlag -> Demand -> Bool -> Floats -> CoreExpr -> Bool
wantFloatNested is_rec :: RecFlag
is_rec dmd :: Demand
dmd is_unlifted :: Bool
is_unlifted floats :: Floats
floats rhs :: CoreExpr
rhs
  =  Floats -> Bool
isEmptyFloats Floats
floats
  Bool -> Bool -> Bool
|| Demand -> Bool
forall s u. JointDmd (Str s) (Use u) -> Bool
isStrictDmd Demand
dmd
  Bool -> Bool -> Bool
|| Bool
is_unlifted
  Bool -> Bool -> Bool
|| (RecFlag -> Floats -> Bool
allLazyNested RecFlag
is_rec Floats
floats Bool -> Bool -> Bool
&& CoreExpr -> Bool
exprIsHNF CoreExpr
rhs)
        -- Why the test for allLazyNested?
        --      v = f (x `divInt#` y)
        -- we don't want to float the case, even if f has arity 2,
        -- because floating the case would make it evaluated too early

allLazyTop :: Floats -> Bool
allLazyTop :: Floats -> Bool
allLazyTop (Floats OkToSpec _) = Bool
True
allLazyTop _                   = Bool
False

allLazyNested :: RecFlag -> Floats -> Bool
allLazyNested :: RecFlag -> Floats -> Bool
allLazyNested _      (Floats OkToSpec    _) = Bool
True
allLazyNested _      (Floats NotOkToSpec _) = Bool
False
allLazyNested is_rec :: RecFlag
is_rec (Floats IfUnboxedOk _) = RecFlag -> Bool
isNonRec RecFlag
is_rec

{-
************************************************************************
*                                                                      *
                Cloning
*                                                                      *
************************************************************************
-}

-- ---------------------------------------------------------------------------
--                      The environment
-- ---------------------------------------------------------------------------

-- Note [Inlining in CorePrep]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- There is a subtle but important invariant that must be upheld in the output
-- of CorePrep: there are no "trivial" updatable thunks.  Thus, this Core
-- is impermissible:
--
--      let x :: ()
--          x = y
--
-- (where y is a reference to a GLOBAL variable).  Thunks like this are silly:
-- they can always be profitably replaced by inlining x with y. Consequently,
-- the code generator/runtime does not bother implementing this properly
-- (specifically, there is no implementation of stg_ap_0_upd_info, which is the
-- stack frame that would be used to update this thunk.  The "0" means it has
-- zero free variables.)
--
-- In general, the inliner is good at eliminating these let-bindings.  However,
-- there is one case where these trivial updatable thunks can arise: when
-- we are optimizing away 'lazy' (see Note [lazyId magic], and also
-- 'cpeRhsE'.)  Then, we could have started with:
--
--      let x :: ()
--          x = lazy @ () y
--
-- which is a perfectly fine, non-trivial thunk, but then CorePrep will
-- drop 'lazy', giving us 'x = y' which is trivial and impermissible.
-- The solution is CorePrep to have a miniature inlining pass which deals
-- with cases like this.  We can then drop the let-binding altogether.
--
-- Why does the removal of 'lazy' have to occur in CorePrep?
-- The gory details are in Note [lazyId magic] in MkId, but the
-- main reason is that lazy must appear in unfoldings (optimizer
-- output) and it must prevent call-by-value for catch# (which
-- is implemented by CorePrep.)
--
-- An alternate strategy for solving this problem is to have the
-- inliner treat 'lazy e' as a trivial expression if 'e' is trivial.
-- We decided not to adopt this solution to keep the definition
-- of 'exprIsTrivial' simple.
--
-- There is ONE caveat however: for top-level bindings we have
-- to preserve the binding so that we float the (hacky) non-recursive
-- binding for data constructors; see Note [Data constructor workers].
--
-- Note [CorePrep inlines trivial CoreExpr not Id]
-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-- Why does cpe_env need to be an IdEnv CoreExpr, as opposed to an
-- IdEnv Id?  Naively, we might conjecture that trivial updatable thunks
-- as per Note [Inlining in CorePrep] always have the form
-- 'lazy @ SomeType gbl_id'.  But this is not true: the following is
-- perfectly reasonable Core:
--
--      let x :: ()
--          x = lazy @ (forall a. a) y @ Bool
--
-- When we inline 'x' after eliminating 'lazy', we need to replace
-- occurrences of 'x' with 'y @ bool', not just 'y'.  Situations like
-- this can easily arise with higher-rank types; thus, cpe_env must
-- map to CoreExprs, not Ids.

data CorePrepEnv
  = CPE { CorePrepEnv -> DynFlags
cpe_dynFlags        :: DynFlags
        , CorePrepEnv -> IdEnv CoreExpr
cpe_env             :: IdEnv CoreExpr   -- Clone local Ids
        -- ^ This environment is used for three operations:
        --
        --      1. To support cloning of local Ids so that they are
        --      all unique (see item (6) of CorePrep overview).
        --
        --      2. To support beta-reduction of runRW, see
        --      Note [runRW magic] and Note [runRW arg].
        --
        --      3. To let us inline trivial RHSs of non top-level let-bindings,
        --      see Note [lazyId magic], Note [Inlining in CorePrep]
        --      and Note [CorePrep inlines trivial CoreExpr not Id] (#12076)
        , CorePrepEnv -> Id
cpe_mkIntegerId     :: Id
        , CorePrepEnv -> Id
cpe_mkNaturalId     :: Id
        , CorePrepEnv -> Maybe DataCon
cpe_integerSDataCon :: Maybe DataCon
        , CorePrepEnv -> Maybe DataCon
cpe_naturalSDataCon :: Maybe DataCon
    }

lookupMkIntegerName :: DynFlags -> HscEnv -> IO Id
lookupMkIntegerName :: DynFlags -> HscEnv -> IO Id
lookupMkIntegerName dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env
    = DynFlags -> IO Id -> IO Id
forall a. DynFlags -> IO a -> IO a
guardIntegerUse DynFlags
dflags (IO Id -> IO Id) -> IO Id -> IO Id
forall a b. (a -> b) -> a -> b
$ (TyThing -> Id) -> IO TyThing -> IO Id
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM TyThing -> Id
tyThingId (IO TyThing -> IO Id) -> IO TyThing -> IO Id
forall a b. (a -> b) -> a -> b
$
      HscEnv -> Name -> IO TyThing
lookupGlobal HscEnv
hsc_env Name
mkIntegerName

lookupMkNaturalName :: DynFlags -> HscEnv -> IO Id
lookupMkNaturalName :: DynFlags -> HscEnv -> IO Id
lookupMkNaturalName dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env
    = DynFlags -> IO Id -> IO Id
forall a. DynFlags -> IO a -> IO a
guardNaturalUse DynFlags
dflags (IO Id -> IO Id) -> IO Id -> IO Id
forall a b. (a -> b) -> a -> b
$ (TyThing -> Id) -> IO TyThing -> IO Id
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM TyThing -> Id
tyThingId (IO TyThing -> IO Id) -> IO TyThing -> IO Id
forall a b. (a -> b) -> a -> b
$
      HscEnv -> Name -> IO TyThing
lookupGlobal HscEnv
hsc_env Name
mkNaturalName

-- See Note [The integer library] in PrelNames
lookupIntegerSDataConName :: DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupIntegerSDataConName :: DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupIntegerSDataConName dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env = case DynFlags -> IntegerLibrary
integerLibrary DynFlags
dflags of
    IntegerGMP -> DynFlags -> IO (Maybe DataCon) -> IO (Maybe DataCon)
forall a. DynFlags -> IO a -> IO a
guardIntegerUse DynFlags
dflags (IO (Maybe DataCon) -> IO (Maybe DataCon))
-> IO (Maybe DataCon) -> IO (Maybe DataCon)
forall a b. (a -> b) -> a -> b
$ (TyThing -> Maybe DataCon) -> IO TyThing -> IO (Maybe DataCon)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (DataCon -> Maybe DataCon
forall a. a -> Maybe a
Just (DataCon -> Maybe DataCon)
-> (TyThing -> DataCon) -> TyThing -> Maybe DataCon
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyThing -> DataCon
tyThingDataCon) (IO TyThing -> IO (Maybe DataCon))
-> IO TyThing -> IO (Maybe DataCon)
forall a b. (a -> b) -> a -> b
$
                  HscEnv -> Name -> IO TyThing
lookupGlobal HscEnv
hsc_env Name
integerSDataConName
    IntegerSimple -> Maybe DataCon -> IO (Maybe DataCon)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe DataCon
forall a. Maybe a
Nothing

lookupNaturalSDataConName :: DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupNaturalSDataConName :: DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupNaturalSDataConName dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env = case DynFlags -> IntegerLibrary
integerLibrary DynFlags
dflags of
    IntegerGMP -> DynFlags -> IO (Maybe DataCon) -> IO (Maybe DataCon)
forall a. DynFlags -> IO a -> IO a
guardNaturalUse DynFlags
dflags (IO (Maybe DataCon) -> IO (Maybe DataCon))
-> IO (Maybe DataCon) -> IO (Maybe DataCon)
forall a b. (a -> b) -> a -> b
$ (TyThing -> Maybe DataCon) -> IO TyThing -> IO (Maybe DataCon)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (DataCon -> Maybe DataCon
forall a. a -> Maybe a
Just (DataCon -> Maybe DataCon)
-> (TyThing -> DataCon) -> TyThing -> Maybe DataCon
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyThing -> DataCon
tyThingDataCon) (IO TyThing -> IO (Maybe DataCon))
-> IO TyThing -> IO (Maybe DataCon)
forall a b. (a -> b) -> a -> b
$
                  HscEnv -> Name -> IO TyThing
lookupGlobal HscEnv
hsc_env Name
naturalSDataConName
    IntegerSimple -> Maybe DataCon -> IO (Maybe DataCon)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe DataCon
forall a. Maybe a
Nothing

-- | Helper for 'lookupMkIntegerName', 'lookupIntegerSDataConName'
guardIntegerUse :: DynFlags -> IO a -> IO a
guardIntegerUse :: DynFlags -> IO a -> IO a
guardIntegerUse dflags :: DynFlags
dflags act :: IO a
act
  | DynFlags -> UnitId
thisPackage DynFlags
dflags UnitId -> UnitId -> Bool
forall a. Eq a => a -> a -> Bool
== UnitId
primUnitId
  = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ String -> a
forall a. String -> a
panic "Can't use Integer in ghc-prim"
  | DynFlags -> UnitId
thisPackage DynFlags
dflags UnitId -> UnitId -> Bool
forall a. Eq a => a -> a -> Bool
== UnitId
integerUnitId
  = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ String -> a
forall a. String -> a
panic "Can't use Integer in integer-*"
  | Bool
otherwise = IO a
act

-- | Helper for 'lookupMkNaturalName', 'lookupNaturalSDataConName'
--
-- Just like we can't use Integer literals in `integer-*`, we can't use Natural
-- literals in `base`. If we do, we get interface loading error for GHC.Natural.
guardNaturalUse :: DynFlags -> IO a -> IO a
guardNaturalUse :: DynFlags -> IO a -> IO a
guardNaturalUse dflags :: DynFlags
dflags act :: IO a
act
  | DynFlags -> UnitId
thisPackage DynFlags
dflags UnitId -> UnitId -> Bool
forall a. Eq a => a -> a -> Bool
== UnitId
primUnitId
  = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ String -> a
forall a. String -> a
panic "Can't use Natural in ghc-prim"
  | DynFlags -> UnitId
thisPackage DynFlags
dflags UnitId -> UnitId -> Bool
forall a. Eq a => a -> a -> Bool
== UnitId
integerUnitId
  = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ String -> a
forall a. String -> a
panic "Can't use Natural in integer-*"
  | DynFlags -> UnitId
thisPackage DynFlags
dflags UnitId -> UnitId -> Bool
forall a. Eq a => a -> a -> Bool
== UnitId
baseUnitId
  = a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ String -> a
forall a. String -> a
panic "Can't use Natural in base"
  | Bool
otherwise = IO a
act

mkInitialCorePrepEnv :: DynFlags -> HscEnv -> IO CorePrepEnv
mkInitialCorePrepEnv :: DynFlags -> HscEnv -> IO CorePrepEnv
mkInitialCorePrepEnv dflags :: DynFlags
dflags hsc_env :: HscEnv
hsc_env
    = do Id
mkIntegerId <- DynFlags -> HscEnv -> IO Id
lookupMkIntegerName DynFlags
dflags HscEnv
hsc_env
         Id
mkNaturalId <- DynFlags -> HscEnv -> IO Id
lookupMkNaturalName DynFlags
dflags HscEnv
hsc_env
         Maybe DataCon
integerSDataCon <- DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupIntegerSDataConName DynFlags
dflags HscEnv
hsc_env
         Maybe DataCon
naturalSDataCon <- DynFlags -> HscEnv -> IO (Maybe DataCon)
lookupNaturalSDataConName DynFlags
dflags HscEnv
hsc_env
         CorePrepEnv -> IO CorePrepEnv
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> IO CorePrepEnv) -> CorePrepEnv -> IO CorePrepEnv
forall a b. (a -> b) -> a -> b
$ CPE :: DynFlags
-> IdEnv CoreExpr
-> Id
-> Id
-> Maybe DataCon
-> Maybe DataCon
-> CorePrepEnv
CPE {
                      cpe_dynFlags :: DynFlags
cpe_dynFlags = DynFlags
dflags,
                      cpe_env :: IdEnv CoreExpr
cpe_env = IdEnv CoreExpr
forall a. VarEnv a
emptyVarEnv,
                      cpe_mkIntegerId :: Id
cpe_mkIntegerId = Id
mkIntegerId,
                      cpe_mkNaturalId :: Id
cpe_mkNaturalId = Id
mkNaturalId,
                      cpe_integerSDataCon :: Maybe DataCon
cpe_integerSDataCon = Maybe DataCon
integerSDataCon,
                      cpe_naturalSDataCon :: Maybe DataCon
cpe_naturalSDataCon = Maybe DataCon
naturalSDataCon
                  }

extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv :: CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv cpe :: CorePrepEnv
cpe id :: Id
id id' :: Id
id'
    = CorePrepEnv
cpe { cpe_env :: IdEnv CoreExpr
cpe_env = IdEnv CoreExpr -> Id -> CoreExpr -> IdEnv CoreExpr
forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv (CorePrepEnv -> IdEnv CoreExpr
cpe_env CorePrepEnv
cpe) Id
id (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
id') }

extendCorePrepEnvExpr :: CorePrepEnv -> Id -> CoreExpr -> CorePrepEnv
extendCorePrepEnvExpr :: CorePrepEnv -> Id -> CoreExpr -> CorePrepEnv
extendCorePrepEnvExpr cpe :: CorePrepEnv
cpe id :: Id
id expr :: CoreExpr
expr
    = CorePrepEnv
cpe { cpe_env :: IdEnv CoreExpr
cpe_env = IdEnv CoreExpr -> Id -> CoreExpr -> IdEnv CoreExpr
forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv (CorePrepEnv -> IdEnv CoreExpr
cpe_env CorePrepEnv
cpe) Id
id CoreExpr
expr }

extendCorePrepEnvList :: CorePrepEnv -> [(Id,Id)] -> CorePrepEnv
extendCorePrepEnvList :: CorePrepEnv -> [(Id, Id)] -> CorePrepEnv
extendCorePrepEnvList cpe :: CorePrepEnv
cpe prs :: [(Id, Id)]
prs
    = CorePrepEnv
cpe { cpe_env :: IdEnv CoreExpr
cpe_env = IdEnv CoreExpr -> [(Id, CoreExpr)] -> IdEnv CoreExpr
forall a. VarEnv a -> [(Id, a)] -> VarEnv a
extendVarEnvList (CorePrepEnv -> IdEnv CoreExpr
cpe_env CorePrepEnv
cpe)
                        (((Id, Id) -> (Id, CoreExpr)) -> [(Id, Id)] -> [(Id, CoreExpr)]
forall a b. (a -> b) -> [a] -> [b]
map (\(id :: Id
id, id' :: Id
id') -> (Id
id, Id -> CoreExpr
forall b. Id -> Expr b
Var Id
id')) [(Id, Id)]
prs) }

lookupCorePrepEnv :: CorePrepEnv -> Id -> CoreExpr
lookupCorePrepEnv :: CorePrepEnv -> Id -> CoreExpr
lookupCorePrepEnv cpe :: CorePrepEnv
cpe id :: Id
id
  = case IdEnv CoreExpr -> Id -> Maybe CoreExpr
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (CorePrepEnv -> IdEnv CoreExpr
cpe_env CorePrepEnv
cpe) Id
id of
        Nothing  -> Id -> CoreExpr
forall b. Id -> Expr b
Var Id
id
        Just exp :: CoreExpr
exp -> CoreExpr
exp

getMkIntegerId :: CorePrepEnv -> Id
getMkIntegerId :: CorePrepEnv -> Id
getMkIntegerId = CorePrepEnv -> Id
cpe_mkIntegerId

getMkNaturalId :: CorePrepEnv -> Id
getMkNaturalId :: CorePrepEnv -> Id
getMkNaturalId = CorePrepEnv -> Id
cpe_mkNaturalId

------------------------------------------------------------------------------
-- Cloning binders
-- ---------------------------------------------------------------------------

cpCloneBndrs :: CorePrepEnv -> [InVar] -> UniqSM (CorePrepEnv, [OutVar])
cpCloneBndrs :: CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
cpCloneBndrs env :: CorePrepEnv
env bs :: [Id]
bs = (CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id))
-> CorePrepEnv -> [Id] -> UniqSM (CorePrepEnv, [Id])
forall (m :: * -> *) acc x y.
Monad m =>
(acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id)
cpCloneBndr CorePrepEnv
env [Id]
bs

cpCloneBndr  :: CorePrepEnv -> InVar -> UniqSM (CorePrepEnv, OutVar)
cpCloneBndr :: CorePrepEnv -> Id -> UniqSM (CorePrepEnv, Id)
cpCloneBndr env :: CorePrepEnv
env bndr :: Id
bndr
  | Bool -> Bool
not (Id -> Bool
isId Id
bndr)
  = (CorePrepEnv, Id) -> UniqSM (CorePrepEnv, Id)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv
env, Id
bndr)

  | Bool
otherwise
  = do { Id
bndr' <- Id -> UniqSM Id
forall (m :: * -> *). MonadUnique m => Id -> m Id
clone_it Id
bndr

       -- Drop (now-useless) rules/unfoldings
       -- See Note [Drop unfoldings and rules]
       -- and Note [Preserve evaluatedness] in CoreTidy
       ; let unfolding' :: Unfolding
unfolding' = Unfolding -> Unfolding
zapUnfolding (Id -> Unfolding
realIdUnfolding Id
bndr)
                          -- Simplifier will set the Id's unfolding

             bndr'' :: Id
bndr'' = Id
bndr' Id -> Unfolding -> Id
`setIdUnfolding`      Unfolding
unfolding'
                            Id -> RuleInfo -> Id
`setIdSpecialisation` RuleInfo
emptyRuleInfo

       ; (CorePrepEnv, Id) -> UniqSM (CorePrepEnv, Id)
forall (m :: * -> *) a. Monad m => a -> m a
return (CorePrepEnv -> Id -> Id -> CorePrepEnv
extendCorePrepEnv CorePrepEnv
env Id
bndr Id
bndr'', Id
bndr'') }
  where
    clone_it :: Id -> m Id
clone_it bndr :: Id
bndr
      | Id -> Bool
isLocalId Id
bndr, Bool -> Bool
not (Id -> Bool
isCoVar Id
bndr)
      = do { Unique
uniq <- m Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM; Id -> m Id
forall (m :: * -> *) a. Monad m => a -> m a
return (Id -> Unique -> Id
setVarUnique Id
bndr Unique
uniq) }
      | Bool
otherwise   -- Top level things, which we don't want
                    -- to clone, have become GlobalIds by now
                    -- And we don't clone tyvars, or coercion variables
      = Id -> m Id
forall (m :: * -> *) a. Monad m => a -> m a
return Id
bndr

{- Note [Drop unfoldings and rules]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to drop the unfolding/rules on every Id:

  - We are now past interface-file generation, and in the
    codegen pipeline, so we really don't need full unfoldings/rules

  - The unfolding/rule may be keeping stuff alive that we'd like
    to discard.  See  Note [Dead code in CorePrep]

  - Getting rid of unnecessary unfoldings reduces heap usage

  - We are changing uniques, so if we didn't discard unfoldings/rules
    we'd have to substitute in them

HOWEVER, we want to preserve evaluated-ness;
see Note [Preserve evaluatedness] in CoreTidy.
-}

------------------------------------------------------------------------------
-- Cloning ccall Ids; each must have a unique name,
-- to give the code generator a handle to hang it on
-- ---------------------------------------------------------------------------

fiddleCCall :: Id -> UniqSM Id
fiddleCCall :: Id -> UniqSM Id
fiddleCCall id :: Id
id
  | Id -> Bool
isFCallId Id
id = (Id
id Id -> Unique -> Id
`setVarUnique`) (Unique -> Id) -> UniqSM Unique -> UniqSM Id
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UniqSM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
  | Bool
otherwise    = Id -> UniqSM Id
forall (m :: * -> *) a. Monad m => a -> m a
return Id
id

------------------------------------------------------------------------------
-- Generating new binders
-- ---------------------------------------------------------------------------

newVar :: Type -> UniqSM Id
newVar :: Type -> UniqSM Id
newVar ty :: Type
ty
 = Type -> ()
seqType Type
ty () -> UniqSM Id -> UniqSM Id
forall a b. a -> b -> b
`seq` do
     Unique
uniq <- UniqSM Unique
forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
     Id -> UniqSM Id
forall (m :: * -> *) a. Monad m => a -> m a
return (FastString -> Unique -> Type -> Id
mkSysLocalOrCoVar (String -> FastString
fsLit "sat") Unique
uniq Type
ty)


------------------------------------------------------------------------------
-- Floating ticks
-- ---------------------------------------------------------------------------
--
-- Note [Floating Ticks in CorePrep]
--
-- It might seem counter-intuitive to float ticks by default, given
-- that we don't actually want to move them if we can help it. On the
-- other hand, nothing gets very far in CorePrep anyway, and we want
-- to preserve the order of let bindings and tick annotations in
-- relation to each other. For example, if we just wrapped let floats
-- when they pass through ticks, we might end up performing the
-- following transformation:
--
--   src<...> let foo = bar in baz
--   ==>  let foo = src<...> bar in src<...> baz
--
-- Because the let-binding would float through the tick, and then
-- immediately materialize, achieving nothing but decreasing tick
-- accuracy. The only special case is the following scenario:
--
--   let foo = src<...> (let a = b in bar) in baz
--   ==>  let foo = src<...> bar; a = src<...> b in baz
--
-- Here we would not want the source tick to end up covering "baz" and
-- therefore refrain from pushing ticks outside. Instead, we copy them
-- into the floating binds (here "a") in cpePair. Note that where "b"
-- or "bar" are (value) lambdas we have to push the annotations
-- further inside in order to uphold our rules.
--
-- All of this is implemented below in @wrapTicks@.

-- | Like wrapFloats, but only wraps tick floats
wrapTicks :: Floats -> CoreExpr -> (Floats, CoreExpr)
wrapTicks :: Floats -> CoreExpr -> (Floats, CoreExpr)
wrapTicks (Floats flag :: OkToSpec
flag floats0 :: OrdList FloatingBind
floats0) expr :: CoreExpr
expr =
    (OkToSpec -> OrdList FloatingBind -> Floats
Floats OkToSpec
flag ([FloatingBind] -> OrdList FloatingBind
forall a. [a] -> OrdList a
toOL ([FloatingBind] -> OrdList FloatingBind)
-> [FloatingBind] -> OrdList FloatingBind
forall a b. (a -> b) -> a -> b
$ [FloatingBind] -> [FloatingBind]
forall a. [a] -> [a]
reverse [FloatingBind]
floats1), (Tickish Id -> CoreExpr -> CoreExpr)
-> CoreExpr -> [Tickish Id] -> CoreExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tickish Id -> CoreExpr -> CoreExpr
mkTick CoreExpr
expr ([Tickish Id] -> [Tickish Id]
forall a. [a] -> [a]
reverse [Tickish Id]
ticks1))
  where (floats1 :: [FloatingBind]
floats1, ticks1 :: [Tickish Id]
ticks1) = (([FloatingBind], [Tickish Id])
 -> FloatingBind -> ([FloatingBind], [Tickish Id]))
-> ([FloatingBind], [Tickish Id])
-> OrdList FloatingBind
-> ([FloatingBind], [Tickish Id])
forall b a. (b -> a -> b) -> b -> OrdList a -> b
foldlOL ([FloatingBind], [Tickish Id])
-> FloatingBind -> ([FloatingBind], [Tickish Id])
go ([], []) (OrdList FloatingBind -> ([FloatingBind], [Tickish Id]))
-> OrdList FloatingBind -> ([FloatingBind], [Tickish Id])
forall a b. (a -> b) -> a -> b
$ OrdList FloatingBind
floats0
        -- Deeply nested constructors will produce long lists of
        -- redundant source note floats here. We need to eliminate
        -- those early, as relying on mkTick to spot it after the fact
        -- can yield O(n^3) complexity [#11095]
        go :: ([FloatingBind], [Tickish Id])
-> FloatingBind -> ([FloatingBind], [Tickish Id])
go (floats :: [FloatingBind]
floats, ticks :: [Tickish Id]
ticks) (FloatTick t :: Tickish Id
t)
          = ASSERT(tickishPlace t == PlaceNonLam)
            ([FloatingBind]
floats, if (Tickish Id -> Bool) -> [Tickish Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Tickish Id -> Tickish Id -> Bool)
-> Tickish Id -> Tickish Id -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip Tickish Id -> Tickish Id -> Bool
forall b. Eq b => Tickish b -> Tickish b -> Bool
tickishContains Tickish Id
t) [Tickish Id]
ticks
                     then [Tickish Id]
ticks else Tickish Id
tTickish Id -> [Tickish Id] -> [Tickish Id]
forall a. a -> [a] -> [a]
:[Tickish Id]
ticks)
        go (floats :: [FloatingBind]
floats, ticks :: [Tickish Id]
ticks) f :: FloatingBind
f
          = ((Tickish Id -> FloatingBind -> FloatingBind)
-> FloatingBind -> [Tickish Id] -> FloatingBind
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tickish Id -> FloatingBind -> FloatingBind
wrap FloatingBind
f ([Tickish Id] -> [Tickish Id]
forall a. [a] -> [a]
reverse [Tickish Id]
ticks)FloatingBind -> [FloatingBind] -> [FloatingBind]
forall a. a -> [a] -> [a]
:[FloatingBind]
floats, [Tickish Id]
ticks)

        wrap :: Tickish Id -> FloatingBind -> FloatingBind
wrap t :: Tickish Id
t (FloatLet bind :: CoreBind
bind)    = CoreBind -> FloatingBind
FloatLet (Tickish Id -> CoreBind -> CoreBind
wrapBind Tickish Id
t CoreBind
bind)
        wrap t :: Tickish Id
t (FloatCase b :: Id
b r :: CoreExpr
r ok :: Bool
ok) = Id -> CoreExpr -> Bool -> FloatingBind
FloatCase Id
b (Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
t CoreExpr
r) Bool
ok
        wrap _ other :: FloatingBind
other              = String -> SDoc -> FloatingBind
forall a. HasCallStack => String -> SDoc -> a
pprPanic "wrapTicks: unexpected float!"
                                             (FloatingBind -> SDoc
forall a. Outputable a => a -> SDoc
ppr FloatingBind
other)
        wrapBind :: Tickish Id -> CoreBind -> CoreBind
wrapBind t :: Tickish Id
t (NonRec binder :: Id
binder rhs :: CoreExpr
rhs) = Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
binder (Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
t CoreExpr
rhs)
        wrapBind t :: Tickish Id
t (Rec pairs :: [(Id, CoreExpr)]
pairs)         = [(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ((CoreExpr -> CoreExpr) -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
forall b c a. (b -> c) -> [(a, b)] -> [(a, c)]
mapSnd (Tickish Id -> CoreExpr -> CoreExpr
mkTick Tickish Id
t) [(Id, CoreExpr)]
pairs)

------------------------------------------------------------------------------
-- Collecting cost centres
-- ---------------------------------------------------------------------------

-- | Collect cost centres defined in the current module, including those in
-- unfoldings.
collectCostCentres :: Module -> CoreProgram -> S.Set CostCentre
collectCostCentres :: Module -> CoreProgram -> Set CostCentre
collectCostCentres mod_name :: Module
mod_name
  = (Set CostCentre -> CoreBind -> Set CostCentre)
-> Set CostCentre -> CoreProgram -> Set CostCentre
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Set CostCentre -> CoreBind -> Set CostCentre
go_bind Set CostCentre
forall a. Set a
S.empty
  where
    go :: Set CostCentre -> CoreExpr -> Set CostCentre
go cs :: Set CostCentre
cs e :: CoreExpr
e = case CoreExpr
e of
      Var{} -> Set CostCentre
cs
      Lit{} -> Set CostCentre
cs
      App e1 :: CoreExpr
e1 e2 :: CoreExpr
e2 -> Set CostCentre -> CoreExpr -> Set CostCentre
go (Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
e1) CoreExpr
e2
      Lam _ e :: CoreExpr
e -> Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
e
      Let b :: CoreBind
b e :: CoreExpr
e -> Set CostCentre -> CoreExpr -> Set CostCentre
go (Set CostCentre -> CoreBind -> Set CostCentre
go_bind Set CostCentre
cs CoreBind
b) CoreExpr
e
      Case scrt :: CoreExpr
scrt _ _ alts :: [Alt Id]
alts -> Set CostCentre -> [Alt Id] -> Set CostCentre
go_alts (Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
scrt) [Alt Id]
alts
      Cast e :: CoreExpr
e _ -> Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
e
      Tick (ProfNote cc :: CostCentre
cc _ _) e :: CoreExpr
e ->
        Set CostCentre -> CoreExpr -> Set CostCentre
go (if CostCentre -> Module -> Bool
ccFromThisModule CostCentre
cc Module
mod_name then CostCentre -> Set CostCentre -> Set CostCentre
forall a. Ord a => a -> Set a -> Set a
S.insert CostCentre
cc Set CostCentre
cs else Set CostCentre
cs) CoreExpr
e
      Tick _ e :: CoreExpr
e -> Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
e
      Type{} -> Set CostCentre
cs
      Coercion{} -> Set CostCentre
cs

    go_alts :: Set CostCentre -> [Alt Id] -> Set CostCentre
go_alts = (Set CostCentre -> Alt Id -> Set CostCentre)
-> Set CostCentre -> [Alt Id] -> Set CostCentre
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\cs :: Set CostCentre
cs (_con :: AltCon
_con, _bndrs :: [Id]
_bndrs, e :: CoreExpr
e) -> Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs CoreExpr
e)

    go_bind :: S.Set CostCentre -> CoreBind -> S.Set CostCentre
    go_bind :: Set CostCentre -> CoreBind -> Set CostCentre
go_bind cs :: Set CostCentre
cs (NonRec b :: Id
b e :: CoreExpr
e) =
      Set CostCentre -> CoreExpr -> Set CostCentre
go (Set CostCentre
-> (CoreExpr -> Set CostCentre) -> Maybe CoreExpr -> Set CostCentre
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set CostCentre
cs (Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs) (Id -> Maybe CoreExpr
get_unf Id
b)) CoreExpr
e
    go_bind cs :: Set CostCentre
cs (Rec bs :: [(Id, CoreExpr)]
bs) =
      (Set CostCentre -> (Id, CoreExpr) -> Set CostCentre)
-> Set CostCentre -> [(Id, CoreExpr)] -> Set CostCentre
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\cs' :: Set CostCentre
cs' (b :: Id
b, e :: CoreExpr
e) -> Set CostCentre -> CoreExpr -> Set CostCentre
go (Set CostCentre
-> (CoreExpr -> Set CostCentre) -> Maybe CoreExpr -> Set CostCentre
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set CostCentre
cs' (Set CostCentre -> CoreExpr -> Set CostCentre
go Set CostCentre
cs') (Id -> Maybe CoreExpr
get_unf Id
b)) CoreExpr
e) Set CostCentre
cs [(Id, CoreExpr)]
bs

    -- Unfoldings may have cost centres that in the original definion are
    -- optimized away, see #5889.
    get_unf :: Id -> Maybe CoreExpr
get_unf = Unfolding -> Maybe CoreExpr
maybeUnfoldingTemplate (Unfolding -> Maybe CoreExpr)
-> (Id -> Unfolding) -> Id -> Maybe CoreExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Unfolding
realIdUnfolding