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}