2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
4 \section[SatStgRhs]{Saturates RHSs when they are partial applications}
8 Subject: arg satis check
9 Date: Wed, 29 Apr 92 13:33:58 +0100
10 From: Simon L Peyton Jones <simonpj>
14 Another transformation to consider. We'd like to avoid
15 argument-satisfaction checks wherever possible. So, whenever we have an
16 STG binding application
18 f = vs \ xs -> g e1 ... en
20 where xs has one or more elements
22 where g is a known function with arity m+n,
26 f = vs \ xs++{x1...xm} -> g e1 ... en x1 .. xm
28 Now g has enough args. One arg-satisfaction check disappears;
29 the one for the closure incorporates the one for g.
31 You might like to consider variants, applying the transformation more
32 widely. I concluded that this was the only instance which made
33 sense, but I could be wrong.
38 The algorithm proceeds as follows:
41 Gather the arity information of the functions defined in this module
42 (as @getIdArity@ only knows about the arity of @ImportedIds@).
45 for every definition of the form
47 v = /\ts -> \vs -> f args
49 we try to match the arity of \tr{f} with the number of arguments.
50 If they do not match we insert extra lambdas to make that application
54 This is done for local definitions as well.
57 #include "HsVersions.h"
59 module SatStgRhs ( satStgRhs ) where
63 import AbsUniType ( splitTypeWithDictsAsArgs, Class,
64 TyVarTemplate, TauType(..)
68 import Id ( mkSysLocal, getIdUniType, getIdArity, addIdArity )
69 import IdInfo -- SIGH: ( arityMaybe, ArityInfo, OptIdInfo(..) )
70 import SrcLoc ( mkUnknownSrcLoc, SrcLoc )
79 type ExprArityInfo = Maybe Int -- Just n => This expression has a guaranteed
81 -- Nothing => Don't know how many args it needs
83 type Id_w_Arity = Id -- An Id with correct arity info pinned on it
84 type SatEnv = IdEnv Id_w_Arity -- Binds only local, let(rec)-bound things
89 \item adds extra args where necessary;
90 \item pins the correct arity on everything.
93 %************************************************************************
95 \subsection{Top-level list of bindings (a ``program'')}
97 %************************************************************************
100 satStgRhs :: PlainStgProgram -> SUniqSM PlainStgProgram
102 satStgRhs p = satProgram nullIdEnv p
104 satProgram :: SatEnv -> PlainStgProgram -> SUniqSM PlainStgProgram
105 satProgram env [] = returnSUs []
107 satProgram env (bind:binds)
108 = satBinding True{-toplevel-} env bind `thenSUs` \ (env2, bind2) ->
109 satProgram env2 binds `thenSUs` \ binds2 ->
110 returnSUs (bind2 : binds2)
113 %************************************************************************
115 \subsection{Bindings}
117 %************************************************************************
120 satBinding :: Bool -- True <=> top-level
123 -> SUniqSM (SatEnv, PlainStgBinding)
125 satBinding top env (StgNonRec b rhs)
126 = satRhs top env (b, rhs) `thenSUs` \ (b2, rhs2) ->
128 env2 = addOneToIdEnv env b b2
130 returnSUs (env2, StgNonRec b2 rhs2)
132 satBinding top env (StgRec pairs)
133 = -- Do it once to get the arities right...
134 mapSUs (satRhs top env) pairs `thenSUs` \ pairs2 ->
136 env2 = growIdEnvList env (map fst pairs `zip` map fst pairs2)
138 -- Do it again to *use* those arities:
139 mapSUs (satRhs top env2) pairs `thenSUs` \ pairs3 ->
141 returnSUs (env2, StgRec pairs3)
143 satRhs :: Bool -> SatEnv -> (Id, PlainStgRhs) -> SUniqSM (Id_w_Arity, PlainStgRhs)
145 satRhs top env (b, StgRhsCon cc con args) -- Nothing much to do here
147 b2 = b `addIdArity` 0 -- bound to a saturated constructor; hence zero.
149 returnSUs (b2, StgRhsCon cc con (lookupArgs env args))
151 satRhs top env (b, StgRhsClosure cc bi fv u args body)
152 = satExpr env body `thenSUs` \ (arity_info, body2) ->
154 num_args = length args
158 returnSUs (num_args, StgRhsClosure cc bi fv u args body2)
161 ASSERT(needed_args >= 1)
163 let -- the arity we're aiming for is: what we already have ("args")
164 -- plus the ones requested in "arity_info"
165 new_arity = num_args + needed_args
167 -- get type info for this function:
168 (_,all_arg_tys,_) = splitTypeWithDictsAsArgs (getIdUniType b)
170 -- now, we already have "args"; we drop that many types
171 args_we_dont_have_tys = drop num_args all_arg_tys
173 -- finally, we take some of those (up to maybe all of them),
174 -- depending on how many "needed_args"
175 args_to_add_tys = take needed_args args_we_dont_have_tys
177 -- make up names for them
178 mapSUs newName args_to_add_tys `thenSUs` \ nns ->
180 -- and do the business
182 body3 = saturate body2 (map StgVarAtom nns)
184 new_cc -- if we're adding args, we'd better not
185 -- keep calling something a CAF! (what about DICTs? ToDo: WDP 95/02)
186 = if not (isCafCC cc)
188 else if top then subsumedCosts else useCurrentCostCentre
190 returnSUs (new_arity, StgRhsClosure new_cc bi fv ReEntrant (args++nns) body3)
192 `thenSUs` \ (arity, rhs2) ->
194 b2 = b `addIdArity` arity
199 %************************************************************************
201 \subsection{Expressions}
203 %************************************************************************
206 satExpr :: SatEnv -> PlainStgExpr -> SUniqSM (ExprArityInfo, PlainStgExpr)
208 satExpr env app@(StgApp (StgLitAtom lit) [] lvs) = returnSUs (Nothing, app)
210 satExpr env app@(StgApp (StgVarAtom f) as lvs)
211 = returnSUs (arity_to_return, StgApp (StgVarAtom f2) as2 lvs)
213 as2 = lookupArgs env as
215 arity_to_return = case arityMaybe (getIdArity f2) of
218 Just f_arity -> if remaining_arity > 0
219 then Just remaining_arity
222 remaining_arity = f_arity - length as
224 satExpr env app@(StgConApp con as lvs)
225 = returnSUs (Nothing, StgConApp con (lookupArgs env as) lvs)
227 satExpr env app@(StgPrimApp op as lvs)
228 = returnSUs (Nothing, StgPrimApp op (lookupArgs env as) lvs)
230 satExpr env (StgSCC ty l e)
231 = satExpr env e `thenSUs` \ (_, e2) ->
232 returnSUs (Nothing, StgSCC ty l e2)
234 {- OMITTED: Let-no-escapery should come *after* saturation
236 satExpr (StgLetNoEscape lvs_whole lvs_rhss binds body)
237 = satBinding binds `thenSUs` \ (binds2, c) ->
238 satExpr body `thenSUs` \ (_, body2, c2) ->
239 returnSUs (Nothing, StgLetNoEscape lvs_whole lvs_rhss binds2 body2, c + c2)
242 satExpr env (StgLet binds body)
243 = satBinding False{-not top-level-} env binds `thenSUs` \ (env2, binds2) ->
244 satExpr env2 body `thenSUs` \ (_, body2) ->
245 returnSUs (Nothing, StgLet binds2 body2)
247 satExpr env (StgCase expr lve lva uniq alts)
248 = satExpr env expr `thenSUs` \ (_, expr2) ->
249 sat_alts alts `thenSUs` \ alts2 ->
250 returnSUs (Nothing, StgCase expr2 lve lva uniq alts2)
252 sat_alts (StgAlgAlts ty alts def)
253 = mapSUs sat_alg_alt alts `thenSUs` \ alts2 ->
254 sat_deflt def `thenSUs` \ def2 ->
255 returnSUs (StgAlgAlts ty alts2 def2)
257 sat_alg_alt (id, bs, use_mask, e)
258 = satExpr env e `thenSUs` \ (_, e2) ->
259 returnSUs (id, bs, use_mask, e2)
261 sat_alts (StgPrimAlts ty alts def)
262 = mapSUs sat_prim_alt alts `thenSUs` \ alts2 ->
263 sat_deflt def `thenSUs` \ def2 ->
264 returnSUs (StgPrimAlts ty alts2 def2)
267 = satExpr env e `thenSUs` \ (_, e2) ->
270 sat_deflt StgNoDefault
271 = returnSUs StgNoDefault
273 sat_deflt (StgBindDefault b u expr)
274 = satExpr env expr `thenSUs` \ (_,expr2) ->
275 returnSUs (StgBindDefault b u expr2)
278 %************************************************************************
280 \subsection{Utility functions}
282 %************************************************************************
285 saturate :: PlainStgExpr -> [PlainStgAtom] -> PlainStgExpr
287 saturate (StgApp f as lvs) ids = StgApp f (as ++ ids) lvs
288 saturate other _ = panic "SatStgRhs: saturate"
292 lookupArgs :: SatEnv -> [PlainStgAtom] -> [PlainStgAtom]
293 lookupArgs env args = map do args
295 do (StgVarAtom v) = StgVarAtom (lookupVar env v)
296 do a@(StgLitAtom lit) = a
298 lookupVar :: SatEnv -> Id -> Id
299 lookupVar env v = case lookupIdEnv env v of
303 newName :: UniType -> SUniqSM Id
305 = getSUnique `thenSUs` \ uniq ->
306 returnSUs (mkSysLocal SLIT("sat") uniq ut mkUnknownSrcLoc)