{-# 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
    }

-- TODO(22292): derive
instance Functor StgArgMap where
    fmap :: forall a b. (a -> b) -> StgArgMap a -> StgArgMap b
fmap 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 = (a -> b) -> DVarEnv a -> DVarEnv b
forall a b. (a -> b) -> UniqDFM Id a -> UniqDFM Id b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f DVarEnv a
varm, sam_lit :: LiteralMap b
sam_lit = (a -> b) -> LiteralMap a -> LiteralMap b
forall a b. (a -> b) -> Map Literal a -> Map Literal b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f LiteralMap a
litm }

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

-- TODO(22292): derive
instance Functor ConAppMap where
    fmap :: forall a b. (a -> b) -> ConAppMap a -> ConAppMap b
fmap a -> b
f = DNameEnv (ListMap StgArgMap b) -> ConAppMap b
forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM (DNameEnv (ListMap StgArgMap b) -> ConAppMap b)
-> (ConAppMap a -> DNameEnv (ListMap StgArgMap b))
-> ConAppMap a
-> ConAppMap b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ListMap StgArgMap a -> ListMap StgArgMap b)
-> UniqDFM Name (ListMap StgArgMap a)
-> DNameEnv (ListMap StgArgMap b)
forall a b. (a -> b) -> UniqDFM Name a -> UniqDFM Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> ListMap StgArgMap a -> ListMap StgArgMap b
forall a b. (a -> b) -> ListMap StgArgMap a -> ListMap StgArgMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) (UniqDFM Name (ListMap StgArgMap a)
 -> DNameEnv (ListMap StgArgMap b))
-> (ConAppMap a -> UniqDFM Name (ListMap StgArgMap a))
-> ConAppMap a
-> DNameEnv (ListMap StgArgMap b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConAppMap a -> UniqDFM Name (ListMap StgArgMap a)
forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam
    {-# INLINE fmap #-}

instance TrieMap ConAppMap where
    type Key ConAppMap = (DataCon, [StgArg])
    emptyTM :: forall a. ConAppMap a
emptyTM  = DNameEnv (ListMap StgArgMap a) -> ConAppMap a
forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM DNameEnv (ListMap StgArgMap a)
forall a. UniqDFM Name a
forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    lookupTM :: forall b. Key ConAppMap -> ConAppMap b -> Maybe b
lookupTM (DataCon
dataCon, [StgArg]
args) = ConAppMap b -> DNameEnv (ListMap StgArgMap b)
forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam (ConAppMap b -> DNameEnv (ListMap StgArgMap b))
-> (DNameEnv (ListMap StgArgMap b) -> Maybe b)
-> ConAppMap b
-> Maybe b
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> DataCon
-> DNameEnv (ListMap StgArgMap b) -> Maybe (ListMap StgArgMap b)
forall n a. NamedThing n => n -> DNameEnv a -> Maybe a
lkDNamed DataCon
dataCon (DNameEnv (ListMap StgArgMap b) -> Maybe (ListMap StgArgMap b))
-> (ListMap StgArgMap b -> Maybe b)
-> DNameEnv (ListMap StgArgMap b)
-> Maybe b
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Key (ListMap StgArgMap) -> ListMap StgArgMap b -> Maybe b
forall b. Key (ListMap StgArgMap) -> ListMap StgArgMap b -> Maybe b
forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM [StgArg]
Key (ListMap StgArgMap)
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 = un_cam m |> xtDNamed dataCon |>> alterTM args f }
    foldTM :: forall a b. (a -> b -> b) -> ConAppMap a -> b -> b
foldTM a -> b -> b
k = ConAppMap a -> DNameEnv (ListMap StgArgMap a)
forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam (ConAppMap a -> DNameEnv (ListMap StgArgMap a))
-> (DNameEnv (ListMap StgArgMap a) -> b -> b)
-> ConAppMap a
-> b
-> b
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> (ListMap StgArgMap a -> b -> b)
-> DNameEnv (ListMap StgArgMap a) -> b -> b
forall a b. (a -> b -> b) -> UniqDFM Name a -> b -> b
forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM ((a -> b -> b) -> ListMap StgArgMap a -> b -> b
forall a b. (a -> b -> b) -> ListMap StgArgMap a -> b -> b
forall (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM a -> b -> b
k)
    filterTM :: forall a. (a -> Bool) -> ConAppMap a -> ConAppMap a
filterTM a -> Bool
f = ConAppMap a -> DNameEnv (ListMap StgArgMap a)
forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam (ConAppMap a -> DNameEnv (ListMap StgArgMap a))
-> (DNameEnv (ListMap StgArgMap a) -> ConAppMap a)
-> ConAppMap a
-> ConAppMap a
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> (ListMap StgArgMap a -> ListMap StgArgMap a)
-> DNameEnv (ListMap StgArgMap a) -> DNameEnv (ListMap StgArgMap a)
forall a b. (a -> b) -> UniqDFM Name a -> UniqDFM Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> Bool) -> ListMap StgArgMap a -> ListMap StgArgMap a
forall a. (a -> Bool) -> ListMap StgArgMap a -> ListMap StgArgMap a
forall (m :: * -> *) a. TrieMap m => (a -> Bool) -> m a -> m a
filterTM a -> Bool
f) (DNameEnv (ListMap StgArgMap a) -> DNameEnv (ListMap StgArgMap a))
-> (DNameEnv (ListMap StgArgMap a) -> ConAppMap a)
-> DNameEnv (ListMap StgArgMap a)
-> ConAppMap a
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> DNameEnv (ListMap StgArgMap a) -> ConAppMap a
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 = ConAppMap Id
forall a. ConAppMap a
forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    , ce_subst :: IdEnv Id
ce_subst     = IdEnv Id
forall a. VarEnv a
emptyVarEnv
    , ce_bndrMap :: IdEnv Id
