UnariseStg.hs 26.9 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3
{-
(c) The GRASP/AQUA Project, Glasgow University, 1992-2012

4 5
Note [Unarisation]
~~~~~~~~~~~~~~~~~~
6 7 8 9 10 11
The idea of this pass is to translate away *all* unboxed-tuple and unboxed-sum
binders. So for example:

  f (x :: (# Int, Bool #)) = f x + f (# 1, True #)

  ==>
12

13
  f (x1 :: Int) (x2 :: Bool) = f x1 x2 + f 1 True
14

15 16 17
It is important that we do this at the STG level and NOT at the Core level
because it would be very hard to make this pass Core-type-preserving. In this
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 25 26 27

Suppose that a variable x : (# t1, t2 #).

  * At the binding site for x, make up fresh vars  x1:t1, x2:t2

28
  * Extend the UnariseEnv   x :-> MultiVal [x1,x2]
29 30

  * Replace the binding with a curried binding for x1,x2
31

32 33 34
       Lambda:   \x.e                ==>   \x1 x2. e
       Case alt: MkT a b x c d -> e  ==>   MkT a b x1 x2 c d -> e

35 36 37
  * Replace argument occurrences with a sequence of args via a lookup in
    UnariseEnv

38 39
       f a b x c d   ==>   f a b x1 x2 c d

40 41 42
  * Replace tail-call occurrences with an unboxed tuple via a lookup in
    UnariseEnv

43
       x  ==>  (# x1, x2 #)
44

45
    So, for example
46

47 48
       f x = x    ==>   f x1 x2 = (# x1, x2 #)

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
  * We /always/ eliminate a case expression when

       - It scrutinises an unboxed tuple or unboxed sum

       - The scrutinee is a variable (or when it is an explicit tuple, but the
         simplifier eliminates those)

    The case alternative (there can be only one) can be one of these two
    things:

      - An unboxed tuple pattern. e.g.

          case v of x { (# x1, x2, x3 #) -> ... }

        Scrutinee has to be in form `(# t1, t2, t3 #)` so we just extend the
        environment with

          x :-> MultiVal [t1,t2,t3]
          x1 :-> UnaryVal t1, x2 :-> UnaryVal t2, x3 :-> UnaryVal t3

      - A DEFAULT alternative. Just the same, without the bindings for x1,x2,x3

By the end of this pass, we only have unboxed tuples in return positions.
Unboxed sums are completely eliminated, see next note.

Note [Translating unboxed sums to unboxed tuples]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Unarise also eliminates unboxed sum binders, and translates unboxed sums in
return positions to unboxed tuples. We want to overlap fields of a sum when
translating it to a tuple to have efficient memory layout. When translating a
sum pattern to a tuple pattern, we need to translate it so that binders of sum
alternatives will be mapped to right arguments after the term translation. So
translation of sum DataCon applications to tuple DataCon applications and
translation of sum patterns to tuple patterns need to be in sync.

These translations work like this. Suppose we have

  (# x1 | | ... #) :: (# t1 | t2 | ... #)

remember that t1, t2 ... can be sums and tuples too. So we first generate
layouts of those. Then we "merge" layouts of each alternative, which gives us a
sum layout with best overlapping possible.

Layout of a flat type 'ty1' is just [ty1].
Layout of a tuple is just concatenation of layouts of its fields.

For layout of a sum type,
96

97 98 99
  - We first get layouts of all alternatives.
  - We sort these layouts based on their "slot types".
  - We merge all the alternatives.
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
For example, say we have (# (# Int#, Char #) | (# Int#, Int# #) | Int# #)

  - Layouts of alternatives: [ [Word, Ptr], [Word, Word], [Word] ]
  - Sorted: [ [Ptr, Word], [Word, Word], [Word] ]
  - Merge all alternatives together: [ Ptr, Word, Word ]

We add a slot for the tag to the first position. So our tuple type is

  (# Tag#, Any, Word#, Word# #)
  (we use Any for pointer slots)

Now, any term of this sum type needs to generate a tuple of this type instead.
The translation works by simply putting arguments to first slots that they fit
in. Suppose we had

  (# (# 42#, 'c' #) | | #)

42# fits in Word#, 'c' fits in Any, so we generate this application:

  (# 1#, 'c', 42#, rubbish #)

Another example using the same type: (# | (# 2#, 3# #) | #). 2# fits in Word#,
3# fits in Word #, so we get:

  (# 2#, rubbish, 2#, 3# #).

Note [Types in StgConApp]
~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have this unboxed sum term:

  (# 123 | #)

What will be the unboxed tuple representation? We can't tell without knowing the
type of this term. For example, these are all valid tuples for this:

  (# 1#, 123 #)          -- when type is (# Int | String #)
  (# 1#, 123, rubbish #) -- when type is (# Int | Float# #)
  (# 1#, 123, rubbish, rubbish #)
                         -- when type is (# Int | (# Int, Int, Int #) #)

So we pass type arguments of the DataCon's TyCon in StgConApp to decide what
layout to use. Note that unlifted values can't be let-bound, so we don't need
types in StgRhsCon.

Note [UnariseEnv can map to literals]
146
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
147 148 149 150
To avoid redundant case expressions when unarising unboxed sums, UnariseEnv
needs to map variables to literals too. Suppose we have this Core:

  f (# x | #)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
151

152
  ==> (CorePrep)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
153

154 155 156
  case (# x | #) of y {
    _ -> f y
  }
Simon Peyton Jones's avatar
Simon Peyton Jones committed
157

158
  ==> (MultiVal)
159

160 161 162
  case (# 1#, x #) of [x1, x2] {
    _ -> f x1 x2
  }
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 171

Note [Unarisation and arity]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
172 173 174 175 176
Because of unarisation, the arity that will be recorded in the generated info
table for an Id may be larger than the idArity. Instead we record what we call
the RepArity, which is the Arity taking into account any expanded arguments, and
corresponds to the number of (possibly-void) *registers* arguments will arrive
in.
177 178 179 180 181 182 183 184 185 186 187 188 189 190

Note [Post-unarisation invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
STG programs after unarisation have these invariants:

  * No unboxed sums at all.

  * No unboxed tuple binders. Tuples only appear in return position.

  * DataCon applications (StgRhsCon and StgConApp) don't have void arguments.
    This means that it's safe to wrap `StgArg`s of DataCon applications with
    `StgCmmEnv.NonVoid`, for example.

  * Alt binders (binders in patterns) are always non-void.
Austin Seipp's avatar
Austin Seipp committed
191
-}
192

193
{-# LANGUAGE CPP, TupleSections #-}
194

195 196 197 198
module UnariseStg (unarise) where

#include "HsVersions.h"

199 200
import GhcPrelude

201
import BasicTypes
202
import CoreSyn
203 204
import DataCon
import FastString (FastString, mkFastString)
205
import Id
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
206
import Literal (Literal (..), literalType)
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
207
import MkCore (aBSENT_SUM_FIELD_ERROR_ID)
208 209 210 211 212
import MkId (voidPrimId, voidArgId)
import MonadUtils (mapAccumLM)
import Outputable
import RepType
import StgSyn
213
import Type
Richard Eisenberg's avatar
Richard Eisenberg committed
214
import TysPrim (intPrimTy)
215
import TysWiredIn
216
import UniqSupply
217
import Util
218
import VarEnv
219

220 221 222
import Data.Bifunctor (second)
import Data.Maybe (mapMaybe)
import qualified Data.IntMap as IM
223

224 225 226 227 228 229
--------------------------------------------------------------------------------

-- | A mapping from binders to the Ids they were expanded/renamed to.
--
--   x :-> MultiVal [a,b,c] in rho
--
Richard Eisenberg's avatar
Richard Eisenberg committed
230
-- iff  x's typePrimRep is not a singleton, or equivalently
231 232 233 234 235 236
--      x's type is an unboxed tuple, sum or void.
--
--    x :-> UnaryVal x'
--
-- iff x's RepType is UnaryRep or equivalently
--     x's type is not unboxed tuple, sum or void.
237
--
238 239 240
-- So
--     x :-> MultiVal [a] in rho
-- means x is represented by singleton tuple.
241
--
242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260
--     x :-> MultiVal [] in rho
-- means x is void.
--
-- INVARIANT: OutStgArgs in the range only have NvUnaryTypes
--            (i.e. no unboxed tuples, sums or voids)
--
type UnariseEnv = VarEnv UnariseVal

data UnariseVal
  = MultiVal [OutStgArg] -- MultiVal to tuple. Can be empty list (void).
  | UnaryVal OutStgArg   -- See NOTE [Renaming during unarisation].

instance Outputable UnariseVal where
  ppr (MultiVal args) = text "MultiVal" <+> ppr args
  ppr (UnaryVal arg)   = text "UnaryVal" <+> ppr arg

-- | Extend the environment, checking the UnariseEnv invariant.
extendRho :: UnariseEnv -> Id -> UnariseVal -> UnariseEnv
extendRho rho x (MultiVal args)
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
261
  = ASSERT(all (isNvUnaryType . stgArgType) args)
262 263
    extendVarEnv rho x (MultiVal args)
extendRho rho x (UnaryVal val)
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
264
  = ASSERT(isNvUnaryType (stgArgType val))
265 266 267 268
    extendVarEnv rho x (UnaryVal val)

--------------------------------------------------------------------------------

269 270 271 272 273 274 275
unarise :: UniqSupply -> [StgTopBinding] -> [StgTopBinding]
unarise us binds = initUs_ us (mapM (unariseTopBinding emptyVarEnv) binds)

unariseTopBinding :: UnariseEnv -> StgTopBinding -> UniqSM StgTopBinding
unariseTopBinding rho (StgTopLifted bind)
  = StgTopLifted <$> unariseBinding rho bind
unariseTopBinding _ bind@StgTopStringLit{} = return bind
276 277 278 279 280 281 282 283 284 285 286 287 288 289 290

unariseBinding :: UnariseEnv -> StgBinding -> UniqSM StgBinding
unariseBinding rho (StgNonRec x rhs)
  = StgNonRec x <$> unariseRhs rho rhs
unariseBinding rho (StgRec xrhss)
  = StgRec <$> mapM (\(x, rhs) -> (x,) <$> unariseRhs rho rhs) xrhss

unariseRhs :: UnariseEnv -> StgRhs -> UniqSM StgRhs
unariseRhs rho (StgRhsClosure ccs b_info fvs update_flag args expr)
  = do (rho', args1) <- unariseFunArgBinders rho args
       expr' <- unariseExpr rho' expr
       let fvs' = unariseFreeVars rho fvs
       return (StgRhsClosure ccs b_info fvs' update_flag args1 expr')

unariseRhs rho (StgRhsCon ccs con args)
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
291
  = ASSERT(not (isUnboxedTupleCon con || isUnboxedSumCon con))
292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310
    return (StgRhsCon ccs con (unariseConArgs rho args))

--------------------------------------------------------------------------------

unariseExpr :: UnariseEnv -> StgExpr -> UniqSM StgExpr

unariseExpr rho e@(StgApp f [])
  = case lookupVarEnv rho f of
      Just (MultiVal args)  -- Including empty tuples
        -> return (mkTuple args)
      Just (UnaryVal (StgVarArg f'))
        -> return (StgApp f' [])
      Just (UnaryVal (StgLitArg f'))
        -> return (StgLit f')
      Nothing
        -> return e

unariseExpr rho e@(StgApp f args)
  = return (StgApp f' (unariseFunArgs rho args))
311
  where
312 313 314 315 316 317 318 319 320 321 322 323 324
    f' = case lookupVarEnv rho f of
           Just (UnaryVal (StgVarArg f')) -> f'
           Nothing -> f
           err -> pprPanic "unariseExpr - app2" (ppr e $$ ppr err)
               -- Can't happen because 'args' is non-empty, and
               -- a tuple or sum cannot be applied to anything

unariseExpr _ (StgLit l)
  = return (StgLit l)

unariseExpr rho (StgConApp dc args ty_args)
  | Just args' <- unariseMulti_maybe rho dc args ty_args
  = return (mkTuple args')
Simon Peyton Jones's avatar
Simon Peyton Jones committed
325 326

  | otherwise
327 328
  , let args' = unariseConArgs rho args
  = return (StgConApp dc args' (map stgArgType args'))
Simon Peyton Jones's avatar
Simon Peyton Jones committed
329

330 331
unariseExpr rho (StgOpApp op args ty)
  = return (StgOpApp op (unariseFunArgs rho args) ty)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
332

333 334
unariseExpr _ e@StgLam{}
  = pprPanic "unariseExpr: found lambda" (ppr e)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
335

336
unariseExpr rho (StgCase scrut bndr alt_ty alts)
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
337
  -- tuple/sum binders in the scrutinee can always be eliminated
338 339 340
  | StgApp v [] <- scrut
  , Just (MultiVal xs) <- lookupVarEnv rho v
  = elimCase rho xs bndr alt_ty alts
Simon Peyton Jones's avatar
Simon Peyton Jones committed
341

342 343 344 345 346 347
  -- Handle strict lets for tuples and sums:
  --   case (# a,b #) of r -> rhs
  -- and analogously for sums
  | StgConApp dc args ty_args <- scrut
  , Just args' <- unariseMulti_maybe rho dc args ty_args
  = elimCase rho args' bndr alt_ty alts
Simon Peyton Jones's avatar
Simon Peyton Jones committed
348

349 350 351 352 353
  -- general case
  | otherwise
  = do scrut' <- unariseExpr rho scrut
       alts'  <- unariseAlts rho alt_ty bndr alts
       return (StgCase scrut' bndr alt_ty alts')
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
354 355
                       -- bndr may have a unboxed sum/tuple type but it will be
                       -- dead after unarise (checked in StgLint)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
356

357 358
unariseExpr rho (StgLet bind e)
  = StgLet <$> unariseBinding rho bind <*> unariseExpr rho e
Simon Peyton Jones's avatar
Simon Peyton Jones committed
359

360 361
unariseExpr rho (StgLetNoEscape bind e)
  = StgLetNoEscape <$> unariseBinding rho bind <*> unariseExpr rho e
Simon Peyton Jones's avatar
Simon Peyton Jones committed
362

363 364
unariseExpr rho (StgTick tick e)
  = StgTick tick <$> unariseExpr rho e
Simon Peyton Jones's avatar
Simon Peyton Jones committed
365

366 367 368 369 370
-- Doesn't return void args.
unariseMulti_maybe :: UnariseEnv -> DataCon -> [InStgArg] -> [Type] -> Maybe [OutStgArg]
unariseMulti_maybe rho dc args ty_args
  | isUnboxedTupleCon dc
  = Just (unariseConArgs rho args)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
371

372
  | isUnboxedSumCon dc
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
373
  , let args1 = ASSERT(isSingleton args) (unariseConArgs rho args)
374
  = Just (mkUbxSum dc ty_args args1)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
375

376 377 378 379 380 381 382 383 384 385 386 387 388 389 390
  | otherwise
  = Nothing

--------------------------------------------------------------------------------

elimCase :: UnariseEnv
         -> [OutStgArg] -- non-void args
         -> InId -> AltType -> [InStgAlt] -> UniqSM OutStgExpr

elimCase rho args bndr (MultiValAlt _) [(_, bndrs, rhs)]
  = do let rho1 = extendRho rho bndr (MultiVal args)
           rho2
             | isUnboxedTupleBndr bndr
             = mapTupleIdBinders bndrs args rho1
             | otherwise
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
391
             = ASSERT(isUnboxedSumBndr bndr)
392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425
               if null bndrs then rho1
                             else mapSumIdBinders bndrs args rho1

       unariseExpr rho2 rhs

elimCase rho args bndr (MultiValAlt _) alts
  | isUnboxedSumBndr bndr
  = do let (tag_arg : real_args) = args
       tag_bndr <- mkId (mkFastString "tag") tagTy
          -- this won't be used but we need a binder anyway
       let rho1 = extendRho rho bndr (MultiVal args)
           scrut' = case tag_arg of
                      StgVarArg v     -> StgApp v []
                      StgLitArg l     -> StgLit l

       alts' <- unariseSumAlts rho1 real_args alts
       return (StgCase scrut' tag_bndr tagAltTy alts')

elimCase _ args bndr alt_ty alts
  = pprPanic "elimCase - unhandled case"
      (ppr args <+> ppr bndr <+> ppr alt_ty $$ ppr alts)

--------------------------------------------------------------------------------

unariseAlts :: UnariseEnv -> AltType -> InId -> [StgAlt] -> UniqSM [StgAlt]
unariseAlts rho (MultiValAlt n) bndr [(DEFAULT, [], e)]
  | isUnboxedTupleBndr bndr
  = do (rho', ys) <- unariseConArgBinder rho bndr
       e' <- unariseExpr rho' e
       return [(DataAlt (tupleDataCon Unboxed n), ys, e')]

unariseAlts rho (MultiValAlt n) bndr [(DataAlt _, ys, e)]
  | isUnboxedTupleBndr bndr
  = do (rho', ys1) <- unariseConArgBinders rho ys
426
       MASSERT(ys1 `lengthIs` n)
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
       let rho'' = extendRho rho' bndr (MultiVal (map StgVarArg ys1))
       e' <- unariseExpr rho'' e
       return [(DataAlt (tupleDataCon Unboxed n), ys1, e')]

unariseAlts _ (MultiValAlt _) bndr alts
  | isUnboxedTupleBndr bndr
  = pprPanic "unariseExpr: strange multi val alts" (ppr alts)

-- In this case we don't need to scrutinize the tag bit
unariseAlts rho (MultiValAlt _) bndr [(DEFAULT, _, rhs)]
  | isUnboxedSumBndr bndr
  = do (rho_sum_bndrs, sum_bndrs) <- unariseConArgBinder rho bndr
       rhs' <- unariseExpr rho_sum_bndrs rhs
       return [(DataAlt (tupleDataCon Unboxed (length sum_bndrs)), sum_bndrs, rhs')]

unariseAlts rho (MultiValAlt _) bndr alts
  | isUnboxedSumBndr bndr
  = do (rho_sum_bndrs, scrt_bndrs@(tag_bndr : real_bndrs)) <- unariseConArgBinder rho bndr
       alts' <- unariseSumAlts rho_sum_bndrs (map StgVarArg real_bndrs) alts
       let inner_case = StgCase (StgApp tag_bndr []) tag_bndr tagAltTy alts'
       return [ (DataAlt (tupleDataCon Unboxed (length scrt_bndrs)),
                 scrt_bndrs,
                 inner_case) ]

unariseAlts rho _ _ alts
  = mapM (\alt -> unariseAlt rho alt) alts

unariseAlt :: UnariseEnv -> StgAlt -> UniqSM StgAlt
unariseAlt rho (con, xs, e)
  = do (rho', xs') <- unariseConArgBinders rho xs
       (con, xs',) <$> unariseExpr rho' e

--------------------------------------------------------------------------------

-- | Make alternatives that match on the tag of a sum
-- (i.e. generate LitAlts for the tag)
unariseSumAlts :: UnariseEnv
               -> [StgArg] -- sum components _excluding_ the tag bit.
               -> [StgAlt] -- original alternative with sum LHS
               -> UniqSM [StgAlt]
unariseSumAlts env args alts
  = do alts' <- mapM (unariseSumAlt env args) alts
       return (mkDefaultLitAlt alts')

unariseSumAlt :: UnariseEnv
              -> [StgArg] -- sum components _excluding_ the tag bit.
              -> StgAlt   -- original alternative with sum LHS
              -> UniqSM StgAlt
unariseSumAlt rho _ (DEFAULT, _, e)
  = ( DEFAULT, [], ) <$> unariseExpr rho e

unariseSumAlt rho args (DataAlt sumCon, bs, e)
  = do let rho' = mapSumIdBinders bs args rho
       e' <- unariseExpr rho' e
       return ( LitAlt (MachInt (fromIntegral (dataConTag sumCon))), [], e' )

unariseSumAlt _ scrt alt
  = pprPanic "unariseSumAlt" (ppr scrt $$ ppr alt)

--------------------------------------------------------------------------------

mapTupleIdBinders
  :: [InId]       -- Un-processed binders of a tuple alternative.
                  -- Can have void binders.
  -> [OutStgArg]  -- Arguments that form the tuple (after unarisation).
                  -- Can't have void args.
  -> UnariseEnv
  -> UnariseEnv
mapTupleIdBinders ids args0 rho0
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
496
  = ASSERT(not (any (isVoidTy . stgArgType) args0))
497
    let
Richard Eisenberg's avatar
Richard Eisenberg committed
498 499
      ids_unarised :: [(Id, [PrimRep])]
      ids_unarised = map (\id -> (id, typePrimRep (idType id))) ids
500

Richard Eisenberg's avatar
Richard Eisenberg committed
501
      map_ids :: UnariseEnv -> [(Id, [PrimRep])] -> [StgArg] -> UnariseEnv
502
      map_ids rho [] _  = rho
Richard Eisenberg's avatar
Richard Eisenberg committed
503
      map_ids rho ((x, x_reps) : xs) args =
504
        let
Richard Eisenberg's avatar
Richard Eisenberg committed
505
          x_arity = length x_reps
506 507 508 509 510
          (x_args, args') =
            ASSERT(args `lengthAtLeast` x_arity)
            splitAt x_arity args

          rho'
Richard Eisenberg's avatar
Richard Eisenberg committed
511
            | x_arity == 1
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
512
            = ASSERT(x_args `lengthIs` 1)
513
              extendRho rho x (UnaryVal (head x_args))
Richard Eisenberg's avatar
Richard Eisenberg committed
514 515
            | otherwise
            = extendRho rho x (MultiVal x_args)
516 517 518 519 520 521 522 523 524 525 526 527 528 529
        in
          map_ids rho' xs args'
    in
      map_ids rho0 ids_unarised args0

mapSumIdBinders
  :: [InId]      -- Binder of a sum alternative (remember that sum patterns
                 -- only have one binder, so this list should be a singleton)
  -> [OutStgArg] -- Arguments that form the sum (NOT including the tag).
                 -- Can't have void args.
  -> UnariseEnv
  -> UnariseEnv

mapSumIdBinders [id] args rho0
Erik de Castro Lopo's avatar
Erik de Castro Lopo committed
530
  = ASSERT(not (any (isVoidTy . stgArgType) args))
531
    let
Richard Eisenberg's avatar
Richard Eisenberg committed
532 533 534
      arg_slots = map primRepSlot $ concatMap (typePrimRep . stgArgType) args
      id_slots  = map primRepSlot $ typePrimRep (idType id)
      layout1   = layoutUbxSum arg_slots id_slots
535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560
    in
      if isMultiValBndr id
        then extendRho rho0 id (MultiVal [ args !! i | i <- layout1 ])
        else ASSERT(layout1 `lengthIs` 1)
             extendRho rho0 id (UnaryVal (args !! head layout1))

mapSumIdBinders ids sum_args _
  = pprPanic "mapSumIdBinders" (ppr ids $$ ppr sum_args)

-- | Build a unboxed sum term from arguments of an alternative.
--
-- Example, for (# x | #) :: (# (# #) | Int #) we call
--
--   mkUbxSum (# _ | #) [ (# #), Int ] [ voidPrimId ]
--
-- which returns
--
--   [ 1#, rubbish ]
--
mkUbxSum
  :: DataCon      -- Sum data con
  -> [Type]       -- Type arguments of the sum data con
  -> [OutStgArg]  -- Actual arguments of the alternative.
  -> [OutStgArg]  -- Final tuple arguments
mkUbxSum dc ty_args args0
  = let
Richard Eisenberg's avatar
Richard Eisenberg committed
561
      (_ : sum_slots) = ubxSumRepType (map typePrimRep ty_args)
562 563 564 565
        -- drop tag slot

      tag = dataConTag dc

Richard Eisenberg's avatar
Richard Eisenberg committed
566
      layout'  = layoutUbxSum sum_slots (mapMaybe (typeSlotTy . stgArgType) args0)
567 568 569 570 571 572 573 574 575 576
      tag_arg  = StgLitArg (MachInt (fromIntegral tag))
      arg_idxs = IM.fromList (zipEqual "mkUbxSum" layout' args0)

      mkTupArgs :: Int -> [SlotTy] -> IM.IntMap StgArg -> [StgArg]
      mkTupArgs _ [] _
        = []
      mkTupArgs arg_idx (slot : slots_left) arg_map
        | Just stg_arg <- IM.lookup arg_idx arg_map
        = stg_arg : mkTupArgs (arg_idx + 1) slots_left arg_map
        | otherwise
577 578 579
        = slotRubbishArg slot : mkTupArgs (arg_idx + 1) slots_left arg_map

      slotRubbishArg :: SlotTy -> StgArg
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
580 581
      slotRubbishArg PtrSlot    = StgVarArg aBSENT_SUM_FIELD_ERROR_ID
                         -- See Note [aBSENT_SUM_FIELD_ERROR_ID] in MkCore
582 583 584 585
      slotRubbishArg WordSlot   = StgLitArg (MachWord 0)
      slotRubbishArg Word64Slot = StgLitArg (MachWord64 0)
      slotRubbishArg FloatSlot  = StgLitArg (MachFloat 0)
      slotRubbishArg DoubleSlot = StgLitArg (MachDouble 0)
586 587 588 589
    in
      tag_arg : mkTupArgs 0 sum_slots arg_idxs

--------------------------------------------------------------------------------
Simon Peyton Jones's avatar
Simon Peyton Jones committed
590

591 592
{-
For arguments (StgArg) and binders (Id) we have two kind of unarisation:
593

594 595
  - When unarising function arg binders and arguments, we don't want to remove
    void binders and arguments. For example,
596

597 598
      f :: (# (# #), (# #) #) -> Void# -> RealWorld# -> ...
      f x y z = <body>
599

600 601
    Here after unarise we should still get a function with arity 3. Similarly
    in the call site we shouldn't remove void arguments:
602

603
      f (# (# #), (# #) #) voidId rw
604

605
    When unarising <body>, we extend the environment with these binders:
Simon Peyton Jones's avatar
Simon Peyton Jones committed
606

607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646
      x :-> MultiVal [], y :-> MultiVal [], z :-> MultiVal []

    Because their rep types are `MultiRep []` (aka. void). This means that when
    we see `x` in a function argument position, we actually replace it with a
    void argument. When we see it in a DataCon argument position, we just get
    rid of it, because DataCon applications in STG are always saturated.

  - When unarising case alternative binders we remove void binders, but we
    still update the environment the same way, because those binders may be
    used in the RHS. Example:

      case x of y {
        (# x1, x2, x3 #) -> <RHS>
      }

    We know that y can't be void, because we don't scrutinize voids, so x will
    be unarised to some number of arguments, and those arguments will have at
    least one non-void thing. So in the rho we will have something like:

      x :-> MultiVal [xu1, xu2]

    Now, after we eliminate void binders in the pattern, we get exactly the same
    number of binders, and extend rho again with these:

      x1 :-> UnaryVal xu1
      x2 :-> MultiVal [] -- x2 is void
      x3 :-> UnaryVal xu2

    Now when we see x2 in a function argument position or in return position, we
    generate void#. In constructor argument position, we just remove it.

So in short, when we have a void id,

  - We keep it if it's a lambda argument binder or
                       in argument position of an application.

  - We remove it if it's a DataCon field binder or
                         in argument position of a DataCon application.
-}

Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675
unariseArgBinder
    :: Bool -- data con arg?
    -> UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
unariseArgBinder is_con_arg rho x =
  case typePrimRep (idType x) of
    []
      | is_con_arg
      -> return (extendRho rho x (MultiVal []), [])
      | otherwise -- fun arg, do not remove void binders
      -> return (extendRho rho x (MultiVal []), [voidArgId])

    [rep]
      -- Arg represented as single variable, but original type may still be an
      -- unboxed sum/tuple, e.g. (# Void# | Void# #).
      --
      -- While not unarising the binder in this case does not break any programs
      -- (because it unarises to a single variable), it triggers StgLint as we
      -- break the the post-unarisation invariant that says unboxed tuple/sum
      -- binders should vanish. See Note [Post-unarisation invariants].
      | isUnboxedSumType (idType x) || isUnboxedTupleType (idType x)
      -> do x' <- mkId (mkFastString "us") (primRepToType rep)
            return (extendRho rho x (MultiVal [StgVarArg x']), [x'])
      | otherwise
      -> return (rho, [x])

    reps -> do
      xs <- mkIds (mkFastString "us") (map primRepToType reps)
      return (extendRho rho x (MultiVal (map StgVarArg xs)), xs)

676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694
--------------------------------------------------------------------------------

-- | MultiVal a function argument. Never returns an empty list.
unariseFunArg :: UnariseEnv -> StgArg -> [StgArg]
unariseFunArg rho (StgVarArg x) =
  case lookupVarEnv rho x of
    Just (MultiVal [])  -> [voidArg]   -- NB: do not remove void args
    Just (MultiVal as)  -> as
    Just (UnaryVal arg) -> [arg]
    Nothing             -> [StgVarArg x]
unariseFunArg _ arg = [arg]

unariseFunArgs :: UnariseEnv -> [StgArg] -> [StgArg]
unariseFunArgs = concatMap . unariseFunArg

unariseFunArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
unariseFunArgBinders rho xs = second concat <$> mapAccumLM unariseFunArgBinder rho xs

-- Result list of binders is never empty
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
695 696
unariseFunArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
unariseFunArgBinder = unariseArgBinder False
697 698 699 700 701 702 703 704

--------------------------------------------------------------------------------

-- | MultiVal a DataCon argument. Returns an empty list when argument is void.
unariseConArg :: UnariseEnv -> InStgArg -> [OutStgArg]
unariseConArg rho (StgVarArg x) =
  case lookupVarEnv rho x of
    Just (UnaryVal arg) -> [arg]
705
    Just (MultiVal as) -> as      -- 'as' can be empty
706 707 708 709 710
    Nothing
      | isVoidTy (idType x) -> [] -- e.g. C realWorld#
                                  -- Here realWorld# is not in the envt, but
                                  -- is a void, and so should be eliminated
      | otherwise -> [StgVarArg x]
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
711 712 713
unariseConArg _ arg@(StgLitArg lit) =
    ASSERT(not (isVoidTy (literalType lit)))  -- We have no void literals
    [arg]
714 715 716 717 718 719 720

unariseConArgs :: UnariseEnv -> [InStgArg] -> [OutStgArg]
unariseConArgs = concatMap . unariseConArg

unariseConArgBinders :: UnariseEnv -> [Id] -> UniqSM (UnariseEnv, [Id])
unariseConArgBinders rho xs = second concat <$> mapAccumLM unariseConArgBinder rho xs

Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
721 722
-- Different from `unariseFunArgBinder`: result list of binders may be empty.
-- See DataCon applications case in Note [Post-unarisation invariants].
723
unariseConArgBinder :: UnariseEnv -> Id -> UniqSM (UnariseEnv, [Id])
Ömer Sinan Ağacan's avatar
Ömer Sinan Ağacan committed
724
unariseConArgBinder = unariseArgBinder True
725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752

unariseFreeVars :: UnariseEnv -> [InId] -> [OutId]
unariseFreeVars rho fvs
 = [ v | fv <- fvs, StgVarArg v <- unariseFreeVar rho fv ]
   -- Notice that we filter out any StgLitArgs
   -- e.g.   case e of (x :: (# Int | Bool #))
   --           (# v | #) ->  ... let {g = \y. ..x...} in ...
   --           (# | w #) -> ...
   --     Here 'x' is free in g's closure, and the env will have
   --       x :-> [1, v]
   --     we want to capture 'v', but not 1, in the free vars

unariseFreeVar :: UnariseEnv -> Id -> [StgArg]
unariseFreeVar rho x =
  case lookupVarEnv rho x of
    Just (MultiVal args) -> args
    Just (UnaryVal arg)  -> [arg]
    Nothing              -> [StgVarArg x]

--------------------------------------------------------------------------------

mkIds :: FastString -> [UnaryType] -> UniqSM [Id]
mkIds fs tys = mapM (mkId fs) tys

mkId :: FastString -> UnaryType -> UniqSM Id
mkId = mkSysLocalOrCoVarM

isMultiValBndr :: Id -> Bool
Richard Eisenberg's avatar
Richard Eisenberg committed
753 754 755 756 757
isMultiValBndr id
  | [_] <- typePrimRep (idType id)
  = False
  | otherwise
  = True
758 759 760 761 762 763 764 765 766 767 768

isUnboxedSumBndr :: Id -> Bool
isUnboxedSumBndr = isUnboxedSumType . idType

isUnboxedTupleBndr :: Id -> Bool
isUnboxedTupleBndr = isUnboxedTupleType . idType

mkTuple :: [StgArg] -> StgExpr
mkTuple args = StgConApp (tupleDataCon Unboxed (length args)) args (map stgArgType args)

tagAltTy :: AltType
Richard Eisenberg's avatar
Richard Eisenberg committed
769
tagAltTy = PrimAlt IntRep
770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786

tagTy :: Type
tagTy = intPrimTy

voidArg :: StgArg
voidArg = StgVarArg voidPrimId

mkDefaultLitAlt :: [StgAlt] -> [StgAlt]
-- We have an exhauseive list of literal alternatives
--    1# -> e1
--    2# -> e2
-- Since they are exhaustive, we can replace one with DEFAULT, to avoid
-- generating a final test. Remember, the DEFAULT comes first if it exists.
mkDefaultLitAlt [] = pprPanic "elimUbxSumExpr.mkDefaultAlt" (text "Empty alts")
mkDefaultLitAlt alts@((DEFAULT, _, _) : _) = alts
mkDefaultLitAlt ((LitAlt{}, [], rhs) : alts) = (DEFAULT, [], rhs) : alts
mkDefaultLitAlt alts = pprPanic "mkDefaultLitAlt" (text "Not a lit alt:" <+> ppr alts)