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