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 IMP_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, SYN_IE(IdEnv)
  71                         )
  72 import SrcLoc           ( noSrcLoc )
  73 import Type             ( splitSigmaTy, splitForAllTy, splitFunTyExpandingDicts )
  74 import UniqSupply       ( returnUs, thenUs, mapUs, getUnique, SYN_IE(UniqSM) )
  75 import Util             ( panic, assertPanic )
  76
  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 :: [StgBinding] -> UniqSM [StgBinding]
 101 satStgRhs = panic "satStgRhs"
 102
 103 {-              NUKED FOR NOW  SLPJ Dec 96
 104
 105
 106 satStgRhs p = satProgram nullIdEnv p
 107
 108 satProgram :: SatEnv -> [StgBinding] -> UniqSM [StgBinding]
 109 satProgram env [] = returnUs []
 110
 111 satProgram env (bind:binds)
 112   = satBinding True{-toplevel-} env bind    `thenUs` \ (env2, bind2) ->
 113     satProgram env2 binds                   `thenUs` \ binds2 ->
 114     returnUs (bind2 : binds2)
 115 \end{code}
 116
 117 %************************************************************************
 118 %*                                                                      *
 119 \subsection{Bindings}
 120 %*                                                                      *
 121 %************************************************************************
 122
 123 \begin{code}
 124 satBinding :: Bool      -- True <=> top-level
 125            -> SatEnv
 126            -> StgBinding
 127            -> UniqSM (SatEnv, StgBinding)
 128
 129 satBinding top env (StgNonRec b rhs)
 130   = satRhs top env (b, rhs)     `thenUs` \ (b2, rhs2) ->
 131     let
 132         env2 = addOneToIdEnv env b b2
 133     in
 134     returnUs (env2, StgNonRec b2 rhs2)
 135
 136 satBinding top env (StgRec pairs)
 137   = -- Do it once to get the arities right...
 138     mapUs (satRhs top env) pairs   `thenUs` \ pairs2 ->
 139     let
 140         env2 = growIdEnvList env (map fst pairs `zip` map fst pairs2)
 141     in
 142     -- Do it again to *use* those arities:
 143     mapUs (satRhs top env2) pairs  `thenUs` \ pairs3 ->
 144
 145     returnUs (env2, StgRec pairs3)
 146
 147 satRhs :: Bool -> SatEnv -> (Id, StgRhs) -> UniqSM (Id_w_Arity, StgRhs)
 148
 149 satRhs top env (b, StgRhsCon cc con args)       -- Nothing much to do here
 150   = let
 151         b2 = b `addIdArity` 0 -- bound to a saturated constructor; hence zero.
 152     in
 153     returnUs (b2, StgRhsCon cc con (lookupArgs env args))
 154
 155 satRhs top env (b, StgRhsClosure cc bi fv u args body)
 156   = satExpr env body    `thenUs` \ (arity_info, body2) ->
 157     let
 158         num_args = length args
 159     in
 160     (case arity_info of
 161       Nothing ->
 162         returnUs (num_args, StgRhsClosure cc bi fv u args body2)
 163
 164       Just needed_args ->
 165         ASSERT(needed_args >= 1)
 166
 167         let  -- the arity we're aiming for is: what we already have ("args")
 168              -- plus the ones requested in "arity_info"
 169             new_arity = num_args + needed_args
 170
 171              -- get type info for this function:
 172             (_, rho_ty) = splitForAllTy (idType b)
 173             (all_arg_tys, _) = splitFunTyExpandingDicts rho_ty
 174
 175              -- now, we already have "args"; we drop that many types
 176             args_we_dont_have_tys = drop num_args all_arg_tys
 177
 178              -- finally, we take some of those (up to maybe all of them),
 179              -- depending on how many "needed_args"
 180             args_to_add_tys = take needed_args args_we_dont_have_tys
 181         in
 182             -- make up names for them
 183         mapUs newName args_to_add_tys   `thenUs` \ nns ->
 184
 185             -- and do the business
 186         let
 187             body3  = saturate body2 (map StgVarArg nns)
 188
 189             new_cc -- if we're adding args, we'd better not
 190                    -- keep calling something a CAF! (what about DICTs? ToDo: WDP 95/02)
 191               = if not (isCafCC cc)
 192                 then cc -- unchanged
 193                 else if top then subsumedCosts else useCurrentCostCentre
 194         in
 195         returnUs (new_arity, StgRhsClosure new_cc bi fv ReEntrant (args++nns) body3)
 196     )
 197                                 `thenUs` \ (arity, rhs2) ->
 198     let
 199         b2 = b `addIdArity` arity
 200     in
 201     returnUs (b2, rhs2)
 202 \end{code}
 203
 204 %************************************************************************
 205 %*                                                                      *
 206 \subsection{Expressions}
 207 %*                                                                      *
 208 %************************************************************************
 209
 210 \begin{code}
 211 satExpr :: SatEnv -> StgExpr -> UniqSM (ExprArityInfo, StgExpr)
 212
 213 satExpr env app@(StgApp (StgLitArg lit) [] lvs) = returnUs (Nothing, app)
 214
 215 satExpr env app@(StgApp (StgVarArg f) as lvs)
 216   = returnUs (arity_to_return, StgApp (StgVarArg f2) as2 lvs)
 217   where
 218     as2 = lookupArgs env as
 219     f2  = lookupVar  env f
 220     arity_to_return = case arityMaybe (getIdArity f2) of
 221                         Nothing      -> Nothing
 222
 223                         Just f_arity -> if remaining_arity > 0
 224                                         then Just remaining_arity
 225                                         else Nothing
 226                                      where
 227                                         remaining_arity = f_arity - length as
 228
 229 satExpr env app@(StgCon con as lvs)
 230   = returnUs (Nothing, StgCon con (lookupArgs env as) lvs)
 231
 232 satExpr env app@(StgPrim op as lvs)
 233   = returnUs (Nothing, StgPrim op (lookupArgs env as) lvs)
 234
 235 satExpr env (StgSCC ty l e)
 236   = satExpr env e        `thenUs` \ (_, e2) ->
 237     returnUs (Nothing, StgSCC ty l e2)
 238
 239 {- OMITTED: Let-no-escapery should come *after* saturation
 240
 241 satExpr (StgLetNoEscape lvs_whole lvs_rhss binds body)
 242   = satBinding binds    `thenUs` \ (binds2, c) ->
 243     satExpr body        `thenUs` \ (_, body2, c2) ->
 244     returnUs (Nothing, StgLetNoEscape lvs_whole lvs_rhss binds2 body2, c + c2)
 245 -}
 246
 247 satExpr env (StgLet binds body)
 248   = satBinding False{-not top-level-} env binds `thenUs` \ (env2, binds2) ->
 249     satExpr env2 body                           `thenUs` \ (_, body2) ->
 250     returnUs (Nothing, StgLet binds2 body2)
 251
 252 satExpr env (StgCase expr lve lva uniq alts)
 253   = satExpr env expr    `thenUs` \ (_, expr2) ->
 254     sat_alts alts       `thenUs` \ alts2 ->
 255     returnUs (Nothing, StgCase expr2 lve lva uniq alts2)
 256     where
 257       sat_alts (StgAlgAlts ty alts def)
 258         = mapUs sat_alg_alt alts        `thenUs` \ alts2 ->
 259           sat_deflt def                 `thenUs` \ def2 ->
 260           returnUs (StgAlgAlts ty alts2 def2)
 261         where
 262           sat_alg_alt (id, bs, use_mask, e)
 263             = satExpr env e `thenUs` \ (_, e2) ->
 264               returnUs (id, bs, use_mask, e2)
 265
 266       sat_alts (StgPrimAlts ty alts def)
 267         = mapUs sat_prim_alt alts       `thenUs` \ alts2 ->
 268           sat_deflt def                 `thenUs` \ def2 ->
 269           returnUs (StgPrimAlts ty alts2 def2)
 270         where
 271           sat_prim_alt (l, e)
 272             = satExpr env e `thenUs` \ (_, e2) ->
 273               returnUs (l, e2)
 274
 275       sat_deflt StgNoDefault
 276         = returnUs StgNoDefault
 277
 278       sat_deflt (StgBindDefault b u expr)
 279         = satExpr env expr      `thenUs` \ (_,expr2) ->
 280           returnUs (StgBindDefault b u expr2)
 281 \end{code}
 282
 283 %************************************************************************
 284 %*                                                                      *
 285 \subsection{Utility functions}
 286 %*                                                                      *
 287 %************************************************************************
 288
 289 \begin{code}
 290 saturate :: StgExpr -> [StgArg] -> StgExpr
 291
 292 saturate (StgApp f as lvs) ids = StgApp f (as ++ ids) lvs
 293 saturate other              _  = panic "SatStgRhs: saturate"
 294 \end{code}
 295
 296 \begin{code}
 297 lookupArgs :: SatEnv -> [StgArg] -> [StgArg]
 298 lookupArgs env args = map doo args
 299   where
 300     doo    (StgVarArg v)  = StgVarArg (lookupVar env v)
 301     doo a@(StgLitArg lit) = a
 302
 303 lookupVar :: SatEnv -> Id -> Id
 304 lookupVar env v = case lookupIdEnv env v of
 305                         Nothing -> v
 306                         Just v2 -> v2
 307
 308 newName :: Type -> UniqSM Id
 309 newName ut
 310   = getUnique   `thenUs` \ uniq ->
 311     returnUs (mkSysLocal SLIT("sat") uniq ut noSrcLoc)
 312
 313 -}
 314 \end{code}