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

4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

Note [Unarisation]
~~~~~~~~~~~~~~~~~~

The idea of this pass is to translate away *all* unboxed-tuple binders. So for example:

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

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.

STG fed to the code generators *must* be unarised because the code generators do
not support unboxed tuple binders natively.


Note [Unarisation and arity]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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.
Austin Seipp's avatar
Austin Seipp committed
28
-}
29

30
31
{-# LANGUAGE CPP #-}

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
module UnariseStg (unarise) where

#include "HsVersions.h"

import CoreSyn
import StgSyn
import VarEnv
import UniqSupply
import Id
import MkId (realWorldPrimId)
import Type
import TysWiredIn
import DataCon
import VarSet
import OccName
import Name
import Util
import Outputable
import BasicTypes


-- | A mapping from unboxed-tuple binders to the Ids they were expanded to.
--
-- INVARIANT: Ids in the range don't have unboxed tuple types.
--
-- Those in-scope variables without unboxed-tuple types are not present in
-- the domain of the mapping at all.
type UnariseEnv = VarEnv [Id]

ubxTupleId0 :: Id
62
ubxTupleId0 = dataConWorkId (tupleDataCon Unboxed 0)
63
64
65

unarise :: UniqSupply -> [StgBinding] -> [StgBinding]
unarise us binds = zipWith (\us -> unariseBinding us init_env) (listSplitUniqSupply us) binds
66
  where -- See Note [Nullary unboxed tuple] in Type.hs
67
68
69
70
71
        init_env = unitVarEnv ubxTupleId0 [realWorldPrimId]

unariseBinding :: UniqSupply -> UnariseEnv -> StgBinding -> StgBinding
unariseBinding us rho bind = case bind of
  StgNonRec x rhs -> StgNonRec x (unariseRhs us rho rhs)
Austin Seipp's avatar
Austin Seipp committed
72
  StgRec xrhss    -> StgRec $ zipWith (\us (x, rhs) -> (x, unariseRhs us rho rhs))
Simon Peyton Jones's avatar
Simon Peyton Jones committed
73
                                      (listSplitUniqSupply us) xrhss
74
75
76
77

unariseRhs :: UniqSupply -> UnariseEnv -> StgRhs -> StgRhs
unariseRhs us rho rhs = case rhs of
  StgRhsClosure ccs b_info fvs update_flag srt args expr
Austin Seipp's avatar
Austin Seipp committed
78
    -> StgRhsClosure ccs b_info (unariseIds rho fvs) update_flag
Simon Peyton Jones's avatar
Simon Peyton Jones committed
79
                     (unariseSRT rho srt) args' (unariseExpr us' rho' expr)
80
81
82
83
    where (us', rho', args') = unariseIdBinders us rho args
  StgRhsCon ccs con args
    -> StgRhsCon ccs con (unariseArgs rho args)

Simon Peyton Jones's avatar
Simon Peyton Jones committed
84
------------------------
85
unariseExpr :: UniqSupply -> UnariseEnv -> StgExpr -> StgExpr
Simon Peyton Jones's avatar
Simon Peyton Jones committed
86
87
88
unariseExpr _ rho (StgApp f args)
  | null args
  , UbxTupleRep tys <- repType (idType f)
Austin Seipp's avatar
Austin Seipp committed
89
  =  -- Particularly important where (##) is concerned
Simon Peyton Jones's avatar
Simon Peyton Jones committed
90
     -- See Note [Nullary unboxed tuple]
91
    StgConApp (tupleDataCon Unboxed (length tys))
Simon Peyton Jones's avatar
Simon Peyton Jones committed
92
93
94
95
96
              (map StgVarArg (unariseId rho f))

  | otherwise
  = StgApp f (unariseArgs rho args)

Austin Seipp's avatar
Austin Seipp committed
97
unariseExpr _ _ (StgLit l)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
98
99
100
  = StgLit l

unariseExpr _ rho (StgConApp dc args)
101
  | isUnboxedTupleCon dc = StgConApp (tupleDataCon Unboxed (length args')) args'
Simon Peyton Jones's avatar
Simon Peyton Jones committed
102
  | otherwise            = StgConApp dc args'
Austin Seipp's avatar
Austin Seipp committed
103
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
104
105
106
107
108
109
110
    args' = unariseArgs rho args

unariseExpr _ rho (StgOpApp op args ty)
  = StgOpApp op (unariseArgs rho args) ty

unariseExpr us rho (StgLam xs e)
  = StgLam xs' (unariseExpr us' rho' e)
Austin Seipp's avatar
Austin Seipp committed
111
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
112
113
114
    (us', rho', xs') = unariseIdBinders us rho xs

unariseExpr us rho (StgCase e case_lives alts_lives bndr srt alt_ty alts)
Austin Seipp's avatar
Austin Seipp committed
115
116
  = StgCase (unariseExpr us1 rho e) (unariseLives rho case_lives)
            (unariseLives rho alts_lives) bndr (unariseSRT rho srt)
117
            alt_ty alts'
Austin Seipp's avatar
Austin Seipp committed
118
 where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
119
    (us1, us2) = splitUniqSupply us
120
    alts'      = unariseAlts us2 rho alt_ty bndr alts
Simon Peyton Jones's avatar
Simon Peyton Jones committed
121
122
123

unariseExpr us rho (StgLet bind e)
  = StgLet (unariseBinding us1 rho bind) (unariseExpr us2 rho e)
Austin Seipp's avatar
Austin Seipp committed
124
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
125
126
127
    (us1, us2) = splitUniqSupply us

unariseExpr us rho (StgLetNoEscape live_in_let live_in_bind bind e)
Austin Seipp's avatar
Austin Seipp committed
128
  = StgLetNoEscape (unariseLives rho live_in_let) (unariseLives rho live_in_bind)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
129
                   (unariseBinding us1 rho bind) (unariseExpr us2 rho e)
Austin Seipp's avatar
Austin Seipp committed
130
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
131
132
    (us1, us2) = splitUniqSupply us

133
134
unariseExpr us rho (StgTick tick e)
  = StgTick tick (unariseExpr us rho e)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
135
136

------------------------
137
138
139
unariseAlts :: UniqSupply -> UnariseEnv -> AltType -> Id -> [StgAlt] -> [StgAlt]
unariseAlts us rho (UbxTupAlt n) bndr [(DEFAULT, [], [], e)]
  = [(DataAlt (tupleDataCon Unboxed n), ys, uses, unariseExpr us2' rho' e)]
Austin Seipp's avatar
Austin Seipp committed
140
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
141
142
143
    (us2', rho', ys) = unariseIdBinder us rho bndr
    uses = replicate (length ys) (not (isDeadBinder bndr))

144
145
unariseAlts us rho (UbxTupAlt n) bndr [(DataAlt _, ys, uses, e)]
  = [(DataAlt (tupleDataCon Unboxed n), ys', uses', unariseExpr us2' rho'' e)]
Austin Seipp's avatar
Austin Seipp committed
146
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
147
148
149
    (us2', rho', ys', uses') = unariseUsedIdBinders us rho ys uses
    rho'' = extendVarEnv rho' bndr ys'

150
unariseAlts _ _ (UbxTupAlt _) _ alts
Simon Peyton Jones's avatar
Simon Peyton Jones committed
151
152
  = pprPanic "unariseExpr: strange unboxed tuple alts" (ppr alts)

153
154
155
unariseAlts us rho _ _ alts
  = zipWith (\us alt -> unariseAlt us rho alt) (listSplitUniqSupply us) alts

Simon Peyton Jones's avatar
Simon Peyton Jones committed
156
--------------------------
157
unariseAlt :: UniqSupply -> UnariseEnv -> StgAlt -> StgAlt
Austin Seipp's avatar
Austin Seipp committed
158
unariseAlt us rho (con, xs, uses, e)
Simon Peyton Jones's avatar
Simon Peyton Jones committed
159
  = (con, xs', uses', unariseExpr us' rho' e)
Austin Seipp's avatar
Austin Seipp committed
160
  where
Simon Peyton Jones's avatar
Simon Peyton Jones committed
161
    (us', rho', xs', uses') = unariseUsedIdBinders us rho xs uses
162

Simon Peyton Jones's avatar
Simon Peyton Jones committed
163
------------------------
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
unariseSRT :: UnariseEnv -> SRT -> SRT
unariseSRT _   NoSRT            = NoSRT
unariseSRT rho (SRTEntries ids) = SRTEntries (concatMapVarSet (unariseId rho) ids)

unariseLives :: UnariseEnv -> StgLiveVars -> StgLiveVars
unariseLives rho ids = concatMapVarSet (unariseId rho) ids

unariseArgs :: UnariseEnv -> [StgArg] -> [StgArg]
unariseArgs rho = concatMap (unariseArg rho)

unariseArg :: UnariseEnv -> StgArg -> [StgArg]
unariseArg rho (StgVarArg x) = map StgVarArg (unariseId rho x)
unariseArg _   (StgLitArg l) = [StgLitArg l]

unariseIds :: UnariseEnv -> [Id] -> [Id]
unariseIds rho = concatMap (unariseId rho)

unariseId :: UnariseEnv -> Id -> [Id]
Austin Seipp's avatar
Austin Seipp committed
182
unariseId rho x
Simon Peyton Jones's avatar
Simon Peyton Jones committed
183
  | Just ys <- lookupVarEnv rho x
Austin Seipp's avatar
Austin Seipp committed
184
  = ASSERT2( case repType (idType x) of UbxTupleRep _ -> True; _ -> x == ubxTupleId0
Simon Peyton Jones's avatar
Simon Peyton Jones committed
185
186
187
188
189
190
191
192
           , text "unariseId: not unboxed tuple" <+> ppr x )
    ys

  | otherwise
  = ASSERT2( case repType (idType x) of UbxTupleRep _ -> False; _ -> True
           , text "unariseId: was unboxed tuple" <+> ppr x )
    [x]

Austin Seipp's avatar
Austin Seipp committed
193
unariseUsedIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> [Bool]
Simon Peyton Jones's avatar
Simon Peyton Jones committed
194
                     -> (UniqSupply, UnariseEnv, [Id], [Bool])
Austin Seipp's avatar
Austin Seipp committed
195
unariseUsedIdBinders us rho xs uses
Simon Peyton Jones's avatar
Simon Peyton Jones committed
196
197
198
199
  = case mapAccumL2 do_one us rho (zipEqual "unariseUsedIdBinders" xs uses) of
      (us', rho', xs_usess) -> uncurry ((,,,) us' rho') (unzip (concat xs_usess))
  where
    do_one us rho (x, use) = third3 (map (flip (,) use)) (unariseIdBinder us rho x)
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

unariseIdBinders :: UniqSupply -> UnariseEnv -> [Id] -> (UniqSupply, UnariseEnv, [Id])
unariseIdBinders us rho xs = third3 concat $ mapAccumL2 unariseIdBinder us rho xs

unariseIdBinder :: UniqSupply -> UnariseEnv -> Id -> (UniqSupply, UnariseEnv, [Id])
unariseIdBinder us rho x = case repType (idType x) of
    UnaryRep _      -> (us, rho, [x])
    UbxTupleRep tys -> let (us0, us1) = splitUniqSupply us
                           ys   = unboxedTupleBindersFrom us0 x tys
                           rho' = extendVarEnv rho x ys
                       in (us1, rho', ys)

unboxedTupleBindersFrom :: UniqSupply -> Id -> [UnaryType] -> [Id]
unboxedTupleBindersFrom us x tys = zipWith (mkSysLocal fs) (uniqsFromSupply us) tys
  where fs = occNameFS (getOccName x)

concatMapVarSet :: (Var -> [Var]) -> VarSet -> VarSet
concatMapVarSet f xs = mkVarSet [x' | x <- varSetElems xs, x' <- f x]