{-# 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 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) XT b
f 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) XT b
f 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 a -> b -> b
k 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 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) XT b
f 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 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 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 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 [StgArg]
args 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 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 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 Var
_ DataCon
_ [] CseEnv
env = CseEnv
env
addDataCon Var
bndr DataCon
dataCon [StgArg]
args 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 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 Var
from Var
to 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 Var
from Var
to 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 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 Var
from) = Var -> StgArg
StgVarArg (CseEnv -> Var -> Var
substVar CseEnv
env Var
from)
substArg CseEnv
_   (StgLitArg Literal
lit)  = Literal -> StgArg
StgLitArg Literal
lit

substVar :: CseEnv -> InId -> OutId
substVar :: CseEnv -> Var -> Var
substVar CseEnv
env 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 CseEnv
env 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'
            | Bool
otherwise = 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 }

substBndrs :: CseEnv -> [InVar] -> (CseEnv, [OutVar])
substBndrs :: CseEnv -> [Var] -> (CseEnv, [Var])
substBndrs CseEnv
env [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 CseEnv
env [(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 CseEnv
env (Var
id, b
x) = let (CseEnv
env', 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 [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 InScopeSet
in_scope t :: InStgTopBinding
t@(StgTopStringLit Var
_ 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 -> Var -> InScopeSet
`extendInScopeSet` Var
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 [ (Var
BinderP 'Vanilla
bndr, InScopeSet -> GenStgRhs 'Vanilla -> GenStgRhs 'Vanilla
stgCseTopLvlRhs InScopeSet
in_scope' GenStgRhs 'Vanilla
rhs) | (Var
bndr, 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 | (Var
bndr, GenStgRhs 'Vanilla
_) <- [(Var, 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 [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 CseEnv
env (StgApp Var
fun [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 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
_ (StgLam NonEmpty (BinderP 'Vanilla)
_ GenStgExpr 'Vanilla
_)
    = String -> SDoc -> GenStgExpr 'Vanilla
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"stgCseExp" (String -> SDoc
text String
"StgLam")
stgCseExpr CseEnv
env (StgTick Tickish Var
tick 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 CseEnv
env (StgCase GenStgExpr 'Vanilla
scrut BinderP 'Vanilla
bndr AltType
ty [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
    (CseEnv
env1, Var
bndr') = CseEnv -> Var -> (CseEnv, Var)
substBndr CseEnv
env Var
BinderP 'Vanilla
bndr
    env2 :: CseEnv
env2 | StgApp 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 CseEnv
env (StgConApp DataCon
dataCon [StgArg]
args [Type]
tys)
    | Just 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 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 -> Var -> GenStgAlt 'Vanilla -> GenStgAlt 'Vanilla
stgCseAlt CseEnv
env AltType
ty Var
case_bndr (DataAlt DataCon
dataCon, [BinderP 'Vanilla]
args, GenStgExpr 'Vanilla
rhs)
    = let (CseEnv
env1, [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 CseEnv
env AltType
_ Var
_ (AltCon
altCon, [BinderP 'Vanilla]
args, GenStgExpr 'Vanilla
rhs)
    = let (CseEnv
env1, [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 CseEnv
env (StgNonRec BinderP 'Vanilla
b GenStgRhs 'Vanilla
e)
    = let (CseEnv
env1, 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
        (Maybe (Var, GenStgRhs 'Vanilla)
Nothing,      CseEnv
env2) -> (Maybe (GenStgBinding 'Vanilla)
forall a. Maybe a
Nothing,                CseEnv
env2)
        (Just (Var
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 Var
BinderP 'Vanilla
b2 GenStgRhs 'Vanilla
e'), CseEnv
env2)
stgCseBind CseEnv
env (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
pairs)
    = let (CseEnv
env1, [(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
        ([],     CseEnv
env2) -> (Maybe (GenStgBinding 'Vanilla)
forall a. Maybe a
Nothing, CseEnv
env2)
        ([(Var, 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 [(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 CseEnv
env [] = ([], CseEnv
env)
stgCsePairs CseEnv
env0 ((Var
b,GenStgRhs 'Vanilla
e):[(Var, GenStgRhs 'Vanilla)]
pairs)
  = let (Maybe (Var, GenStgRhs 'Vanilla)
pairMB, CseEnv
env1) = CseEnv
-> Var
-> GenStgRhs 'Vanilla
-> (Maybe (Var, GenStgRhs 'Vanilla), CseEnv)
stgCseRhs CseEnv
env0 Var
b GenStgRhs 'Vanilla
e
        ([(Var, GenStgRhs 'Vanilla)]
pairs', 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 CseEnv
env Var
bndr (StgRhsCon CostCentreStack
ccs DataCon
dataCon [StgArg]
args)
    | Just 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 CseEnv
env Var
bndr (StgRhsClosure XRhsClosure 'Vanilla
ext CostCentreStack
ccs UpdateFlag
upd [BinderP 'Vanilla]
args GenStgExpr 'Vanilla
body)
    = let (CseEnv
env1, [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 GenStgExpr 'Vanilla
scrut Var
bndr AltType
ty [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 (AltCon
_, [Var]
_, StgApp Var
f []) = Var
f Var -> Var -> Bool
forall a. Eq a => a -> a -> Bool
== Var
bndr
    isBndr (AltCon, [Var], GenStgExpr 'Vanilla)
_                   = 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 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.
-}