[project @ 1996-01-08 20:28:12 by partain]
[ghc.git] / ghc / compiler / simplStg / SatStgRhs.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1995
3 %
4 \section[SatStgRhs]{Saturates RHSs when they are partial applications}
5
6
7 \begin{display}
8 Subject: arg satis check
9 Date: Wed, 29 Apr 92 13:33:58 +0100
10 From: Simon L Peyton Jones <simonpj>
11
12 Andre
13
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
17
18         f = vs \ xs -> g e1 ... en
19
20 where xs has one or more elements
21 and
22 where g is a known function with arity m+n,
23
24 then: change it to
25
26         f = vs \ xs++{x1...xm} -> g e1 ... en x1 .. xm
27
28 Now g has enough args.   One arg-satisfaction check disappears;
29 the one for the closure incorporates the one for g.
30
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.
34
35 Simon
36 \end{display}
37
38 The algorithm proceeds as follows:
39 \begin{enumerate}
40 \item
41 Gather the arity information of the functions defined in this module
42 (as @getIdArity@ only knows about the arity of @ImportedIds@).
43
44 \item
45 for every definition of the form
46 \begin{verbatim}
47     v = /\ts -> \vs -> f args
48 \end{verbatim}
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
51 saturated.
52 \end{enumerate}
53
54 This is done for local definitions as well.
55
56 \begin{code}
57 #include "HsVersions.h"
58
59 module SatStgRhs ( satStgRhs ) where
60
61 import StgSyn
62
63 import AbsUniType       ( splitTypeWithDictsAsArgs, Class,
64                           TyVarTemplate, TauType(..)
65                         )
66 import CostCentre
67 import IdEnv
68 import Id               ( mkSysLocal, getIdUniType, getIdArity, addIdArity )
69 import IdInfo           -- SIGH: ( arityMaybe, ArityInfo, OptIdInfo(..) )
70 import SrcLoc           ( mkUnknownSrcLoc, SrcLoc )
71 import SplitUniq
72 import Unique
73 import Util
74 import Maybes
75
76 type Arity = Int
77 type Count = Int
78
79 type ExprArityInfo = Maybe Int      -- Just n  => This expression has a guaranteed
80                                     --            arity of n
81                                     -- Nothing => Don't know how many args it needs
82                                     
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
85 \end{code}
86
87 This pass 
88 \begin{itemize}
89 \item adds extra args where necessary;
90 \item pins the correct arity on everything.
91 \end{itemize}
92
93 %************************************************************************
94 %*                                                                      *
95 \subsection{Top-level list of bindings (a ``program'')}
96 %*                                                                      *
97 %************************************************************************
98
99 \begin{code}
100 satStgRhs :: PlainStgProgram -> SUniqSM PlainStgProgram
101
102 satStgRhs p = satProgram nullIdEnv p
103
104 satProgram :: SatEnv -> PlainStgProgram -> SUniqSM PlainStgProgram
105 satProgram env [] = returnSUs []
106
107 satProgram env (bind:binds) 
108   = satBinding True{-toplevel-} env bind    `thenSUs` \ (env2, bind2) ->
109     satProgram env2 binds                   `thenSUs` \ binds2 ->
110     returnSUs (bind2 : binds2)
111 \end{code}
112
113 %************************************************************************
114 %*                                                                      *
115 \subsection{Bindings}
116 %*                                                                      *
117 %************************************************************************
118
119 \begin{code}
120 satBinding :: Bool      -- True <=> top-level
121            -> SatEnv 
122            -> PlainStgBinding 
123            -> SUniqSM (SatEnv, PlainStgBinding)
124
125 satBinding top env (StgNonRec b rhs)
126   = satRhs top env (b, rhs)     `thenSUs` \ (b2, rhs2) ->
127     let
128         env2 = addOneToIdEnv env b b2
129     in
130     returnSUs (env2, StgNonRec b2 rhs2)
131
132 satBinding top env (StgRec pairs)
133   = -- Do it once to get the arities right...
134     mapSUs (satRhs top env) pairs   `thenSUs` \ pairs2 ->
135     let
136         env2 = growIdEnvList env (map fst pairs `zip` map fst pairs2)
137     in
138     -- Do it again to *use* those arities:
139     mapSUs (satRhs top env2) pairs  `thenSUs` \ pairs3 ->
140
141     returnSUs (env2, StgRec pairs3)
142
143 satRhs :: Bool -> SatEnv -> (Id, PlainStgRhs) -> SUniqSM (Id_w_Arity, PlainStgRhs)
144
145 satRhs top env (b, StgRhsCon cc con args)       -- Nothing much to do here
146   = let 
147         b2 = b `addIdArity` 0 -- bound to a saturated constructor; hence zero.
148     in
149     returnSUs (b2, StgRhsCon cc con (lookupArgs env args))
150
151 satRhs top env (b, StgRhsClosure cc bi fv u args body)
152   = satExpr env body    `thenSUs` \ (arity_info, body2) ->
153     let
154         num_args = length args
155     in
156     (case arity_info of
157       Nothing ->
158         returnSUs (num_args, StgRhsClosure cc bi fv u args body2)
159
160       Just needed_args ->
161         ASSERT(needed_args >= 1)
162
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
166
167              -- get type info for this function:
168             (_,all_arg_tys,_) = splitTypeWithDictsAsArgs (getIdUniType b)
169
170              -- now, we already have "args"; we drop that many types
171             args_we_dont_have_tys = drop num_args all_arg_tys
172
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
176         in
177             -- make up names for them
178         mapSUs newName args_to_add_tys  `thenSUs` \ nns ->
179
180             -- and do the business
181         let
182             body3  = saturate body2 (map StgVarAtom nns)
183
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) 
187                 then cc -- unchanged
188                 else if top then subsumedCosts else useCurrentCostCentre
189         in
190         returnSUs (new_arity, StgRhsClosure new_cc bi fv ReEntrant (args++nns) body3)
191     )
192                                 `thenSUs` \ (arity, rhs2) ->
193     let 
194         b2 = b `addIdArity` arity
195     in
196     returnSUs (b2, rhs2)
197 \end{code}
198
199 %************************************************************************
200 %*                                                                      *
201 \subsection{Expressions}
202 %*                                                                      *
203 %************************************************************************
204
205 \begin{code}    
206 satExpr :: SatEnv -> PlainStgExpr -> SUniqSM (ExprArityInfo, PlainStgExpr)
207
208 satExpr env app@(StgApp (StgLitAtom lit) [] lvs) = returnSUs (Nothing, app)
209
210 satExpr env app@(StgApp (StgVarAtom f) as lvs)
211   = returnSUs (arity_to_return, StgApp (StgVarAtom f2) as2 lvs)
212   where
213     as2 = lookupArgs env as
214     f2  = lookupVar  env f
215     arity_to_return = case arityMaybe (getIdArity f2) of
216                         Nothing      -> Nothing
217
218                         Just f_arity -> if remaining_arity > 0 
219                                         then Just remaining_arity
220                                         else Nothing
221                                      where
222                                         remaining_arity = f_arity - length as
223                                 
224 satExpr env app@(StgConApp con as lvs)
225   = returnSUs (Nothing, StgConApp con (lookupArgs env as) lvs)
226
227 satExpr env app@(StgPrimApp op as lvs)
228   = returnSUs (Nothing, StgPrimApp op (lookupArgs env as) lvs)
229
230 satExpr env (StgSCC ty l e)
231   = satExpr env e        `thenSUs` \ (_, e2) ->
232     returnSUs (Nothing, StgSCC ty l e2)
233
234 {- OMITTED: Let-no-escapery should come *after* saturation
235
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)
240 -}
241
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)
246
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)
251     where
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)
256         where
257           sat_alg_alt (id, bs, use_mask, e)
258             = satExpr env e `thenSUs` \ (_, e2) ->
259               returnSUs (id, bs, use_mask, e2)
260
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)
265         where
266           sat_prim_alt (l, e)
267             = satExpr env e `thenSUs` \ (_, e2) ->
268               returnSUs (l, e2)
269
270       sat_deflt StgNoDefault
271         = returnSUs StgNoDefault
272
273       sat_deflt (StgBindDefault b u expr)
274         = satExpr env expr      `thenSUs` \ (_,expr2) ->
275           returnSUs (StgBindDefault b u expr2)
276 \end{code}
277
278 %************************************************************************
279 %*                                                                      *
280 \subsection{Utility functions}
281 %*                                                                      *
282 %************************************************************************
283
284 \begin{code}
285 saturate :: PlainStgExpr -> [PlainStgAtom] -> PlainStgExpr
286
287 saturate (StgApp f as lvs) ids = StgApp f (as ++ ids) lvs
288 saturate other              _  = panic "SatStgRhs: saturate"
289 \end{code}
290
291 \begin{code}
292 lookupArgs :: SatEnv -> [PlainStgAtom] -> [PlainStgAtom]
293 lookupArgs env args = map do args
294   where 
295     do    (StgVarAtom v)  = StgVarAtom (lookupVar env v)
296     do a@(StgLitAtom lit) = a
297
298 lookupVar :: SatEnv -> Id -> Id
299 lookupVar env v = case lookupIdEnv env v of
300                         Nothing -> v
301                         Just v2 -> v2
302
303 newName :: UniType -> SUniqSM Id
304 newName ut
305   = getSUnique  `thenSUs` \ uniq ->
306     returnSUs (mkSysLocal SLIT("sat") uniq ut mkUnknownSrcLoc)
307 \end{code}