ce_bndrMap   = IdEnv Id
forall a. VarEnv a
emptyVarEnv
    , ce_in_scope :: InScopeSet
ce_in_scope  = InScopeSet
in_scope
    }

-------------------
normaliseConArgs :: CseEnv -> [OutStgArg] -> [OutStgArg]
-- See Note [Trivial case scrutinee]
normaliseConArgs :: CseEnv -> [StgArg] -> [StgArg]
normaliseConArgs CseEnv
env [StgArg]
args
  = (StgArg -> StgArg) -> [StgArg] -> [StgArg]
forall a b. (a -> b) -> [a] -> [b]
map StgArg -> StgArg
go [StgArg]
args
  where
    bndr_map :: IdEnv Id
bndr_map = CseEnv -> IdEnv Id
ce_bndrMap CseEnv
env
    go :: StgArg -> StgArg
go (StgVarArg Id
v  ) = Id -> StgArg
StgVarArg (IdEnv Id -> Id -> Id
normaliseId IdEnv Id
bndr_map Id
v)
    go (StgLitArg Literal
lit) = Literal -> StgArg
StgLitArg Literal
lit

normaliseId :: IdEnv OutId -> OutId -> OutId
normaliseId :: IdEnv Id -> Id -> Id
normaliseId IdEnv Id
bndr_map Id
v = case IdEnv Id -> Id -> Maybe Id
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv IdEnv Id
bndr_map Id
v of
                           Just Id
v' -> Id
v'
                           Maybe Id
Nothing -> Id
v

addTrivCaseBndr :: OutId -> OutId -> CseEnv -> CseEnv
-- See Note [Trivial case scrutinee]
addTrivCaseBndr :: Id -> Id -> CseEnv -> CseEnv
addTrivCaseBndr Id
from Id
to CseEnv
env
    = CseEnv
env { ce_bndrMap = extendVarEnv bndr_map from norm_to }
    where
      bndr_map :: IdEnv Id
