    {-# LANGUAGE TypeFamilies #-}

    {- |
    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 (
        annTopBindingsFreeVars,
        annBindingFreeVars
      ) where

    import GHC.Prelude

    import GHC.Stg.Syntax
    import GHC.Types.Id
    import GHC.Types.Var.Set
    import GHC.Types.Tickish ( GenTickish(Breakpoint) )
    import GHC.Utils.Misc

    import Data.Maybe ( mapMaybe )

    newtype Env
      = Env
      { Env -> IdSet
    locals :: IdSet
      }

    emptyEnv :: Env
    emptyEnv :: Env
    emptyEnv = IdSet -> Env
    Env IdSet
    emptyVarSet

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

    -- | Annotates a top-level STG binding group with its free variables.
    annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
    annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
    annTopBindingsFreeVars = (StgTopBinding -> CgStgTopBinding)
    -> [StgTopBinding] -> [CgStgTopBinding]
    forall a b. (a -> b) -> [a] -> [b]
    map StgTopBinding -> CgStgTopBinding
    go
      where
        go :: StgTopBinding -> CgStgTopBinding
    go (StgTopStringLit Var
    id ByteString
    bs) = Var -> ByteString -> CgStgTopBinding
    forall (pass :: StgPass).
    Var -> ByteString -> GenStgTopBinding pass
    StgTopStringLit Var
    id ByteString
    bs
        go (StgTopLifted GenStgBinding 'Vanilla
    bind)
          = GenStgBinding 'CodeGen -> CgStgTopBinding
    forall (pass :: StgPass).
    GenStgBinding pass -> GenStgTopBinding pass
    StgTopLifted (GenStgBinding 'Vanilla -> GenStgBinding 'CodeGen
    annBindingFreeVars GenStgBinding 'Vanilla
    bind)

    -- | Annotates an STG binding with its free variables.
    annBindingFreeVars :: StgBinding -> CgStgBinding
    annBindingFreeVars :: GenStgBinding 'Vanilla -> GenStgBinding 'CodeGen
    annBindingFreeVars = (GenStgBinding 'CodeGen, DIdSet) -> GenStgBinding 'CodeGen
    forall a b. (a, b) -> a
    fst ((GenStgBinding 'CodeGen, DIdSet) -> GenStgBinding 'CodeGen)
    -> (GenStgBinding 'Vanilla -> (GenStgBinding 'CodeGen, DIdSet))
    -> GenStgBinding 'Vanilla
    -> GenStgBinding 'CodeGen
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Env
    -> DIdSet
    -> GenStgBinding 'Vanilla
    -> (GenStgBinding 'CodeGen, DIdSet)
    binding Env
    emptyEnv DIdSet
    emptyDVarSet

    boundIds :: StgBinding -> [Id]
    boundIds :: GenStgBinding 'Vanilla -> [Var]
    boundIds (StgNonRec BinderP 'Vanilla
    b GenStgRhs 'Vanilla
    _) = [Var
    BinderP 'Vanilla
    b]
    boundIds (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
    pairs)  = ((Var, GenStgRhs 'Vanilla) -> Var)
    -> [(Var, GenStgRhs 'Vanilla)] -> [Var]
    forall a b. (a -> b) -> [a] -> [b]
    map (Var, GenStgRhs 'Vanilla) -> Var
    forall a b. (a, b) -> a
    fst [(Var, GenStgRhs 'Vanilla)]
    [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
    pairs

    -- Note [Tracking local binders]
    -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    -- 'locals' contains non-toplevel, non-imported binders.
    -- We maintain the set in 'expr', 'alt' and 'rhs', which are the only
    -- places where new local binders are introduced.
    -- Why do it there rather than in 'binding'? Two reasons:
    --
    --   1. We call 'binding' from 'annTopBindingsFreeVars', which would
    --      add top-level bindings to the 'locals' set.
    --   2. In the let(-no-escape) case, we need to extend the environment
    --      prior to analysing the body, but we also need the fvs from the
    --      body to analyse the RHSs. No way to do this without some
    --      knot-tying.

    -- | This makes sure that only local, non-global free vars make it into the set.
    mkFreeVarSet :: Env -> [Id] -> DIdSet
    mkFreeVarSet :: Env -> [Var] -> DIdSet
    mkFreeVarSet Env
    env = [Var] -> DIdSet
    mkDVarSet ([Var] -> DIdSet) -> ([Var] -> [Var]) -> [Var] -> DIdSet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Var -> Bool) -> [Var] -> [Var]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Var -> IdSet -> Bool
    `elemVarSet` Env -> IdSet
    locals Env
    env)

    args :: Env -> [StgArg] -> DIdSet
    args :: Env -> [StgArg] -> DIdSet
    args Env
    env = Env -> [Var] -> DIdSet
    mkFreeVarSet Env
    env ([Var] -> DIdSet) -> ([StgArg] -> [Var]) -> [StgArg] -> DIdSet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (StgArg -> Maybe Var) -> [StgArg] -> [Var]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe StgArg -> Maybe Var
    f
      where
        f :: StgArg -> Maybe Var
    f (StgVarArg Var
    occ) = Var -> Maybe Var
    forall a. a -> Maybe a
    Just Var
    occ
        f StgArg
    _               = Maybe Var
    forall a. Maybe a
    Nothing

    binding :: Env -> DIdSet -> StgBinding -> (CgStgBinding, DIdSet)
    binding :: Env
    -> DIdSet
    -> GenStgBinding 'Vanilla
    -> (GenStgBinding 'CodeGen, DIdSet)
    binding Env
    env DIdSet
    body_fv (StgNonRec BinderP 'Vanilla
    bndr GenStgRhs 'Vanilla
    r) = (BinderP 'CodeGen -> GenStgRhs 'CodeGen -> GenStgBinding 'CodeGen
    forall (pass :: StgPass).
    BinderP pass -> GenStgRhs pass -> GenStgBinding pass
    StgNonRec BinderP 'CodeGen
    BinderP 'Vanilla
    bndr GenStgRhs 'CodeGen
    r', DIdSet
    fvs)
      where
        -- See Note [Tracking local binders]
        (GenStgRhs 'CodeGen
    r', DIdSet
    rhs_fvs) = Env -> GenStgRhs 'Vanilla -> (GenStgRhs 'CodeGen, DIdSet)
    rhs Env
    env GenStgRhs 'Vanilla
    r
        fvs :: DIdSet
    fvs = DIdSet -> Var -> DIdSet
    delDVarSet DIdSet
    body_fv Var
    BinderP 'Vanilla
    bndr DIdSet -> DIdSet -> DIdSet
    `unionDVarSet` DIdSet
    rhs_fvs
    binding Env
    env DIdSet
    body_fv (StgRec [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
    pairs) = ([(BinderP 'CodeGen, GenStgRhs 'CodeGen)] -> GenStgBinding 'CodeGen
    forall (pass :: StgPass).
    [(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
    StgRec [(Var, GenStgRhs 'CodeGen)]
    [(BinderP 'CodeGen, GenStgRhs 'CodeGen)]
    pairs', DIdSet
    fvs)
      where
        -- See Note [Tracking local binders]
        bndrs :: [Var]
    bndrs = ((Var, GenStgRhs 'Vanilla) -> Var)
    -> [(Var, GenStgRhs 'Vanilla)] -> [Var]
    forall a b. (a -> b) -> [a] -> [b]
    map (Var, GenStgRhs 'Vanilla) -> Var
    forall a b. (a, b) -> a
    fst [(Var, GenStgRhs 'Vanilla)]
    [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
    pairs
        ([GenStgRhs 'CodeGen]
    rhss, [DIdSet]
    rhs_fvss) = ((Var, GenStgRhs 'Vanilla) -> (GenStgRhs 'CodeGen, DIdSet))
    -> [(Var, GenStgRhs 'Vanilla)] -> ([GenStgRhs 'CodeGen], [DIdSet])
    forall a b c. (a -> (b, c)) -> [a] -> ([b], [c])
    mapAndUnzip (Env -> GenStgRhs 'Vanilla -> (GenStgRhs 'CodeGen, DIdSet)
    rhs Env
    env (GenStgRhs 'Vanilla -> (GenStgRhs 'CodeGen, DIdSet))
    -> ((Var, GenStgRhs 'Vanilla) -> GenStgRhs 'Vanilla)
    -> (Var, GenStgRhs 'Vanilla)
    -> (GenStgRhs 'CodeGen, DIdSet)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Var, GenStgRhs 'Vanilla) -> GenStgRhs 'Vanilla
    forall a b. (a, b) -> b
    snd) [(Var, GenStgRhs 'Vanilla)]
    [(BinderP 'Vanilla, GenStgRhs 'Vanilla)]
    pairs
        pairs' :: [(Var, GenStgRhs 'CodeGen)]
    pairs' = [Var] -> [GenStgRhs 'CodeGen] -> [(Var, GenStgRhs 'CodeGen)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Var]
    bndrs [GenStgRhs 'CodeGen]
    rhss
        fvs :: DIdSet
    fvs = DIdSet -> [Var] -> DIdSet
    delDVarSetList ([DIdSet] -> DIdSet
    unionDVarSets (DIdSet
    body_fvDIdSet -> [DIdSet] -> [DIdSet]
    forall a. a -> [a] -> [a]
    :[DIdSet]
    rhs_fvss)) [Var]
    bndrs

    expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
    expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
    expr Env
    env = StgExpr -> (CgStgExpr, DIdSet)
    go
      where
        go :: StgExpr -> (CgStgExpr, DIdSet)
    go (StgApp Var
    occ [StgArg]
    as)
          = (Var -> [StgArg] -> CgStgExpr
    forall (pass :: StgPass). Var -> [StgArg] -> GenStgExpr pass
    StgApp Var
    occ [StgArg]
    as, DIdSet -> DIdSet -> DIdSet
    unionDVarSet (Env -> [StgArg] -> DIdSet
    args Env
    env [StgArg]
    as) (Env -> [Var] -> DIdSet
    mkFreeVarSet Env
    env [Var
    occ]))
        go (StgLit Literal
    lit) = (Literal -> CgStgExpr
    forall (pass :: StgPass). Literal -> GenStgExpr pass
    StgLit Literal
    lit, DIdSet
    emptyDVarSet)
        go (StgConApp DataCon
    dc XConApp 'Vanilla
    n [StgArg]
    as [Type]
    tys) = (DataCon -> XConApp 'CodeGen -> [StgArg] -> [Type] -> CgStgExpr
    forall (pass :: StgPass).
    DataCon -> XConApp pass -> [StgArg] -> [Type] -> GenStgExpr pass
    StgConApp DataCon
    dc XConApp 'CodeGen
    XConApp 'Vanilla
    n [StgArg]
    as [Type]
    tys, Env -> [StgArg] -> DIdSet
    args Env
    env [StgArg]
    as)
        go (StgOpApp StgOp
    op [StgArg]
    as Type
    ty) = (StgOp -> [StgArg] -> Type -> CgStgExpr
    forall (pass :: StgPass).
    StgOp -> [StgArg] -> Type -> GenStgExpr pass
    StgOpApp StgOp
    op [StgArg]
    as Type
    ty, Env -> [StgArg] -> DIdSet
    args Env
    env [StgArg]
    as)
        go (StgCase StgExpr
    scrut BinderP 'Vanilla
    bndr AltType
    ty [GenStgAlt 'Vanilla]
    alts) = (CgStgExpr
    -> BinderP 'CodeGen -> AltType -> [GenStgAlt 'CodeGen] -> CgStgExpr
    forall (pass :: StgPass).
    GenStgExpr pass
    -> BinderP pass -> AltType -> [GenStgAlt pass] -> GenStgExpr pass
    StgCase CgStgExpr
    scrut' BinderP 'CodeGen
    BinderP 'Vanilla
    bndr AltType
    ty [(AltCon, [Var], CgStgExpr)]
    [GenStgAlt 'CodeGen]
    alts', DIdSet
    fvs)
          where
            (CgStgExpr
    scrut', DIdSet
    scrut_fvs) = StgExpr -> (CgStgExpr, DIdSet)
    go StgExpr
    scrut
            -- See Note [Tracking local binders]
            ([(AltCon, [Var], CgStgExpr)]
    alts', [DIdSet]
    alt_fvss) = ((AltCon, [Var], StgExpr) -> ((AltCon, [Var], CgStgExpr), DIdSet))
    -> [(AltCon, [Var], StgExpr)]
    -> ([(AltCon, [Var], CgStgExpr)], [DIdSet])
    forall a b c. (a -> (b, c)) -> [a] -> ([b], [c])
    mapAndUnzip (Env -> GenStgAlt 'Vanilla -> (GenStgAlt 'CodeGen, DIdSet)
    alt ([Var] -> Env -> Env
    addLocals [Var
    BinderP 'Vanilla
    bndr] Env
    env)) [(AltCon, [Var], StgExpr)]
    [GenStgAlt 'Vanilla]
    alts
            alt_fvs :: DIdSet
    alt_fvs = [DIdSet] -> DIdSet
    unionDVarSets [DIdSet]
    alt_fvss
            fvs :: DIdSet
    fvs = DIdSet -> Var -> DIdSet
    delDVarSet (DIdSet -> DIdSet -> DIdSet
    unionDVarSet DIdSet
    scrut_fvs DIdSet
    alt_fvs) Var
    BinderP 'Vanilla
    bndr
        go (StgLet XLet 'Vanilla
    ext GenStgBinding 'Vanilla
    bind StgExpr
    body) = (GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr)
    -> GenStgBinding 'Vanilla -> StgExpr -> (CgStgExpr, DIdSet)
    go_bind (XLet 'CodeGen -> GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr
    forall (pass :: StgPass).
    XLet pass
    -> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
    StgLet XLet 'CodeGen
    XLet 'Vanilla
    ext) GenStgBinding 'Vanilla
    bind StgExpr
    body
        go (StgLetNoEscape XLetNoEscape 'Vanilla
    ext GenStgBinding 'Vanilla
    bind StgExpr
    body) = (GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr)
    -> GenStgBinding 'Vanilla -> StgExpr -> (CgStgExpr, DIdSet)
    go_bind (XLetNoEscape 'CodeGen
    -> GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr
    forall (pass :: StgPass).
    XLetNoEscape pass
    -> GenStgBinding pass -> GenStgExpr pass -> GenStgExpr pass
    StgLetNoEscape XLetNoEscape 'CodeGen
    XLetNoEscape 'Vanilla
    ext) GenStgBinding 'Vanilla
    bind StgExpr
    body
        go (StgTick StgTickish
    tick StgExpr
    e) = (StgTickish -> CgStgExpr -> CgStgExpr
    forall (pass :: StgPass).
    StgTickish -> GenStgExpr pass -> GenStgExpr pass
    StgTick StgTickish
    tick CgStgExpr
    e', DIdSet
    fvs')
          where
            (CgStgExpr
    e', DIdSet
    fvs) = StgExpr -> (CgStgExpr, DIdSet)
    go StgExpr
    e
            fvs' :: DIdSet
    fvs' = DIdSet -> DIdSet -> DIdSet
    unionDVarSet (StgTickish -> DIdSet
    forall {pass :: TickishPass}.
    (XTickishId pass ~ Var) =>
    GenTickish pass -> DIdSet
    tickish StgTickish
    tick) DIdSet
    fvs
            tickish :: GenTickish pass -> DIdSet
    tickish (Breakpoint XBreakpoint pass
    _ Int
    _ [XTickishId pass]
    ids) = [Var] -> DIdSet
    mkDVarSet [Var]
    [XTickishId pass]
    ids
            tickish GenTickish pass
    _                    = DIdSet
    emptyDVarSet

        go_bind :: (GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr)
    -> GenStgBinding 'Vanilla -> StgExpr -> (CgStgExpr, DIdSet)
    go_bind GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr
    dc GenStgBinding 'Vanilla
    bind StgExpr
    body = (GenStgBinding 'CodeGen -> CgStgExpr -> CgStgExpr
    dc GenStgBinding 'CodeGen
    bind' CgStgExpr
    body', DIdSet
    fvs)
          where
            -- See Note [Tracking local binders]
            env' :: Env
    env' = [Var] -> Env -> Env
    addLocals (GenStgBinding 'Vanilla -> [Var]
    boundIds GenStgBinding 'Vanilla
    bind) Env
    env
            (CgStgExpr
    body', DIdSet
    body_fvs) = Env -> StgExpr -> (CgStgExpr, DIdSet)
    expr Env
    env' StgExpr
    body
            (GenStgBinding 'CodeGen
    bind', DIdSet
    fvs) = Env
    -> DIdSet
    -> GenStgBinding 'Vanilla
    -> (GenStgBinding 'CodeGen, DIdSet)
    binding Env
    env' DIdSet
    body_fvs GenStgBinding 'Vanilla
    bind

    rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
    rhs :: Env -> GenStgRhs 'Vanilla -> (GenStgRhs 'CodeGen, DIdSet)
    rhs Env
    env (StgRhsClosure XRhsClosure 'Vanilla
    _ CostCentreStack
    ccs UpdateFlag
    uf [BinderP 'Vanilla]
    bndrs StgExpr
    body)
      = (XRhsClosure 'CodeGen
    -> CostCentreStack
    -> UpdateFlag
    -> [BinderP 'CodeGen]
    -> CgStgExpr
    -> GenStgRhs 'CodeGen
    forall (pass :: StgPass).
    XRhsClosure pass
    -> CostCentreStack
    -> UpdateFlag
    -> [BinderP pass]
    -> GenStgExpr pass
    -> GenStgRhs pass
    StgRhsClosure DIdSet
    XRhsClosure 'CodeGen
    fvs CostCentreStack
    ccs UpdateFlag
    uf [BinderP 'CodeGen]
    [BinderP 'Vanilla]
    bndrs CgStgExpr
    body', DIdSet
    fvs)
      where
        -- See Note [Tracking local binders]
        (CgStgExpr
    body', DIdSet
    body_fvs) = Env -> StgExpr -> (CgStgExpr, DIdSet)
    expr ([Var] -> Env -> Env
    addLocals [Var]
    [BinderP 'Vanilla]
    bndrs Env
    env) StgExpr
    body
        fvs :: DIdSet
    fvs = DIdSet -> [Var] -> DIdSet
    delDVarSetList DIdSet
    body_fvs [Var]
    [BinderP 'Vanilla]
    bndrs
    rhs Env
    env (StgRhsCon CostCentreStack
    ccs DataCon
    dc ConstructorNumber
    mu [StgTickish]
    ts [StgArg]
    as) = (CostCentreStack
    -> DataCon
    -> ConstructorNumber
    -> [StgTickish]
    -> [StgArg]
    -> GenStgRhs 'CodeGen
    forall (pass :: StgPass).
    CostCentreStack
    -> DataCon
    -> ConstructorNumber
    -> [StgTickish]
    -> [StgArg]
    -> GenStgRhs pass
    StgRhsCon CostCentreStack
    ccs DataCon
    dc ConstructorNumber
    mu [StgTickish]
    ts [StgArg]
    as, Env -> [StgArg] -> DIdSet
    args Env
    env [StgArg]
    as)

    alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
    alt :: Env -> GenStgAlt 'Vanilla -> (GenStgAlt 'CodeGen, DIdSet)
    alt Env
    env (AltCon
    con, [BinderP 'Vanilla]
    bndrs, StgExpr
    e) = ((AltCon
    con, [BinderP 'CodeGen]
    [BinderP 'Vanilla]
    bndrs, CgStgExpr
    e'), DIdSet
    fvs)
      where
        -- See Note [Tracking local binders]
        (CgStgExpr
    e', DIdSet
    rhs_fvs) = Env -> StgExpr -> (CgStgExpr, DIdSet)
    expr ([Var] -> Env -> Env
    addLocals [Var]
    [BinderP 'Vanilla]
    bndrs Env
    env) StgExpr
    e
        fvs :: DIdSet
    fvs = DIdSet -> [Var] -> DIdSet
    delDVarSetList DIdSet
    rhs_fvs [Var]
    [BinderP 'Vanilla]
    bndrs