fc308599801cc57257c29fec88d792be44f24a91
[ghc.git] / compiler / simplStg / UnariseStg.hs
1 {-
2 (c) The GRASP/AQUA Project, Glasgow University, 1992-2012
3
4 Note [Unarisation]
5 ~~~~~~~~~~~~~~~~~~
6 The idea of this pass is to translate away *all* unboxed-tuple and unboxed-sum
7 binders. So for example:
8
9 f (x :: (# Int, Bool #)) = f x + f (# 1, True #)
10
11 ==>
12
13 f (x1 :: Int) (x2 :: Bool) = f x1 x2 + f 1 True
14
15 It is important that we do this at the STG level and NOT at the Core level
16 because it would be very hard to make this pass Core-type-preserving. In this
17 example the type of 'f' changes, for example.
18
19 STG fed to the code generators *must* be unarised because the code generators do
20 not support unboxed tuple and unboxed sum binders natively.
21
22 In more detail: (see next note for unboxed sums)
23
24 Suppose that a variable x : (# t1, t2 #).
25
26 * At the binding site for x, make up fresh vars x1:t1, x2:t2
27
28 * Extend the UnariseEnv x :-> MultiVal [x1,x2]
29
30 * Replace the binding with a curried binding for x1,x2
31
32 Lambda: \x.e ==> \x1 x2. e
33 Case alt: MkT a b x c d -> e ==> MkT a b x1 x2 c d -> e
34
35 * Replace argument occurrences with a sequence of args via a lookup in
36 UnariseEnv
37
38 f a b x c d ==> f a b x1 x2 c d
39
40 * Replace tail-call occurrences with an unboxed tuple via a lookup in
41 UnariseEnv
42
43 x ==> (# x1, x2 #)
44
45 So, for example
46
47 f x = x ==> f x1 x2 = (# x1, x2 #)
48
49 * We /always/ eliminate a case expression when
50
51 - It scrutinises an unboxed tuple or unboxed sum
52
53 - The scrutinee is a variable (or when it is an explicit tuple, but the
54 simplifier eliminates those)
55
56 The case alternative (there can be only one) can be one of these two
57 things:
58
59 - An unboxed tuple pattern. e.g.
60
61 case v of x { (# x1, x2, x3 #) -> ... }
62
63 Scrutinee has to be in form `(# t1, t2, t3 #)` so we just extend the
64 environment with
65
66 x :-> MultiVal [t1,t2,t3]
67 x1 :-> UnaryVal t1, x2 :-> UnaryVal t2, x3 :-> UnaryVal t3
68
69 - A DEFAULT alternative. Just the same, without the bindings for x1,x2,x3
70
71 By the end of this pass, we only have unboxed tuples in return positions.
72 Unboxed sums are completely eliminated, see next note.
73
74 Note [Translating unboxed sums to unboxed tuples]
75 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
76 Unarise also eliminates unboxed sum binders, and translates unboxed sums in
77 return positions to unboxed tuples. We want to overlap fields of a sum when
78 translating it to a tuple to have efficient memory layout. When translating a
79 sum pattern to a tuple pattern, we need to translate it so that binders of sum
80 alternatives will be mapped to right arguments after the term translation. So
81 translation of sum DataCon applications to tuple DataCon applications and
82 translation of sum patterns to tuple patterns need to be in sync.
83
84 These translations work like this. Suppose we have
85
86 (# x1 | | ... #) :: (# t1 | t2 | ... #)
87
88 remember that t1, t2 ... can be sums and tuples too. So we first generate
89 layouts of those. Then we "merge" layouts of each alternative, which gives us a
90 sum layout with best overlapping possible.
91
92 Layout of a flat type 'ty1' is just [ty1].
93 Layout of a tuple is just concatenation of layouts of its fields.
94
95 For layout of a sum type,
96
97 - We first get layouts of all alternatives.
98 - We sort these layouts based on their "slot types".
99 - We merge all the alternatives.
100
101 For example, say we have (# (# Int#, Char #) | (# Int#, Int# #) | Int# #)
102
103 - Layouts of alternatives: [ [Word, Ptr], [Word, Word], [Word] ]
104 - Sorted: [ [Ptr, Word], [Word, Word], [Word] ]
105 - Merge all alternatives together: [ Ptr, Word, Word ]
106
107 We add a slot for the tag to the first position. So our tuple type is
108
109 (# Tag#, Any, Word#, Word# #)
110 (we use Any for pointer slots)
111
112 Now, any term of this sum type needs to generate a tuple of this type instead.
113 The translation works by simply putting arguments to first slots that they fit
114 in. Suppose we had
115
116 (# (# 42#, 'c' #) | | #)
117
118 42# fits in Word#, 'c' fits in Any, so we generate this application:
119
120 (# 1#, 'c', 42#, rubbish #)
121
122 Another example using the same type: (# | (# 2#, 3# #) | #). 2# fits in Word#,
123 3# fits in Word #, so we get:
124
125 (# 2#, rubbish, 2#, 3# #).
126
127 Note [Types in StgConApp]
128 ~~~~~~~~~~~~~~~~~~~~~~~~~
129 Suppose we have this unboxed sum term:
130
131 (# 123 | #)
132
133 What will be the unboxed tuple representation? We can't tell without knowing the
134 type of this term. For example, these are all valid tuples for this:
135
136 (# 1#, 123 #) -- when type is (# Int | String #)
137 (# 1#, 123, rubbish #) -- when type is (# Int | Float# #)
138 (# 1#, 123, rubbish, rubbish #)
139 -- when type is (# Int | (# Int, Int, Int #) #)
140
141 So we pass type arguments of the DataCon's TyCon in StgConApp to decide what
142 layout to use. Note that unlifted values can't be let-bound, so we don't need
143 types in StgRhsCon.
144
145 Note [UnariseEnv can map to literals]
146 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
147 To avoid redundant case expressions when unarising unboxed sums, UnariseEnv
148 needs to map variables to literals too. Suppose we have this Core:
149
150 f (# x | #)
151
152 ==> (CorePrep)
153
154 case (# x | #) of y {
155 _ -> f y
156 }
157
158 ==> (MultiVal)
159
160 case (# 1#, x #) of [x1, x2] {
161 _ -> f x1 x2
162 }
163
164 To eliminate this case expression we need to map x1 to 1# in UnariseEnv:
165
166 x1 :-> UnaryVal 1#, x2 :-> UnaryVal x
167
168 so that `f x1 x2` becomes `f 1# x`.
169
170 Note [Unarisation and arity]
171 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
172 Because of unarisation, the arity that will be recorded in the generated info
173 table for an Id may be larger than the idArity. Instead we record what we call
174 the RepArity, which is the Arity taking into account any expanded arguments, and
175 corresponds to the number of (possibly-void) *registers* arguments will arrive
176 in.
177
178 Note [Post-unarisation invariants]
179 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
180 STG programs after unarisation have these invariants:
181
182 * No unboxed sums at all.
183
184 * No unboxed tuple binders. Tuples only appear in return position.
185
186 * DataCon applications (StgRhsCon and StgConApp) don't have void arguments.
187 This means that it's safe to wrap `StgArg`s of DataCon applications with
188 `StgCmmEnv.NonVoid`, for example.
189
190 * Alt binders (binders in patterns) are always non-void.
191 -}
192
193 {-# LANGUAGE CPP, TupleSections #-}
194
195 module UnariseStg (unarise) where
196
197 #include "HsVersions.h"
198
199 import BasicTypes
200 import CoreSyn
201 import DataCon
202 import FastString (FastString, mkFastString)
203 import Id
204 import Literal (Literal (..))
205 import MkCore (aBSENT_ERROR_ID)
206 import MkId (voidPrimId, voidArgId)
207 import MonadUtils (mapAccumLM)
208 import Outputable
209 import RepType
210 import StgSyn
211 import Type
212 import TysPrim (intPrimTyCon, intPrimTy)
213 import TysWiredIn
214 import UniqSupply
215 import Util
216 import VarEnv
217
218 import Data.Bifunctor (second)
219 import Data.Maybe (mapMaybe)
220 import qualified Data.IntMap as IM
221
222 --------------------------------------------------------------------------------
223
224 -- | A mapping from binders to the Ids they were expanded/renamed to.
225 --
226 -- x :-> MultiVal [a,b,c] in rho
227 --
228 -- iff x's repType is a MultiRep, or equivalently
229 -- x's type is an unboxed tuple, sum or void.
230 --
231 -- x :-> UnaryVal x'
232 --
233 -- iff x's RepType is UnaryRep or equivalently
234 -- x's type is not unboxed tuple, sum or void.
235 --
236 -- So
237 -- x :-> MultiVal [a] in rho
238 -- means x is represented by singleton tuple.
239 --
240 -- x :-> MultiVal [] in rho
241 -- means x is void.
242 --
243 -- INVARIANT: OutStgArgs in the range only have NvUnaryTypes
244 -- (i.e. no unboxed tuples, sums or voids)
245 --
246 type UnariseEnv = VarEnv UnariseVal
247
248 data UnariseVal
249 = MultiVal [OutStgArg] -- MultiVal to tuple. Can be empty list (void).
250 | UnaryVal OutStgArg -- See NOTE [Renaming during unarisation].
251
252 instance Outputable UnariseVal where
253 ppr (MultiVal args) = text "MultiVal" <+> ppr args
254 ppr (UnaryVal arg) = text "UnaryVal" <+> ppr arg
255
256 -- | Extend the environment, checking the UnariseEnv invariant.
257 extendRho :: UnariseEnv -> Id -> UnariseVal -> UnariseEnv
258 extendRho rho x (MultiVal args)
259 = ASSERT(all (isNvUnaryType . stgArgType) args)
260 extendVarEnv rho x (MultiVal args)
261 extendRho rho x (UnaryVal val)
262 = ASSERT(isNvUnaryType (stgArgType val))
263 extendVarEnv rho x (UnaryVal val)
264
265 --------------------------------------------------------------------------------
266
267 type OutStgExpr = StgExpr
268 type InStgAlt = StgAlt
269 type InStgArg = StgArg
270 type OutStgArg = StgArg
271
272 unarise :: UniqSupply -> [StgBinding] -> [StgBinding]
273 unarise us binds = initUs_ us (mapM (unariseBinding emptyVarEnv) binds)
274
275 unariseBinding :: UnariseEnv -> StgBinding -> UniqSM StgBinding
276 unariseBinding rho (StgNonRec x rhs)
277 = StgNonRec x <$> unariseRhs rho rhs
278 unariseBinding rho (StgRec xrhss)
279 = StgRec <$> mapM (\(x, rhs) -> (x,) <$> unariseRhs rho rhs) xrhss
280
281 unariseRhs :: UnariseEnv -> StgRhs -> UniqSM StgRhs
282 unariseRhs rho (StgRhsClosure ccs b_info fvs update_flag args expr)
283 = do (rho', args1) <- unariseFunArgBinders rho args
284 expr' <- unariseExpr rho' expr
285 let fvs' = unariseFreeVars rho fvs
286 return (StgRhsClosure ccs b_info fvs' update_flag args1 expr')
287
288 unariseRhs rho (StgRhsCon ccs con args)
289 = ASSERT(not (isUnboxedTupleCon con || isUnboxedSumCon con))
290 return (StgRhsCon ccs con (unariseConArgs rho args))
291
292 --------------------------------------------------------------------------------
293
294 unariseExpr :: UnariseEnv -> StgExpr -> UniqSM StgExpr
295
296 unariseExpr rho e@(StgApp f [])
297 = case lookupVarEnv rho f of
298 Just (MultiVal args) -- Including empty tuples
299 -> return (mkTuple args)
300 Just (UnaryVal (StgVarArg f'))
301 -> return (StgApp f' [])
302 Just (UnaryVal (StgLitArg f'))
303 -> return (StgLit f')
304 Nothing
305 -> return e
306
307 unariseExpr rho e@(StgApp f args)
308 = return (StgApp f' (unariseFunArgs rho args))
309 where
310 f' = case lookupVarEnv rho f of
311 Just (UnaryVal (StgVarArg f')) -> f'
312 Nothing -> f
313 err -> pprPanic "unariseExpr - app2" (ppr e $$ ppr err)
314 -- Can't happen because 'args' is non-empty, and
315 -- a tuple or sum cannot be applied to anything
316
317 unariseExpr _ (StgLit l)
318 = return (StgLit l)
319
320 unariseExpr rho (StgConApp dc args ty_args)
321 | Just args' <- unariseMulti_maybe rho dc args ty_args
322 = return (mkTuple args')
323
324 | otherwise
325 , let args' = unariseConArgs rho args
326 = return (StgConApp dc args' (map stgArgType args'))
327
328 unariseExpr rho (StgOpApp op args ty)
329 = return (StgOpApp op (unariseFunArgs rho args) ty)
330
331 unariseExpr _ e@StgLam{}
332 = pprPanic "unariseExpr: found lambda" (ppr e)
333
334 unariseExpr rho (StgCase scrut bndr alt_ty alts)
335 -- a tuple/sum binders in the scrutinee can always be eliminated
336 | StgApp v [] <- scrut
337 , Just (MultiVal xs) <- lookupVarEnv rho v
338 = elimCase rho xs bndr alt_ty alts
339
340 -- Handle strict lets for tuples and sums:
341 -- case (# a,b #) of r -> rhs
342 -- and analogously for sums
343 | StgConApp dc args ty_args <- scrut
344 , Just args' <- unariseMulti_maybe rho dc args ty_args
345 = elimCase rho args' bndr alt_ty alts
346
347 -- general case
348 | otherwise
349 = do scrut' <- unariseExpr rho scrut
350 alts' <- unariseAlts rho alt_ty bndr alts
351 return (StgCase scrut' bndr alt_ty alts')
352 -- bndr will be dead after unarise
353
354 unariseExpr rho (StgLet bind e)
355 = StgLet <$> unariseBinding rho bind <*> unariseExpr rho e
356
357 unariseExpr rho (StgLetNoEscape bind e)
358 = StgLetNoEscape <$> unariseBinding rho bind <*> unariseExpr rho e
359
360 unariseExpr rho (StgTick tick e)
361 = StgTick tick <$> unariseExpr rho e
362
363 -- Doesn't return void args.
364 unariseMulti_maybe :: UnariseEnv -> DataCon -> [InStgArg] -> [Type] -> Maybe [OutStgArg]
365 unariseMulti_maybe rho dc args ty_args
366 | isUnboxedTupleCon dc
367 = Just (unariseConArgs rho args)
368
369 | isUnboxedSumCon dc
370 , let args1 = ASSERT(isSingleton args) (unariseConArgs rho args)
371 = Just (mkUbxSum dc ty_args args1)
372
373 | otherwise
374 = Nothing
375
376 --------------------------------------------------------------------------------
377
378 elimCase :: UnariseEnv
379 -> [OutStgArg] -- non-void args
380 -> InId -> AltType -> [InStgAlt] -> UniqSM OutStgExpr
381
382 elimCase rho args bndr (MultiValAlt _) [(_, bndrs, rhs)]
383 = do let rho1 = extendRho rho bndr (MultiVal args)
384 rho2
385 | isUnboxedTupleBndr bndr
386 = mapTupleIdBinders bndrs args rho1
387 | otherwise
388 = ASSERT(isUnboxedSumBndr bndr)
389 if null bndrs then rho1
390 else mapSumIdBinders bndrs args rho1
391
392 unariseExpr rho2 rhs
393
394 elimCase rho args bndr (MultiValAlt _) alts
395 | isUnboxedSumBndr bndr
396 = do let (tag_arg : real_args) = args
397 tag_bndr <- mkId (mkFastString "tag") tagTy
398 -- this won't be used but we need a binder anyway
399 let rho1 = extendRho rho bndr (MultiVal args)
400 scrut' = case tag_arg of
401 StgVarArg v -> StgApp v []
402 StgLitArg l -> StgLit l
403
404 alts' <- unariseSumAlts rho1 real_args alts
405 return (StgCase scrut' tag_bndr tagAltTy alts')
406
407 elimCase _ args bndr alt_ty alts
408 = pprPanic "elimCase - unhandled case"
409 (ppr args <+> ppr bndr <+> ppr alt_ty $$ ppr alts)
410
411 --------------------------------------------------------------------------------
412
413 unariseAlts :: UnariseEnv -> AltType -> InId -> [StgAlt] -> UniqSM [StgAlt]
414 unariseAlts rho (MultiValAlt n) bndr [(DEFAULT, [], e)]
415 | isUnboxedTupleBndr bndr
416 = do (rho', ys) <- unariseConArgBinder rho bndr
417 e' <- unariseExpr rho' e
418 return [(DataAlt (tupleDataCon Unboxed n), ys, e')]
419
420 unariseAlts rho (MultiValAlt n) bndr [(DataAlt _, ys, e)]
421 | isUnboxedTupleBndr bndr
422 = do (rho', ys1) <- unariseConArgBinders rho ys
423 MASSERT(n == length ys1)
424 let rho'' = extendRho rho' bndr (MultiVal (map StgVarArg ys1))
425 e' <- unariseExpr rho'' e
426 return [(DataAlt (tupleDataCon Unboxed n), ys1, e')]
427
428 unariseAlts _ (MultiValAlt _) bndr alts
429 | isUnboxedTupleBndr bndr
430 = pprPanic "unariseExpr: strange multi val alts" (ppr alts)
431
432 -- In this case we don't need to scrutinize the tag bit
433 unariseAlts rho (MultiValAlt _) bndr [(DEFAULT, _, rhs)]
434 | isUnboxedSumBndr bndr
435 = do (rho_sum_bndrs, sum_bndrs) <- unariseConArgBinder rho bndr
436 rhs' <- unariseExpr rho_sum_bndrs rhs
437 return [(DataAlt (tupleDataCon Unboxed (length sum_bndrs)), sum_bndrs, rhs')]
438
439 unariseAlts rho (MultiValAlt _) bndr alts
440 | isUnboxedSumBndr bndr
441 = do (rho_sum_bndrs, scrt_bndrs@(tag_bndr : real_bndrs)) <- unariseConArgBinder rho bndr
442 alts' <- unariseSumAlts rho_sum_bndrs (map StgVarArg real_bndrs) alts
443 let inner_case = StgCase (StgApp tag_bndr []) tag_bndr tagAltTy alts'
444 return [ (DataAlt (tupleDataCon Unboxed (length scrt_bndrs)),
445 scrt_bndrs,
446 inner_case) ]
447
448 unariseAlts rho _ _ alts
449 = mapM (\alt -> unariseAlt rho alt) alts
450
451 unariseAlt :: UnariseEnv -> StgAlt -> UniqSM StgAlt
452 unariseAlt rho (con, xs, e)
453 = do (rho', xs') <- unariseConArgBinders rho xs
454 (con, xs',) <$> unariseExpr rho' e
455
456 --------------------------------------------------------------------------------
457
458 -- | Make alternatives that match on the tag of a sum
459 -- (i.e. generate LitAlts for the tag)
460 unariseSumAlts :: UnariseEnv
461 -> [StgArg] -- sum components _excluding_ the tag bit.
462 -> [StgAlt] -- original alternative with sum LHS
463 -> UniqSM [StgAlt]
464 unariseSumAlts env args alts
465 = do alts' <- mapM (unariseSumAlt env args) alts
466 return (mkDefaultLitAlt alts')
467
468 unariseSumAlt :: UnariseEnv
469 -> [StgArg] -- sum components _excluding_ the tag bit.
470 -> StgAlt -- original alternative with sum LHS
471 -> UniqSM StgAlt
472 unariseSumAlt rho _ (DEFAULT, _, e)
473 = ( DEFAULT, [], ) <$> unariseExpr rho e
474
475 unariseSumAlt rho args (DataAlt sumCon, bs, e)
476 = do let rho' = mapSumIdBinders bs args rho
477 e' <- unariseExpr rho' e
478 return ( LitAlt (MachInt (fromIntegral (dataConTag sumCon))), [], e' )
479
480 unariseSumAlt _ scrt alt
481 = pprPanic "unariseSumAlt" (ppr scrt $$ ppr alt)
482
483 --------------------------------------------------------------------------------
484
485 mapTupleIdBinders
486 :: [InId] -- Un-processed binders of a tuple alternative.
487 -- Can have void binders.
488 -> [OutStgArg] -- Arguments that form the tuple (after unarisation).
489 -- Can't have void args.
490 -> UnariseEnv
491 -> UnariseEnv
492 mapTupleIdBinders ids args0 rho0
493 = ASSERT(not (any (isVoidTy . stgArgType) args0))
494 let
495 ids_unarised :: [(Id, RepType)]
496 ids_unarised = map (\id -> (id, repType (idType id))) ids
497
498 map_ids :: UnariseEnv -> [(Id, RepType)] -> [StgArg] -> UnariseEnv
499 map_ids rho [] _ = rho
500 map_ids rho ((x, x_rep) : xs) args =
501 let
502 x_arity = length (repTypeSlots x_rep)
503 (x_args, args') =
504 ASSERT(args `lengthAtLeast` x_arity)
505 splitAt x_arity args
506
507 rho'
508 | isMultiRep x_rep
509 = extendRho rho x (MultiVal x_args)
510 | otherwise
511 = ASSERT(x_args `lengthIs` 1)
512 extendRho rho x (UnaryVal (head x_args))
513 in
514 map_ids rho' xs args'
515 in
516 map_ids rho0 ids_unarised args0
517
518 mapSumIdBinders
519 :: [InId] -- Binder of a sum alternative (remember that sum patterns
520 -- only have one binder, so this list should be a singleton)
521 -> [OutStgArg] -- Arguments that form the sum (NOT including the tag).
522 -- Can't have void args.
523 -> UnariseEnv
524 -> UnariseEnv
525
526 mapSumIdBinders [id] args rho0
527 = ASSERT(not (any (isVoidTy . stgArgType) args))
528 let
529 arg_slots = concatMap (repTypeSlots . repType . stgArgType) args
530 id_slots = repTypeSlots (repType (idType id))
531 layout1 = layout arg_slots id_slots
532 in
533 if isMultiValBndr id
534 then extendRho rho0 id (MultiVal [ args !! i | i <- layout1 ])
535 else ASSERT(layout1 `lengthIs` 1)
536 extendRho rho0 id (UnaryVal (args !! head layout1))
537
538 mapSumIdBinders ids sum_args _
539 = pprPanic "mapSumIdBinders" (ppr ids $$ ppr sum_args)
540
541 -- | Build a unboxed sum term from arguments of an alternative.
542 --
543 -- Example, for (# x | #) :: (# (# #) | Int #) we call
544 --
545 -- mkUbxSum (# _ | #) [ (# #), Int ] [ voidPrimId ]
546 --
547 -- which returns
548 --
549 -- [ 1#, rubbish ]
550 --
551 mkUbxSum
552 :: DataCon -- Sum data con
553 -> [Type] -- Type arguments of the sum data con
554 -> [OutStgArg] -- Actual arguments of the alternative.
555 -> [OutStgArg] -- Final tuple arguments
556 mkUbxSum dc ty_args args0
557 = let
558 (_ : sum_slots) = ubxSumRepType ty_args
559 -- drop tag slot
560
561 tag = dataConTag dc
562
563 layout' = layout sum_slots (mapMaybe (typeSlotTy . stgArgType) args0)
564 tag_arg = StgLitArg (MachInt (fromIntegral tag))
565 arg_idxs = IM.fromList (zipEqual "mkUbxSum" layout' args0)
566
567 mkTupArgs :: Int -> [SlotTy] -> IM.IntMap StgArg -> [StgArg]
568 mkTupArgs _ [] _
569 = []
570 mkTupArgs arg_idx (slot : slots_left) arg_map
571 | Just stg_arg <- IM.lookup arg_idx arg_map
572 = stg_arg : mkTupArgs (arg_idx + 1) slots_left arg_map
573 | otherwise
574 = slotRubbishArg slot : mkTupArgs (arg_idx + 1) slots_left arg_map
575
576 slotRubbishArg :: SlotTy -> StgArg
577 slotRubbishArg PtrSlot = StgVarArg aBSENT_ERROR_ID
578 slotRubbishArg WordSlot = StgLitArg (MachWord 0)
579 slotRubbishArg Word64Slot = StgLitArg (MachWord64 0)
580 slotRubbishArg FloatSlot = StgLitArg (MachFloat 0)
581 slotRubbishArg DoubleSlot = StgLitArg (MachDouble 0)
582 in
583 tag_arg : mkTupArgs 0 sum_slots arg_idxs
584
585 --------------------------------------------------------------------------------
586
587 {-
588 For arguments (StgArg) and binders (Id) we have two kind of unarisation:
589
590 - When unarising function arg binders and arguments, we don't want to remove
591 void binders and arguments. For example,
592
593 f :: (# (# #), (# #) #) -> Void# -> RealWorld# -> ...
594 f x y z = <body>
595
596 Here after unarise we should still get a function with arity 3. Similarly
597 in the call site we shouldn't remove void arguments:
598
599 f (# (# #), (# #) #) voidId rw
600
601 When unarising <body>, we extend the environment with these binders:
602
603 x :-> MultiVal [], y :-> MultiVal [], z :-> MultiVal []
604
605 Because their rep types are `MultiRep []` (aka. void). This means that when
606 we see `x` in a function argument position, we actually replace it with a
607 void argument. When we see it in a DataCon argument position, we just get
608 rid of it, because DataCon applications in STG are always saturated.
609
610 - When unarising case alternative binders we remove void binders, but we
611 still update the environment the same way, because those binders may be
612 used in the RHS. Example:
613
614 case x of y {
615 (# x1, x2, x3 #) -> <RHS>
616 }
617
618 We know that y can't be void, because we don't scrutinize voids, so x will
619 be unarised to some number of arguments, and those arguments will have at
620 least one non-void thing. So in the rho we will have something like:
621
622 x :-> MultiVal [xu1, xu2]
623
624 Now, after we eliminate void binders in the pattern, we get exactly the same
625 number of binders, and extend rho again with these:
626
627 x1 :-> UnaryVal xu1
628 x2 :-> MultiVal [] -- x2 is void
629 x3 :-> UnaryVal xu2
630
631 Now when we see x2 in a function argument position or in return position, we
632 generate void#. In constructor argument position, we just remove it.
633
634 So in short, when we have a void id,
635
636 - We keep it if it's a lambda argument binder or
637 in argument position of an application.
638
639 - We remove it if it's a DataCon field binder or
640 in argument position of a DataCon application.
641 -}
642
643 --------------------------------------------------------------------------------
644
645 -- | MultiVal a function argument. Never returns an empty list.
646 unariseFunArg :: UnariseEnv -> StgArg -> [StgArg]
647 unariseFunArg rho (StgVarArg x) =
648 case lookupVarEnv rho x of
649 Just (MultiVal []) -> [voidArg] -- NB: do not remove void args
650 Just (MultiVal as) -> as
651 Just (UnaryVal arg) -> [arg]
652 Nothing -> [StgVarArg x]
653 unariseFunArg _ arg = [arg]
654
655 unariseFunArgs :: UnariseEnv -> [StgArg] -> [StgArg]
656 unariseFunArgs = concatMap . unariseFunArg
657
658 unariseFunArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
659 unariseFunArgBinders rho xs = second concat <$> mapAccumLM unariseFunArgBinder rho xs
660
661 unariseFunArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
662 -- Result list of binders is never empty
663 unariseFunArgBinder rho x =
664 case repType (idType x) of
665 UnaryRep _ -> return (rho, [x])
666 MultiRep [] -> return (extendRho rho x (MultiVal []), [voidArgId])
667 -- NB: do not remove void binders
668 MultiRep slots -> do
669 xs <- mkIds (mkFastString "us") (map slotTyToType slots)
670 return (extendRho rho x (MultiVal (map StgVarArg xs)), xs)
671
672 --------------------------------------------------------------------------------
673
674 -- | MultiVal a DataCon argument. Returns an empty list when argument is void.
675 unariseConArg :: UnariseEnv -> InStgArg -> [OutStgArg]
676 unariseConArg rho (StgVarArg x) =
677 case lookupVarEnv rho x of
678 Just (UnaryVal arg) -> [arg]
679 Just (MultiVal as) -> as -- 'as' can be empty
680 Nothing
681 | isVoidTy (idType x) -> [] -- e.g. C realWorld#
682 -- Here realWorld# is not in the envt, but
683 -- is a void, and so should be eliminated
684 | otherwise -> [StgVarArg x]
685 unariseConArg _ arg = [arg] -- We have no void literals
686
687 unariseConArgs :: UnariseEnv -> [InStgArg] -> [OutStgArg]
688 unariseConArgs = concatMap . unariseConArg
689
690 unariseConArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
691 unariseConArgBinders rho xs = second concat <$> mapAccumLM unariseConArgBinder rho xs
692
693 unariseConArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
694 unariseConArgBinder rho x =
695 case repType (idType x) of
696 UnaryRep _ -> return (rho, [x])
697 MultiRep slots -> do
698 xs <- mkIds (mkFastString "us") (map slotTyToType slots)
699 return (extendRho rho x (MultiVal (map StgVarArg xs)), xs)
700
701 unariseFreeVars :: UnariseEnv -> [InId] -> [OutId]
702 unariseFreeVars rho fvs
703 = [ v | fv <- fvs, StgVarArg v <- unariseFreeVar rho fv ]
704 -- Notice that we filter out any StgLitArgs
705 -- e.g. case e of (x :: (# Int | Bool #))
706 -- (# v | #) -> ... let {g = \y. ..x...} in ...
707 -- (# | w #) -> ...
708 -- Here 'x' is free in g's closure, and the env will have
709 -- x :-> [1, v]
710 -- we want to capture 'v', but not 1, in the free vars
711
712 unariseFreeVar :: UnariseEnv -> Id -> [StgArg]
713 unariseFreeVar rho x =
714 case lookupVarEnv rho x of
715 Just (MultiVal args) -> args
716 Just (UnaryVal arg) -> [arg]
717 Nothing -> [StgVarArg x]
718
719 --------------------------------------------------------------------------------
720
721 mkIds :: FastString -> [UnaryType] -> UniqSM [Id]
722 mkIds fs tys = mapM (mkId fs) tys
723
724 mkId :: FastString -> UnaryType -> UniqSM Id
725 mkId = mkSysLocalOrCoVarM
726
727 isMultiValBndr :: Id -> Bool
728 isMultiValBndr = isMultiRep . repType . idType
729
730 isUnboxedSumBndr :: Id -> Bool
731 isUnboxedSumBndr = isUnboxedSumType . idType
732
733 isUnboxedTupleBndr :: Id -> Bool
734 isUnboxedTupleBndr = isUnboxedTupleType . idType
735
736 mkTuple :: [StgArg] -> StgExpr
737 mkTuple args = StgConApp (tupleDataCon Unboxed (length args)) args (map stgArgType args)
738
739 tagAltTy :: AltType
740 tagAltTy = PrimAlt intPrimTyCon
741
742 tagTy :: Type
743 tagTy = intPrimTy
744
745 voidArg :: StgArg
746 voidArg = StgVarArg voidPrimId
747
748 mkDefaultLitAlt :: [StgAlt] -> [StgAlt]
749 -- We have an exhauseive list of literal alternatives
750 -- 1# -> e1
751 -- 2# -> e2
752 -- Since they are exhaustive, we can replace one with DEFAULT, to avoid
753 -- generating a final test. Remember, the DEFAULT comes first if it exists.
754 mkDefaultLitAlt [] = pprPanic "elimUbxSumExpr.mkDefaultAlt" (text "Empty alts")
755 mkDefaultLitAlt alts@((DEFAULT, _, _) : _) = alts
756 mkDefaultLitAlt ((LitAlt{}, [], rhs) : alts) = (DEFAULT, [], rhs) : alts
757 mkDefaultLitAlt alts = pprPanic "mkDefaultLitAlt" (text "Not a lit alt:" <+> ppr alts)