simplStg/StgLiftLams/Transformation.hs

    {-# LANGUAGE CPP #-}

    -- | (Mostly) textbook instance of the lambda lifting transformation,
    -- selecting which bindings to lambda lift by consulting 'goodToLift'.
    module StgLiftLams.Transformation (stgLiftLams) where

    #include "HsVersions.h"

    import GhcPrelude

    import BasicTypes
    import DynFlags
    import Id
    import IdInfo
    import StgFVs ( annBindingFreeVars )
    import StgLiftLams.Analysis
    import StgLiftLams.LiftM
    import StgSyn
    import Outputable
    import UniqSupply
    import Util
    import VarSet
    import Control.Monad ( when )
    import Data.Maybe ( isNothing )

    -- | Lambda lifts bindings to top-level deemed worth lifting (see 'goodToLift').
    stgLiftLams :: DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
    stgLiftLams dflags us = runLiftM dflags us . foldr liftTopLvl (pure ())

    liftTopLvl :: InStgTopBinding -> LiftM () -> LiftM ()
    liftTopLvl (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do
      addTopStringLit bndr' lit
      rest
    liftTopLvl (StgTopLifted bind) rest = do
      let is_rec = isRec $ fst $ decomposeStgBinding bind
      when is_rec startBindingGroup
      let bind_w_fvs = annBindingFreeVars bind
      withLiftedBind TopLevel (tagSkeletonTopBind bind_w_fvs) NilSk $ \mb_bind' -> do
        -- We signal lifting of a binding through returning Nothing.
        -- Should never happen for a top-level binding, though, since we are already
        -- at top-level.
        case mb_bind' of
          Nothing -> pprPanic "StgLiftLams" (text "Lifted top-level binding")
          Just bind' -> addLiftedBinding bind'
        when is_rec endBindingGroup
        rest

    withLiftedBind
      :: TopLevelFlag
      -> LlStgBinding
      -> Skeleton
      -> (Maybe OutStgBinding -> LiftM a)
      -> LiftM a
    withLiftedBind top_lvl bind scope k
      | isTopLevel top_lvl
      = withCaffyness (is_caffy pairs) go
      | otherwise
      = go
      where
        (rec, pairs) = decomposeStgBinding bind
        is_caffy = any (mayHaveCafRefs . idCafInfo . binderInfoBndr . fst)
        go = withLiftedBindPairs top_lvl rec pairs scope (k . fmap (mkStgBinding rec))

    withLiftedBindPairs
      :: TopLevelFlag
      -> RecFlag
      -> [(BinderInfo, LlStgRhs)]
      -> Skeleton
      -> (Maybe [(Id, OutStgRhs)] -> LiftM a)
      -> LiftM a
    withLiftedBindPairs top rec pairs scope k = do
      let (infos, rhss) = unzip pairs
      let bndrs = map binderInfoBndr infos
      expander <- liftedIdsExpander
      dflags <- getDynFlags
      case goodToLift dflags top rec expander pairs scope of
        -- @abs_ids@ is the set of all variables that need to become parameters.
        Just abs_ids -> withLiftedBndrs abs_ids bndrs $ \bndrs' -> do
          -- Within this block, all binders in @bndrs@ will be noted as lifted, so
          -- that the return value of @liftedIdsExpander@ in this context will also
          -- expand the bindings in @bndrs@ to their free variables.
          -- Now we can recurse into the RHSs and see if we can lift any further
          -- bindings. We pass the set of expanded free variables (thus OutIds) on
          -- to @liftRhs@ so that it can add them as parameter binders.
          when (isRec rec) startBindingGroup
          rhss' <- traverse (liftRhs (Just abs_ids)) rhss
          let pairs' = zip bndrs' rhss'
          addLiftedBinding (mkStgBinding rec pairs')
          when (isRec rec) endBindingGroup
          k Nothing
        Nothing -> withSubstBndrs bndrs $ \bndrs' -> do
          -- Don't lift the current binding, but possibly some bindings in their
          -- RHSs.
          rhss' <- traverse (liftRhs Nothing) rhss
          let pairs' = zip bndrs' rhss'
          k (Just pairs')

    liftRhs
      :: Maybe (DIdSet)
      -- ^ @Just former_fvs@ <=> this RHS was lifted and we have to add @former_fvs@
      -- as lambda binders, discarding all free vars.
      -> LlStgRhs
      -> LiftM OutStgRhs
    liftRhs mb_former_fvs rhs@(StgRhsCon ccs con args)
      = ASSERT2(isNothing mb_former_fvs, text "Should never lift a constructor" $$ ppr rhs)
        StgRhsCon ccs con <$> traverse liftArgs args
    liftRhs Nothing (StgRhsClosure _ ccs upd infos body) = do
      -- This RHS wasn't lifted.
      withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
        StgRhsClosure noExtFieldSilent ccs upd bndrs' <$> liftExpr body
    liftRhs (Just former_fvs) (StgRhsClosure _ ccs upd infos body) = do
      -- This RHS was lifted. Insert extra binders for @former_fvs@.
      withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> do
        let bndrs'' = dVarSetElems former_fvs ++ bndrs'
        StgRhsClosure noExtFieldSilent ccs upd bndrs'' <$> liftExpr body

    liftArgs :: InStgArg -> LiftM OutStgArg
    liftArgs a@(StgLitArg _) = pure a
    liftArgs (StgVarArg occ) = do
      ASSERTM2( not <$> isLifted occ, text "StgArgs should never be lifted" $$ ppr occ )
      StgVarArg <$> substOcc occ

    liftExpr :: LlStgExpr -> LiftM OutStgExpr
    liftExpr (StgLit lit) = pure (StgLit lit)
    liftExpr (StgTick t e) = StgTick t <$> liftExpr e
    liftExpr (StgApp f args) = do
      f' <- substOcc f
      args' <- traverse liftArgs args
      fvs' <- formerFreeVars f
      let top_lvl_args = map StgVarArg fvs' ++ args'
      pure (StgApp f' top_lvl_args)
    liftExpr (StgConApp con args tys) = StgConApp con <$> traverse liftArgs args <*> pure tys
    liftExpr (StgOpApp op args ty) = StgOpApp op <$> traverse liftArgs args <*> pure ty
    liftExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam")
    liftExpr (StgCase scrut info ty alts) = do
      scrut' <- liftExpr scrut
      withSubstBndr (binderInfoBndr info) $ \bndr' -> do
        alts' <- traverse liftAlt alts
        pure (StgCase scrut' bndr' ty alts')
    liftExpr (StgLet scope bind body)
      = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
          body' <- liftExpr body
          case mb_bind' of
            Nothing -> pure body' -- withLiftedBindPairs decided to lift it and already added floats
            Just bind' -> pure (StgLet noExtFieldSilent bind' body')
    liftExpr (StgLetNoEscape scope bind body)
      = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
          body' <- liftExpr body
          case mb_bind' of
            Nothing -> pprPanic "stgLiftLams" (text "Should never decide to lift LNEs")
            Just bind' -> pure (StgLetNoEscape noExtFieldSilent bind' body')

    liftAlt :: LlStgAlt -> LiftM OutStgAlt
    liftAlt (con, infos, rhs) = withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
      (,,) con bndrs' <$> liftExpr rhs