{-# 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 -- Trac #15300.

Solution: do unarise first.

-}

module StgCse (stgCse) where

import GhcPrelude

import DataCon
import Id
import StgSyn
import Outputable
import VarEnv
import CoreSyn (AltCon(..))
import Data.List (mapAccumL)
import Data.Maybe (fromMaybe)
import CoreMap
import NameEnv
import Control.Monad( (>=>) )

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

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

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

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

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

instance TrieMap ConAppMap where
    type Key ConAppMap = (DataCon, [StgArg])
    emptyTM :: ConAppMap a
emptyTM  = DNameEnv (ListMap StgArgMap a) -> ConAppMap a
forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM DNameEnv (ListMap StgArgMap a)
forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    lookupTM :: Key ConAppMap -> ConAppMap b -> Maybe b
lookupTM (dataCon, 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 (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM [StgArg]
Key (ListMap StgArgMap)
args
    alterTM :: Key ConAppMap -> XT b -> ConAppMap b -> ConAppMap b
alterTM  (dataCon, args) f :: XT b
f m :: ConAppMap b
m =
        ConAppMap b
m { un_cam :: DNameEnv (ListMap StgArgMap b)
un_cam = ConAppMap b -> DNameEnv (ListMap StgArgMap b)
forall a. ConAppMap a -> DNameEnv (ListMap StgArgMap a)
un_cam ConAppMap b
m DNameEnv (ListMap StgArgMap b)
-> (DNameEnv (ListMap StgArgMap b)
    -> DNameEnv (ListMap StgArgMap b))
-> DNameEnv (ListMap StgArgMap b)
forall a b. a -> (a -> b) -> b
|> DataCon
-> XT (ListMap StgArgMap b)
-> DNameEnv (ListMap StgArgMap b)
-> DNameEnv (ListMap StgArgMap b)
forall n a. NamedThing n => n -> XT a -> DNameEnv a -> DNameEnv a
xtDNamed DataCon
dataCon (XT (ListMap StgArgMap b)
 -> DNameEnv (ListMap StgArgMap b)
 -> DNameEnv (ListMap StgArgMap b))
-> (ListMap StgArgMap b -> ListMap StgArgMap b)
-> DNameEnv (ListMap StgArgMap b)
-> DNameEnv (ListMap StgArgMap b)
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)
|>> Key (ListMap StgArgMap)
-> XT b -> ListMap StgArgMap b -> ListMap StgArgMap b
forall (m :: * -> *) b. TrieMap m => Key m -> XT b -> m b -> m b
alterTM [StgArg]
Key (ListMap StgArgMap)
args XT b
f }
    foldTM :: (a -> b -> b) -> ConAppMap a -> b -> b
foldTM k :: 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 (m :: * -> *) a b.
TrieMap m =>
(a -> b -> b) -> m a -> b -> b
foldTM ((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)
    mapTM :: (a -> b) -> ConAppMap a -> ConAppMap b
mapTM f :: a -> b
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 b)
-> ConAppMap a
-> ConAppMap b
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> (ListMap StgArgMap a -> ListMap StgArgMap b)
-> DNameEnv (ListMap StgArgMap a) -> UniqDFM (ListMap StgArgMap b)
forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM ((a -> b) -> ListMap StgArgMap a -> ListMap StgArgMap b
forall (m :: * -> *) a b. TrieMap m => (a -> b) -> m a -> m b
mapTM a -> b
f) (DNameEnv (ListMap StgArgMap a) -> UniqDFM (ListMap StgArgMap b))
-> (UniqDFM (ListMap StgArgMap b) -> ConAppMap b)
-> DNameEnv (ListMap StgArgMap a)
-> ConAppMap b
forall a b c. (a -> b) -> (b -> c) -> a -> c
>.> UniqDFM (ListMap StgArgMap b) -> ConAppMap b
forall a. DNameEnv (ListMap StgArgMap a) -> ConAppMap a
CAM

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

-- | The CSE environment. See note [CseEnv Example]
data CseEnv = CseEnv
    { CseEnv -> ConAppMap Var
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 Var
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 Var
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 in_scope :: InScopeSet
in_scope = CseEnv :: ConAppMap Var -> IdEnv Var -> IdEnv Var -> InScopeSet -> CseEnv
CseEnv
    { ce_conAppMap :: ConAppMap Var
ce_conAppMap = ConAppMap Var
forall (m :: * -> *) a. TrieMap m => m a
emptyTM
    , ce_subst :: IdEnv Var
ce_subst     = IdEnv Var
forall a. VarEnv a
emptyVarEnv
    , ce_bndrMap :: IdEnv Var
ce_bndrMap   = IdEnv Var
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 Var
envLookup dataCon :: DataCon
dataCon args :: [StgArg]
args env :: CseEnv
env = Key ConAppMap -> ConAppMap Var -> Maybe Var
forall (m :: * -> *) b. TrieMap m => Key m -> m b -> Maybe b
lookupTM (DataCon
dataCon, [StgArg]
args') (CseEnv -> ConAppMap Var
ce_conAppMap CseEnv
env)
  where args' :: [StgArg]
args' = (StgArg -> StgArg) -> [StgArg] -> [StgArg]
forall a b. (a -> b) -> [a] -> [b]
map StgArg -> StgArg
go [StgArg]
args -- See Note [Trivial case scrutinee]
        go :: StgArg -> StgArg
go (StgVarArg v :: Var
v  ) = Var -> StgArg
StgVarArg (Var -> Maybe Var -> Var
forall a. a -> Maybe a -> a
fromMaybe Var
v (Maybe Var -> Var) -> Maybe Var -> Var
forall a b. (a -> b) -> a -> b
$ IdEnv Var -> Var -> Maybe Var
forall a. VarEnv a -> Var -> Maybe a
lookupVarEnv (CseEnv -> IdEnv Var
ce_bndrMap CseEnv
env) Var
v)
        go (StgLitArg lit :: 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 :: Var -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon _ _ [] env :: CseEnv
env = CseEnv
env
addDataCon bndr :: Var
bndr dataCon :: DataCon
dataCon args :: [StgArg]
args env :: CseEnv
env = CseEnv
env { ce_conAppMap :: ConAppMap Var
ce_conAppMap = ConAppMap Var
new_env }
  where
    new_env :: ConAppMap Var
new_env = Key ConAppMap -> Var -> ConAppMap Var -> ConAppMap Var
forall (m :: * -> *) a. TrieMap m => Key m -> a -> m a -> m a
insertTM (DataCon
dataCon, [StgArg]
args) Var
bndr (CseEnv -> ConAppMap Var
ce_conAppMap CseEnv
env)

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

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

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

substArgs :: CseEnv -> [InStgArg] -> [OutStgArg]
substArgs :: CseEnv -> [StgArg] -> [StgArg]
substArgs env :: 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 env :: CseEnv
env (StgVarArg from :: Var
from) = Var -> StgArg
StgVarArg (CseEnv -> Var -> Var
substVar CseEnv
env Var
from)
substArg _   (StgLitArg lit :: Literal
lit)  = Literal -> StgArg
StgLitArg Literal
lit

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

-- Functions to enter binders

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

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

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

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

-- Main entry point

stgCse :: [InStgTopBinding] -> [OutStgTopBinding]
stgCse :: [InStgTopBinding] -> [InStgTopBinding]
stgCse binds :: [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 :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
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 in_scope :: InScopeSet
in_scope t :: InStgTopBinding
t@(StgTopStringLit _ _) = (InScopeSet
in_scope, InStgTopBinding
t)
stgCseTopLvl in_scope :: InScopeSet
in_scope (StgTopLifted (StgNonRec bndr :: BinderP 'Vanilla
bndr rhs :: 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 -> Var -> InScopeSet
`extendInScopeSet` Var
BinderP 'Vanilla
bndr

stgCseTopLvl in_scope :: InScopeSet
in_scope (StgTopLifted (StgRec eqs :: [(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 [ (Var
BinderP 'Vanilla
bndr, InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope' GenStgRhs 'Vanilla
rhs) | (bndr :: Var
bndr, rhs :: GenStgRhs 'Vanilla
rhs) <- [(Var, GenStgRhs 'Vanilla)]
[(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs ]))
  where in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> [Var] -> InScopeSet
`extendInScopeSetList` [ Var
bndr | (bndr :: Var
bndr, _) <- [(Var, GenStgRhs 'Vanilla)]
[(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
eqs ]

stgCseTopLvlRhs :: InScopeSet -> InStgRhs -> OutStgRhs
stgCseTopLvlRhs :: InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs in_scope :: InScopeSet
in_scope (StgRhsClosure ext :: XRhsClosure 'Vanilla
ext ccs :: CostCentreStack
ccs upd :: UpdateFlag
upd args :: [BinderP 'Vanilla]
args body :: 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 _ (StgRhsCon ccs :: CostCentreStack
ccs dataCon :: DataCon
dataCon args :: [StgArg]
args)
    = CostCentreStack -> DataCon -> [StgArg] -> GenStgRhs 'Vanilla
forall (pass :: StgPass).
CostCentreStack -> DataCon -> [StgArg] -> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon [StgArg]
args

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

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


-- A constructor application.
-- To be removed by a variable use when found in the CSE environment
stgCseExpr env :: CseEnv
env (StgConApp dataCon :: DataCon
dataCon args :: [StgArg]
args tys :: [Type]
tys)
    | Just bndr' :: Var
bndr' <- DataCon -> [StgArg] -> CseEnv -> Maybe Var
envLookup DataCon
dataCon [StgArg]
args' CseEnv
env
    = Var -> [StgArg] -> GenStgExpr 'Vanilla
forall (pass :: StgPass). Var -> [StgArg] -> GenStgExpr pass
StgApp Var
bndr' []
    | Bool
otherwise
    = DataCon -> [StgArg] -> [Type] -> GenStgExpr 'Vanilla
forall (pass :: StgPass).
DataCon -> [StgArg] -> [Type] -> GenStgExpr pass
StgConApp DataCon
dataCon [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 env :: CseEnv
env (StgLet ext :: XLet 'Vanilla
ext binds :: GenStgBinding 'Vanilla
binds body :: GenStgExpr 'Vanilla
body)
    = let (binds' :: Maybe (GenStgBinding 'Vanilla)
binds', env' :: 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 env :: CseEnv
env (StgLetNoEscape ext :: XLetNoEscape 'Vanilla
ext binds :: GenStgBinding 'Vanilla
binds body :: GenStgExpr 'Vanilla
body)
    = let (binds' :: Maybe (GenStgBinding 'Vanilla)
binds', env' :: 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 -> Var -> GenStgAlt 'Vanilla -> GenStgAlt 'Vanilla
stgCseAlt env :: CseEnv
env ty :: AltType
ty case_bndr :: Var
case_bndr (DataAlt dataCon :: DataCon
dataCon, args :: [BinderP 'Vanilla]
args, rhs :: GenStgExpr 'Vanilla
rhs)
    = let (env1 :: CseEnv
env1, args' :: [Var]
args') = CseEnv -> [Var] -> (CseEnv, [Var])
substBndrs CseEnv
env [Var]
[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
            = Var -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon Var
case_bndr DataCon
dataCon ((Var -> StgArg) -> [Var] -> [StgArg]
forall a b. (a -> b) -> [a] -> [b]
map Var -> StgArg
StgVarArg [Var]
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, [Var]
[BinderP 'Vanilla]
args', GenStgExpr 'Vanilla
rhs')
stgCseAlt env :: CseEnv
env _ _ (altCon :: AltCon
altCon, args :: [BinderP 'Vanilla]
args, rhs :: GenStgExpr 'Vanilla
rhs)
    = let (env1 :: CseEnv
env1, args' :: [Var]
args') = CseEnv -> [Var] -> (CseEnv, [Var])
substBndrs CseEnv
env [Var]
[BinderP 'Vanilla]
args
          rhs' :: GenStgExpr 'Vanilla
rhs' = CseEnv -> GenStgExpr 'Vanilla -> GenStgExpr 'Vanilla
stgCseExpr CseEnv
env1 GenStgExpr 'Vanilla
rhs
      in (AltCon
altCon, [Var]
[BinderP 'Vanilla]
args', GenStgExpr 'Vanilla
rhs')

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

stgCsePairs :: CseEnv -> [(OutId, InStgRhs)] -> ([(OutId, OutStgRhs)], CseEnv)
stgCsePairs :: CseEnv
-> [(Var, GenStgRhs 'Vanilla)]
-> ([(Var, GenStgRhs 'Vanilla)], CseEnv)
stgCsePairs env :: CseEnv
env [] = ([], CseEnv
env)
stgCsePairs env0 :: CseEnv
env0 ((b :: Var
b,e :: GenStgRhs 'Vanilla
e):pairs :: [(Var, GenStgRhs 'Vanilla)]
pairs)
  = let (pairMB :: Maybe (Var, GenStgRhs 'Vanilla)
pairMB, env1 :: CseEnv
env1) = CseEnv
-> Var
-> GenStgRhs 'Vanilla
-> (Maybe (Var, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs CseEnv
env0 Var
b GenStgRhs 'Vanilla
e
        (pairs' :: [(Var, GenStgRhs 'Vanilla)]
pairs', env2 :: CseEnv
env2) = CseEnv
-> [(Var, GenStgRhs 'Vanilla)]
-> ([(Var, GenStgRhs 'Vanilla)], CseEnv)
stgCsePairs CseEnv
env1 [(Var, GenStgRhs 'Vanilla)]
pairs
    in (Maybe (Var, GenStgRhs 'Vanilla)
pairMB Maybe (Var, GenStgRhs 'Vanilla)
-> [(Var, GenStgRhs 'Vanilla)] -> [(Var, GenStgRhs 'Vanilla)]
forall a. Maybe a -> [a] -> [a]
`mbCons` [(Var, 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
-> Var
-> GenStgRhs 'Vanilla
-> (Maybe (Var, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs env :: CseEnv
env bndr :: Var
bndr (StgRhsCon ccs :: CostCentreStack
ccs dataCon :: DataCon
dataCon args :: [StgArg]
args)
    | Just other_bndr :: Var
other_bndr <- DataCon -> [StgArg] -> CseEnv -> Maybe Var
envLookup DataCon
dataCon [StgArg]
args' CseEnv
env
    = let env' :: CseEnv
env' = Var -> Var -> CseEnv -> CseEnv
addSubst Var
bndr Var
other_bndr CseEnv
env
      in (Maybe (Var, GenStgRhs 'Vanilla)
forall a. Maybe a
Nothing, CseEnv
env')
    | Bool
otherwise
    = let env' :: CseEnv
env' = Var -> DataCon -> [StgArg] -> CseEnv -> CseEnv
addDataCon Var
bndr DataCon
dataCon [StgArg]
args' CseEnv
env
            -- see note [Case 1: CSEing allocated closures]
          pair :: (Var, GenStgRhs 'Vanilla)
pair = (Var
bndr, CostCentreStack -> DataCon -> [StgArg] -> GenStgRhs 'Vanilla
forall (pass :: StgPass).
CostCentreStack -> DataCon -> [StgArg] -> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dataCon [StgArg]
args')
      in ((Var, GenStgRhs 'Vanilla) -> Maybe (Var, GenStgRhs 'Vanilla)
forall a. a -> Maybe a
Just (Var, GenStgRhs 'Vanilla)
pair, CseEnv
env')
  where args' :: [StgArg]
args' = CseEnv -> [StgArg] -> [StgArg]
substArgs CseEnv
env [StgArg]
args
stgCseRhs env :: CseEnv
env bndr :: Var
bndr (StgRhsClosure ext :: XRhsClosure 'Vanilla
ext ccs :: CostCentreStack
ccs upd :: UpdateFlag
upd args :: [BinderP 'Vanilla]
args body :: GenStgExpr 'Vanilla
body)
    = let (env1 :: CseEnv
env1, args' :: [Var]
args') = CseEnv -> [Var] -> (CseEnv, [Var])
substBndrs CseEnv
env [Var]
[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 ((Var, GenStgRhs 'Vanilla) -> Maybe (Var, GenStgRhs 'Vanilla)
forall a. a -> Maybe a
Just (CseEnv -> Var -> Var
substVar CseEnv
env Var
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 [Var]
[BinderP 'Vanilla]
args' GenStgExpr 'Vanilla
body'), CseEnv
env)


mkStgCase :: StgExpr -> OutId -> AltType -> [StgAlt] -> StgExpr
mkStgCase :: GenStgExpr 'Vanilla
-> Var -> AltType -> [GenStgAlt 'Vanilla] -> GenStgExpr 'Vanilla
mkStgCase scrut :: GenStgExpr 'Vanilla
scrut bndr :: Var
bndr ty :: AltType
ty alts :: [GenStgAlt 'Vanilla]
alts | ((AltCon, [Var], GenStgExpr 'Vanilla) -> Bool)
-> [(AltCon, [Var], GenStgExpr 'Vanilla)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (AltCon, [Var], GenStgExpr 'Vanilla) -> Bool
isBndr [(AltCon, [Var], GenStgExpr 'Vanilla)]
[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 Var
BinderP 'Vanilla
bndr AltType
ty [GenStgAlt 'Vanilla]
alts

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


-- Utilities

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