never executed always true always false
    1 
    2 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
    3 {-# LANGUAGE TypeFamilies #-}
    4 
    5 -- | Hides away distracting bookkeeping while lambda lifting into a 'LiftM'
    6 -- monad.
    7 module GHC.Stg.Lift.Monad (
    8     decomposeStgBinding, mkStgBinding,
    9     Env (..),
   10     -- * #floats# Handling floats
   11     -- $floats
   12     FloatLang (..), collectFloats, -- Exported just for the docs
   13     -- * Transformation monad
   14     LiftM, runLiftM,
   15     -- ** Adding bindings
   16     startBindingGroup, endBindingGroup, addTopStringLit, addLiftedBinding,
   17     -- ** Substitution and binders
   18     withSubstBndr, withSubstBndrs, withLiftedBndr, withLiftedBndrs,
   19     -- ** Occurrences
   20     substOcc, isLifted, formerFreeVars, liftedIdsExpander
   21   ) where
   22 
   23 import GHC.Prelude
   24 
   25 import GHC.Types.Basic
   26 import GHC.Types.CostCentre ( isCurrentCCS, dontCareCCS )
   27 import GHC.Driver.Session
   28 import GHC.Data.FastString
   29 import GHC.Types.Id
   30 import GHC.Types.Name
   31 import GHC.Utils.Outputable
   32 import GHC.Data.OrdList
   33 import GHC.Stg.Subst
   34 import GHC.Stg.Syntax
   35 import GHC.Core.Utils
   36 import GHC.Types.Unique.Supply
   37 import GHC.Utils.Panic
   38 import GHC.Utils.Panic.Plain
   39 import GHC.Types.Var.Env
   40 import GHC.Types.Var.Set
   41 import GHC.Core.Multiplicity
   42 
   43 import Control.Arrow ( second )
   44 import Control.Monad.Trans.Class
   45 import Control.Monad.Trans.RWS.Strict ( RWST, runRWST )
   46 import qualified Control.Monad.Trans.RWS.Strict as RWS
   47 import Control.Monad.Trans.Cont ( ContT (..) )
   48 import Data.ByteString ( ByteString )
   49 
   50 -- | @uncurry 'mkStgBinding' . 'decomposeStgBinding' = id@
   51 decomposeStgBinding :: GenStgBinding pass -> (RecFlag, [(BinderP pass, GenStgRhs pass)])
   52 decomposeStgBinding (StgRec pairs) = (Recursive, pairs)
   53 decomposeStgBinding (StgNonRec bndr rhs) = (NonRecursive, [(bndr, rhs)])
   54 
   55 mkStgBinding :: RecFlag -> [(BinderP pass, GenStgRhs pass)] -> GenStgBinding pass
   56 mkStgBinding Recursive = StgRec
   57 mkStgBinding NonRecursive = uncurry StgNonRec . head
   58 
   59 -- | Environment threaded around in a scoped, @Reader@-like fashion.
   60 data Env
   61   = Env
   62   { e_dflags     :: !DynFlags
   63   -- ^ Read-only.
   64   , e_subst      :: !Subst
   65   -- ^ We need to track the renamings of local 'InId's to their lifted 'OutId',
   66   -- because shadowing might make a closure's free variables unavailable at its
   67   -- call sites. Consider:
   68   -- @
   69   --    let f y = x + y in let x = 4 in f x
   70   -- @
   71   -- Here, @f@ can't be lifted to top-level, because its free variable @x@ isn't
   72   -- available at its call site.
   73   , e_expansions :: !(IdEnv DIdSet)
   74   -- ^ Lifted 'Id's don't occur as free variables in any closure anymore, because
   75   -- they are bound at the top-level. Every occurrence must supply the formerly
   76   -- free variables of the lifted 'Id', so they in turn become free variables of
   77   -- the call sites. This environment tracks this expansion from lifted 'Id's to
   78   -- their free variables.
   79   --
   80   -- 'InId's to 'OutId's.
   81   --
   82   -- Invariant: 'Id's not present in this map won't be substituted.
   83   }
   84 
   85 emptyEnv :: DynFlags -> Env
   86 emptyEnv dflags = Env dflags emptySubst emptyVarEnv
   87 
   88 
   89 -- Note [Handling floats]
   90 -- ~~~~~~~~~~~~~~~~~~~~~~
   91 -- $floats
   92 -- Consider the following expression:
   93 --
   94 -- @
   95 --     f x =
   96 --       let g y = ... f y ...
   97 --       in g x
   98 -- @
   99 --
  100 -- What happens when we want to lift @g@? Normally, we'd put the lifted @l_g@
  101 -- binding above the binding for @f@:
  102 --
  103 -- @
  104 --     g f y = ... f y ...
  105 --     f x = g f x
  106 -- @
  107 --
  108 -- But this very unnecessarily turns a known call to @f@ into an unknown one, in
  109 -- addition to complicating matters for the analysis.
  110 -- Instead, we'd really like to put both functions in the same recursive group,
  111 -- thereby preserving the known call:
  112 --
  113 -- @
  114 --     Rec {
  115 --       g y = ... f y ...
  116 --       f x = g x
  117 --     }
  118 -- @
  119 --
  120 -- But we don't want this to happen for just /any/ binding. That would create
  121 -- possibly huge recursive groups in the process, calling for an occurrence
  122 -- analyser on STG.
  123 -- So, we need to track when we lift a binding out of a recursive RHS and add
  124 -- the binding to the same recursive group as the enclosing recursive binding
  125 -- (which must have either already been at the top-level or decided to be
  126 -- lifted itself in order to preserve the known call).
  127 --
  128 -- This is done by expressing this kind of nesting structure as a 'Writer' over
  129 -- @['FloatLang']@ and flattening this expression in 'runLiftM' by a call to
  130 -- 'collectFloats'.
  131 -- API-wise, the analysis will not need to know about the whole 'FloatLang'
  132 -- business and will just manipulate it indirectly through actions in 'LiftM'.
  133 
  134 -- | We need to detect when we are lifting something out of the RHS of a
  135 -- recursive binding (c.f. "GHC.Stg.Lift.Monad#floats"), in which case that
  136 -- binding needs to be added to the same top-level recursive group. This
  137 -- requires we detect a certain nesting structure, which is encoded by
  138 -- 'StartBindingGroup' and 'EndBindingGroup'.
  139 --
  140 -- Although 'collectFloats' will only ever care if the current binding to be
  141 -- lifted (through 'LiftedBinding') will occur inside such a binding group or
  142 -- not, e.g. doesn't care about the nesting level as long as its greater than 0.
  143 data FloatLang
  144   = StartBindingGroup
  145   | EndBindingGroup
  146   | PlainTopBinding OutStgTopBinding
  147   | LiftedBinding OutStgBinding
  148 
  149 instance Outputable FloatLang where
  150   ppr StartBindingGroup = char '('
  151   ppr EndBindingGroup = char ')'
  152   ppr (PlainTopBinding StgTopStringLit{}) = text "<str>"
  153   ppr (PlainTopBinding (StgTopLifted b)) = ppr (LiftedBinding b)
  154   ppr (LiftedBinding bind) = (if isRec rec then char 'r' else char 'n') <+> ppr (map fst pairs)
  155     where
  156       (rec, pairs) = decomposeStgBinding bind
  157 
  158 -- | Flattens an expression in @['FloatLang']@ into an STG program, see "GHC.Stg.Lift.Monad#floats".
  159 -- Important pre-conditions: The nesting of opening 'StartBindinGroup's and
  160 -- closing 'EndBindinGroup's is balanced. Also, it is crucial that every binding
  161 -- group has at least one recursive binding inside. Otherwise there's no point
  162 -- in announcing the binding group in the first place and an @ASSERT@ will
  163 -- trigger.
  164 collectFloats :: [FloatLang] -> [OutStgTopBinding]
  165 collectFloats = go (0 :: Int) []
  166   where
  167     go 0 [] [] = []
  168     go _ _ [] = pprPanic "collectFloats" (text "unterminated group")
  169     go n binds (f:rest) = case f of
  170       StartBindingGroup -> go (n+1) binds rest
  171       EndBindingGroup
  172         | n == 0 -> pprPanic "collectFloats" (text "no group to end")
  173         | n == 1 -> StgTopLifted (merge_binds binds) : go 0 [] rest
  174         | otherwise -> go (n-1) binds rest
  175       PlainTopBinding top_bind
  176         | n == 0 -> top_bind : go n binds rest
  177         | otherwise -> pprPanic "collectFloats" (text "plain top binding inside group")
  178       LiftedBinding bind
  179         | n == 0 -> StgTopLifted (rm_cccs bind) : go n binds rest
  180         | otherwise -> go n (bind:binds) rest
  181 
  182     map_rhss f = uncurry mkStgBinding . second (map (second f)) . decomposeStgBinding
  183     rm_cccs = map_rhss removeRhsCCCS
  184     merge_binds binds = assert (any is_rec binds) $
  185                         StgRec (concatMap (snd . decomposeStgBinding . rm_cccs) binds)
  186     is_rec StgRec{} = True
  187     is_rec _ = False
  188 
  189 -- | Omitting this makes for strange closure allocation schemes that crash the
  190 -- GC.
  191 removeRhsCCCS :: GenStgRhs pass -> GenStgRhs pass
  192 removeRhsCCCS (StgRhsClosure ext ccs upd bndrs body)
  193   | isCurrentCCS ccs
  194   = StgRhsClosure ext dontCareCCS upd bndrs body
  195 removeRhsCCCS (StgRhsCon ccs con mu ts args)
  196   | isCurrentCCS ccs
  197   = StgRhsCon dontCareCCS con mu ts args
  198 removeRhsCCCS rhs = rhs
  199 
  200 -- | The analysis monad consists of the following 'RWST' components:
  201 --
  202 --     * 'Env': Reader-like context. Contains a substitution, info about how
  203 --       how lifted identifiers are to be expanded into applications and details
  204 --       such as 'DynFlags'.
  205 --
  206 --     * @'OrdList' 'FloatLang'@: Writer output for the resulting STG program.
  207 --
  208 --     * No pure state component
  209 --
  210 --     * But wrapping around 'UniqSM' for generating fresh lifted binders.
  211 --       (The @uniqAway@ approach could give the same name to two different
  212 --       lifted binders, so this is necessary.)
  213 newtype LiftM a
  214   = LiftM { unwrapLiftM :: RWST Env (OrdList FloatLang) () UniqSM a }
  215   deriving (Functor, Applicative, Monad)
  216 
  217 instance HasDynFlags LiftM where
  218   getDynFlags = LiftM (RWS.asks e_dflags)
  219 
  220 instance MonadUnique LiftM where
  221   getUniqueSupplyM = LiftM (lift getUniqueSupplyM)
  222   getUniqueM = LiftM (lift getUniqueM)
  223   getUniquesM = LiftM (lift getUniquesM)
  224 
  225 runLiftM :: DynFlags -> UniqSupply -> LiftM () -> [OutStgTopBinding]
  226 runLiftM dflags us (LiftM m) = collectFloats (fromOL floats)
  227   where
  228     (_, _, floats) = initUs_ us (runRWST m (emptyEnv dflags) ())
  229 
  230 -- | Writes a plain 'StgTopStringLit' to the output.
  231 addTopStringLit :: OutId -> ByteString -> LiftM ()
  232 addTopStringLit id = LiftM . RWS.tell . unitOL . PlainTopBinding . StgTopStringLit id
  233 
  234 -- | Starts a recursive binding group. See "GHC.Stg.Lift.Monad#floats" and 'collectFloats'.
  235 startBindingGroup :: LiftM ()
  236 startBindingGroup = LiftM $ RWS.tell $ unitOL $ StartBindingGroup
  237 
  238 -- | Ends a recursive binding group. See "GHC.Stg.Lift.Monad#floats" and 'collectFloats'.
  239 endBindingGroup :: LiftM ()
  240 endBindingGroup = LiftM $ RWS.tell $ unitOL $ EndBindingGroup
  241 
  242 -- | Lifts a binding to top-level. Depending on whether it's declared inside
  243 -- a recursive RHS (see "GHC.Stg.Lift.Monad#floats" and 'collectFloats'), this might be added to
  244 -- an existing recursive top-level binding group.
  245 addLiftedBinding :: OutStgBinding -> LiftM ()
  246 addLiftedBinding = LiftM . RWS.tell . unitOL . LiftedBinding
  247 
  248 -- | Takes a binder and a continuation which is called with the substituted
  249 -- binder. The continuation will be evaluated in a 'LiftM' context in which that
  250 -- binder is deemed in scope. Think of it as a 'RWS.local' computation: After
  251 -- the continuation finishes, the new binding won't be in scope anymore.
  252 withSubstBndr :: Id -> (Id -> LiftM a) -> LiftM a
  253 withSubstBndr bndr inner = LiftM $ do
  254   subst <- RWS.asks e_subst
  255   let (bndr', subst') = substBndr bndr subst
  256   RWS.local (\e -> e { e_subst = subst' }) (unwrapLiftM (inner bndr'))
  257 
  258 -- | See 'withSubstBndr'.
  259 withSubstBndrs :: Traversable f => f Id -> (f Id -> LiftM a) -> LiftM a
  260 withSubstBndrs = runContT . traverse (ContT . withSubstBndr)
  261 
  262 -- | Similarly to 'withSubstBndr', this function takes a set of variables to
  263 -- abstract over, the binder to lift (and generate a fresh, substituted name
  264 -- for) and a continuation in which that fresh, lifted binder is in scope.
  265 --
  266 -- It takes care of all the details involved with copying and adjusting the
  267 -- binder and fresh name generation.
  268 withLiftedBndr :: DIdSet -> Id -> (Id -> LiftM a) -> LiftM a
  269 withLiftedBndr abs_ids bndr inner = do
  270   uniq <- getUniqueM
  271   let str = "$l" ++ occNameString (getOccName bndr)
  272   let ty = mkLamTypes (dVarSetElems abs_ids) (idType bndr)
  273   let bndr'
  274         -- See Note [transferPolyIdInfo] in GHC.Types.Id. We need to do this at least
  275         -- for arity information.
  276         = transferPolyIdInfo bndr (dVarSetElems abs_ids)
  277         . mkSysLocal (mkFastString str) uniq Many
  278         $ ty
  279   LiftM $ RWS.local
  280     (\e -> e
  281       { e_subst = extendSubst bndr bndr' $ extendInScope bndr' $ e_subst e
  282       , e_expansions = extendVarEnv (e_expansions e) bndr abs_ids
  283       })
  284     (unwrapLiftM (inner bndr'))
  285 
  286 -- | See 'withLiftedBndr'.
  287 withLiftedBndrs :: Traversable f => DIdSet -> f Id -> (f Id -> LiftM a) -> LiftM a
  288 withLiftedBndrs abs_ids = runContT . traverse (ContT . withLiftedBndr abs_ids)
  289 
  290 -- | Substitutes a binder /occurrence/, which was brought in scope earlier by
  291 -- 'withSubstBndr' \/ 'withLiftedBndr'.
  292 substOcc :: Id -> LiftM Id
  293 substOcc id = LiftM (RWS.asks (lookupIdSubst id . e_subst))
  294 
  295 -- | Whether the given binding was decided to be lambda lifted.
  296 isLifted :: InId -> LiftM Bool
  297 isLifted bndr = LiftM (RWS.asks (elemVarEnv bndr . e_expansions))
  298 
  299 -- | Returns an empty list for a binding that was not lifted and the list of all
  300 -- local variables the binding abstracts over (so, exactly the additional
  301 -- arguments at adjusted call sites) otherwise.
  302 formerFreeVars :: InId -> LiftM [OutId]
  303 formerFreeVars f = LiftM $ do
  304   expansions <- RWS.asks e_expansions
  305   pure $ case lookupVarEnv expansions f of
  306     Nothing -> []
  307     Just fvs -> dVarSetElems fvs
  308 
  309 -- | Creates an /expander function/ for the current set of lifted binders.
  310 -- This expander function will replace any 'InId' by their corresponding 'OutId'
  311 -- and, in addition, will expand any lifted binders by the former free variables
  312 -- it abstracts over.
  313 liftedIdsExpander :: LiftM (DIdSet -> DIdSet)
  314 liftedIdsExpander = LiftM $ do
  315   expansions <- RWS.asks e_expansions
  316   subst <- RWS.asks e_subst
  317   -- We use @noWarnLookupIdSubst@ here in order to suppress "not in scope"
  318   -- warnings generated by 'lookupIdSubst' due to local bindings within RHS.
  319   -- These are not in the InScopeSet of @subst@ and extending the InScopeSet in
  320   -- @goodToLift@/@closureGrowth@ before passing it on to @expander@ is too much
  321   -- trouble.
  322   let go set fv = case lookupVarEnv expansions fv of
  323         Nothing -> extendDVarSet set (noWarnLookupIdSubst fv subst) -- Not lifted
  324         Just fvs' -> unionDVarSet set fvs'
  325   let expander fvs = foldl' go emptyDVarSet (dVarSetElems fvs)
  326   pure expander