[project @ 1996-04-05 08:26:04 by partain]
[ghc-hetmet.git] / ghc / compiler / simplStg / SatStgRhs.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
3 %
4 \section[SatStgRhs]{Saturates RHSs when they are partial applications}
5
6 96/03: This is actually an essential module, as it sets arity info
7 for the code generator.
8
9 \begin{display}
10 Subject: arg satis check
11 Date: Wed, 29 Apr 92 13:33:58 +0100
12 From: Simon L Peyton Jones <simonpj>
13
14 Andre
15
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
19
20         f = vs \ xs -> g e1 ... en
21
22 where xs has one or more elements
23 and
24 where g is a known function with arity m+n,
25
26 then: change it to
27
28         f = vs \ xs++{x1...xm} -> g e1 ... en x1 .. xm
29
30 Now g has enough args.   One arg-satisfaction check disappears;
31 the one for the closure incorporates the one for g.
32
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.
36
37 Simon
38 \end{display}
39
40 The algorithm proceeds as follows:
41 \begin{enumerate}
42 \item
43 Gather the arity information of the functions defined in this module
44 (as @getIdArity@ only knows about the arity of @ImportedIds@).
45
46 \item
47 for every definition of the form
48 \begin{verbatim}
49     v = /\ts -> \vs -> f args
50 \end{verbatim}
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
53 saturated.
54 \end{enumerate}
55
56 This is done for local definitions as well.
57
58 \begin{code}
59 #include "HsVersions.h"
60
61 module SatStgRhs ( satStgRhs ) where
62
63 import Ubiq{-uitous-}
64
65 import StgSyn
66
67 import CostCentre       ( isCafCC, subsumedCosts, useCurrentCostCentre )
68 import Id               ( idType, getIdArity, addIdArity, mkSysLocal,
69                           nullIdEnv, addOneToIdEnv, growIdEnvList,
70                           lookupIdEnv, IdEnv(..)
71                         )
72 import IdInfo           ( arityMaybe )
73 import SrcLoc           ( mkUnknownSrcLoc )
74 import UniqSupply       ( returnUs, thenUs, mapUs, getUnique, UniqSM(..) )
75 import Util             ( panic, assertPanic )
76
77 splitTypeWithDictsAsArgs = panic "SatStgRhs.splitTypeWithDictsAsArgs (ToDo)"
78
79 type Count = Int
80
81 type ExprArityInfo = Maybe Int      -- Just n  => This expression has a guaranteed
82                                     --            arity of n
83                                     -- Nothing => Don't know how many args it needs
84
85 type Id_w_Arity = Id                -- An Id with correct arity info pinned on it
86 type SatEnv     = IdEnv Id_w_Arity  -- Binds only local, let(rec)-bound things
87 \end{code}
88
89 This pass
90 \begin{itemize}
91 \item adds extra args where necessary;
92 \item pins the correct arity on everything.
93 \end{itemize}
94
95 %************************************************************************
96 %*                                                                      *
97 \subsection{Top-level list of bindings (a ``program'')}
98 %*                                                                      *
99 %************************************************************************
100
101 \begin{code}
102 satStgRhs :: [StgBinding] -> UniqSM [StgBinding]
103
104 satStgRhs p = satProgram nullIdEnv p
105
106 satProgram :: SatEnv -> [StgBinding] -> UniqSM [StgBinding]
107 satProgram env [] = returnUs []
108
109 satProgram env (bind:binds)
110   = satBinding True{-toplevel-} env bind    `thenUs` \ (env2, bind2) ->
111     satProgram env2 binds                   `thenUs` \ binds2 ->
112     returnUs (bind2 : binds2)
113 \end{code}
114
115 %************************************************************************
116 %*                                                                      *
117 \subsection{Bindings}
118 %*                                                                      *
119 %************************************************************************
120
121 \begin{code}
122 satBinding :: Bool      -- True <=> top-level
123            -> SatEnv
124            -> StgBinding
125            -> UniqSM (SatEnv, StgBinding)
126
127 satBinding top env (StgNonRec b rhs)
128   = satRhs top env (b, rhs)     `thenUs` \ (b2, rhs2) ->
129     let
130         env2 = addOneToIdEnv env b b2
131     in
132     returnUs (env2, StgNonRec b2 rhs2)
133
134 satBinding top env (StgRec pairs)
135   = -- Do it once to get the arities right...
136     mapUs (satRhs top env) pairs   `thenUs` \ pairs2 ->
137     let
138         env2 = growIdEnvList env (map fst pairs `zip` map fst pairs2)
139     in
140     -- Do it again to *use* those arities:
141     mapUs (satRhs top env2) pairs  `thenUs` \ pairs3 ->
142
143     returnUs (env2, StgRec pairs3)
144
145 satRhs :: Bool -> SatEnv -> (Id, StgRhs) -> UniqSM (Id_w_Arity, StgRhs)
146
147 satRhs top env (b, StgRhsCon cc con args)       -- Nothing much to do here
148   = let
149         b2 = b `addIdArity` 0 -- bound to a saturated constructor; hence zero.
150     in
151     returnUs (b2, StgRhsCon cc con (lookupArgs env args))
152
153 satRhs top env (b, StgRhsClosure cc bi fv u args body)
154   = satExpr env body    `thenUs` \ (arity_info, body2) ->
155     let
156         num_args = length args
157     in
158     (case arity_info of
159       Nothing ->
160         returnUs (num_args, StgRhsClosure cc bi fv u args body2)
161
162       Just needed_args ->
163         ASSERT(needed_args >= 1)
164
165         let  -- the arity we're aiming for is: what we already have ("args")
166              -- plus the ones requested in "arity_info"
167             new_arity = num_args + needed_args
168
169              -- get type info for this function:
170             (_,all_arg_tys,_) = splitTypeWithDictsAsArgs (idType b)
171
172              -- now, we already have "args"; we drop that many types
173             args_we_dont_have_tys = drop num_args all_arg_tys
174
175              -- finally, we take some of those (up to maybe all of them),
176              -- depending on how many "needed_args"
177             args_to_add_tys = take needed_args args_we_dont_have_tys
178         in
179             -- make up names for them
180         mapUs newName args_to_add_tys   `thenUs` \ nns ->
181
182             -- and do the business
183         let
184             body3  = saturate body2 (map StgVarArg nns)
185
186             new_cc -- if we're adding args, we'd better not
187                    -- keep calling something a CAF! (what about DICTs? ToDo: WDP 95/02)
188               = if not (isCafCC cc)
189                 then cc -- unchanged
190                 else if top then subsumedCosts else useCurrentCostCentre
191         in
192         returnUs (new_arity, StgRhsClosure new_cc bi fv ReEntrant (args++nns) body3)
193     )
194                                 `thenUs` \ (arity, rhs2) ->
195     let
196         b2 = b `addIdArity` arity
197     in
198     returnUs (b2, rhs2)
199 \end{code}
200
201 %************************************************************************
202 %*                                                                      *
203 \subsection{Expressions}
204 %*                                                                      *
205 %************************************************************************
206
207 \begin{code}
208 satExpr :: SatEnv -> StgExpr -> UniqSM (ExprArityInfo, StgExpr)
209
210 satExpr env app@(StgApp (StgLitArg lit) [] lvs) = returnUs (Nothing, app)
211
212 satExpr env app@(StgApp (StgVarArg f) as lvs)
213   = returnUs (arity_to_return, StgApp (StgVarArg f2) as2 lvs)
214   where
215     as2 = lookupArgs env as
216     f2  = lookupVar  env f
217     arity_to_return = case arityMaybe (getIdArity f2) of
218                         Nothing      -> Nothing
219
220                         Just f_arity -> if remaining_arity > 0
221                                         then Just remaining_arity
222                                         else Nothing
223                                      where
224                                         remaining_arity = f_arity - length as
225
226 satExpr env app@(StgCon con as lvs)
227   = returnUs (Nothing, StgCon con (lookupArgs env as) lvs)
228
229 satExpr env app@(StgPrim op as lvs)
230   = returnUs (Nothing, StgPrim op (lookupArgs env as) lvs)
231
232 satExpr env (StgSCC ty l e)
233   = satExpr env e        `thenUs` \ (_, e2) ->
234     returnUs (Nothing, StgSCC ty l e2)
235
236 {- OMITTED: Let-no-escapery should come *after* saturation
237
238 satExpr (StgLetNoEscape lvs_whole lvs_rhss binds body)
239   = satBinding binds    `thenUs` \ (binds2, c) ->
240     satExpr body        `thenUs` \ (_, body2, c2) ->
241     returnUs (Nothing, StgLetNoEscape lvs_whole lvs_rhss binds2 body2, c + c2)
242 -}
243
244 satExpr env (StgLet binds body)
245   = satBinding False{-not top-level-} env binds `thenUs` \ (env2, binds2) ->
246     satExpr env2 body                           `thenUs` \ (_, body2) ->
247     returnUs (Nothing, StgLet binds2 body2)
248
249 satExpr env (StgCase expr lve lva uniq alts)
250   = satExpr env expr    `thenUs` \ (_, expr2) ->
251     sat_alts alts       `thenUs` \ alts2 ->
252     returnUs (Nothing, StgCase expr2 lve lva uniq alts2)
253     where
254       sat_alts (StgAlgAlts ty alts def)
255         = mapUs sat_alg_alt alts        `thenUs` \ alts2 ->
256           sat_deflt def                 `thenUs` \ def2 ->
257           returnUs (StgAlgAlts ty alts2 def2)
258         where
259           sat_alg_alt (id, bs, use_mask, e)
260             = satExpr env e `thenUs` \ (_, e2) ->
261               returnUs (id, bs, use_mask, e2)
262
263       sat_alts (StgPrimAlts ty alts def)
264         = mapUs sat_prim_alt alts       `thenUs` \ alts2 ->
265           sat_deflt def                 `thenUs` \ def2 ->
266           returnUs (StgPrimAlts ty alts2 def2)
267         where
268           sat_prim_alt (l, e)
269             = satExpr env e `thenUs` \ (_, e2) ->
270               returnUs (l, e2)
271
272       sat_deflt StgNoDefault
273         = returnUs StgNoDefault
274
275       sat_deflt (StgBindDefault b u expr)
276         = satExpr env expr      `thenUs` \ (_,expr2) ->
277           returnUs (StgBindDefault b u expr2)
278 \end{code}
279
280 %************************************************************************
281 %*                                                                      *
282 \subsection{Utility functions}
283 %*                                                                      *
284 %************************************************************************
285
286 \begin{code}
287 saturate :: StgExpr -> [StgArg] -> StgExpr
288
289 saturate (StgApp f as lvs) ids = StgApp f (as ++ ids) lvs
290 saturate other              _  = panic "SatStgRhs: saturate"
291 \end{code}
292
293 \begin{code}
294 lookupArgs :: SatEnv -> [StgArg] -> [StgArg]
295 lookupArgs env args = map do args
296   where
297     do    (StgVarArg v)  = StgVarArg (lookupVar env v)
298     do a@(StgLitArg lit) = a
299
300 lookupVar :: SatEnv -> Id -> Id
301 lookupVar env v = case lookupIdEnv env v of
302                         Nothing -> v
303                         Just v2 -> v2
304
305 newName :: Type -> UniqSM Id
306 newName ut
307   = getUnique   `thenUs` \ uniq ->
308     returnUs (mkSysLocal SLIT("sat") uniq ut mkUnknownSrcLoc)
309 \end{code}