Replace .lhs with .hs in compiler comments
[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 (tupleCon UnboxedTuple 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 (tupleCon UnboxedTuple (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 (tupleCon UnboxedTuple (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 (alt_ty', alts') = unariseAlts us2 rho alt_ty bndr (repType (idType 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 -> RepType -> [StgAlt] -> (AltType, [StgAlt])
138 unariseAlts us rho alt_ty _ (UnaryRep _) alts
139 = (alt_ty, zipWith (\us alt -> unariseAlt us rho alt) (listSplitUniqSupply us) alts)
140
141 unariseAlts us rho _ bndr (UbxTupleRep tys) ((DEFAULT, [], [], e) : _)
142 = (UbxTupAlt n, [(DataAlt (tupleCon UnboxedTuple n), ys, uses, unariseExpr us2' rho' e)])
143 where
144 (us2', rho', ys) = unariseIdBinder us rho bndr
145 uses = replicate (length ys) (not (isDeadBinder bndr))
146 n = length tys
147
148 unariseAlts us rho _ bndr (UbxTupleRep _) [(DataAlt _, ys, uses, e)]
149 = (UbxTupAlt n, [(DataAlt (tupleCon UnboxedTuple n), ys', uses', unariseExpr us2' rho'' e)])
150 where
151 (us2', rho', ys', uses') = unariseUsedIdBinders us rho ys uses
152 rho'' = extendVarEnv rho' bndr ys'
153 n = length ys'
154
155 unariseAlts _ _ _ _ (UbxTupleRep _) alts
156 = pprPanic "unariseExpr: strange unboxed tuple alts" (ppr alts)
157
158 --------------------------
159 unariseAlt :: UniqSupply -> UnariseEnv -> StgAlt -> StgAlt
160 unariseAlt us rho (con, xs, uses, e)
161 = (con, xs', uses', unariseExpr us' rho' e)
162 where
163 (us', rho', xs', uses') = unariseUsedIdBinders us rho xs uses
164
165 ------------------------
166 unariseSRT :: UnariseEnv -> SRT -> SRT
167 unariseSRT _ NoSRT = NoSRT
168 unariseSRT rho (SRTEntries ids) = SRTEntries (concatMapVarSet (unariseId rho) ids)
169 unariseSRT _ (SRT {}) = panic "unariseSRT"
170
171 unariseLives :: UnariseEnv -> StgLiveVars -> StgLiveVars
172 unariseLives rho ids = concatMapVarSet (unariseId rho) ids
173
174 unariseArgs :: UnariseEnv -> [StgArg] -> [StgArg]
175 unariseArgs rho = concatMap (unariseArg rho)
176
177 unariseArg :: UnariseEnv -> StgArg -> [StgArg]
178 unariseArg rho (StgVarArg x) = map StgVarArg (unariseId rho x)
179 unariseArg _ (StgLitArg l) = [StgLitArg l]
180
181 unariseIds :: UnariseEnv -> [Id] -> [Id]
182 unariseIds rho = concatMap (unariseId rho)
183
184 unariseId :: UnariseEnv -> Id -> [Id]
185 unariseId rho x
186 | Just ys <- lookupVarEnv rho x
187 = ASSERT2( case repType (idType x) of UbxTupleRep _ -> True; _ -> x == ubxTupleId0
188 , text "unariseId: not unboxed tuple" <+> ppr x )
189 ys
190
191 | otherwise
192 = ASSERT2( case repType (idType x) of UbxTupleRep _ -> False; _ -> True
193 , text "unariseId: was unboxed tuple" <+> ppr x )
194 [x]
195
196 unariseUsedIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> [Bool]
197 -> (UniqSupply, UnariseEnv, [Id], [Bool])
198 unariseUsedIdBinders us rho xs uses
199 = case mapAccumL2 do_one us rho (zipEqual "unariseUsedIdBinders" xs uses) of
200 (us', rho', xs_usess) -> uncurry ((,,,) us' rho') (unzip (concat xs_usess))
201 where
202 do_one us rho (x, use) = third3 (map (flip (,) use)) (unariseIdBinder us rho x)
203
204 unariseIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> (UniqSupply, UnariseEnv, [Id])
205 unariseIdBinders us rho xs = third3 concat $ mapAccumL2 unariseIdBinder us rho xs
206
207 unariseIdBinder :: UniqSupply -> UnariseEnv -> Id -> (UniqSupply, UnariseEnv, [Id])
208 unariseIdBinder us rho x = case repType (idType x) of
209 UnaryRep _ -> (us, rho, [x])
210 UbxTupleRep tys -> let (us0, us1) = splitUniqSupply us
211 ys = unboxedTupleBindersFrom us0 x tys
212 rho' = extendVarEnv rho x ys
213 in (us1, rho', ys)
214
215 unboxedTupleBindersFrom :: UniqSupply -> Id -> [UnaryType] -> [Id]
216 unboxedTupleBindersFrom us x tys = zipWith (mkSysLocal fs) (uniqsFromSupply us) tys
217 where fs = occNameFS (getOccName x)
218
219 concatMapVarSet :: (Var -> [Var]) -> VarSet -> VarSet
220 concatMapVarSet f xs = mkVarSet [x' | x <- varSetElems xs, x' <- f x]