    module GHC.Core.Opt.Exitify ( exitifyProgram ) where

    {-
    Note [Exitification]
    ~~~~~~~~~~~~~~~~~~~~

    This module implements Exitification. The goal is to pull as much code out of
    recursive functions as possible, as the simplifier is better at inlining into
    call-sites that are not in recursive functions.

    Example:

      let t = foo bar
      joinrec go 0     x y = t (x*x)
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    We’d like to inline `t`, but that does not happen: Because t is a thunk and is
    used in a recursive function, doing so might lose sharing in general. In
    this case, however, `t` is on the _exit path_ of `go`, so called at most once.
    How do we make this clearly visible to the simplifier?

    A code path (i.e., an expression in a tail-recursive position) in a recursive
    function is an exit path if it does not contain a recursive call. We can bind
    this expression outside the recursive function, as a join-point.

    Example result:

      let t = foo bar
      join exit x = t (x*x)
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    Now `t` is no longer in a recursive function, and good things happen!
    -}

    import GHC.Prelude
    import GHC.Types.Var
    import GHC.Types.Id
    import GHC.Types.Id.Info
    import GHC.Core
    import GHC.Core.Utils
    import GHC.Utils.Monad.State.Strict
    import GHC.Builtin.Uniques
    import GHC.Types.Var.Set
    import GHC.Types.Var.Env
    import GHC.Core.FVs
    import GHC.Data.FastString
    import GHC.Core.Type
    import GHC.Utils.Misc( mapSnd )

    import Data.Bifunctor
    import Control.Monad

    -- | Traverses the AST, simply to find all joinrecs and call 'exitify' on them.
    -- The really interesting function is exitifyRec
    exitifyProgram :: CoreProgram -> CoreProgram
    exitifyProgram :: CoreProgram -> CoreProgram
    exitifyProgram CoreProgram
    binds = (Bind JoinId -> Bind JoinId) -> CoreProgram -> CoreProgram
    forall a b. (a -> b) -> [a] -> [b]
    map Bind JoinId -> Bind JoinId
    goTopLvl CoreProgram
    binds
      where
        goTopLvl :: Bind JoinId -> Bind JoinId
    goTopLvl (NonRec JoinId
    v CoreExpr
    e) = JoinId -> CoreExpr -> Bind JoinId
    forall b. b -> Expr b -> Bind b
    NonRec JoinId
    v (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope_toplvl CoreExpr
    e)
        goTopLvl (Rec [(JoinId, CoreExpr)]
    pairs) = [(JoinId, CoreExpr)] -> Bind JoinId
    forall b. [(b, Expr b)] -> Bind b
    Rec (((JoinId, CoreExpr) -> (JoinId, CoreExpr))
    -> [(JoinId, CoreExpr)] -> [(JoinId, CoreExpr)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((CoreExpr -> CoreExpr) -> (JoinId, CoreExpr) -> (JoinId, CoreExpr)
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope_toplvl)) [(JoinId, CoreExpr)]
    pairs)
          -- Top-level bindings are never join points

        in_scope_toplvl :: InScopeSet
    in_scope_toplvl = InScopeSet
    emptyInScopeSet InScopeSet -> [JoinId] -> InScopeSet
    `extendInScopeSetList` CoreProgram -> [JoinId]
    forall b. [Bind b] -> [b]
    bindersOfBinds CoreProgram
    binds

        go :: InScopeSet -> CoreExpr -> CoreExpr
        go :: InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    _    e :: CoreExpr
    e@(Var{})       = CoreExpr
    e
        go InScopeSet
    _    e :: CoreExpr
    e@(Lit {})      = CoreExpr
    e
        go InScopeSet
    _    e :: CoreExpr
    e@(Type {})     = CoreExpr
    e
        go InScopeSet
    _    e :: CoreExpr
    e@(Coercion {}) = CoreExpr
    e
        go InScopeSet
    in_scope (Cast CoreExpr
    e' CoercionR
    c) = CoreExpr -> CoercionR -> CoreExpr
    forall b. Expr b -> CoercionR -> Expr b
    Cast (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    e') CoercionR
    c
        go InScopeSet
    in_scope (Tick CoreTickish
    t CoreExpr
    e') = CoreTickish -> CoreExpr -> CoreExpr
    forall b. CoreTickish -> Expr b -> Expr b
    Tick CoreTickish
    t (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    e')
        go InScopeSet
    in_scope (App CoreExpr
    e1 CoreExpr
    e2) = CoreExpr -> CoreExpr -> CoreExpr
    forall b. Expr b -> Expr b -> Expr b
    App (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    e1) (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    e2)

        go InScopeSet
    in_scope (Lam JoinId
    v CoreExpr
    e')
          = JoinId -> CoreExpr -> CoreExpr
    forall b. b -> Expr b -> Expr b
    Lam JoinId
    v (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope' CoreExpr
    e')
          where in_scope' :: InScopeSet
    in_scope' = InScopeSet
    in_scope InScopeSet -> JoinId -> InScopeSet
    `extendInScopeSet` JoinId
    v

        go InScopeSet
    in_scope (Case CoreExpr
    scrut JoinId
    bndr Type
    ty [Alt JoinId]
    alts)
          = CoreExpr -> JoinId -> Type -> [Alt JoinId] -> CoreExpr
    forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
    Case (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    scrut) JoinId
    bndr Type
    ty ((Alt JoinId -> Alt JoinId) -> [Alt JoinId] -> [Alt JoinId]
    forall a b. (a -> b) -> [a] -> [b]
    map Alt JoinId -> Alt JoinId
    go_alt [Alt JoinId]
    alts)
          where
            in_scope1 :: InScopeSet
    in_scope1 = InScopeSet
    in_scope InScopeSet -> JoinId -> InScopeSet
    `extendInScopeSet` JoinId
    bndr
            go_alt :: Alt JoinId -> Alt JoinId
    go_alt (Alt AltCon
    dc [JoinId]
    pats CoreExpr
    rhs) = AltCon -> [JoinId] -> CoreExpr -> Alt JoinId
    forall b. AltCon -> [b] -> Expr b -> Alt b
    Alt AltCon
    dc [JoinId]
    pats (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope' CoreExpr
    rhs)
               where in_scope' :: InScopeSet
    in_scope' = InScopeSet
    in_scope1 InScopeSet -> [JoinId] -> InScopeSet
    `extendInScopeSetList` [JoinId]
    pats

        go InScopeSet
    in_scope (Let (NonRec JoinId
    bndr CoreExpr
    rhs) CoreExpr
    body)
          = Bind JoinId -> CoreExpr -> CoreExpr
    forall b. Bind b -> Expr b -> Expr b
    Let (JoinId -> CoreExpr -> Bind JoinId
    forall b. b -> Expr b -> Bind b
    NonRec JoinId
    bndr (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope CoreExpr
    rhs)) (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope' CoreExpr
    body)
          where
            in_scope' :: InScopeSet
    in_scope' = InScopeSet
    in_scope InScopeSet -> JoinId -> InScopeSet
    `extendInScopeSet` JoinId
    bndr

        go InScopeSet
    in_scope (Let (Rec [(JoinId, CoreExpr)]
    pairs) CoreExpr
    body)
          | Bool
    is_join_rec = CoreProgram -> CoreExpr -> CoreExpr
    forall b. [Bind b] -> Expr b -> Expr b
    mkLets (InScopeSet -> [(JoinId, CoreExpr)] -> CoreProgram
    exitifyRec InScopeSet
    in_scope' [(JoinId, CoreExpr)]
    pairs') CoreExpr
    body'
          | Bool
    otherwise   = Bind JoinId -> CoreExpr -> CoreExpr
    forall b. Bind b -> Expr b -> Expr b
    Let ([(JoinId, CoreExpr)] -> Bind JoinId
    forall b. [(b, Expr b)] -> Bind b
    Rec [(JoinId, CoreExpr)]
    pairs') CoreExpr
    body'
          where
            is_join_rec :: Bool
    is_join_rec = ((JoinId, CoreExpr) -> Bool) -> [(JoinId, CoreExpr)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (JoinId -> Bool
    isJoinId (JoinId -> Bool)
    -> ((JoinId, CoreExpr) -> JoinId) -> (JoinId, CoreExpr) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (JoinId, CoreExpr) -> JoinId
    forall a b. (a, b) -> a
    fst) [(JoinId, CoreExpr)]
    pairs
            in_scope' :: InScopeSet
    in_scope'   = InScopeSet
    in_scope InScopeSet -> [JoinId] -> InScopeSet
    `extendInScopeSetList` Bind JoinId -> [JoinId]
    forall b. Bind b -> [b]
    bindersOf ([(JoinId, CoreExpr)] -> Bind JoinId
    forall b. [(b, Expr b)] -> Bind b
    Rec [(JoinId, CoreExpr)]
    pairs)
            pairs' :: [(JoinId, CoreExpr)]
    pairs'      = (CoreExpr -> CoreExpr)
    -> [(JoinId, CoreExpr)] -> [(JoinId, CoreExpr)]
    forall b c a. (b -> c) -> [(a, b)] -> [(a, c)]
    mapSnd (InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope') [(JoinId, CoreExpr)]
    pairs
            body' :: CoreExpr
    body'       = InScopeSet -> CoreExpr -> CoreExpr
    go InScopeSet
    in_scope' CoreExpr
    body

    -- | State Monad used inside `exitify`
    type ExitifyM =  State [(JoinId, CoreExpr)]

    -- | Given a recursive group of a joinrec, identifies “exit paths” and binds them as
    --   join-points outside the joinrec.
    exitifyRec :: InScopeSet -> [(Var,CoreExpr)] -> [CoreBind]
    exitifyRec :: InScopeSet -> [(JoinId, CoreExpr)] -> CoreProgram
    exitifyRec InScopeSet
    in_scope [(JoinId, CoreExpr)]
    pairs
      = [ JoinId -> CoreExpr -> Bind JoinId
    forall b. b -> Expr b -> Bind b
    NonRec JoinId
    xid CoreExpr
    rhs | (JoinId
    xid,CoreExpr
    rhs) <- [(JoinId, CoreExpr)]
    exits ] CoreProgram -> CoreProgram -> CoreProgram
    forall a. [a] -> [a] -> [a]
    ++ [[(JoinId, CoreExpr)] -> Bind JoinId
    forall b. [(b, Expr b)] -> Bind b
    Rec [(JoinId, CoreExpr)]
    pairs']
      where
        -- We need the set of free variables of many subexpressions here, so
        -- annotate the AST with them
        -- see Note [Calculating free variables]
        ann_pairs :: [(JoinId, CoreExprWithFVs)]
    ann_pairs = ((JoinId, CoreExpr) -> (JoinId, CoreExprWithFVs))
    -> [(JoinId, CoreExpr)] -> [(JoinId, CoreExprWithFVs)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((CoreExpr -> CoreExprWithFVs)
    -> (JoinId, CoreExpr) -> (JoinId, CoreExprWithFVs)
    forall b c a. (b -> c) -> (a, b) -> (a, c)
    forall (p :: * -> * -> *) b c a.
    Bifunctor p =>
    (b -> c) -> p a b -> p a c
    second CoreExpr -> CoreExprWithFVs
    freeVars) [(JoinId, CoreExpr)]
    pairs

        -- Which are the recursive calls?
        recursive_calls :: VarSet
    recursive_calls = [JoinId] -> VarSet
    mkVarSet ([JoinId] -> VarSet) -> [JoinId] -> VarSet
    forall a b. (a -> b) -> a -> b
    $ ((JoinId, CoreExpr) -> JoinId) -> [(JoinId, CoreExpr)] -> [JoinId]
    forall a b. (a -> b) -> [a] -> [b]
    map (JoinId, CoreExpr) -> JoinId
    forall a b. (a, b) -> a
    fst [(JoinId, CoreExpr)]
    pairs

        ([(JoinId, CoreExpr)]
    pairs',[(JoinId, CoreExpr)]
    exits) = (State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    -> [(JoinId, CoreExpr)]
    -> ([(JoinId, CoreExpr)], [(JoinId, CoreExpr)])
    forall s a. State s a -> s -> (a, s)
    `runState` []) (State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
     -> ([(JoinId, CoreExpr)], [(JoinId, CoreExpr)]))
    -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    -> ([(JoinId, CoreExpr)], [(JoinId, CoreExpr)])
    forall a b. (a -> b) -> a -> b
    $
            [(JoinId, CoreExprWithFVs)]
    -> ((JoinId, CoreExprWithFVs)
        -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
    -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [(JoinId, CoreExprWithFVs)]
    ann_pairs (((JoinId, CoreExprWithFVs)
      -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
     -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)])
    -> ((JoinId, CoreExprWithFVs)
        -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
    -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall a b. (a -> b) -> a -> b
    $ \(JoinId
    x,CoreExprWithFVs
    rhs) -> do
                -- go past the lambdas of the join point
                let ([JoinId]
    args, CoreExprWithFVs
    body) = Int -> CoreExprWithFVs -> ([JoinId], CoreExprWithFVs)
    forall bndr annot.
    Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
    collectNAnnBndrs (JoinId -> Int
    idJoinArity JoinId
    x) CoreExprWithFVs
    rhs
                CoreExpr
    body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go [JoinId]
    args CoreExprWithFVs
    body
                let rhs' :: CoreExpr
    rhs' = [JoinId] -> CoreExpr -> CoreExpr
    forall b. [b] -> Expr b -> Expr b
    mkLams [JoinId]
    args CoreExpr
    body'
                (JoinId, CoreExpr) -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr)
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (JoinId
    x, CoreExpr
    rhs')

        ---------------------
        -- 'go' is the main working function.
        -- It goes through the RHS (tail-call positions only),
        -- checks if there are no more recursive calls, if so, abstracts over
        -- variables bound on the way and lifts it out as a join point.
        --
        -- ExitifyM is a state monad to keep track of floated binds
        go :: [Var]           -- Variables that are in-scope here, but
                              -- not in scope at the joinrec; that is,
                              -- we must potentially abstract over them.
                              -- Invariant: they are kept in dependency order
           -> CoreExprWithFVs -- Current expression in tail position
           -> ExitifyM CoreExpr

        -- We first look at the expression (no matter what it shape is)
        -- and determine if we can turn it into a exit join point
        go :: [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go [JoinId]
    captured CoreExprWithFVs
    ann_e
            | -- An exit expression has no recursive calls
              let fvs :: VarSet
    fvs = DVarSet -> VarSet
    dVarSetToVarSet (CoreExprWithFVs -> DVarSet
    freeVarsOf CoreExprWithFVs
    ann_e)
            , VarSet -> VarSet -> Bool
    disjointVarSet VarSet
    fvs VarSet
    recursive_calls
            = [JoinId] -> CoreExpr -> VarSet -> ExitifyM CoreExpr
    go_exit [JoinId]
    captured (CoreExprWithFVs -> CoreExpr
    forall bndr annot. AnnExpr bndr annot -> Expr bndr
    deAnnotate CoreExprWithFVs
    ann_e) VarSet
    fvs

        -- We could not turn it into a exit join point. So now recurse
        -- into all expression where eligible exit join points might sit,
        -- i.e. into all tail-call positions:

        -- Case right hand sides are in tail-call position
        go [JoinId]
    captured (DVarSet
    _, AnnCase CoreExprWithFVs
    scrut JoinId
    bndr Type
    ty [AnnAlt JoinId DVarSet]
    alts) = do
            [Alt JoinId]
    alts' <- [AnnAlt JoinId DVarSet]
    -> (AnnAlt JoinId DVarSet
        -> State [(JoinId, CoreExpr)] (Alt JoinId))
    -> State [(JoinId, CoreExpr)] [Alt JoinId]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [AnnAlt JoinId DVarSet]
    alts ((AnnAlt JoinId DVarSet -> State [(JoinId, CoreExpr)] (Alt JoinId))
     -> State [(JoinId, CoreExpr)] [Alt JoinId])
    -> (AnnAlt JoinId DVarSet
        -> State [(JoinId, CoreExpr)] (Alt JoinId))
    -> State [(JoinId, CoreExpr)] [Alt JoinId]
    forall a b. (a -> b) -> a -> b
    $ \(AnnAlt AltCon
    dc [JoinId]
    pats CoreExprWithFVs
    rhs) -> do
                CoreExpr
    rhs' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId
    bndr] [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId]
    pats) CoreExprWithFVs
    rhs
                Alt JoinId -> State [(JoinId, CoreExpr)] (Alt JoinId)
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (AltCon -> [JoinId] -> CoreExpr -> Alt JoinId
    forall b. AltCon -> [b] -> Expr b -> Alt b
    Alt AltCon
    dc [JoinId]
    pats CoreExpr
    rhs')
            CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExpr -> ExitifyM CoreExpr) -> CoreExpr -> ExitifyM CoreExpr
    forall a b. (a -> b) -> a -> b
    $ CoreExpr -> JoinId -> Type -> [Alt JoinId] -> CoreExpr
    forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
    Case (CoreExprWithFVs -> CoreExpr
    forall bndr annot. AnnExpr bndr annot -> Expr bndr
    deAnnotate CoreExprWithFVs
    scrut) JoinId
    bndr Type
    ty [Alt JoinId]
    alts'

        go [JoinId]
    captured (DVarSet
    _, AnnLet AnnBind JoinId DVarSet
    ann_bind CoreExprWithFVs
    body)
            -- join point, RHS and body are in tail-call position
            | AnnNonRec JoinId
    j CoreExprWithFVs
    rhs <- AnnBind JoinId DVarSet
    ann_bind
            , Just Int
    join_arity <- JoinId -> Maybe Int
    isJoinId_maybe JoinId
    j
            = do let ([JoinId]
    params, CoreExprWithFVs
    join_body) = Int -> CoreExprWithFVs -> ([JoinId], CoreExprWithFVs)
    forall bndr annot.
    Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
    collectNAnnBndrs Int
    join_arity CoreExprWithFVs
    rhs
                 CoreExpr
    join_body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId]
    params) CoreExprWithFVs
    join_body
                 let rhs' :: CoreExpr
    rhs' = [JoinId] -> CoreExpr -> CoreExpr
    forall b. [b] -> Expr b -> Expr b
    mkLams [JoinId]
    params CoreExpr
    join_body'
                 CoreExpr
    body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId
    j]) CoreExprWithFVs
    body
                 CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExpr -> ExitifyM CoreExpr) -> CoreExpr -> ExitifyM CoreExpr
    forall a b. (a -> b) -> a -> b
    $ Bind JoinId -> CoreExpr -> CoreExpr
    forall b. Bind b -> Expr b -> Expr b
    Let (JoinId -> CoreExpr -> Bind JoinId
    forall b. b -> Expr b -> Bind b
    NonRec JoinId
    j CoreExpr
    rhs') CoreExpr
    body'

            -- rec join point, RHSs and body are in tail-call position
            | AnnRec [(JoinId, CoreExprWithFVs)]
    pairs <- AnnBind JoinId DVarSet
    ann_bind
            , JoinId -> Bool
    isJoinId ((JoinId, CoreExprWithFVs) -> JoinId
    forall a b. (a, b) -> a
    fst ([(JoinId, CoreExprWithFVs)] -> (JoinId, CoreExprWithFVs)
    forall a. HasCallStack => [a] -> a
    head [(JoinId, CoreExprWithFVs)]
    pairs))
            = do let js :: [JoinId]
    js = ((JoinId, CoreExprWithFVs) -> JoinId)
    -> [(JoinId, CoreExprWithFVs)] -> [JoinId]
    forall a b. (a -> b) -> [a] -> [b]
    map (JoinId, CoreExprWithFVs) -> JoinId
    forall a b. (a, b) -> a
    fst [(JoinId, CoreExprWithFVs)]
    pairs
                 [(JoinId, CoreExpr)]
    pairs' <- [(JoinId, CoreExprWithFVs)]
    -> ((JoinId, CoreExprWithFVs)
        -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
    -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [(JoinId, CoreExprWithFVs)]
    pairs (((JoinId, CoreExprWithFVs)
      -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
     -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)])
    -> ((JoinId, CoreExprWithFVs)
        -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr))
    -> State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall a b. (a -> b) -> a -> b
    $ \(JoinId
    j,CoreExprWithFVs
    rhs) -> do
                     let join_arity :: Int
    join_arity = JoinId -> Int
    idJoinArity JoinId
    j
                         ([JoinId]
    params, CoreExprWithFVs
    join_body) = Int -> CoreExprWithFVs -> ([JoinId], CoreExprWithFVs)
    forall bndr annot.
    Int -> AnnExpr bndr annot -> ([bndr], AnnExpr bndr annot)
    collectNAnnBndrs Int
    join_arity CoreExprWithFVs
    rhs
                     CoreExpr
    join_body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId]
    js [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId]
    params) CoreExprWithFVs
    join_body
                     let rhs' :: CoreExpr
    rhs' = [JoinId] -> CoreExpr -> CoreExpr
    forall b. [b] -> Expr b -> Expr b
    mkLams [JoinId]
    params CoreExpr
    join_body'
                     (JoinId, CoreExpr) -> State [(JoinId, CoreExpr)] (JoinId, CoreExpr)
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (JoinId
    j, CoreExpr
    rhs')
                 CoreExpr
    body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ [JoinId]
    js) CoreExprWithFVs
    body
                 CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExpr -> ExitifyM CoreExpr) -> CoreExpr -> ExitifyM CoreExpr
    forall a b. (a -> b) -> a -> b
    $ Bind JoinId -> CoreExpr -> CoreExpr
    forall b. Bind b -> Expr b -> Expr b
    Let ([(JoinId, CoreExpr)] -> Bind JoinId
    forall b. [(b, Expr b)] -> Bind b
    Rec [(JoinId, CoreExpr)]
    pairs') CoreExpr
    body'

            -- normal Let, only the body is in tail-call position
            | Bool
    otherwise
            = do CoreExpr
    body' <- [JoinId] -> CoreExprWithFVs -> ExitifyM CoreExpr
    go ([JoinId]
    captured [JoinId] -> [JoinId] -> [JoinId]
    forall a. [a] -> [a] -> [a]
    ++ Bind JoinId -> [JoinId]
    forall b. Bind b -> [b]
    bindersOf Bind JoinId
    bind ) CoreExprWithFVs
    body
                 CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExpr -> ExitifyM CoreExpr) -> CoreExpr -> ExitifyM CoreExpr
    forall a b. (a -> b) -> a -> b
    $ Bind JoinId -> CoreExpr -> CoreExpr
    forall b. Bind b -> Expr b -> Expr b
    Let Bind JoinId
    bind CoreExpr
    body'
          where bind :: Bind JoinId
    bind = AnnBind JoinId DVarSet -> Bind JoinId
    forall b annot. AnnBind b annot -> Bind b
    deAnnBind AnnBind JoinId DVarSet
    ann_bind

        -- Cannot be turned into an exit join point, but also has no
        -- tail-call subexpression. Nothing to do here.
        go [JoinId]
    _ CoreExprWithFVs
    ann_e = CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExprWithFVs -> CoreExpr
    forall bndr annot. AnnExpr bndr annot -> Expr bndr
    deAnnotate CoreExprWithFVs
    ann_e)

        ---------------------
        go_exit :: [Var]      -- Variables captured locally
                -> CoreExpr   -- An exit expression
                -> VarSet     -- Free vars of the expression
                -> ExitifyM CoreExpr
        -- go_exit deals with a tail expression that is floatable
        -- out as an exit point; that is, it mentions no recursive calls
        go_exit :: [JoinId] -> CoreExpr -> VarSet -> ExitifyM CoreExpr
    go_exit [JoinId]
    captured CoreExpr
    e VarSet
    fvs
          -- Do not touch an expression that is already a join jump where all arguments
          -- are captured variables. See Note [Idempotency]
          -- But _do_ float join jumps with interesting arguments.
          -- See Note [Jumps can be interesting]
          | (Var JoinId
    f, [CoreExpr]
    args) <- CoreExpr -> (CoreExpr, [CoreExpr])
    forall b. Expr b -> (Expr b, [Expr b])
    collectArgs CoreExpr
    e
          , JoinId -> Bool
    isJoinId JoinId
    f
          , (CoreExpr -> Bool) -> [CoreExpr] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all CoreExpr -> Bool
    isCapturedVarArg [CoreExpr]
    args
          = CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return CoreExpr
    e

          -- Do not touch a boring expression (see Note [Interesting expression])
          | Bool -> Bool
    not Bool
    is_interesting
          = CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return CoreExpr
    e

          -- Cannot float out if local join points are used, as
          -- we cannot abstract over them
          | Bool
    captures_join_points
          = CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return CoreExpr
    e

          -- We have something to float out!
          | Bool
    otherwise
          = do { -- Assemble the RHS of the exit join point
                 let rhs :: CoreExpr
    rhs   = [JoinId] -> CoreExpr -> CoreExpr
    forall b. [b] -> Expr b -> Expr b
    mkLams [JoinId]
    abs_vars CoreExpr
    e
                     avoid :: InScopeSet
    avoid = InScopeSet
    in_scope InScopeSet -> [JoinId] -> InScopeSet
    `extendInScopeSetList` [JoinId]
    captured
                 -- Remember this binding under a suitable name
               ; JoinId
    v <- InScopeSet -> Int -> CoreExpr -> ExitifyM JoinId
    addExit InScopeSet
    avoid ([JoinId] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [JoinId]
    abs_vars) CoreExpr
    rhs
                 -- And jump to it from here
               ; CoreExpr -> ExitifyM CoreExpr
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (CoreExpr -> ExitifyM CoreExpr) -> CoreExpr -> ExitifyM CoreExpr
    forall a b. (a -> b) -> a -> b
    $ CoreExpr -> [JoinId] -> CoreExpr
    forall b. Expr b -> [JoinId] -> Expr b
    mkVarApps (JoinId -> CoreExpr
    forall b. JoinId -> Expr b
    Var JoinId
    v) [JoinId]
    abs_vars }

          where
            -- Used to detect exit expressions that are already proper exit jumps
            isCapturedVarArg :: CoreExpr -> Bool
    isCapturedVarArg (Var JoinId
    v) = JoinId
    v JoinId -> [JoinId] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [JoinId]
    captured
            isCapturedVarArg CoreExpr
    _ = Bool
    False

            -- An interesting exit expression has free, non-imported
            -- variables from outside the recursive group
            -- See Note [Interesting expression]
            is_interesting :: Bool
    is_interesting = (JoinId -> Bool) -> VarSet -> Bool
    anyVarSet JoinId -> Bool
    isLocalId (VarSet -> Bool) -> VarSet -> Bool
    forall a b. (a -> b) -> a -> b
    $
                             VarSet
    fvs VarSet -> VarSet -> VarSet
    `minusVarSet` [JoinId] -> VarSet
    mkVarSet [JoinId]
    captured

            -- The arguments of this exit join point
            -- See Note [Picking arguments to abstract over]
            abs_vars :: [JoinId]
    abs_vars = (VarSet, [JoinId]) -> [JoinId]
    forall a b. (a, b) -> b
    snd ((VarSet, [JoinId]) -> [JoinId]) -> (VarSet, [JoinId]) -> [JoinId]
    forall a b. (a -> b) -> a -> b
    $ (JoinId -> (VarSet, [JoinId]) -> (VarSet, [JoinId]))
    -> (VarSet, [JoinId]) -> [JoinId] -> (VarSet, [JoinId])
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr JoinId -> (VarSet, [JoinId]) -> (VarSet, [JoinId])
    pick (VarSet
    fvs, []) [JoinId]
    captured
              where
                pick :: JoinId -> (VarSet, [JoinId]) -> (VarSet, [JoinId])
    pick JoinId
    v (VarSet
    fvs', [JoinId]
    acc) | JoinId
    v JoinId -> VarSet -> Bool
    `elemVarSet` VarSet
    fvs' = (VarSet
    fvs' VarSet -> JoinId -> VarSet
    `delVarSet` JoinId
    v, JoinId -> JoinId
    zap JoinId
    v JoinId -> [JoinId] -> [JoinId]
    forall a. a -> [a] -> [a]
    : [JoinId]
    acc)
                                   | Bool
    otherwise           = (VarSet
    fvs',               [JoinId]
    acc)

            -- We are going to abstract over these variables, so we must
            -- zap any IdInfo they have; see #15005
            -- cf. GHC.Core.Opt.SetLevels.abstractVars
            zap :: JoinId -> JoinId
    zap JoinId
    v | JoinId -> Bool
    isId JoinId
    v = JoinId -> IdInfo -> JoinId
    setIdInfo JoinId
    v IdInfo
    vanillaIdInfo
                  | Bool
    otherwise = JoinId
    v

            -- We cannot abstract over join points
            captures_join_points :: Bool
    captures_join_points = (JoinId -> Bool) -> [JoinId] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any JoinId -> Bool
    isJoinId [JoinId]
    abs_vars

    -- Picks a new unique, which is disjoint from
    --  * the free variables of the whole joinrec
    --  * any bound variables (captured)
    --  * any exit join points created so far.
    mkExitJoinId :: InScopeSet -> Type -> JoinArity -> ExitifyM JoinId
    mkExitJoinId :: InScopeSet -> Type -> Int -> ExitifyM JoinId
    mkExitJoinId InScopeSet
    in_scope Type
    ty Int
    join_arity = do
        [(JoinId, CoreExpr)]
    fs <- State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall s. State s s
    get
        let avoid :: InScopeSet
    avoid = InScopeSet
    in_scope InScopeSet -> [JoinId] -> InScopeSet
    `extendInScopeSetList` (((JoinId, CoreExpr) -> JoinId) -> [(JoinId, CoreExpr)] -> [JoinId]
    forall a b. (a -> b) -> [a] -> [b]
    map (JoinId, CoreExpr) -> JoinId
    forall a b. (a, b) -> a
    fst [(JoinId, CoreExpr)]
    fs)
                             InScopeSet -> JoinId -> InScopeSet
    `extendInScopeSet` JoinId
    exit_id_tmpl -- just cosmetics
        JoinId -> ExitifyM JoinId
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (InScopeSet -> JoinId -> JoinId
    uniqAway InScopeSet
    avoid JoinId
    exit_id_tmpl)
      where
        exit_id_tmpl :: JoinId
    exit_id_tmpl = FastString -> Unique -> Type -> Type -> JoinId
    mkSysLocal (String -> FastString
    fsLit String
    "exit") Unique
    initExitJoinUnique Type
    Many Type
    ty
                        JoinId -> Int -> JoinId
    `asJoinId` Int
    join_arity

    addExit :: InScopeSet -> JoinArity -> CoreExpr -> ExitifyM JoinId
    addExit :: InScopeSet -> Int -> CoreExpr -> ExitifyM JoinId
    addExit InScopeSet
    in_scope Int
    join_arity CoreExpr
    rhs = do
        -- Pick a suitable name
        let ty :: Type
    ty = (() :: Constraint) => CoreExpr -> Type
    CoreExpr -> Type
    exprType CoreExpr
    rhs
        JoinId
    v <- InScopeSet -> Type -> Int -> ExitifyM JoinId
    mkExitJoinId InScopeSet
    in_scope Type
    ty Int
    join_arity
        [(JoinId, CoreExpr)]
    fs <- State [(JoinId, CoreExpr)] [(JoinId, CoreExpr)]
    forall s. State s s
    get
        [(JoinId, CoreExpr)] -> State [(JoinId, CoreExpr)] ()
    forall s. s -> State s ()
    put ((JoinId
    v,CoreExpr
    rhs)(JoinId, CoreExpr) -> [(JoinId, CoreExpr)] -> [(JoinId, CoreExpr)]
    forall a. a -> [a] -> [a]
    :[(JoinId, CoreExpr)]
    fs)
        JoinId -> ExitifyM JoinId
    forall a. a -> State [(JoinId, CoreExpr)] a
    forall (m :: * -> *) a. Monad m => a -> m a
    return JoinId
    v

    {-
    Note [Interesting expression]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    We do not want this to happen:

      joinrec go 0     x y = x
              go (n-1) x y = jump go (n-1) (x+y)
      in …
    ==>
      join exit x = x
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    because the floated exit path (`x`) is simply a parameter of `go`; there are
    not useful interactions exposed this way.

    Neither do we want this to happen

      joinrec go 0     x y = x+x
              go (n-1) x y = jump go (n-1) (x+y)
      in …
    ==>
      join exit x = x+x
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    where the floated expression `x+x` is a bit more complicated, but still not
    intersting.

    Expressions are interesting when they move an occurrence of a variable outside
    the recursive `go` that can benefit from being obviously called once, for example:
     * a local thunk that can then be inlined (see example in Note [Exitification])
     * the parameter of a function, where the demand analyzer then can then
       see that it is called at most once, and hence improve the function’s
       strictness signature

    So we only hoist an exit expression out if it mentiones at least one free,
    non-imported variable.

    Note [Jumps can be interesting]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    A jump to a join point can be interesting, if its arguments contain free
    non-exported variables (z in the following example):

      joinrec go 0     x y = jump j (x+z)
              go (n-1) x y = jump go (n-1) (x+y)
      in …
    ==>
      join exit x y = jump j (x+z)
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)

    The join point itself can be interesting, even if none if its
    arguments have free variables free in the joinrec.  For example

      join j p = case p of (x,y) -> x+y
      joinrec go 0     x y = jump j (x,y)
              go (n-1) x y = jump go (n-1) (x+y) y
      in …

    Here, `j` would not be inlined because we do not inline something that looks
    like an exit join point (see Note [Do not inline exit join points]). But
    if we exitify the 'jump j (x,y)' we get

      join j p = case p of (x,y) -> x+y
      join exit x y = jump j (x,y)
      joinrec go 0     x y = jump exit x y
              go (n-1) x y = jump go (n-1) (x+y) y
      in …

    and now 'j' can inline, and we get rid of the pair. Here's another
    example (assume `g` to be an imported function that, on its own,
    does not make this interesting):

      join j y = map f y
      joinrec go 0     x y = jump j (map g x)
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    Again, `j` would not be inlined because we do not inline something that looks
    like an exit join point (see Note [Do not inline exit join points]).

    But after exitification we have

      join j y = map f y
      join exit x = jump j (map g x)
      joinrec go 0     x y = jump j (map g x)
                  go (n-1) x y = jump go (n-1) (x+y)
      in …

    and now we can inline `j` and this will allow `map/map` to fire.

    Note [Idempotency]
    ~~~~~~~~~~~~~~~~~~

    We do not want this to happen, where we replace the floated expression with
    essentially the same expression:

      join exit x = t (x*x)
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)
      in …
    ==>
      join exit x = t (x*x)
      join exit' x = jump exit x
      joinrec go 0     x y = jump exit' x
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    So when the RHS is a join jump, and all of its arguments are captured variables,
    then we leave it in place.

    Note that `jump exit x` in this example looks interesting, as `exit` is a free
    variable. Therefore, idempotency does not simply follow from floating only
    interesting expressions.

    Note [Calculating free variables]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    We have two options where to annotate the tree with free variables:

     A) The whole tree.
     B) Each individual joinrec as we come across it.

    Downside of A: We pay the price on the whole module, even outside any joinrecs.
    Downside of B: We pay the price per joinrec, possibly multiple times when
    joinrecs are nested.

    Further downside of A: If the exitify function returns annotated expressions,
    it would have to ensure that the annotations are correct.

    We therefore choose B, and calculate the free variables in `exitify`.

    Note [Do not inline exit join points]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    When we have

      let t = foo bar
      join exit x = t (x*x)
      joinrec go 0     x y = jump exit x
              go (n-1) x y = jump go (n-1) (x+y)
      in …

    we do not want the simplifier to simply inline `exit` back in (which it happily
    would).

    To prevent this, we need to recognize exit join points, and then disable
    inlining.

    Exit join points, recognizeable using `isExitJoinId` are join points with an
    occurrence in a recursive group, and can be recognized (after the occurrence
    analyzer ran!) using `isExitJoinId`.
    This function detects joinpoints with `occ_in_lam (idOccinfo id) == True`,
    because the lambdas of a non-recursive join point are not considered for
    `occ_in_lam`.  For example, in the following code, `j1` is /not/ marked
    occ_in_lam, because `j2` is called only once.

      join j1 x = x+1
      join j2 y = join j1 (y+2)

    To prevent inlining, we check for isExitJoinId
    * In `preInlineUnconditionally` directly.
    * In `simplLetUnfolding` we simply give exit join points no unfolding, which
      prevents inlining in `postInlineUnconditionally` and call sites.

    Note [Placement of the exitification pass]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    I (Joachim) experimented with multiple positions for the Exitification pass in
    the Core2Core pipeline:

     A) Before the `simpl_phases`
     B) Between the `simpl_phases` and the "main" simplifier pass
     C) After demand_analyser
     D) Before the final simplification phase

    Here is the table (this is without inlining join exit points in the final
    simplifier run):

            Program |                       Allocs                      |                      Instrs
                    | ABCD.log     A.log     B.log     C.log     D.log  | ABCD.log     A.log     B.log     C.log     D.log
    ----------------|---------------------------------------------------|-------------------------------------------------
     fannkuch-redux |   -99.9%     +0.0%    -99.9%    -99.9%    -99.9%  |    -3.9%     +0.5%     -3.0%     -3.9%     -3.9%
              fasta |    -0.0%     +0.0%     +0.0%     -0.0%     -0.0%  |    -8.5%     +0.0%     +0.0%     -0.0%     -8.5%
                fem |     0.0%      0.0%      0.0%      0.0%     +0.0%  |    -2.2%     -0.1%     -0.1%     -2.1%     -2.1%
               fish |     0.0%      0.0%      0.0%      0.0%     +0.0%  |    -3.1%     +0.0%     -1.1%     -1.1%     -0.0%
       k-nucleotide |   -91.3%    -91.0%    -91.0%    -91.3%    -91.3%  |    -6.3%    +11.4%    +11.4%     -6.3%     -6.2%
                scs |    -0.0%     -0.0%     -0.0%     -0.0%     -0.0%  |    -3.4%     -3.0%     -3.1%     -3.3%     -3.3%
             simple |    -6.0%      0.0%     -6.0%     -6.0%     +0.0%  |    -3.4%     +0.0%     -5.2%     -3.4%     -0.1%
      spectral-norm |    -0.0%      0.0%      0.0%     -0.0%     +0.0%  |    -2.7%     +0.0%     -2.7%     -5.4%     -5.4%
    ----------------|---------------------------------------------------|-------------------------------------------------
                Min |   -95.0%    -91.0%    -95.0%    -95.0%    -95.0%  |    -8.5%     -3.0%     -5.2%     -6.3%     -8.5%
                Max |    +0.2%     +0.2%     +0.2%     +0.2%     +1.5%  |    +0.4%    +11.4%    +11.4%     +0.4%     +1.5%
     Geometric Mean |    -4.7%     -2.1%     -4.7%     -4.7%     -4.6%  |    -0.4%     +0.1%     -0.1%     -0.3%     -0.2%

    Position A is disqualified, as it does not get rid of the allocations in
    fannkuch-redux.
    Position A and B are disqualified because it increases instructions in k-nucleotide.
    Positions C and D have their advantages: C decreases allocations in simpl, but D instructions in fasta.

    Assuming we have a budget of _one_ run of Exitification, then C wins (but we
    could get more from running it multiple times, as seen in fish).

    Note [Picking arguments to abstract over]
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    When we create an exit join point, so we need to abstract over those of its
    free variables that are be out-of-scope at the destination of the exit join
    point. So we go through the list `captured` and pick those that are actually
    free variables of the join point.

    We do not just `filter (`elemVarSet` fvs) captured`, as there might be
    shadowing, and `captured` may contain multiple variables with the same Unique. I
    these cases we want to abstract only over the last occurrence, hence the `foldr`
    (with emphasis on the `r`). This is #15110.

    -}