{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE LambdaCase #-}

{- |
Non-global free variable analysis on STG terms. This pass annotates
non-top-level closure bindings with captured variables. Global variables are not
captured. For example, in a top-level binding like (pseudo-STG)

    f = \[x,y] .
      let g = \[p] . reverse (x ++ p)
      in g y

In g, `reverse` and `(++)` are global variables so they're not considered free.
`p` is an argument, so `x` is the only actual free variable here. The annotated
version is thus:

    f = \[x,y] .
      let g = [x] \[p] . reverse (x ++ p)
      in g y

Note that non-top-level recursive bindings are also considered free within the
group:

    map = {} \r [f xs0]
      let {
        Rec {
          go = {f, go} \r [xs1]
            case xs1 of {
              [] -> [] [];
              : x xs2 ->
                  let { xs' = {go, xs2} \u [] go xs2; } in
                  let { x' = {f, x} \u [] f x; } in
                  : [x' xs'];
            };
        end Rec }
      } in go xs0;

Here go is free in its RHS.

Top-level closure bindings never capture variables as all of their free
variables are global.
-}
module GHC.Stg.FVs (
    depSortWithAnnotStgPgm,
    annBindingFreeVars
  ) where

import GHC.Prelude hiding (mod)

import GHC.Stg.Syntax
import GHC.Stg.Utils (bindersOf)
import GHC.Types.Id
import GHC.Types.Name (Name, nameIsLocalOrFrom)
import GHC.Types.Tickish ( GenTickish(Breakpoint) )
import GHC.Types.Unique.Set (nonDetEltsUniqSet)
import GHC.Types.Var.Set
import GHC.Unit.Module (Module)
import GHC.Utils.Misc

import Data.Graph (SCC (..))
import GHC.Data.Graph.Directed( Node(..), stronglyConnCompFromEdgedVerticesUniq )

{- Note [Why do we need dependency analysis?]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The program needs to be in dependency order for the SRT algorithm to
work (see CmmBuildInfoTables, which also includes a detailed
description of the algorithm).

But isn't it in correct dependency order already?  No:

* The simplifier does not guarantee to produce programs in dependency
  order (see #16192 and Note [Glomming] in GHC.Core.Opt.OccurAnal).
  This could be solved by a final run of the occurrence analyser, but
  that's more work

* We also don't guarantee that StgLiftLams will preserve the order or
  only create minimal recursive groups.
-}

--------------------------------------------------------------------------------
-- | Dependency sort a STG program, and annotate it with free variables
-- The returned bindings:
--   * Are in dependency order
--   * Each StgRhsClosure is correctly annotated (in its extension field)
--     with the free variables needed in the closure
--   * Each StgCase is correctly annotated (in its extension field) with
--     the variables that must be saved across the case
depSortWithAnnotStgPgm :: Module -> [StgTopBinding] -> [(CgStgTopBinding,ImpFVs)]
depSortWithAnnotStgPgm :: Module -> [StgTopBinding] -> [(CgStgTopBinding, ImpFVs)]
depSortWithAnnotStgPgm Module
this_mod [StgTopBinding]
binds
  = {-# SCC "STG.depSortWithAnnotStgPgm" #-}
    [CgStgTopBinding] -> [ImpFVs] -> [(CgStgTopBinding, ImpFVs)]
forall a b. [a] -> [b] -> [(a, b)]
zip [CgStgTopBinding]
lit_binds (ImpFVs -> [ImpFVs]
forall a. a -> [a]
repeat ImpFVs
emptyVarSet) [(CgStgTopBinding, ImpFVs)]
-> [(CgStgTopBinding, ImpFVs)] -> [(CgStgTopBinding, ImpFVs)]
forall a. [a] -> [a] -> [a]
++ (SCC (Id, CgStgRhs, ImpFVs) -> (CgStgTopBinding, ImpFVs))
-> [SCC (Id, CgStgRhs, ImpFVs)] -> [(CgStgTopBinding, ImpFVs)]
forall a b. (a -> b) -> [a] -> [b]
map SCC (Id, CgStgRhs, ImpFVs) -> (CgStgTopBinding, ImpFVs)
SCC (BinderP 'CodeGen, CgStgRhs, ImpFVs)
-> (CgStgTopBinding, ImpFVs)
forall {pass :: StgPass}.
SCC (BinderP pass, GenStgRhs pass, ImpFVs)
-> (GenStgTopBinding pass, ImpFVs)
from_scc [SCC (Id, CgStgRhs, ImpFVs)]
sccs
  where
    lit_binds :: [CgStgTopBinding]
    pairs     :: [(Id, StgRhs)]
    ([CgStgTopBinding]
lit_binds, [(Id, StgRhs)]
pairs) = [StgTopBinding] -> ([CgStgTopBinding], [(Id, StgRhs)])
flattenTopStgBindings [StgTopBinding]
binds

    nodes :: [Node Name (Id, CgStgRhs, ImpFVs)]
    nodes :: [Node Name (Id, CgStgRhs, ImpFVs)]
nodes = ((Id, StgRhs) -> Node Name (Id, CgStgRhs, ImpFVs))
-> [(Id, StgRhs)] -> [Node Name (Id, CgStgRhs, ImpFVs)]
forall a b. (a -> b) -> [a] -> [b]
map (Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs, ImpFVs)
annotateTopPair Env
env0) [(Id, StgRhs)]
pairs
    env0 :: Env
env0 = Env { locals :: ImpFVs
locals = ImpFVs
emptyVarSet, mod :: Module
mod = Module
this_mod }

    -- Do strongly connected component analysis.  Why?
    -- See Note [Why do we need dependency analysis?]
    sccs :: [SCC (Id,CgStgRhs,ImpFVs)]
    sccs :: [SCC (Id, CgStgRhs, ImpFVs)]
sccs  = [Node Name (Id, CgStgRhs, ImpFVs)] -> [SCC (Id, CgStgRhs, ImpFVs)]
forall key payload.
Uniquable key =>
[Node key payload] -> [SCC payload]
stronglyConnCompFromEdgedVerticesUniq [Node Name (Id, CgStgRhs, ImpFVs)]
nodes

    from_scc :: SCC (BinderP pass, GenStgRhs pass, ImpFVs)
-> (GenStgTopBinding pass, ImpFVs)
from_scc = \case
      AcyclicSCC (BinderP pass
bndr,GenStgRhs pass
rhs,ImpFVs
imp_fvs) -> (GenStgBinding pass -> GenStgTopBinding pass
forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted (BinderP pass -> GenStgRhs pass -> GenStgBinding pass
forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec BinderP pass
bndr GenStgRhs pass
rhs), ImpFVs
imp_fvs)
      CyclicSCC [(BinderP pass, GenStgRhs pass, ImpFVs)]
triples             -> (GenStgBinding pass -> GenStgTopBinding pass
forall (pass :: StgPass).
GenStgBinding pass -> GenStgTopBinding pass
StgTopLifted ([(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [(BinderP pass, GenStgRhs pass)]
pairs), ImpFVs
imp_fvs)
        where
          ([BinderP pass]
ids,[GenStgRhs pass]
rhss,[ImpFVs]
imp_fvss) = [(BinderP pass, GenStgRhs pass, ImpFVs)]
-> ([BinderP pass], [GenStgRhs pass], [ImpFVs])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 [(BinderP pass, GenStgRhs pass, ImpFVs)]
triples
          pairs :: [(BinderP pass, GenStgRhs pass)]
pairs = [BinderP pass]
-> [GenStgRhs pass] -> [(BinderP pass, GenStgRhs pass)]
forall a b. [a] -> [b] -> [(a, b)]
zip [BinderP pass]
ids [GenStgRhs pass]
rhss
          imp_fvs :: ImpFVs
imp_fvs = [ImpFVs] -> ImpFVs
unionVarSets [ImpFVs]
imp_fvss


flattenTopStgBindings :: [StgTopBinding] -> ([CgStgTopBinding], [(Id,StgRhs)])
flattenTopStgBindings :: [StgTopBinding] -> ([CgStgTopBinding], [(Id, StgRhs)])
flattenTopStgBindings [StgTopBinding]
binds
  = [CgStgTopBinding]
-> [(BinderP 'Vanilla, StgRhs)]
-> [StgTopBinding]
-> ([CgStgTopBinding], [(BinderP 'Vanilla, StgRhs)])
forall {pass :: StgPass} {pass :: StgPass}.
[GenStgTopBinding pass]
-> [(BinderP pass, GenStgRhs pass)]
-> [GenStgTopBinding pass]
-> ([GenStgTopBinding pass], [(BinderP pass, GenStgRhs pass)])
go [] [] [StgTopBinding]
binds
  where
    go :: [GenStgTopBinding pass]
-> [(BinderP pass, GenStgRhs pass)]
-> [GenStgTopBinding pass]
-> ([GenStgTopBinding pass], [(BinderP pass, GenStgRhs pass)])
go [GenStgTopBinding pass]
lits [(BinderP pass, GenStgRhs pass)]
pairs [] = ([GenStgTopBinding pass]
lits, [(BinderP pass, GenStgRhs pass)]
pairs)
    go [GenStgTopBinding pass]
lits [(BinderP pass, GenStgRhs pass)]
pairs (GenStgTopBinding pass
bind:[GenStgTopBinding pass]
binds)
      = case GenStgTopBinding pass
bind of
          StgTopStringLit Id
bndr ByteString
rhs -> [GenStgTopBinding pass]
-> [(BinderP pass, GenStgRhs pass)]
-> [GenStgTopBinding pass]
-> ([GenStgTopBinding pass], [(BinderP pass, GenStgRhs pass)])
go (Id -> ByteString -> GenStgTopBinding pass
forall (pass :: StgPass). Id -> ByteString -> GenStgTopBinding pass
StgTopStringLit Id
bndr ByteString
rhsGenStgTopBinding pass
-> [GenStgTopBinding pass] -> [GenStgTopBinding pass]
forall a. a -> [a] -> [a]
:[GenStgTopBinding pass]
lits) [(BinderP pass, GenStgRhs pass)]
pairs [GenStgTopBinding pass]
binds
          StgTopLifted GenStgBinding pass
stg_bind -> [GenStgTopBinding pass]
-> [(BinderP pass, GenStgRhs pass)]
-> [GenStgTopBinding pass]
-> ([GenStgTopBinding pass], [(BinderP pass, GenStgRhs pass)])
go [GenStgTopBinding pass]
lits (GenStgBinding pass -> [(BinderP pass, GenStgRhs pass)]
forall {pass :: StgPass}.
GenStgBinding pass -> [(BinderP pass, GenStgRhs pass)]
flatten_one GenStgBinding pass
stg_bind [(BinderP pass, GenStgRhs pass)]
-> [(BinderP pass, GenStgRhs pass)]
-> [(BinderP pass, GenStgRhs pass)]
forall a. [a] -> [a] -> [a]
++ [(BinderP pass, GenStgRhs pass)]
pairs) [GenStgTopBinding pass]
binds

    flatten_one :: GenStgBinding pass -> [(BinderP pass, GenStgRhs pass)]
flatten_one (StgNonRec BinderP pass
b GenStgRhs pass
r) = [(BinderP pass
b,GenStgRhs pass
r)]
    flatten_one (StgRec [(BinderP pass, GenStgRhs pass)]
pairs)  = [(BinderP pass, GenStgRhs pass)]
pairs

annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs, ImpFVs)
annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs, ImpFVs)
annotateTopPair Env
env0 (Id
bndr, StgRhs
rhs)
  = DigraphNode { node_key :: Name
node_key          = Id -> Name
idName Id
bndr
                , node_payload :: (Id, CgStgRhs, ImpFVs)
node_payload      = (Id
bndr, CgStgRhs
rhs', ImpFVs
imp_fvs)
                , node_dependencies :: [Name]
node_dependencies = (Id -> Name) -> [Id] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map Id -> Name
idName (ImpFVs -> [Id]
forall elt. UniqSet elt -> [elt]
nonDetEltsUniqSet ImpFVs
top_fvs) }
  where
    (CgStgRhs
rhs', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
_) = Env -> StgRhs -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs)
rhsFVs Env
env0 StgRhs
rhs

--------------------------------------------------------------------------------
-- * Non-global free variable analysis

data Env
  = Env
  { -- | Set of locally-bound, not-top-level binders in scope.
    -- That is, variables bound by a let (but not let-no-escape), a lambda
    -- (in a StgRhsClsoure), a case binder, or a case alternative.  These
    -- are the variables that must be captured in a function closure, if they
    -- are free in the RHS. Example
    --   f = \x. let g = \y. x+1
    --           let h = \z. g z + 1
    --           in h x
    -- In the body of h we have locals = {x, g, z}.  Note that f is top level
    -- and does not appear in locals.
    Env -> ImpFVs
locals :: IdSet
  , Env -> Module
mod    :: Module
  }

addLocals :: [Id] -> Env -> Env
addLocals :: [Id] -> Env -> Env
addLocals [Id]
bndrs Env
env
  = Env
env { locals = extendVarSetList (locals env) bndrs }

--------------------------------------------------------------------------------
-- | TopFVs: set of variables that are:
--    (a) bound at the top level of this module, and
--    (b) appear free in the expression
-- It is a /non-deterministic/ set because we use it only to perform dependency
-- analysis on the top-level bindings.
type TopFVs   = IdSet

-- | ImpFVs: set of variables that are imported
--
-- It is a /non-deterministic/ set because we use it only to perform module
-- dependency analysis.
type ImpFVs   = IdSet

-- | LocalFVs: set of variable that are:
--     (a) bound locally (by a lambda, non-top-level let, or case); that is,
--         it appears in the 'locals' field of 'Env'
--     (b) appear free in the expression
-- It is a /deterministic/ set because it is used to annotate closures with
-- their free variables, and we want closure layout to be deterministic.
--
-- Invariant: the LocalFVs returned is a subset of the 'locals' field of Env
type LocalFVs = DIdSet

-- | Dependency analysis on STG terms.
--
-- Dependencies of a binding are just free variables in the binding. This
-- includes imported ids and ids in the current module. For recursive groups we
-- just return one set of free variables which is just the union of dependencies
-- of all bindings in the group.
--
-- Implementation: pass bound variables (NestedIds) to recursive calls, get free
-- variables (TopFVs) back. We ignore imported TopFVs as they do not change the
-- ordering but it improves performance (see `nameIsExternalFrom` call in `vars_fvs`).
--

annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
annBindingFreeVars Module
this_mod = (CgStgBinding, ImpFVs, ImpFVs, LocalFVs) -> CgStgBinding
forall a b c d. (a, b, c, d) -> a
fstOf4 ((CgStgBinding, ImpFVs, ImpFVs, LocalFVs) -> CgStgBinding)
-> (StgBinding -> (CgStgBinding, ImpFVs, ImpFVs, LocalFVs))
-> StgBinding
-> CgStgBinding
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env
-> LocalFVs
-> StgBinding
-> (CgStgBinding, ImpFVs, ImpFVs, LocalFVs)
bindingFVs (ImpFVs -> Module -> Env
Env ImpFVs
emptyVarSet Module
this_mod) LocalFVs
emptyDVarSet

bindingFVs :: Env -> LocalFVs -> StgBinding -> (CgStgBinding, ImpFVs, TopFVs, LocalFVs)
bindingFVs :: Env
-> LocalFVs
-> StgBinding
-> (CgStgBinding, ImpFVs, ImpFVs, LocalFVs)
bindingFVs Env
env LocalFVs
body_fv StgBinding
b =
  case StgBinding
b of
    StgNonRec BinderP 'Vanilla
bndr StgRhs
r -> (BinderP 'CodeGen -> CgStgRhs -> CgStgBinding
forall (pass :: StgPass).
BinderP pass -> GenStgRhs pass -> GenStgBinding pass
StgNonRec BinderP 'Vanilla
BinderP 'CodeGen
bndr CgStgRhs
r', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)
      where
        (CgStgRhs
r', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
rhs_lcl_fvs) = Env -> StgRhs -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs)
rhsFVs Env
env StgRhs
r
        lcl_fvs :: LocalFVs
lcl_fvs = LocalFVs -> Id -> LocalFVs
delDVarSet LocalFVs
body_fv Id
BinderP 'Vanilla
bndr LocalFVs -> LocalFVs -> LocalFVs
`unionDVarSet` LocalFVs
rhs_lcl_fvs

    StgRec [(BinderP 'Vanilla, StgRhs)]
pairs -> ([(BinderP 'CodeGen, CgStgRhs)] -> CgStgBinding
forall (pass :: StgPass).
[(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
StgRec [(Id, CgStgRhs)]
[(BinderP 'CodeGen, CgStgRhs)]
pairs', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvss)
      where
        bndrs :: [Id]
bndrs = ((Id, StgRhs) -> Id) -> [(Id, StgRhs)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, StgRhs) -> Id
forall a b. (a, b) -> a
fst [(Id, StgRhs)]
[(BinderP 'Vanilla, StgRhs)]
pairs
        env' :: Env
env' = [Id] -> Env -> Env
addLocals [Id]
bndrs Env
env
        ([CgStgRhs]
rhss, [ImpFVs]
rhs_imp_fvss, [ImpFVs]
rhs_top_fvss, [LocalFVs]
rhs_lcl_fvss) = ((Id, StgRhs) -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs))
-> [(Id, StgRhs)] -> ([CgStgRhs], [ImpFVs], [ImpFVs], [LocalFVs])
forall a b c d e.
(a -> (b, c, d, e)) -> [a] -> ([b], [c], [d], [e])
mapAndUnzip4 (Env -> StgRhs -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs)
rhsFVs Env
env' (StgRhs -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs))
-> ((Id, StgRhs) -> StgRhs)
-> (Id, StgRhs)
-> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, StgRhs) -> StgRhs
forall a b. (a, b) -> b
snd) [(Id, StgRhs)]
[(BinderP 'Vanilla, StgRhs)]
pairs
        top_fvs :: ImpFVs
top_fvs = [ImpFVs] -> ImpFVs
unionVarSets [ImpFVs]
rhs_top_fvss
        imp_fvs :: ImpFVs
imp_fvs = [ImpFVs] -> ImpFVs
unionVarSets [ImpFVs]
rhs_imp_fvss
        pairs' :: [(Id, CgStgRhs)]
pairs' = [Id] -> [CgStgRhs] -> [(Id, CgStgRhs)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
bndrs [CgStgRhs]
rhss
        lcl_fvss :: LocalFVs
lcl_fvss = LocalFVs -> [Id] -> LocalFVs
delDVarSetList ([LocalFVs] -> LocalFVs
unionDVarSets (LocalFVs
body_fvLocalFVs -> [LocalFVs] -> [LocalFVs]
forall a. a -> [a] -> [a]
:[LocalFVs]
rhs_lcl_fvss)) [Id]
bndrs

varFVs :: Env -> Id -> (ImpFVs, TopFVs, LocalFVs) -> (ImpFVs, TopFVs, LocalFVs)
varFVs :: Env
-> Id -> (ImpFVs, ImpFVs, LocalFVs) -> (ImpFVs, ImpFVs, LocalFVs)
varFVs Env
env Id
v (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)
  | Id
v Id -> ImpFVs -> Bool
`elemVarSet` Env -> ImpFVs
locals Env
env                -- v is locally bound
  = (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs LocalFVs -> Id -> LocalFVs
`extendDVarSet` Id
v)
  | Module -> Name -> Bool
nameIsLocalOrFrom (Env -> Module
mod Env
env) (Id -> Name
idName Id
v)   -- v is bound at top level
  = (ImpFVs
imp_fvs, ImpFVs
top_fvs ImpFVs -> Id -> ImpFVs
`extendVarSet` Id
v, LocalFVs
lcl_fvs)
  | Bool
otherwise                                -- v is imported
  = (ImpFVs
imp_fvs ImpFVs -> Id -> ImpFVs
`extendVarSet` Id
v, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

exprFVs :: Env -> StgExpr -> (CgStgExpr, ImpFVs, TopFVs, LocalFVs)
exprFVs :: Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs Env
env = StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
go
  where
    go :: StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
go (StgApp Id
f [StgArg]
as)
      | (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env
-> Id -> (ImpFVs, ImpFVs, LocalFVs) -> (ImpFVs, ImpFVs, LocalFVs)
varFVs Env
env Id
f (Env -> [StgArg] -> (ImpFVs, ImpFVs, LocalFVs)
argsFVs Env
env [StgArg]
as)
      = (Id -> [StgArg] -> CgStgExpr
forall (pass :: StgPass). Id -> [StgArg] -> GenStgExpr pass
StgApp Id
f [StgArg]
as, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

    go (StgLit Literal
lit) = (Literal -> CgStgExpr
forall (pass :: StgPass). Literal -> GenStgExpr pass
StgLit Literal
lit, ImpFVs
emptyVarSet, ImpFVs
emptyVarSet, LocalFVs
emptyDVarSet)

    go (StgConApp DataCon
dc ConstructorNumber
n [StgArg]
as [[PrimRep]]
tys)
      | (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> [StgArg] -> (ImpFVs, ImpFVs, LocalFVs)
argsFVs Env
env [StgArg]
as
      = (DataCon
-> ConstructorNumber -> [StgArg] -> [[PrimRep]] -> CgStgExpr
forall (pass :: StgPass).
DataCon
-> ConstructorNumber -> [StgArg] -> [[PrimRep]] -> GenStgExpr pass
StgConApp DataCon
dc ConstructorNumber
n [StgArg]
as [[PrimRep]]
tys, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

    go (StgOpApp StgOp
op [StgArg]
as Type
ty)
      | (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> [StgArg] -> (ImpFVs, ImpFVs, LocalFVs)
argsFVs Env
env [StgArg]
as
      = (StgOp -> [StgArg] -> Type -> CgStgExpr
forall (pass :: StgPass).
StgOp -> [StgArg] -> Type -> GenStgExpr pass
StgOpApp StgOp
op [StgArg]
as Type
ty, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

    go (StgCase StgExpr
scrut BinderP 'Vanilla
bndr AltType
ty [GenStgAlt 'Vanilla]
alts)
      | (CgStgExpr
scrut',ImpFVs
scrut_imp_fvs,ImpFVs
scrut_top_fvs,LocalFVs
scrut_lcl_fvs) <- Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs Env
env StgExpr
scrut
      , ([CgStgAlt]
alts',[ImpFVs]
alts_imp_fvss,[ImpFVs]
alts_top_fvss,[LocalFVs]
alts_lcl_fvss)
          <- (GenStgAlt 'Vanilla -> (CgStgAlt, ImpFVs, ImpFVs, LocalFVs))
-> [GenStgAlt 'Vanilla]
-> ([CgStgAlt], [ImpFVs], [ImpFVs], [LocalFVs])
forall a b c d e.
(a -> (b, c, d, e)) -> [a] -> ([b], [c], [d], [e])
mapAndUnzip4 (Env -> GenStgAlt 'Vanilla -> (CgStgAlt, ImpFVs, ImpFVs, LocalFVs)
altFVs ([Id] -> Env -> Env
addLocals [Id
BinderP 'Vanilla
bndr] Env
env)) [GenStgAlt 'Vanilla]
alts
      , let top_fvs :: ImpFVs
top_fvs = ImpFVs
scrut_top_fvs ImpFVs -> ImpFVs -> ImpFVs
`unionVarSet` [ImpFVs] -> ImpFVs
unionVarSets [ImpFVs]
alts_top_fvss
            imp_fvs :: ImpFVs
imp_fvs = ImpFVs
scrut_imp_fvs ImpFVs -> ImpFVs -> ImpFVs
`unionVarSet` [ImpFVs] -> ImpFVs
unionVarSets [ImpFVs]
alts_imp_fvss
            alts_lcl_fvs :: LocalFVs
alts_lcl_fvs = [LocalFVs] -> LocalFVs
unionDVarSets [LocalFVs]
alts_lcl_fvss
            lcl_fvs :: LocalFVs
lcl_fvs = LocalFVs -> Id -> LocalFVs
delDVarSet (LocalFVs -> LocalFVs -> LocalFVs
unionDVarSet LocalFVs
scrut_lcl_fvs LocalFVs
alts_lcl_fvs) Id
BinderP 'Vanilla
bndr
      = (CgStgExpr -> BinderP 'CodeGen -> AltType -> [CgStgAlt] -> CgStgExpr
forall (pass :: StgPass).
GenStgExpr pass
-> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
StgCase CgStgExpr
scrut' BinderP 'Vanilla
BinderP 'CodeGen
bndr AltType
ty [CgStgAlt]
alts', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

    go (StgLet XLet 'Vanilla
ext         StgBinding
bind StgExpr
body) = (CgStgBinding -> CgStgExpr -> CgStgExpr)
-> StgBinding -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
go_bind (XLet 'CodeGen -> CgStgBinding -> CgStgExpr -> CgStgExpr
forall (pass :: StgPass).
XLet pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLet XLet 'Vanilla
XLet 'CodeGen
ext) StgBinding
bind StgExpr
body
    go (StgLetNoEscape XLetNoEscape 'Vanilla
ext StgBinding
bind StgExpr
body) = (CgStgBinding -> CgStgExpr -> CgStgExpr)
-> StgBinding -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
go_bind (XLetNoEscape 'CodeGen -> CgStgBinding -> CgStgExpr -> CgStgExpr
forall (pass :: StgPass).
XLetNoEscape pass
-> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
StgLetNoEscape XLetNoEscape 'Vanilla
XLetNoEscape 'CodeGen
ext) StgBinding
bind StgExpr
body

    go (StgTick StgTickish
tick StgExpr
e)
      | (CgStgExpr
e', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs Env
env StgExpr
e
      , let lcl_fvs' :: LocalFVs
lcl_fvs' = LocalFVs -> LocalFVs -> LocalFVs
unionDVarSet (StgTickish -> LocalFVs
forall {pass :: TickishPass}.
(XTickishId pass ~ Id) =>
GenTickish pass -> LocalFVs
tickish StgTickish
tick) LocalFVs
lcl_fvs
      = (StgTickish -> CgStgExpr -> CgStgExpr
forall (pass :: StgPass).
StgTickish -> GenStgExpr pass -> GenStgExpr pass
StgTick StgTickish
tick CgStgExpr
e', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs')
        where
          tickish :: GenTickish pass -> LocalFVs
tickish (Breakpoint XBreakpoint pass
_ Int
_ [XTickishId pass]
ids Module
_) = [Id] -> LocalFVs
mkDVarSet [Id]
[XTickishId pass]
ids
          tickish GenTickish pass
_                      = LocalFVs
emptyDVarSet

    go_bind :: (CgStgBinding -> CgStgExpr -> CgStgExpr)
-> StgBinding -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
go_bind CgStgBinding -> CgStgExpr -> CgStgExpr
dc StgBinding
bind StgExpr
body = (CgStgBinding -> CgStgExpr -> CgStgExpr
dc CgStgBinding
bind' CgStgExpr
body', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)
      where
        env' :: Env
env' = [Id] -> Env -> Env
addLocals (StgBinding -> [Id]
forall (a :: StgPass). (BinderP a ~ Id) => GenStgBinding a -> [Id]
bindersOf StgBinding
bind) Env
env
        (CgStgExpr
body', ImpFVs
body_imp_fvs, ImpFVs
body_top_fvs, LocalFVs
body_lcl_fvs) = Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs Env
env' StgExpr
body
        (CgStgBinding
bind', ImpFVs
bind_imp_fvs, ImpFVs
bind_top_fvs, LocalFVs
lcl_fvs)      = Env
-> LocalFVs
-> StgBinding
-> (CgStgBinding, ImpFVs, ImpFVs, LocalFVs)
bindingFVs Env
env' LocalFVs
body_lcl_fvs StgBinding
bind
        top_fvs :: ImpFVs
top_fvs = ImpFVs
bind_top_fvs ImpFVs -> ImpFVs -> ImpFVs
`unionVarSet` ImpFVs
body_top_fvs
        imp_fvs :: ImpFVs
imp_fvs = ImpFVs
bind_imp_fvs ImpFVs -> ImpFVs -> ImpFVs
`unionVarSet` ImpFVs
body_imp_fvs


rhsFVs :: Env -> StgRhs -> (CgStgRhs, ImpFVs, TopFVs, LocalFVs)
rhsFVs :: Env -> StgRhs -> (CgStgRhs, ImpFVs, ImpFVs, LocalFVs)
rhsFVs Env
env (StgRhsClosure XRhsClosure 'Vanilla
_ CostCentreStack
ccs UpdateFlag
uf [BinderP 'Vanilla]
bs StgExpr
body Type
typ)
  | (CgStgExpr
body', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs ([Id] -> Env -> Env
addLocals [Id]
[BinderP 'Vanilla]
bs Env
env) StgExpr
body
  , let lcl_fvs' :: LocalFVs
lcl_fvs' = LocalFVs -> [Id] -> LocalFVs
delDVarSetList LocalFVs
lcl_fvs [Id]
[BinderP 'Vanilla]
bs
  = (XRhsClosure 'CodeGen
-> CostCentreStack
-> UpdateFlag
-> [BinderP 'CodeGen]
-> CgStgExpr
-> Type
-> CgStgRhs
forall (pass :: StgPass).
XRhsClosure pass
-> CostCentreStack
-> UpdateFlag
-> [BinderP pass]
-> GenStgExpr pass
-> Type
-> GenStgRhs pass
StgRhsClosure LocalFVs
XRhsClosure 'CodeGen
lcl_fvs' CostCentreStack
ccs UpdateFlag
uf [BinderP 'Vanilla]
[BinderP 'CodeGen]
bs CgStgExpr
body' Type
typ, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs')
rhsFVs Env
env (StgRhsCon CostCentreStack
ccs DataCon
dc ConstructorNumber
mu [StgTickish]
ts [StgArg]
bs Type
typ)
  | (ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> [StgArg] -> (ImpFVs, ImpFVs, LocalFVs)
argsFVs Env
env [StgArg]
bs
  = (CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> Type
-> CgStgRhs
forall (pass :: StgPass).
CostCentreStack
-> DataCon
-> ConstructorNumber
-> [StgTickish]
-> [StgArg]
-> Type
-> GenStgRhs pass
StgRhsCon CostCentreStack
ccs DataCon
dc ConstructorNumber
mu [StgTickish]
ts [StgArg]
bs Type
typ, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs)

argsFVs :: Env -> [StgArg] -> (ImpFVs, TopFVs, LocalFVs)
argsFVs :: Env -> [StgArg] -> (ImpFVs, ImpFVs, LocalFVs)
argsFVs Env
env = ((ImpFVs, ImpFVs, LocalFVs)
 -> StgArg -> (ImpFVs, ImpFVs, LocalFVs))
-> (ImpFVs, ImpFVs, LocalFVs)
-> [StgArg]
-> (ImpFVs, ImpFVs, LocalFVs)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (ImpFVs, ImpFVs, LocalFVs) -> StgArg -> (ImpFVs, ImpFVs, LocalFVs)
f (ImpFVs
emptyVarSet, ImpFVs
emptyVarSet, LocalFVs
emptyDVarSet)
  where
    f :: (ImpFVs, ImpFVs, LocalFVs) -> StgArg -> (ImpFVs, ImpFVs, LocalFVs)
f (ImpFVs
imp_fvs,ImpFVs
fvs,LocalFVs
ids) StgLitArg{}   = (ImpFVs
imp_fvs, ImpFVs
fvs, LocalFVs
ids)
    f (ImpFVs
imp_fvs,ImpFVs
fvs,LocalFVs
ids) (StgVarArg Id
v) = Env
-> Id -> (ImpFVs, ImpFVs, LocalFVs) -> (ImpFVs, ImpFVs, LocalFVs)
varFVs Env
env Id
v (ImpFVs
imp_fvs, ImpFVs
fvs, LocalFVs
ids)

altFVs :: Env -> StgAlt -> (CgStgAlt, ImpFVs, TopFVs, LocalFVs)
altFVs :: Env -> GenStgAlt 'Vanilla -> (CgStgAlt, ImpFVs, ImpFVs, LocalFVs)
altFVs Env
env GenStgAlt{alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con=AltCon
con, alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs=[BinderP 'Vanilla]
bndrs, alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs=StgExpr
e}
  | (CgStgExpr
e', ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs) <- Env -> StgExpr -> (CgStgExpr, ImpFVs, ImpFVs, LocalFVs)
exprFVs ([Id] -> Env -> Env
addLocals [Id]
[BinderP 'Vanilla]
bndrs Env
env) StgExpr
e
  , let lcl_fvs' :: LocalFVs
lcl_fvs' = LocalFVs -> [Id] -> LocalFVs
delDVarSetList LocalFVs
lcl_fvs [Id]
[BinderP 'Vanilla]
bndrs
  , let newAlt :: CgStgAlt
newAlt   = GenStgAlt{alt_con :: AltCon
alt_con=AltCon
con, alt_bndrs :: [BinderP 'CodeGen]
alt_bndrs=[BinderP 'Vanilla]
[BinderP 'CodeGen]
bndrs, alt_rhs :: CgStgExpr
alt_rhs=CgStgExpr
e'}
  = (CgStgAlt
newAlt, ImpFVs
imp_fvs, ImpFVs
top_fvs, LocalFVs
lcl_fvs')