bndr_map = CseEnv -> IdEnv Id
ce_bndrMap CseEnv
env
      norm_to :: Id
norm_to = IdEnv Id -> Id -> Id
normaliseId IdEnv Id
bndr_map Id
to

envLookup :: DataCon -> [OutStgArg] -> CseEnv -> Maybe OutId
envLookup :: DataCon -> [StgArg] -> CseEnv -> Maybe Id
envLookup DataCon
dataCon [StgArg]
args CseEnv
env
  = Key ConAppMap -> ConAppMap Id -> Maybe Id
forall b. Key ConAppMap -> ConAppMap b -> Maybe b
forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM (DataCon
dataCon, CseEnv -> [StgArg] -> [StgArg]
normaliseConArgs CseEnv
env [StgArg]
args)
             (CseEnv -> ConAppMap Id
ce_conAppMap CseEnv
env)
    -- normaliseConArgs: See Note [Trivial case scrutinee]

addDataCon :: OutId -> DataCon -> [OutStgArg] -> CseEnv -> CseEnv
-- Do not bother with nullary data constructors; they are static anyway
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 = new_env }
  where
    new_env :: ConAppMap Id
new_env = Key ConAppMap -> Id -> ConAppMap Id -> ConAppMap Id
forall (m :: * -> *) a. TrieMap m => Key m -> a -> m a -> m a
insertTM (DataCon
dataCon, CseEnv -> [StgArg] -> [StgArg]
normaliseConArgs CseEnv
env [StgArg]
args)
                       Id
bndr (CseEnv -> ConAppMap Id
ce_conAppMap CseEnv
env)
    -- normaliseConArgs: See Note [Trivial case scrutinee]

-------------------
forgetCse :: CseEnv -> CseEnv
forgetCse :: CseEnv -> CseEnv
forgetCse CseEnv
env = CseEnv
env { ce_conAppMap = 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 = extendVarEnv (ce_subst env) from to }

substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
substArgs :: CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env = (StgArg -> StgArg) -> [StgArg] -> [StgArg]
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 = Id -> Maybe Id -> Id
forall a. a -> Maybe a -> a
fromMaybe Id
id (Maybe Id -> Id) -> Maybe Id -> Id
forall a b. (a -> b) -> a -> b
$ IdEnv Id -> Id -> Maybe Id
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 Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
old_id
    env' :: CseEnv
env' = CseEnv
env { ce_in_scope = ce_in_scope env `extendInScopeSet` new_id }
    new_env :: CseEnv
new_env | Bool
no_change = CseEnv
env'
            | Bool
otherwise = CseEnv
env' { ce_subst = extendVarEnv (ce_subst env) old_id new_id }

substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
substBndrs :: CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [Id]
bndrs = (CseEnv -> Id -> (CseEnv, Id)) -> CseEnv -> [Id] -> (CseEnv, [Id])
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 = (CseEnv -> (Id, a) -> (CseEnv, (Id, a)))
-> CseEnv -> [(Id, a)] -> (CseEnv, [(Id, a)])
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL CseEnv -> (Id, a) -> (CseEnv, (Id, a))
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 = (InScopeSet, [InStgTopBinding]) -> [InStgTopBinding]
forall a b. (a, b) -> b
snd ((InScopeSet, [InStgTopBinding]) -> [InStgTopBinding])
-> (InScopeSet, [InStgTopBinding]) -> [InStgTopBinding]
forall a b. (a -> b) -> a -> b
$ (InScopeSet -> InStgTopBinding -> (InScopeSet, InStgTopBinding))
-> InScopeSet
-> [InStgTopBinding]
-> (InScopeSet, [InStgTopBinding])
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'
      , GenStgBinding 'Vanilla -> InStgTopBinding
forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted (BinderP 'Vanilla -> GenStgRhs 'Vanilla -> GenStgBinding 'Vanilla
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` Id
BinderP 'Vanilla
bndr

stgCseTopLvl InScopeSet
in_scope (StgTopLifted (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs))
    = ( InScopeSet
in_scope'
      , GenStgBinding 'Vanilla -> InStgTopBinding
forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted ([(BinderP 'Vanilla, GenStgRhs 'Vanilla)] -> GenStgBinding 'Vanilla
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [ (Id
BinderP 'Vanilla
bndr, InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope' GenStgRhs 'Vanilla
rhs) | (Id
bndr, GenStgRhs 'Vanilla
rhs) <- [(Id, GenStgRhs 'Vanilla)]
[(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs ]))
  where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> [Id] -> InScopeSet
`extendInScopeSetList` [ Id
bndr | (Id
bndr, GenStgRhs 'Vanilla
_) <- [(Id, 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  XRhsClosure 'Vanilla
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'Vanilla]
-> GenStgExpr 'Vanilla
-> GenStgRhs 'Vanilla
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)
    = CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs 'Vanilla
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)
    = Id -> [StgArg] -> GenStgExpr 'Vanilla
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)
    = Literal -> GenStgExpr 'Vanilla
forall (pass :: StgPass). Literal -> GenStgExpr pass
StgLit Literal
lit
stgCseExpr CseEnv
env (StgOpApp StgOp
op [StgArg]
args Type
tys)
    = StgOp -> [StgArg] -> Type -> GenStgExpr 'Vanilla
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 StgTickish -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
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 [GenStgAlt '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 Id
BinderP 'Vanilla
bndr
    env2 :: CseEnv
env2 | StgApp Id
trivial_scrut [] <- GenStgExpr 'Vanilla
scrut'
         = Id -> Id -> CseEnv -> CseEnv
addTrivCaseBndr Id
BinderP 'Vanilla
bndr Id
trivial_scrut CseEnv
env1
                 -- See Note [Trivial case scrutinee]
         | Bool
otherwise
         = CseEnv
env1
    alts' :: [GenStgAlt 'Vanilla]
alts' = (GenStgAlt 'Vanilla -> GenStgAlt 'Vanilla)
-> [GenStgAlt 'Vanilla] -> [GenStgAlt 'Vanilla]
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 ConstructorNumber
n [StgArg]
args [Type]
tys)
    | Just Id
bndr' <- DataCon -> [StgArg] -> CseEnv -> Maybe Id
envLookup DataCon
dataCon [StgArg]
args' CseEnv
env
    = Id -> [StgArg] -> GenStgExpr 'Vanilla
forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
bndr' []
    | Bool
otherwise
    = DataCon
-> ConstructorNumber -> [StgArg] -> [Type] -> GenStgExpr 'Vanilla
forall (pass :: StgPass).
DataCon
-> ConstructorNumber -> [StgArg] -> [Type] -> GenStgExpr pass
StgConApp DataCon
dataCon ConstructorNumber
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 (GenStgBinding 'Vanilla
 -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla)
-> Maybe (GenStgBinding 'Vanilla)
-> GenStgExpr 'Vanilla
-> GenStgExpr 'Vanilla
forall a b. (a -> b -> b) -> Maybe a -> b -> b
mkStgLet (XLet 'Vanilla
-> GenStgBinding 'Vanilla
-> GenStgExpr 'Vanilla
-> GenStgExpr 'Vanilla
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 (GenStgBinding 'Vanilla
 -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla)
-> Maybe (GenStgBinding 'Vanilla)
-> GenStgExpr 'Vanilla
-> GenStgExpr 'Vanilla
forall a b. (a -> b -> b) -> Maybe a -> b -> b
mkStgLet (XLetNoEscape 'Vanilla
-> GenStgBinding 'Vanilla
-> GenStgExpr 'Vanilla
-> GenStgExpr 'Vanilla
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 GenStgAlt{alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con=DataAlt DataCon
dataCon, alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs=[BinderP 'Vanilla]
args, alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs=GenStgExpr 'Vanilla
rhs}
    = let (CseEnv
env1, [Id]
args') = CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [Id]
[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 ((Id -> StgArg) -> [Id] -> [StgArg]
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 AltCon
-> [BinderP 'Vanilla] -> GenStgExpr 'Vanilla -> GenStgAlt 'Vanilla
forall (pass :: StgPass).
AltCon -> [BinderP pass] -> GenStgExpr pass -> GenStgAlt pass
GenStgAlt (DataCon -> AltCon
DataAlt DataCon
dataCon) [Id]
[BinderP 'Vanilla]
args' GenStgExpr 'Vanilla
rhs'
stgCseAlt CseEnv
env AltType
_ Id
_ g :: GenStgAlt 'Vanilla
g@GenStgAlt{alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con=AltCon
_, alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs=[BinderP 'Vanilla]
args, alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs=GenStgExpr 'Vanilla
rhs}
    = let (CseEnv
env1, [Id]
args') = CseEnv -> [Id] -> (CseEnv, [Id])
substBndrs CseEnv
env [Id]
[BinderP 'Vanilla]
args
          rhs' :: GenStgExpr 'Vanilla
rhs' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env1 GenStgExpr 'Vanilla
rhs
      in GenStgAlt 'Vanilla
g {alt_bndrs=args', alt_rhs=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 Id
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) -> (Maybe (GenStgBinding 'Vanilla)
forall a. Maybe a
Nothing,                CseEnv
env2)
        (Just (Id
b2,GenStgRhs 'Vanilla
e'), CseEnv
env2) -> (GenStgBinding 'Vanilla -> Maybe (GenStgBinding 'Vanilla)
forall a. a -> Maybe a
Just (BinderP 'Vanilla -> GenStgRhs 'Vanilla -> GenStgBinding 'Vanilla
forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec Id
BinderP 'Vanilla
b2 GenStgRhs 'Vanilla
e'), CseEnv
env2)
stgCseBind CseEnv
env (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
pairs)
    = let (CseEnv
env1, [(Id, GenStgRhs 'Vanilla)]
pairs1) = CseEnv
-> [(Id, GenStgRhs 'Vanilla)]
-> (CseEnv, [(Id, GenStgRhs 'Vanilla)])
forall a. CseEnv -> [(Id, a)] -> (CseEnv, [(Id, a)])
substPairs CseEnv
env [(Id, GenStgRhs 'Vanilla)]
[(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) -> (Maybe (GenStgBinding 'Vanilla)
forall a. Maybe a
Nothing, CseEnv
env2)
        ([(Id, GenStgRhs 'Vanilla)]
pairs2, CseEnv
env2) -> (GenStgBinding 'Vanilla -> Maybe (GenStgBinding 'Vanilla)
forall a. a -> Maybe a
Just ([(BinderP 'Vanilla, GenStgRhs 'Vanilla)] -> GenStgBinding 'Vanilla
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [(Id, GenStgRhs 'Vanilla)]
[(BinderP 'Vanilla, 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 Maybe (Id, GenStgRhs 'Vanilla)
-> [(Id, GenStgRhs 'Vanilla)] -> [(Id, GenStgRhs 'Vanilla)]
forall {a}. Maybe a -> [a] -> [a]
`mbCons` [(Id, GenStgRhs 'Vanilla)]
pairs', CseEnv
env2)
  where
    mbCons :: Maybe a -> [a] -> [a]
mbCons = ([a] -> [a]) -> (a -> [a] -> [a]) -> Maybe a -> [a] -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [a] -> [a]
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 (Maybe (Id, GenStgRhs 'Vanilla)
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, CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs 'Vanilla
forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon ConstructorNumber
mu [StgTickish]
ticks [StgArg]
args')
      in ((Id, GenStgRhs 'Vanilla) -> Maybe (Id, GenStgRhs 'Vanilla)
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 [Id]
[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 ((Id, GenStgRhs 'Vanilla) -> Maybe (Id, GenStgRhs 'Vanilla)
forall a. a -> Maybe a
Just (CseEnv -> Id -> Id
substVar CseEnv
env Id
bndr, XRhsClosure 'Vanilla
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'Vanilla]
-> GenStgExpr 'Vanilla
-> GenStgRhs 'Vanilla
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> GenStgRhs pass
StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [Id]
[BinderP 'Vanilla]
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 | (GenStgAlt 'Vanilla -> Bool) -> [GenStgAlt 'Vanilla] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all GenStgAlt 'Vanilla -> Bool
isBndr [GenStgAlt 'Vanilla]
alts = GenStgExpr 'Vanilla
scrut
                             | Bool
otherwise       = GenStgExpr 'Vanilla
-> BinderP 'Vanilla
-> AltType
-> [GenStgAlt 'Vanilla]
-> GenStgExpr 'Vanilla
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase GenStgExpr 'Vanilla
scrut Id
BinderP 'Vanilla
bndr AltType
ty [GenStgAlt 'Vanilla]
alts

  where
    -- see Note [All alternatives are the binder]
    isBndr :: GenStgAlt 'Vanilla -> Bool
isBndr GenStgAlt{alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con=AltCon
_,alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs=[BinderP 'Vanilla]
_,alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs=StgApp Id
f []} = Id
f Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
bndr
    isBndr GenStgAlt '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 CSE 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

We want the RHS of the first branch to be just the original argument.
The RHS of 'nested' will look like
    case x of r1
      Right a -> case a of r2
              Right b -> let v = Right b
                         in Right v
Then:
* We create the ce_conAppMap [Right a :-> r1, Right b :-> r2].
* When we encounter v = Right b, we'll drop the binding and extend
  the substitution with [v :-> r2]
* But now when we see (Right v), we'll substitute to get (Right r2)...and
  fail to find that in the ce_conAppMap!

Solution:

* When passing (case x of bndr { alts }), where 'x' is a variable, we
  add [bndr :-> x] to the ce_bndrMap.  In our example the ce_bndrMap will
  be [r1 :-> x, r2 :-> a]. This is done in addTrivCaseBndr.

* Before doing the /lookup/ in ce_conAppMap, we "normalise" the
  arguments with the ce_bndrMap.  In our example, we normalise
  (Right r2) to (Right a), and then find it in the map.  Normalisation
  is done by normaliseConArgs.

* Similarly before /inserting/ in ce_conAppMap, we normalise the arguments.
  This is a bit more subtle. Suppose we have
       case x of y
         DEFAULT -> let a = Just y
                    let b = Just y
                    in ...
  We'll have [y :-> x] in the ce_bndrMap.  When looking up (Just y) in
  the map, we'll normalise it to (Just x).  So we'd better normalise
  the (Just y) in the defn of 'a', before inserting it!

* When inserting into cs_bndrMap, we must normalise that too!
      case x of y
        DEFAULT -> case y of z
                      DEFAULT -> ...
  We want the cs_bndrMap to be [y :-> x, z :-> x]!
  Hence the call to normaliseId in addTrivCaseBinder.

All this is a bit tricky.  Why does it not occur for the Core version
of CSE?  See Note [CSE for bindings] in GHC.Core.Opt.CSE.  The reason
is this: in Core CSE we augment the /main substitution/ with [y :-> x]
etc, so as a side consequence we transform
    case x of y       ===>    case x of y
      pat -> ...y...             pat -> ...x...
That is, the /exact reverse/ of the binder-swap transformation done by
the occurrence analyser.  However, it's easy for CSE to do on-the-fly,
and it completely solves the above tricky problem, using only two maps:
the main reverse-map, and the substitution.  The occurrence analyser
puts it back the way it should be, the next time it runs.

However in STG there is no occurrence analyser, and we don't want to
require another pass.  So the ce_bndrMap is a little swizzle that we
apply just when manipulating the ce_conAppMap, but that does not
affect the output program.


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.
-}