{-# LANGUAGE TypeFamilies #-}

{-|
Note [CSE for Stg]
~~~~~~~~~~~~~~~~~~

This module implements a simple common subexpression elimination pass for STG.
This is useful because there are expressions that we want to common up (because
they are operationally equivalent), but that we cannot common up in Core, because
their types differ.
This was originally reported as #9291.

There are two types of common code occurrences that we aim for, see
note [Case 1: CSEing allocated closures] and
note [Case 2: CSEing case binders] below.


Note [Case 1: CSEing allocated closures]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The first kind of CSE opportunity we aim for is generated by this Haskell code:

    bar :: a -> (Either Int a, Either Bool a)
    bar x = (Right x, Right x)

which produces this Core:

    bar :: forall a. a -> (Either Int a, Either Bool a)
    bar @a x = (Right @Int @a x, Right @Bool @a x)

where the two components of the tuple are different terms, and cannot be
commoned up (easily). On the STG level we have

    bar [x] = let c1 = Right [x]
                  c2 = Right [x]
              in (c1,c2)

and now it is obvious that we can write

    bar [x] = let c1 = Right [x]
              in (c1,c1)

instead.


Note [Case 2: CSEing case binders]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The second kind of CSE opportunity we aim for is more interesting, and
came up in #9291 and #5344: The Haskell code

    foo :: Either Int a -> Either Bool a
    foo (Right x) = Right x
    foo _         = Left False

produces this Core

    foo :: forall a. Either Int a -> Either Bool a
    foo @a e = case e of b { Left n -> …
                           , Right x -> Right @Bool @a x }

where we cannot CSE `Right @Bool @a x` with the case binder `b` as they have
different types. But in STG we have

    foo [e] = case e of b { Left [n] -> …
                          , Right [x] -> Right [x] }

and nothing stops us from transforming that to

    foo [e] = case e of b { Left [n] -> …
                          , Right [x] -> b}


Note [StgCse after unarisation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Consider two unboxed sum terms:

    (# 1 | #) :: (# Int | Int# #)
    (# 1 | #) :: (# Int | Int  #)

These two terms are not equal as they unarise to different unboxed
tuples. However if we run StgCse before Unarise, it'll think the two
terms (# 1 | #) are equal, and replace one of these with a binder to
the other. That's bad -- #15300.

Solution: do unarise first.

-}

module GHC.Stg.CSE (stgCse) where

import GHC.Prelude

import GHC.Core.DataCon
import GHC.Types.Id
import GHC.Stg.Syntax
import GHC.Types.Basic (isWeakLoopBreaker)
import GHC.Types.Var.Env
import GHC.Core (AltCon(..))
import Data.List (mapAccumL)
import Data.Maybe (fromMaybe)
import GHC.Core.Map.Expr
import GHC.Data.TrieMap
import GHC.Types.Name.Env
import Control.Monad( (>=>) )

--------------
-- The Trie --
--------------

-- A lookup trie for data constructor applications, i.e.
-- keys of type `(DataCon, [StgArg])`, following the patterns in GHC.Data.TrieMap.

data StgArgMap a = SAM
    { forall a. StgArgMap a -> DVarEnv a
sam_var :: DVarEnv a
    , forall a. StgArgMap a -> LiteralMap a
sam_lit :: LiteralMap a
    }

instance TrieMap StgArgMap where
    type Key StgArgMap = StgArg
    emptyTM :: forall a. StgArgMap a
emptyTM  = SAM { sam_var :: DVarEnv a
sam_var = forall (m :: * -> *) a. TrieMap m => m a
emptyTM
                   , sam_lit :: LiteralMap a
sam_lit = forall (m :: * -> *) a. TrieMap m => m a
emptyTM }
    lookupTM :: forall b. Key StgArgMap -> StgArgMap b -> Maybe b
lookupTM (StgVarArg Id
var) = forall a. StgArgMap a -> DVarEnv a
sam_var forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall a. Id -> DVarEnv a -> Maybe a
lkDFreeVar Id
var
    lookupTM (StgLitArg Literal
lit) = forall a. StgArgMap a -> LiteralMap a
sam_lit forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM Literal
lit
    alterTM :: forall b. Key StgArgMap -> XT b -> StgArgMap b -> StgArgMap b
alterTM  (StgVarArg Id
var) XT b
f StgArgMap b
m = StgArgMap b
m { sam_var :: DVarEnv b
sam_var = forall a. StgArgMap a -> DVarEnv a
sam_var StgArgMap b
m forall a b. a -> (a -> b) -> b
|> forall a. Id -> XT a -> DVarEnv a -> DVarEnv a
xtDFreeVar Id
var XT b
f }
    alterTM  (StgLitArg Literal
lit) XT b
f StgArgMap b
m = StgArgMap b
m { sam_lit :: LiteralMap b
sam_lit = forall a. StgArgMap a -> LiteralMap a
sam_lit StgArgMap b
m forall a b. a -> (a -> b) -> b
|> forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM Literal
lit XT b
f }
    foldTM :: forall a b. (a -> b -> b) -> StgArgMap a -> b -> b
foldTM a -> b -> b
k StgArgMap a
m = forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k (forall a. StgArgMap a -> DVarEnv a
sam_var StgArgMap a
m) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k (forall a. StgArgMap a -> LiteralMap a
sam_lit StgArgMap a
m)
    mapTM :: forall a b. (a -> b) -> StgArgMap a -> StgArgMap b
mapTM a -> b
f (SAM {sam_var :: forall a. StgArgMap a -> DVarEnv a
sam_var = DVarEnv a
varm, sam_lit :: forall a. StgArgMap a -> LiteralMap a
sam_lit = LiteralMap a
litm}) =
        SAM { sam_var :: DVarEnv b
sam_var = forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM a -> b
f DVarEnv a
varm, sam_lit :: LiteralMap b
sam_lit = forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM a -> b
f LiteralMap a
litm }
    filterTM :: forall a. (a -> Bool) -> StgArgMap a -> StgArgMap a
filterTM a -> Bool
f (SAM {sam_var :: forall a. StgArgMap a -> DVarEnv a
sam_var = DVarEnv a
varm, sam_lit :: forall a. StgArgMap a -> LiteralMap a
sam_lit = LiteralMap a
litm}) =
        SAM { sam_var :: DVarEnv a
sam_var = forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f DVarEnv a
varm, sam_lit :: LiteralMap a
sam_lit = forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f LiteralMap a
litm }

newtype ConAppMap a = CAM { forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam :: DNameEnv (ListMap StgArgMap a) }

instance TrieMap ConAppMap where
    type Key ConAppMap = (DataCon, [StgArg])
    emptyTM :: forall a. ConAppMap a
emptyTM  = forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    lookupTM :: forall b. Key ConAppMap -> ConAppMap b -> Maybe b
lookupTM (DataCon
dataCon, [StgArg]
args) = forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall n a. NamedThing n => n -> DNameEnv a -> Maybe a
lkDNamed DataCon
dataCon forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM [StgArg]
args
    alterTM :: forall b. Key ConAppMap -> XT b -> ConAppMap b -> ConAppMap b
alterTM  (DataCon
dataCon, [StgArg]
args) XT b
f ConAppMap b
m =
        ConAppMap b
m { un_cam :: DNameEnv (ListMap StgArgMap b)
un_cam = forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam ConAppMap b
m forall a b. a -> (a -> b) -> b
|> forall n a. NamedThing n => n -> XT a -> DNameEnv a -> DNameEnv a
xtDNamed DataCon
dataCon forall (m2 :: * -> *) a (m1 :: * -> *).
TrieMap m2 =>
(XT (m2 a) -> m1 (m2 a) -> m1 (m2 a))
-> (m2 a -> m2 a) -> m1 (m2 a) -> m1 (m2 a)
|>> forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM [StgArg]
args XT b
f }
    foldTM :: forall a b. (a -> b -> b) -> ConAppMap a -> b -> b
foldTM a -> b -> b
k = forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM (forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k)
    mapTM :: forall a b. (a -> b) -> ConAppMap a -> ConAppMap b
mapTM a -> b
f  = forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM (forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM a -> b
f) forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM
    filterTM :: forall a. (a -> Bool) -> ConAppMap a -> ConAppMap a
filterTM a -> Bool
f = forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM (forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f) forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM

-----------------
-- The CSE Env --
-----------------

-- | The CSE environment. See note [CseEnv Example]
data CseEnv = CseEnv
    { CseEnv -> ConAppMap Id
ce_conAppMap :: ConAppMap OutId
        -- ^ The main component of the environment is the trie that maps
        --   data constructor applications (with their `OutId` arguments)
        --   to an in-scope name that can be used instead.
        --   This name is always either a let-bound variable or a case binder.
    , CseEnv -> IdEnv Id
ce_subst     :: IdEnv OutId
        -- ^ This substitution is applied to the code as we traverse it.
        --   Entries have one of two reasons:
        --
        --   * The input might have shadowing (see Note [Shadowing]), so we have
        --     to rename some binders as we traverse the tree.
        --   * If we remove `let x = Con z` because  `let y = Con z` is in scope,
        --     we note this here as x ↦ y.
    , CseEnv -> IdEnv Id
ce_bndrMap     :: IdEnv OutId
        -- ^ If we come across a case expression case x as b of … with a trivial
        --   binder, we add b ↦ x to this.
        --   This map is *only* used when looking something up in the ce_conAppMap.
        --   See Note [Trivial case scrutinee]
    , CseEnv -> InScopeSet
ce_in_scope  :: InScopeSet
        -- ^ The third component is an in-scope set, to rename away any
        --   shadowing binders
    }

{-|
Note [CseEnv Example]
~~~~~~~~~~~~~~~~~~~~~
The following tables shows how the CseEnvironment changes as code is traversed,
as well as the changes to that code.

  InExpr                         OutExpr
     conAppMap                   subst          in_scope
  ───────────────────────────────────────────────────────────
  -- empty                       {}             {}
  case … as a of {Con x y ->     case … as a of {Con x y ->
  -- Con x y ↦ a                 {}             {a,x,y}
  let b = Con x y                (removed)
  -- Con x y ↦ a                 b↦a            {a,x,y,b}
  let c = Bar a                  let c = Bar a
  -- Con x y ↦ a, Bar a ↦ c      b↦a            {a,x,y,b,c}
  let c = some expression        let c' = some expression
  -- Con x y ↦ a, Bar a ↦ c      b↦a, c↦c',     {a,x,y,b,c,c'}
  let d = Bar b                  (removed)
  -- Con x y ↦ a, Bar a ↦ c      b↦a, c↦c', d↦c {a,x,y,b,c,c',d}
  (a, b, c d)                    (a, a, c' c)
-}

initEnv :: InScopeSet -> CseEnv
initEnv :: InScopeSet -> CseEnv
initEnv InScopeSet
in_scope = CseEnv
    { ce_conAppMap :: ConAppMap Id
ce_conAppMap = forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    , ce_subst :: IdEnv Id
ce_subst     = forall a. VarEnv a
emptyVarEnv
    , ce_bndrMap :: IdEnv Id
ce_bndrMap   = forall a. VarEnv a
emptyVarEnv
    , ce_in_scope :: InScopeSet
ce_in_scope  = InScopeSet
in_scope
    }

envLookup :: DataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
envLookup :: DataCon -> [StgArg] -> CseEnv -> Maybe Id
envLookup DataCon
dataCon [StgArg]
args CseEnv
env = forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM (DataCon
dataCon, [StgArg]
args') (CseEnv -> ConAppMap Id
ce_conAppMap CseEnv
env)
  where args' :: [StgArg]
args' = forall a b. (a -> b) -> [a] -> [b]
map StgArg -> StgArg
go [StgArg]
args -- See Note [Trivial case scrutinee]
        go :: StgArg -> StgArg
go (StgVarArg Id
v  ) = Id -> StgArg
StgVarArg (forall a. a -> Maybe a -> a
fromMaybe Id
v forall a b. (a -> b) -> a -> b
$ forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (CseEnv -> IdEnv Id
ce_bndrMap CseEnv
env) Id
v)
        go (StgLitArg Literal
lit) = Literal -> StgArg
StgLitArg Literal
lit

addDataCon :: OutId -> DataCon -> [OutStgArg] -> CseEnv -> CseEnv
-- do not bother with nullary data constructors, they are static anyways
addDataCon :: Id -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon Id
_ DataCon
_ [] CseEnv
env = CseEnv
env
addDataCon Id
bndr DataCon
dataCon [StgArg]
args CseEnv
env = CseEnv
env { ce_conAppMap :: ConAppMap Id
ce_conAppMap = ConAppMap Id
new_env }
  where
    new_env :: ConAppMap Id
new_env = forall (m :: * -> *) a. TrieMap m => Key m -> a -> m a -> m a
insertTM (DataCon
dataCon, [StgArg]
args) Id
bndr (CseEnv -> ConAppMap Id
ce_conAppMap CseEnv
env)

forgetCse :: CseEnv -> CseEnv
forgetCse :: CseEnv -> CseEnv
forgetCse CseEnv
env = CseEnv
env { ce_conAppMap :: ConAppMap Id
ce_conAppMap = forall (m :: * -> *) a. TrieMap m => m a
emptyTM }
    -- See note [Free variables of an StgClosure]

addSubst :: OutId -> OutId -> CseEnv -> CseEnv
addSubst :: Id -> Id -> CseEnv -> CseEnv
addSubst Id
from Id
to CseEnv
env
    = CseEnv
env { ce_subst :: IdEnv Id
ce_subst = forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv (CseEnv -> IdEnv Id
ce_subst CseEnv
env) Id
from Id
to }

addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
addTrivCaseBndr :: Id -> Id -> CseEnv -> CseEnv
addTrivCaseBndr Id
from Id
to CseEnv
env
    = CseEnv
env { ce_bndrMap :: IdEnv Id
ce_bndrMap = forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv (CseEnv -> IdEnv Id
ce_bndrMap CseEnv
env) Id
from Id
to }

substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
substArgs :: CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env = forall a b. (a -> b) -> [a] -> [b]
map (CseEnv -> StgArg -> StgArg
substArg CseEnv
env)

substArg :: CseEnv -> InStgArg -> OutStgArg
substArg :: CseEnv -> StgArg -> StgArg
substArg CseEnv
env (StgVarArg Id
from) = Id -> StgArg
StgVarArg (CseEnv -> Id -> Id
substVar CseEnv
env Id
from)
substArg CseEnv
_   (StgLitArg Literal
lit)  = Literal -> StgArg
StgLitArg Literal
lit

substVar :: CseEnv -> InId -> OutId
substVar :: CseEnv -> Id -> Id
substVar CseEnv
env Id
id = forall a. a -> Maybe a -> a
fromMaybe Id
id forall a b. (a -> b) -> a -> b
$ forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv (CseEnv -> IdEnv Id
ce_subst CseEnv
env) Id
id

-- Functions to enter binders

-- This is much simpler than the equivalent code in GHC.Core.Subst:
--  * We do not substitute type variables, and
--  * There is nothing relevant in GHC.Types.Id.Info at this stage
--    that needs substitutions.
-- Therefore, no special treatment for a recursive group is required.

substBndr :: CseEnv -> InId -> (CseEnv, OutId)
substBndr :: CseEnv -> Id -> (CseEnv, Id)
substBndr CseEnv
env Id
old_id
  = (CseEnv
new_env, Id
new_id)
  where
    new_id :: Id
new_id = InScopeSet -> Id -> Id
uniqAway (CseEnv -> InScopeSet
ce_in_scope CseEnv
env) Id
old_id
    no_change :: Bool
no_change = Id
new_id forall a. Eq a => a -> a -> Bool
== Id
old_id
    env' :: CseEnv
env' = CseEnv
env { ce_in_scope :: InScopeSet
ce_in_scope = CseEnv -> InScopeSet
ce_in_scope CseEnv
env InScopeSet -> Id -> InScopeSet
`extendInScopeSet` Id
new_id }
    new_env :: CseEnv
new_env | Bool
no_change = CseEnv
env'
            | Bool
otherwise = CseEnv
env' { ce_subst :: IdEnv Id
ce_subst = forall a. VarEnv a -> Id -> a -> VarEnv a
extendVarEnv (CseEnv -> IdEnv Id
ce_subst CseEnv
env) Id
old_id Id
new_id }

substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
substBndrs :: CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [Id]
bndrs = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL CseEnv -> Id -> (CseEnv, Id)
substBndr CseEnv
env [Id]
bndrs

substPairs :: CseEnv -> [(InVar, a)] -> (CseEnv, [(OutVar, a)])
substPairs :: forall a. CseEnv -> [(Id, a)] -> (CseEnv, [(Id, a)])
substPairs CseEnv
env [(Id, a)]
bndrs = forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL forall {b}. CseEnv -> (Id, b) -> (CseEnv, (Id, b))
go CseEnv
env [(Id, a)]
bndrs
  where go :: CseEnv -> (Id, b) -> (CseEnv, (Id, b))
go CseEnv
env (Id
id, b
x) = let (CseEnv
env', Id
id') = CseEnv -> Id -> (CseEnv, Id)
substBndr CseEnv
env Id
id
                         in (CseEnv
env', (Id
id', b
x))

-- Main entry point

stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
stgCse :: [InStgTopBinding] -> [InStgTopBinding]
stgCse [InStgTopBinding]
binds = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL InScopeSet -> InStgTopBinding -> (InScopeSet, InStgTopBinding)
stgCseTopLvl InScopeSet
emptyInScopeSet [InStgTopBinding]
binds

-- Top level bindings.
--
-- We do not CSE these, as top-level closures are allocated statically anyways.
-- Also, they might be exported.
-- But we still have to collect the set of in-scope variables, otherwise
-- uniqAway might shadow a top-level closure.

stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, OutStgTopBinding)
stgCseTopLvl :: InScopeSet -> InStgTopBinding -> (InScopeSet, InStgTopBinding)
stgCseTopLvl InScopeSet
in_scope t :: InStgTopBinding
t@(StgTopStringLit Id
_ ByteString
_) = (InScopeSet
in_scope, InStgTopBinding
t)
stgCseTopLvl InScopeSet
in_scope (StgTopLifted (StgNonRec BinderP 'Vanilla
bndr GenStgRhs 'Vanilla
rhs))
    = (InScopeSet
in_scope'
      , forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted (forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec BinderP 'Vanilla
bndr (InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope GenStgRhs 'Vanilla
rhs)))
  where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> Id -> InScopeSet
`extendInScopeSet` BinderP 'Vanilla
bndr

stgCseTopLvl InScopeSet
in_scope (StgTopLifted (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs))
    = ( InScopeSet
in_scope'
      , forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted (forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [ (Id
bndr, InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope' GenStgRhs 'Vanilla
rhs) | (Id
bndr, GenStgRhs 'Vanilla
rhs) <- [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs ]))
  where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> [Id] -> InScopeSet
`extendInScopeSetList` [ Id
bndr | (Id
bndr, GenStgRhs 'Vanilla
_) <- [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs ]

stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
stgCseTopLvlRhs :: InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope (StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [BinderP 'Vanilla]
args GenStgExpr 'Vanilla
body)
    = let body' :: GenStgExpr 'Vanilla
body' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr (InScopeSet -> CseEnv
initEnv InScopeSet
in_scope) GenStgExpr 'Vanilla
body
      in  forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> GenStgRhs pass
StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [BinderP 'Vanilla]
args GenStgExpr 'Vanilla
body'
stgCseTopLvlRhs InScopeSet
_ (StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args)
    = forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args

------------------------------
-- The actual AST traversal --
------------------------------

-- Trivial cases
stgCseExpr :: CseEnv -> InStgExpr -> OutStgExpr
stgCseExpr :: CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env (StgApp Id
fun [StgArg]
args)
    = forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
fun' [StgArg]
args'
  where fun' :: Id
fun' = CseEnv -> Id -> Id
substVar CseEnv
env Id
fun
        args' :: [StgArg]
args' = CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env [StgArg]
args
stgCseExpr CseEnv
_ (StgLit Literal
lit)
    = forall (pass :: StgPass). Literal -> GenStgExpr pass
StgLit Literal
lit
stgCseExpr CseEnv
env (StgOpApp StgOp
op [StgArg]
args Type
tys)
    = forall (pass :: StgPass).
StgOp -> [StgArg] -> Type -> GenStgExpr pass
StgOpApp StgOp
op [StgArg]
args' Type
tys
  where args' :: [StgArg]
args' = CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env [StgArg]
args
stgCseExpr CseEnv
env (StgTick StgTickish
tick GenStgExpr 'Vanilla
body)
    = let body' :: GenStgExpr 'Vanilla
body' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env GenStgExpr 'Vanilla
body
      in forall (pass :: StgPass).
StgTickish -> GenStgExpr pass -> GenStgExpr pass
StgTick StgTickish
tick GenStgExpr 'Vanilla
body'
stgCseExpr CseEnv
env (StgCase GenStgExpr 'Vanilla
scrut BinderP 'Vanilla
bndr AltType
ty [GenStgAlt 'Vanilla]
alts)
    = GenStgExpr 'Vanilla
-> Id -> AltType -> [GenStgAlt 'Vanilla] -> GenStgExpr 'Vanilla
mkStgCase GenStgExpr 'Vanilla
scrut' Id
bndr' AltType
ty [(AltCon, [Id], GenStgExpr 'Vanilla)]
alts'
  where
    scrut' :: GenStgExpr 'Vanilla
scrut' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env GenStgExpr 'Vanilla
scrut
    (CseEnv
env1, Id
bndr') = CseEnv -> Id -> (CseEnv, Id)
substBndr CseEnv
env BinderP 'Vanilla
bndr
    env2 :: CseEnv
env2 | StgApp Id
trivial_scrut [] <- GenStgExpr 'Vanilla
scrut' = Id -> Id -> CseEnv -> CseEnv
addTrivCaseBndr BinderP 'Vanilla
bndr Id
trivial_scrut CseEnv
env1
                 -- See Note [Trivial case scrutinee]
         | Bool
otherwise                         = CseEnv
env1
    alts' :: [(AltCon, [Id], GenStgExpr 'Vanilla)]
alts' = forall a b. (a -> b) -> [a] -> [b]
map (CseEnv -> AltType -> Id -> GenStgAlt 'Vanilla -> GenStgAlt 'Vanilla
stgCseAlt CseEnv
env2 AltType
ty Id
bndr') [GenStgAlt 'Vanilla]
alts


-- A constructor application.
-- To be removed by a variable use when found in the CSE environment
stgCseExpr CseEnv
env (StgConApp DataCon
dataCon XConApp 'Vanilla
n [StgArg]
args [Type]
tys)
    | Just Id
bndr' <- DataCon -> [StgArg] -> CseEnv -> Maybe Id
envLookup DataCon
dataCon [StgArg]
args' CseEnv
env
    = forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
bndr' []
    | Bool
otherwise
    = forall (pass :: StgPass).
DataCon -> XConApp pass -> [StgArg] -> [Type] -> GenStgExpr pass
StgConApp DataCon
dataCon XConApp 'Vanilla
n [StgArg]
args' [Type]
tys
  where args' :: [StgArg]
args' = CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env [StgArg]
args

-- Let bindings
-- The binding might be removed due to CSE (we do not want trivial bindings on
-- the STG level), so use the smart constructor `mkStgLet` to remove the binding
-- if empty.
stgCseExpr CseEnv
env (StgLet XLet 'Vanilla
ext GenStgBinding 'Vanilla
binds GenStgExpr 'Vanilla
body)
    = let (Maybe (GenStgBinding 'Vanilla)
binds', CseEnv
env') = CseEnv
-> GenStgBinding 'Vanilla
-> (Maybe (GenStgBinding 'Vanilla), CseEnv)
stgCseBind CseEnv
env GenStgBinding 'Vanilla
binds
          body' :: GenStgExpr 'Vanilla
body' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env' GenStgExpr 'Vanilla
body
      in forall a b. (a -> b -> b) -> Maybe a -> b -> b
mkStgLet (forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet 'Vanilla
ext) Maybe (GenStgBinding 'Vanilla)
binds' GenStgExpr 'Vanilla
body'
stgCseExpr CseEnv
env (StgLetNoEscape XLetNoEscape 'Vanilla
ext GenStgBinding 'Vanilla
binds GenStgExpr 'Vanilla
body)
    = let (Maybe (GenStgBinding 'Vanilla)
binds', CseEnv
env') = CseEnv
-> GenStgBinding 'Vanilla
-> (Maybe (GenStgBinding 'Vanilla), CseEnv)
stgCseBind CseEnv
env GenStgBinding 'Vanilla
binds
          body' :: GenStgExpr 'Vanilla
body' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env' GenStgExpr 'Vanilla
body
      in forall a b. (a -> b -> b) -> Maybe a -> b -> b
mkStgLet (forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape 'Vanilla
ext) Maybe (GenStgBinding 'Vanilla)
binds' GenStgExpr 'Vanilla
body'

-- Case alternatives
-- Extend the CSE environment
stgCseAlt :: CseEnv -> AltType -> OutId -> InStgAlt -> OutStgAlt
stgCseAlt :: CseEnv -> AltType -> Id -> GenStgAlt 'Vanilla -> GenStgAlt 'Vanilla
stgCseAlt CseEnv
env AltType
ty Id
case_bndr (DataAlt DataCon
dataCon, [BinderP 'Vanilla]
args, GenStgExpr 'Vanilla
rhs)
    = let (CseEnv
env1, [Id]
args') = CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [BinderP 'Vanilla]
args
          env2 :: CseEnv
env2
            -- To avoid dealing with unboxed sums StgCse runs after unarise and
            -- should maintain invariants listed in Note [Post-unarisation
            -- invariants]. One of the invariants is that some binders are not
            -- used (unboxed tuple case binders) which is what we check with
            -- `stgCaseBndrInScope` here. If the case binder is not in scope we
            -- don't add it to the CSE env. See also #15300.
            | AltType -> Bool -> Bool
stgCaseBndrInScope AltType
ty Bool
True -- CSE runs after unarise
            = Id -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon Id
case_bndr DataCon
dataCon (forall a b. (a -> b) -> [a] -> [b]
map Id -> StgArg
StgVarArg [Id]
args') CseEnv
env1
            | Bool
otherwise
            = CseEnv
env1
            -- see note [Case 2: CSEing case binders]
          rhs' :: GenStgExpr 'Vanilla
rhs' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env2 GenStgExpr 'Vanilla
rhs
      in (DataCon -> AltCon
DataAlt DataCon
dataCon, [Id]
args', GenStgExpr 'Vanilla
rhs')
stgCseAlt CseEnv
env AltType
_ Id
_ (AltCon
altCon, [BinderP 'Vanilla]
args, GenStgExpr 'Vanilla
rhs)
    = let (CseEnv
env1, [Id]
args') = CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [BinderP 'Vanilla]
args
          rhs' :: GenStgExpr 'Vanilla
rhs' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env1 GenStgExpr 'Vanilla
rhs
      in (AltCon
altCon, [Id]
args', GenStgExpr 'Vanilla
rhs')

-- Bindings
stgCseBind :: CseEnv -> InStgBinding -> (Maybe OutStgBinding, CseEnv)
stgCseBind :: CseEnv
-> GenStgBinding 'Vanilla
-> (Maybe (GenStgBinding 'Vanilla), CseEnv)
stgCseBind CseEnv
env (StgNonRec BinderP 'Vanilla
b GenStgRhs 'Vanilla
e)
    = let (CseEnv
env1, Id
b') = CseEnv -> Id -> (CseEnv, Id)
substBndr CseEnv
env BinderP 'Vanilla
b
      in case CseEnv
-> Id
-> GenStgRhs 'Vanilla
-> (Maybe (Id, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs CseEnv
env1 Id
b' GenStgRhs 'Vanilla
e of
        (Maybe (Id, GenStgRhs 'Vanilla)
Nothing,      CseEnv
env2) -> (forall a. Maybe a
Nothing,                CseEnv
env2)
        (Just (Id
b2,GenStgRhs 'Vanilla
e'), CseEnv
env2) -> (forall a. a -> Maybe a
Just (forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec Id
b2 GenStgRhs 'Vanilla
e'), CseEnv
env2)
stgCseBind CseEnv
env (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
pairs)
    = let (CseEnv
env1, [(Id, GenStgRhs 'Vanilla)]
pairs1) = forall a. CseEnv -> [(Id, a)] -> (CseEnv, [(Id, a)])
substPairs CseEnv
env [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
pairs
      in case CseEnv
-> [(Id, GenStgRhs 'Vanilla)]
-> ([(Id, GenStgRhs 'Vanilla)], CseEnv)
stgCsePairs CseEnv
env1 [(Id, GenStgRhs 'Vanilla)]
pairs1 of
        ([],     CseEnv
env2) -> (forall a. Maybe a
Nothing, CseEnv
env2)
        ([(Id, GenStgRhs 'Vanilla)]
pairs2, CseEnv
env2) -> (forall a. a -> Maybe a
Just (forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [(Id, GenStgRhs 'Vanilla)]
pairs2), CseEnv
env2)

stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
stgCsePairs :: CseEnv
-> [(Id, GenStgRhs 'Vanilla)]
-> ([(Id, GenStgRhs 'Vanilla)], CseEnv)
stgCsePairs CseEnv
env [] = ([], CseEnv
env)
stgCsePairs CseEnv
env0 ((Id
b,GenStgRhs 'Vanilla
e):[(Id, GenStgRhs 'Vanilla)]
pairs)
  = let (Maybe (Id, GenStgRhs 'Vanilla)
pairMB, CseEnv
env1) = CseEnv
-> Id
-> GenStgRhs 'Vanilla
-> (Maybe (Id, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs CseEnv
env0 Id
b GenStgRhs 'Vanilla
e
        ([(Id, GenStgRhs 'Vanilla)]
pairs', CseEnv
env2) = CseEnv
-> [(Id, GenStgRhs 'Vanilla)]
-> ([(Id, GenStgRhs 'Vanilla)], CseEnv)
stgCsePairs CseEnv
env1 [(Id, GenStgRhs 'Vanilla)]
pairs
    in (Maybe (Id, GenStgRhs 'Vanilla)
pairMB forall {a}. Maybe a -> [a] -> [a]
`mbCons` [(Id, GenStgRhs 'Vanilla)]
pairs', CseEnv
env2)
  where
    mbCons :: Maybe a -> [a] -> [a]
mbCons = forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id (:)

-- The RHS of a binding.
-- If it is a constructor application, either short-cut it or extend the environment
stgCseRhs :: CseEnv -> OutId -> InStgRhs -> (Maybe (OutId, OutStgRhs), CseEnv)
stgCseRhs :: CseEnv
-> Id
-> GenStgRhs 'Vanilla
-> (Maybe (Id, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs CseEnv
env Id
bndr (StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args)
    | Just Id
other_bndr <- DataCon -> [StgArg] -> CseEnv -> Maybe Id
envLookup DataCon
dataCon [StgArg]
args' CseEnv
env
    , Bool -> Bool
not (OccInfo -> Bool
isWeakLoopBreaker (Id -> OccInfo
idOccInfo Id
bndr)) -- See Note [Care with loop breakers]
    = let env' :: CseEnv
env' = Id -> Id -> CseEnv -> CseEnv
addSubst Id
bndr Id
other_bndr CseEnv
env
      in (forall a. Maybe a
Nothing, CseEnv
env')
    | Bool
otherwise
    = let env' :: CseEnv
env' = Id -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon Id
bndr DataCon
dataCon [StgArg]
args' CseEnv
env
            -- see note [Case 1: CSEing allocated closures]
          pair :: (Id, GenStgRhs 'Vanilla)
pair = (Id
bndr, forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args')
      in (forall a. a -> Maybe a
Just (Id, GenStgRhs 'Vanilla)
pair, CseEnv
env')
  where args' :: [StgArg]
args' = CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env [StgArg]
args

stgCseRhs CseEnv
env Id
bndr (StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [BinderP 'Vanilla]
args GenStgExpr 'Vanilla
body)
    = let (CseEnv
env1, [Id]
args') = CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [BinderP 'Vanilla]
args
          env2 :: CseEnv
env2 = CseEnv -> CseEnv
forgetCse CseEnv
env1 -- See note [Free variables of an StgClosure]
          body' :: GenStgExpr 'Vanilla
body' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env2 GenStgExpr 'Vanilla
body
      in (forall a. a -> Maybe a
Just (CseEnv -> Id -> Id
substVar CseEnv
env Id
bndr, forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> GenStgRhs pass
StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [Id]
args' GenStgExpr 'Vanilla
body'), CseEnv
env)


mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
mkStgCase :: GenStgExpr 'Vanilla
-> Id -> AltType -> [GenStgAlt 'Vanilla] -> GenStgExpr 'Vanilla
mkStgCase GenStgExpr 'Vanilla
scrut Id
bndr AltType
ty [GenStgAlt 'Vanilla]
alts | forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (AltCon, [Id], GenStgExpr 'Vanilla) -> Bool
isBndr [GenStgAlt 'Vanilla]
alts = GenStgExpr 'Vanilla
scrut
                             | Bool
otherwise       = forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'Vanilla
scrut Id
bndr AltType
ty [GenStgAlt 'Vanilla]
alts

  where
    -- see Note [All alternatives are the binder]
    isBndr :: (AltCon, [Id], GenStgExpr 'Vanilla) -> Bool
isBndr (AltCon
_, [Id]
_, StgApp Id
f []) = Id
f forall a. Eq a => a -> a -> Bool
== Id
bndr
    isBndr (AltCon, [Id], GenStgExpr 'Vanilla)
_                   = Bool
False


{- Note [Care with loop breakers]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When doing CSE on a letrec we must be careful about loop
breakers.  Consider
  rec { y = K z
      ; z = K z }
Now if, somehow (and wrongly)), y and z are both marked as
loop-breakers, we do *not* want to drop the (z = K z) binding
in favour of a substitution (z :-> y).

I think this bug will only show up if the loop-breaker-ness is done
wrongly (itself a bug), but it still seems better to do the right
thing regardless.
-}

-- Utilities

-- | This function short-cuts let-bindings that are now obsolete
mkStgLet :: (a -> b -> b) -> Maybe a -> b -> b
mkStgLet :: forall a b. (a -> b -> b) -> Maybe a -> b -> b
mkStgLet a -> b -> b
_      Maybe a
Nothing      b
body = b
body
mkStgLet a -> b -> b
stgLet (Just a
binds) b
body = a -> b -> b
stgLet a
binds b
body


{-
Note [All alternatives are the binder]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

When all alternatives simply refer to the case binder, then we do not have
to bother with the case expression at all (#13588). CoreSTG does this as well,
but sometimes, types get into the way:

    newtype T = MkT Int
    f :: (Int, Int) -> (T, Int)
    f (x, y) = (MkT x, y)

Core cannot just turn this into

    f p = p

as this would not be well-typed. But to STG, where MkT is no longer in the way,
we can.

Note [Trivial case scrutinee]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to be able to handle nested reconstruction of constructors as in

    nested :: Either Int (Either Int a) -> Either Bool (Either Bool a)
    nested (Right (Right v)) = Right (Right v)
    nested _ = Left True

So if we come across

    case x of r1
      Right a -> case a of r2
              Right b -> let v = Right b
                         in Right v

we first replace v with r2. Next we want to replace Right r2 with r1. But the
ce_conAppMap contains Right a!

Therefore, we add r1 ↦ x to ce_bndrMap when analysing the outer case, and use
this substitution before looking Right r2 up in ce_conAppMap, and everything
works out.

Note [Free variables of an StgClosure]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
StgClosures (function and thunks) have an explicit list of free variables:

foo [x] =
    let not_a_free_var = Left [x]
    let a_free_var = Right [x]
    let closure = \[x a_free_var] -> \[y] -> bar y (Left [x]) a_free_var
    in closure

If we were to CSE `Left [x]` in the body of `closure` with `not_a_free_var`,
then the list of free variables would be wrong, so for now, we do not CSE
across such a closure, simply because I (Joachim) was not sure about possible
knock-on effects. If deemed safe and worth the slight code complication of
re-calculating this list during or after this pass, this can surely be done.
-}