1{-# LANGUAGE CPP #-}
2
3-- | (Mostly) textbook instance of the lambda lifting transformation,
4-- selecting which bindings to lambda lift by consulting 'goodToLift'.
5module StgLiftLams.Transformation (stgLiftLams) where
6
7#include "HsVersions.h"
8
9import GhcPrelude
10
11import BasicTypes
12import DynFlags
13import Id
14import IdInfo
15import StgFVs ( annBindingFreeVars )
16import StgLiftLams.Analysis
17import StgLiftLams.LiftM
18import StgSyn
19import Outputable
20import UniqSupply
21import Util
22import VarSet
23import Control.Monad ( when )
24import Data.Maybe ( isNothing )
25
26-- | Lambda lifts bindings to top-level deemed worth lifting (see 'goodToLift').
27stgLiftLams :: DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
28stgLiftLams dflags us = runLiftM dflags us . foldr liftTopLvl (pure ())
29
30liftTopLvl :: InStgTopBinding -> LiftM () -> LiftM ()
31liftTopLvl (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do
32  addTopStringLit bndr' lit
33  rest
34liftTopLvl (StgTopLifted bind) rest = do
35  let is_rec = isRec $ fst $ decomposeStgBinding bind
36  when is_rec startBindingGroup
37  let bind_w_fvs = annBindingFreeVars bind
38  withLiftedBind TopLevel (tagSkeletonTopBind bind_w_fvs) NilSk $ \mb_bind' -> do
39    -- We signal lifting of a binding through returning Nothing.
40    -- Should never happen for a top-level binding, though, since we are already
41    -- at top-level.
42    case mb_bind' of
43      Nothing -> pprPanic "StgLiftLams" (text "Lifted top-level binding")
44      Just bind' -> addLiftedBinding bind'
45    when is_rec endBindingGroup
46    rest
47
48withLiftedBind
49  :: TopLevelFlag
50  -> LlStgBinding
51  -> Skeleton
52  -> (Maybe OutStgBinding -> LiftM a)
53  -> LiftM a
54withLiftedBind top_lvl bind scope k
55  | isTopLevel top_lvl
56  = withCaffyness (is_caffy pairs) go
57  | otherwise
58  = go
59  where
60    (rec, pairs) = decomposeStgBinding bind
61    is_caffy = any (mayHaveCafRefs . idCafInfo . binderInfoBndr . fst)
62    go = withLiftedBindPairs top_lvl rec pairs scope (k . fmap (mkStgBinding rec))
63
64withLiftedBindPairs
65  :: TopLevelFlag
66  -> RecFlag
67  -> [(BinderInfo, LlStgRhs)]
68  -> Skeleton
69  -> (Maybe [(Id, OutStgRhs)] -> LiftM a)
70  -> LiftM a
71withLiftedBindPairs top rec pairs scope k = do
72  let (infos, rhss) = unzip pairs
73  let bndrs = map binderInfoBndr infos
74  expander <- liftedIdsExpander
75  dflags <- getDynFlags
76  case goodToLift dflags top rec expander pairs scope of
77    -- @abs_ids@ is the set of all variables that need to become parameters.
78    Just abs_ids -> withLiftedBndrs abs_ids bndrs $ \bndrs' -> do
79      -- Within this block, all binders in @bndrs@ will be noted as lifted, so
80      -- that the return value of @liftedIdsExpander@ in this context will also
81      -- expand the bindings in @bndrs@ to their free variables.
82      -- Now we can recurse into the RHSs and see if we can lift any further
83      -- bindings. We pass the set of expanded free variables (thus OutIds) on
84      -- to @liftRhs@ so that it can add them as parameter binders.
85      when (isRec rec) startBindingGroup
86      rhss' <- traverse (liftRhs (Just abs_ids)) rhss
87      let pairs' = zip bndrs' rhss'
88      addLiftedBinding (mkStgBinding rec pairs')
89      when (isRec rec) endBindingGroup
90      k Nothing
91    Nothing -> withSubstBndrs bndrs $ \bndrs' -> do
92      -- Don't lift the current binding, but possibly some bindings in their
93      -- RHSs.
94      rhss' <- traverse (liftRhs Nothing) rhss
95      let pairs' = zip bndrs' rhss'
96      k (Just pairs')
97
98liftRhs
99  :: Maybe (DIdSet)
100  -- ^ @Just former_fvs@ <=> this RHS was lifted and we have to add @former_fvs@
101  -- as lambda binders, discarding all free vars.
102  -> LlStgRhs
103  -> LiftM OutStgRhs
104liftRhs mb_former_fvs rhs@(StgRhsCon ccs con args)
105  = ASSERT2(isNothing mb_former_fvs, text "Should never lift a constructor" $$ ppr rhs)
106    StgRhsCon ccs con <$> traverse liftArgs args
107liftRhs Nothing (StgRhsClosure _ ccs upd infos body) = do
108  -- This RHS wasn't lifted.
109  withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
110    StgRhsClosure noExtFieldSilent ccs upd bndrs' <$> liftExpr body
111liftRhs (Just former_fvs) (StgRhsClosure _ ccs upd infos body) = do
112  -- This RHS was lifted. Insert extra binders for @former_fvs@.
113  withSubstBndrs (map binderInfoBndr infos) $ \bndrs' -> do
114    let bndrs'' = dVarSetElems former_fvs ++ bndrs'
115    StgRhsClosure noExtFieldSilent ccs upd bndrs'' <$> liftExpr body
116
117liftArgs :: InStgArg -> LiftM OutStgArg
118liftArgs a@(StgLitArg _) = pure a
119liftArgs (StgVarArg occ) = do
120  ASSERTM2( not <$> isLifted occ, text "StgArgs should never be lifted" $$ ppr occ )
121  StgVarArg <$> substOcc occ
122
123liftExpr :: LlStgExpr -> LiftM OutStgExpr
124liftExpr (StgLit lit) = pure (StgLit lit)
125liftExpr (StgTick t e) = StgTick t <$> liftExpr e
126liftExpr (StgApp f args) = do
127  f' <- substOcc f
128  args' <- traverse liftArgs args
129  fvs' <- formerFreeVars f
130  let top_lvl_args = map StgVarArg fvs' ++ args'
131  pure (StgApp f' top_lvl_args)
132liftExpr (StgConApp con args tys) = StgConApp con <$> traverse liftArgs args <*> pure tys
133liftExpr (StgOpApp op args ty) = StgOpApp op <$> traverse liftArgs args <*> pure ty
134liftExpr (StgLam _ _) = pprPanic "stgLiftLams" (text "StgLam")
135liftExpr (StgCase scrut info ty alts) = do
136  scrut' <- liftExpr scrut
137  withSubstBndr (binderInfoBndr info) $ \bndr' -> do
138    alts' <- traverse liftAlt alts
139    pure (StgCase scrut' bndr' ty alts')
140liftExpr (StgLet scope bind body)
141  = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
142      body' <- liftExpr body
143      case mb_bind' of
144        Nothing -> pure body' -- withLiftedBindPairs decided to lift it and already added floats
145        Just bind' -> pure (StgLet noExtFieldSilent bind' body')
146liftExpr (StgLetNoEscape scope bind body)
147  = withLiftedBind NotTopLevel bind scope $ \mb_bind' -> do
148      body' <- liftExpr body
149      case mb_bind' of
150        Nothing -> pprPanic "stgLiftLams" (text "Should never decide to lift LNEs")
151        Just bind' -> pure (StgLetNoEscape noExtFieldSilent bind' body')
152
153liftAlt :: LlStgAlt -> LiftM OutStgAlt
154liftAlt (con, infos, rhs) = withSubstBndrs (map binderInfoBndr infos) $ \bndrs' ->
155  (,,) con bndrs' <$> liftExpr rhs
156