b16220134df7945c7fd6df1aadaa1a4c5b6c3068
[ghc.git] / compiler / simplStg / UnariseStg.hs
1 {-
2 (c) The GRASP/AQUA Project, Glasgow University, 1992-2012
3
4
5 Note [Unarisation]
6 ~~~~~~~~~~~~~~~~~~
7
8 The idea of this pass is to translate away *all* unboxed-tuple binders. So for example:
9
10 f (x :: (# Int, Bool #)) = f x + f (# 1, True #)
11 ==>
12 f (x1 :: Int) (x2 :: Bool) = f x1 x2 + f 1 True
13
14 It is important that we do this at the STG level and NOT at the core level
15 because it would be very hard to make this pass Core-type-preserving.
16
17 STG fed to the code generators *must* be unarised because the code generators do
18 not support unboxed tuple binders natively.
19
20
21 Note [Unarisation and arity]
22 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
23
24 Because of unarisation, the arity that will be recorded in the generated info table
25 for an Id may be larger than the idArity. Instead we record what we call the RepArity,
26 which is the Arity taking into account any expanded arguments, and corresponds to
27 the number of (possibly-void) *registers* arguments will arrive in.
28 -}
29
30 {-# LANGUAGE CPP #-}
31
32 module UnariseStg (unarise) where
33
34 #include "HsVersions.h"
35
36 import CoreSyn
37 import StgSyn
38 import VarEnv
39 import UniqSupply
40 import Id
41 import MkId (realWorldPrimId)
42 import Type
43 import TysWiredIn
44 import DataCon
45 import VarSet
46 import OccName
47 import Name
48 import Util
49 import Outputable
50 import BasicTypes
51
52
53 -- | A mapping from unboxed-tuple binders to the Ids they were expanded to.
54 --
55 -- INVARIANT: Ids in the range don't have unboxed tuple types.
56 --
57 -- Those in-scope variables without unboxed-tuple types are not present in
58 -- the domain of the mapping at all.
59 type UnariseEnv = VarEnv [Id]
60
61 ubxTupleId0 :: Id
62 ubxTupleId0 = dataConWorkId (tupleDataCon Unboxed 0)
63
64 unarise :: UniqSupply -> [StgBinding] -> [StgBinding]
65 unarise us binds = zipWith (\us -> unariseBinding us init_env) (listSplitUniqSupply us) binds
66 where -- See Note [Nullary unboxed tuple] in Type.hs
67 init_env = unitVarEnv ubxTupleId0 [realWorldPrimId]
68
69 unariseBinding :: UniqSupply -> UnariseEnv -> StgBinding -> StgBinding
70 unariseBinding us rho bind = case bind of
71 StgNonRec x rhs -> StgNonRec x (unariseRhs us rho rhs)
72 StgRec xrhss -> StgRec $ zipWith (\us (x, rhs) -> (x, unariseRhs us rho rhs))
73 (listSplitUniqSupply us) xrhss
74
75 unariseRhs :: UniqSupply -> UnariseEnv -> StgRhs -> StgRhs
76 unariseRhs us rho rhs = case rhs of
77 StgRhsClosure ccs b_info fvs update_flag srt args expr
78 -> StgRhsClosure ccs b_info (unariseIds rho fvs) update_flag
79 (unariseSRT rho srt) args' (unariseExpr us' rho' expr)
80 where (us', rho', args') = unariseIdBinders us rho args
81 StgRhsCon ccs con args
82 -> StgRhsCon ccs con (unariseArgs rho args)
83
84 ------------------------
85 unariseExpr :: UniqSupply -> UnariseEnv -> StgExpr -> StgExpr
86 unariseExpr _ rho (StgApp f args)
87 | null args
88 , UbxTupleRep tys <- repType (idType f)
89 = -- Particularly important where (##) is concerned
90 -- See Note [Nullary unboxed tuple]
91 StgConApp (tupleDataCon Unboxed (length tys))
92 (map StgVarArg (unariseId rho f))
93
94 | otherwise
95 = StgApp f (unariseArgs rho args)
96
97 unariseExpr _ _ (StgLit l)
98 = StgLit l
99
100 unariseExpr _ rho (StgConApp dc args)
101 | isUnboxedTupleCon dc = StgConApp (tupleDataCon Unboxed (length args')) args'
102 | otherwise = StgConApp dc args'
103 where
104 args' = unariseArgs rho args
105
106 unariseExpr _ rho (StgOpApp op args ty)
107 = StgOpApp op (unariseArgs rho args) ty
108
109 unariseExpr us rho (StgLam xs e)
110 = StgLam xs' (unariseExpr us' rho' e)
111 where
112 (us', rho', xs') = unariseIdBinders us rho xs
113
114 unariseExpr us rho (StgCase e case_lives alts_lives bndr srt alt_ty alts)
115 = StgCase (unariseExpr us1 rho e) (unariseLives rho case_lives)
116 (unariseLives rho alts_lives) bndr (unariseSRT rho srt)
117 alt_ty alts'
118 where
119 (us1, us2) = splitUniqSupply us
120 alts' = unariseAlts us2 rho alt_ty bndr alts
121
122 unariseExpr us rho (StgLet bind e)
123 = StgLet (unariseBinding us1 rho bind) (unariseExpr us2 rho e)
124 where
125 (us1, us2) = splitUniqSupply us
126
127 unariseExpr us rho (StgLetNoEscape live_in_let live_in_bind bind e)
128 = StgLetNoEscape (unariseLives rho live_in_let) (unariseLives rho live_in_bind)
129 (unariseBinding us1 rho bind) (unariseExpr us2 rho e)
130 where
131 (us1, us2) = splitUniqSupply us
132
133 unariseExpr us rho (StgTick tick e)
134 = StgTick tick (unariseExpr us rho e)
135
136 ------------------------
137 unariseAlts :: UniqSupply -> UnariseEnv -> AltType -> Id -> [StgAlt] -> [StgAlt]
138 unariseAlts us rho (UbxTupAlt n) bndr [(DEFAULT, [], [], e)]
139 = [(DataAlt (tupleDataCon Unboxed n), ys, uses, unariseExpr us2' rho' e)]
140 where
141 (us2', rho', ys) = unariseIdBinder us rho bndr
142 uses = replicate (length ys) (not (isDeadBinder bndr))
143
144 unariseAlts us rho (UbxTupAlt n) bndr [(DataAlt _, ys, uses, e)]
145 = [(DataAlt (tupleDataCon Unboxed n), ys', uses', unariseExpr us2' rho'' e)]
146 where
147 (us2', rho', ys', uses') = unariseUsedIdBinders us rho ys uses
148 rho'' = extendVarEnv rho' bndr ys'
149
150 unariseAlts _ _ (UbxTupAlt _) _ alts
151 = pprPanic "unariseExpr: strange unboxed tuple alts" (ppr alts)
152
153 unariseAlts us rho _ _ alts
154 = zipWith (\us alt -> unariseAlt us rho alt) (listSplitUniqSupply us) alts
155
156 --------------------------
157 unariseAlt :: UniqSupply -> UnariseEnv -> StgAlt -> StgAlt
158 unariseAlt us rho (con, xs, uses, e)
159 = (con, xs', uses', unariseExpr us' rho' e)
160 where
161 (us', rho', xs', uses') = unariseUsedIdBinders us rho xs uses
162
163 ------------------------
164 unariseSRT :: UnariseEnv -> SRT -> SRT
165 unariseSRT _ NoSRT = NoSRT
166 unariseSRT rho (SRTEntries ids) = SRTEntries (concatMapVarSet (unariseId rho) ids)
167
168 unariseLives :: UnariseEnv -> StgLiveVars -> StgLiveVars
169 unariseLives rho ids = concatMapVarSet (unariseId rho) ids
170
171 unariseArgs :: UnariseEnv -> [StgArg] -> [StgArg]
172 unariseArgs rho = concatMap (unariseArg rho)
173
174 unariseArg :: UnariseEnv -> StgArg -> [StgArg]
175 unariseArg rho (StgVarArg x) = map StgVarArg (unariseId rho x)
176 unariseArg _ (StgLitArg l) = [StgLitArg l]
177
178 unariseIds :: UnariseEnv -> [Id] -> [Id]
179 unariseIds rho = concatMap (unariseId rho)
180
181 unariseId :: UnariseEnv -> Id -> [Id]
182 unariseId rho x
183 | Just ys <- lookupVarEnv rho x
184 = ASSERT2( case repType (idType x) of UbxTupleRep _ -> True; _ -> x == ubxTupleId0
185 , text "unariseId: not unboxed tuple" <+> ppr x )
186 ys
187
188 | otherwise
189 = ASSERT2( case repType (idType x) of UbxTupleRep _ -> False; _ -> True
190 , text "unariseId: was unboxed tuple" <+> ppr x )
191 [x]
192
193 unariseUsedIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> [Bool]
194 -> (UniqSupply, UnariseEnv, [Id], [Bool])
195 unariseUsedIdBinders us rho xs uses
196 = case mapAccumL2 do_one us rho (zipEqual "unariseUsedIdBinders" xs uses) of
197 (us', rho', xs_usess) -> uncurry ((,,,) us' rho') (unzip (concat xs_usess))
198 where
199 do_one us rho (x, use) = third3 (map (flip (,) use)) (unariseIdBinder us rho x)
200
201 unariseIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> (UniqSupply, UnariseEnv, [Id])
202 unariseIdBinders us rho xs = third3 concat $ mapAccumL2 unariseIdBinder us rho xs
203
204 unariseIdBinder :: UniqSupply -> UnariseEnv -> Id -> (UniqSupply, UnariseEnv, [Id])
205 unariseIdBinder us rho x = case repType (idType x) of
206 UnaryRep _ -> (us, rho, [x])
207 UbxTupleRep tys -> let (us0, us1) = splitUniqSupply us
208 ys = unboxedTupleBindersFrom us0 x tys
209 rho' = extendVarEnv rho x ys
210 in (us1, rho', ys)
211
212 unboxedTupleBindersFrom :: UniqSupply -> Id -> [UnaryType] -> [Id]
213 unboxedTupleBindersFrom us x tys = zipWith (mkSysLocalOrCoVar fs) (uniqsFromSupply us) tys
214 where fs = occNameFS (getOccName x)
215
216 concatMapVarSet :: (Var -> [Var]) -> VarSet -> VarSet
217 concatMapVarSet f xs = mkVarSet [x' | x <- varSetElems xs, x' <- f x]