1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
{-# 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
|
