2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
4 \section[SatStgRhs]{Saturates RHSs when they are partial applications}
6 96/03: This is actually an essential module, as it sets arity info
7 for the code generator.
10 Subject: arg satis check
11 Date: Wed, 29 Apr 92 13:33:58 +0100
12 From: Simon L Peyton Jones <simonpj>
16 Another transformation to consider. We'd like to avoid
17 argument-satisfaction checks wherever possible. So, whenever we have an
18 STG binding application
20 f = vs \ xs -> g e1 ... en
22 where xs has one or more elements
24 where g is a known function with arity m+n,
28 f = vs \ xs++{x1...xm} -> g e1 ... en x1 .. xm
30 Now g has enough args. One arg-satisfaction check disappears;
31 the one for the closure incorporates the one for g.
33 You might like to consider variants, applying the transformation more
34 widely. I concluded that this was the only instance which made
35 sense, but I could be wrong.
40 The algorithm proceeds as follows:
43 Gather the arity information of the functions defined in this module
44 (as @getIdArity@ only knows about the arity of @ImportedIds@).
47 for every definition of the form
49 v = /\ts -> \vs -> f args
51 we try to match the arity of \tr{f} with the number of arguments.
52 If they do not match we insert extra lambdas to make that application
56 This is done for local definitions as well.
59 #include "HsVersions.h"
61 module SatStgRhs ( satStgRhs ) where
67 import CostCentre ( isCafCC, subsumedCosts, useCurrentCostCentre )
68 import Id ( idType, getIdArity, addIdArity, mkSysLocal,
69 nullIdEnv, addOneToIdEnv, growIdEnvList,
70 lookupIdEnv, IdEnv(..)
72 import IdInfo ( arityMaybe )
73 import SrcLoc ( mkUnknownSrcLoc )
74 import Type ( splitSigmaTy, splitFunTy )
75 import UniqSupply ( returnUs, thenUs, mapUs, getUnique, UniqSM(..) )
76 import Util ( panic, assertPanic )
80 type ExprArityInfo = Maybe Int -- Just n => This expression has a guaranteed
82 -- Nothing => Don't know how many args it needs
84 type Id_w_Arity = Id -- An Id with correct arity info pinned on it
85 type SatEnv = IdEnv Id_w_Arity -- Binds only local, let(rec)-bound things
90 \item adds extra args where necessary;
91 \item pins the correct arity on everything.
94 %************************************************************************
96 \subsection{Top-level list of bindings (a ``program'')}
98 %************************************************************************
101 satStgRhs :: [StgBinding] -> UniqSM [StgBinding]
103 satStgRhs p = satProgram nullIdEnv p
105 satProgram :: SatEnv -> [StgBinding] -> UniqSM [StgBinding]
106 satProgram env [] = returnUs []
108 satProgram env (bind:binds)
109 = satBinding True{-toplevel-} env bind `thenUs` \ (env2, bind2) ->
110 satProgram env2 binds `thenUs` \ binds2 ->
111 returnUs (bind2 : binds2)
114 %************************************************************************
116 \subsection{Bindings}
118 %************************************************************************
121 satBinding :: Bool -- True <=> top-level
124 -> UniqSM (SatEnv, StgBinding)
126 satBinding top env (StgNonRec b rhs)
127 = satRhs top env (b, rhs) `thenUs` \ (b2, rhs2) ->
129 env2 = addOneToIdEnv env b b2
131 returnUs (env2, StgNonRec b2 rhs2)
133 satBinding top env (StgRec pairs)
134 = -- Do it once to get the arities right...
135 mapUs (satRhs top env) pairs `thenUs` \ pairs2 ->
137 env2 = growIdEnvList env (map fst pairs `zip` map fst pairs2)
139 -- Do it again to *use* those arities:
140 mapUs (satRhs top env2) pairs `thenUs` \ pairs3 ->
142 returnUs (env2, StgRec pairs3)
144 satRhs :: Bool -> SatEnv -> (Id, StgRhs) -> UniqSM (Id_w_Arity, StgRhs)
146 satRhs top env (b, StgRhsCon cc con args) -- Nothing much to do here
148 b2 = b `addIdArity` 0 -- bound to a saturated constructor; hence zero.
150 returnUs (b2, StgRhsCon cc con (lookupArgs env args))
152 satRhs top env (b, StgRhsClosure cc bi fv u args body)
153 = satExpr env body `thenUs` \ (arity_info, body2) ->
155 num_args = length args
159 returnUs (num_args, StgRhsClosure cc bi fv u args body2)
162 ASSERT(needed_args >= 1)
164 let -- the arity we're aiming for is: what we already have ("args")
165 -- plus the ones requested in "arity_info"
166 new_arity = num_args + needed_args
168 -- get type info for this function:
169 (_,rho_arg_tys,tau_ty) = splitSigmaTy (idType b)
170 (tau_arg_tys, _) = splitFunTy tau_ty
171 all_arg_tys = ASSERT(null rho_arg_tys) {-rho_arg_tys ++-} tau_arg_tys
173 -- now, we already have "args"; we drop that many types
174 args_we_dont_have_tys = drop num_args all_arg_tys
176 -- finally, we take some of those (up to maybe all of them),
177 -- depending on how many "needed_args"
178 args_to_add_tys = take needed_args args_we_dont_have_tys
180 -- make up names for them
181 mapUs newName args_to_add_tys `thenUs` \ nns ->
183 -- and do the business
185 body3 = saturate body2 (map StgVarArg nns)
187 new_cc -- if we're adding args, we'd better not
188 -- keep calling something a CAF! (what about DICTs? ToDo: WDP 95/02)
189 = if not (isCafCC cc)
191 else if top then subsumedCosts else useCurrentCostCentre
193 returnUs (new_arity, StgRhsClosure new_cc bi fv ReEntrant (args++nns) body3)
195 `thenUs` \ (arity, rhs2) ->
197 b2 = b `addIdArity` arity
202 %************************************************************************
204 \subsection{Expressions}
206 %************************************************************************
209 satExpr :: SatEnv -> StgExpr -> UniqSM (ExprArityInfo, StgExpr)
211 satExpr env app@(StgApp (StgLitArg lit) [] lvs) = returnUs (Nothing, app)
213 satExpr env app@(StgApp (StgVarArg f) as lvs)
214 = returnUs (arity_to_return, StgApp (StgVarArg f2) as2 lvs)
216 as2 = lookupArgs env as
218 arity_to_return = case arityMaybe (getIdArity f2) of
221 Just f_arity -> if remaining_arity > 0
222 then Just remaining_arity
225 remaining_arity = f_arity - length as
227 satExpr env app@(StgCon con as lvs)
228 = returnUs (Nothing, StgCon con (lookupArgs env as) lvs)
230 satExpr env app@(StgPrim op as lvs)
231 = returnUs (Nothing, StgPrim op (lookupArgs env as) lvs)
233 satExpr env (StgSCC ty l e)
234 = satExpr env e `thenUs` \ (_, e2) ->
235 returnUs (Nothing, StgSCC ty l e2)
237 {- OMITTED: Let-no-escapery should come *after* saturation
239 satExpr (StgLetNoEscape lvs_whole lvs_rhss binds body)
240 = satBinding binds `thenUs` \ (binds2, c) ->
241 satExpr body `thenUs` \ (_, body2, c2) ->
242 returnUs (Nothing, StgLetNoEscape lvs_whole lvs_rhss binds2 body2, c + c2)
245 satExpr env (StgLet binds body)
246 = satBinding False{-not top-level-} env binds `thenUs` \ (env2, binds2) ->
247 satExpr env2 body `thenUs` \ (_, body2) ->
248 returnUs (Nothing, StgLet binds2 body2)
250 satExpr env (StgCase expr lve lva uniq alts)
251 = satExpr env expr `thenUs` \ (_, expr2) ->
252 sat_alts alts `thenUs` \ alts2 ->
253 returnUs (Nothing, StgCase expr2 lve lva uniq alts2)
255 sat_alts (StgAlgAlts ty alts def)
256 = mapUs sat_alg_alt alts `thenUs` \ alts2 ->
257 sat_deflt def `thenUs` \ def2 ->
258 returnUs (StgAlgAlts ty alts2 def2)
260 sat_alg_alt (id, bs, use_mask, e)
261 = satExpr env e `thenUs` \ (_, e2) ->
262 returnUs (id, bs, use_mask, e2)
264 sat_alts (StgPrimAlts ty alts def)
265 = mapUs sat_prim_alt alts `thenUs` \ alts2 ->
266 sat_deflt def `thenUs` \ def2 ->
267 returnUs (StgPrimAlts ty alts2 def2)
270 = satExpr env e `thenUs` \ (_, e2) ->
273 sat_deflt StgNoDefault
274 = returnUs StgNoDefault
276 sat_deflt (StgBindDefault b u expr)
277 = satExpr env expr `thenUs` \ (_,expr2) ->
278 returnUs (StgBindDefault b u expr2)
281 %************************************************************************
283 \subsection{Utility functions}
285 %************************************************************************
288 saturate :: StgExpr -> [StgArg] -> StgExpr
290 saturate (StgApp f as lvs) ids = StgApp f (as ++ ids) lvs
291 saturate other _ = panic "SatStgRhs: saturate"
295 lookupArgs :: SatEnv -> [StgArg] -> [StgArg]
296 lookupArgs env args = map do args
298 do (StgVarArg v) = StgVarArg (lookupVar env v)
299 do a@(StgLitArg lit) = a
301 lookupVar :: SatEnv -> Id -> Id
302 lookupVar env v = case lookupIdEnv env v of
306 newName :: Type -> UniqSM Id
308 = getUnique `thenUs` \ uniq ->
309 returnUs (mkSysLocal SLIT("sat") uniq ut mkUnknownSrcLoc)