ghc/compiler/specialise/SpecConstr.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[SpecConstr]{Specialise over constructors}
   5
   6 \begin{code}
   7 module SpecConstr(
   8         specConstrProgram
   9     ) where
  10
  11 #include "HsVersions.h"
  12
  13 import CoreSyn
  14 import CoreLint         ( showPass, endPass )
  15 import CoreUtils        ( exprType, exprIsConApp_maybe, eqExpr )
  16 import CoreFVs          ( exprsFreeVars )
  17 import DataCon          ( isExistentialDataCon )
  18 import PprCore          ( pprCoreRules )
  19 import Id               ( Id, idName, idSpecialisation, mkUserLocal, mkSysLocal )
  20 import Var              ( Var )
  21 import VarEnv
  22 import VarSet
  23 import Name             ( nameOccName, nameSrcLoc )
  24 import Rules            ( addIdSpecialisations )
  25 import OccName          ( mkSpecOcc )
  26 import ErrUtils         ( dumpIfSet_dyn )
  27 import CmdLineOpts      ( DynFlags, DynFlag(..) )
  28 import Outputable
  29
  30 import Maybes           ( orElse )
  31 import Util             ( mapAccumL )
  32 import List             ( nubBy, partition )
  33 import UniqSupply
  34 import Outputable
  35 \end{code}
  36
  37 -----------------------------------------------------
  38                         Game plan
  39 -----------------------------------------------------
  40
  41 Consider
  42         drop n []     = []
  43         drop 0 xs     = []
  44         drop n (x:xs) = drop (n-1) xs
  45
  46 After the first time round, we could pass n unboxed.  This happens in
  47 numerical code too.  Here's what it looks like in Core:
  48
  49         drop n xs = case xs of
  50                       []     -> []
  51                       (y:ys) -> case n of
  52                                   I# n# -> case n# of
  53                                              0 -> []
  54                                              _ -> drop (I# (n# -# 1#)) xs
  55
  56 Notice that the recursive call has an explicit constructor as argument.
  57 Noticing this, we can make a specialised version of drop
  58
  59         RULE: drop (I# n#) xs ==> drop' n# xs
  60
  61         drop' n# xs = let n = I# n# in ...orig RHS...
  62
  63 Now the simplifier will apply the specialisation in the rhs of drop', giving
  64
  65         drop' n# xs = case xs of
  66                       []     -> []
  67                       (y:ys) -> case n# of
  68                                   0 -> []
  69                                   _ -> drop (n# -# 1#) xs
  70
  71 Much better!
  72
  73 We'd also like to catch cases where a parameter is carried along unchanged,
  74 but evaluated each time round the loop:
  75
  76         f i n = if i>0 || i>n then i else f (i*2) n
  77
  78 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  79 In Core, by the time we've w/wd (f is strict in i) we get
  80
  81         f i# n = case i# ># 0 of
  82                    False -> I# i#
  83                    True  -> case n of n' { I# n# ->
  84                             case i# ># n# of
  85                                 False -> I# i#
  86                                 True  -> f (i# *# 2#) n'
  87
  88 At the call to f, we see that the argument, n is know to be (I# n#),
  89 and n is evaluated elsewhere in the body of f, so we can play the same
  90 trick as above.  However we don't want to do that if the boxed version
  91 of n is needed (else we'd avoid the eval but pay more for re-boxing n).
  92 So in this case we want that the *only* uses of n are in case statements.
  93
  94
  95 So we look for
  96
  97 * A self-recursive function.  Ignore mutual recursion for now,
  98   because it's less common, and the code is simpler for self-recursion.
  99
 100 * EITHER
 101
 102    a) At a recursive call, one or more parameters is an explicit
 103       constructor application
 104         AND
 105       That same parameter is scrutinised by a case somewhere in
 106       the RHS of the function
 107
 108   OR
 109
 110     b) At a recursive call, one or more parameters has an unfolding
 111        that is an explicit constructor application
 112         AND
 113       That same parameter is scrutinised by a case somewhere in
 114       the RHS of the function
 115         AND
 116       Those are the only uses of the parameter
 117
 118
 119 There's a bit of a complication with type arguments.  If the call
 120 site looks like
 121
 122         f p = ...f ((:) [a] x xs)...
 123
 124 then our specialised function look like
 125
 126         f_spec x xs = let p = (:) [a] x xs in ....as before....
 127
 128 This only makes sense if either
 129   a) the type variable 'a' is in scope at the top of f, or
 130   b) the type variable 'a' is an argument to f (and hence fs)
 131
 132 Actually, (a) may hold for value arguments too, in which case
 133 we may not want to pass them.  Supose 'x' is in scope at f's
 134 defn, but xs is not.  Then we'd like
 135
 136         f_spec xs = let p = (:) [a] x xs in ....as before....
 137
 138 Similarly (b) may hold too.  If x is already an argument at the
 139 call, no need to pass it again.
 140
 141 Finally, if 'a' is not in scope at the call site, we could abstract
 142 it as we do the term variables:
 143
 144         f_spec a x xs = let p = (:) [a] x xs in ...as before...
 145
 146 So the grand plan is:
 147
 148         * abstract the call site to a constructor-only pattern
 149           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
 150
 151         * Find the free variables of the abstracted pattern
 152
 153         * Pass these variables, less any that are in scope at
 154           the fn defn.
 155
 156
 157 NOTICE that we only abstract over variables that are not in scope,
 158 so we're in no danger of shadowing variables used in "higher up"
 159 in f_spec's RHS.
 160
 161
 162 %************************************************************************
 163 %*                                                                      *
 164 \subsection{Top level wrapper stuff}
 165 %*                                                                      *
 166 %************************************************************************
 167
 168 \begin{code}
 169 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
 170 specConstrProgram dflags us binds
 171   = do
 172         showPass dflags "SpecConstr"
 173
 174         let (binds', _) = initUs us (go emptyScEnv binds)
 175
 176         endPass dflags "SpecConstr" Opt_D_dump_spec binds'
 177
 178         dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
 179                   (vcat (map dump_specs (concat (map bindersOf binds'))))
 180
 181         return binds'
 182   where
 183     go env []           = returnUs []
 184     go env (bind:binds) = scBind env bind       `thenUs` \ (env', _, bind') ->
 185                           go env' binds         `thenUs` \ binds' ->
 186                           returnUs (bind' : binds')
 187
 188 dump_specs var = pprCoreRules var (idSpecialisation var)
 189 \end{code}
 190
 191
 192 %************************************************************************
 193 %*                                                                      *
 194 \subsection{Environments and such}
 195 %*                                                                      *
 196 %************************************************************************
 197
 198 \begin{code}
 199 type ScEnv = VarEnv HowBound
 200
 201 emptyScEnv = emptyVarEnv
 202
 203 data HowBound = RecFun          -- These are the recursive functions for which
 204                                 -- we seek interesting call patterns
 205
 206               | RecArg          -- These are those functions' arguments; we are
 207                                 -- interested to see if those arguments are scrutinised
 208
 209               | Other           -- We track all others so we know what's in scope
 210                                 -- This is used in spec_one to check what needs to be
 211                                 -- passed as a parameter and what is in scope at the
 212                                 -- function definition site
 213
 214 extendBndrs env bndrs = extendVarEnvList env [(b,Other) | b <- bndrs]
 215 extendBndr  env bndr  = extendVarEnv env bndr Other
 216
 217 data ScUsage
 218    = SCU {
 219         calls :: !(IdEnv ([[CoreArg]])),        -- Calls
 220                                                 -- The functions are a subset of the
 221                                                 --      RecFuns in the ScEnv
 222
 223         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 224      }                                  -- The variables are a subset of the
 225                                         --      RecArg in the ScEnv
 226
 227 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 228
 229 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 230                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 231
 232 combineUsages [] = nullUsage
 233 combineUsages us = foldr1 combineUsage us
 234
 235 data ArgOcc = CaseScrut
 236             | OtherOcc
 237             | Both
 238
 239 instance Outputable ArgOcc where
 240   ppr CaseScrut = ptext SLIT("case-scrut")
 241   ppr OtherOcc  = ptext SLIT("other-occ")
 242   ppr Both      = ptext SLIT("case-scrut and other")
 243
 244 combineOcc CaseScrut CaseScrut = CaseScrut
 245 combineOcc OtherOcc  OtherOcc  = OtherOcc
 246 combineOcc _         _         = Both
 247 \end{code}
 248
 249
 250 %************************************************************************
 251 %*                                                                      *
 252 \subsection{The main recursive function}
 253 %*                                                                      *
 254 %************************************************************************
 255
 256 \begin{code}
 257 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 258         -- The unique supply is needed when we invent
 259         -- a new name for the specialised function and its args
 260
 261 scExpr env e@(Type t) = returnUs (nullUsage, e)
 262 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 263 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
 264 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 265                         returnUs (usg, Note n e')
 266 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 267                         returnUs (usg, Lam b e')
 268
 269 scExpr env (Case scrut b alts)
 270   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
 271     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
 272     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
 273               Case scrut' b alts')
 274   where
 275     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
 276     sc_scrut e         = scExpr env e
 277
 278     sc_alt (con,bs,rhs) = scExpr env rhs        `thenUs` \ (usg,rhs') ->
 279                           returnUs (usg, (con,bs,rhs'))
 280                         where
 281                           env1 = extendBndrs env (b:bs)
 282
 283 scExpr env (Let bind body)
 284   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 285     scExpr env' body    `thenUs` \ (body_usg, body') ->
 286     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 287
 288 scExpr env e@(App _ _)
 289   = let
 290         (fn, args) = collectArgs e
 291     in
 292     mapAndUnzipUs (scExpr env) args     `thenUs` \ (usgs, args') ->
 293     let
 294         arg_usg = combineUsages usgs
 295         fn_usg  | Var f <- fn,
 296                   Just RecFun <- lookupVarEnv env f
 297                 = SCU { calls = unitVarEnv f [args], occs = emptyVarEnv }
 298                 | otherwise
 299                 = nullUsage
 300     in
 301     returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
 302         -- Don't bother to look inside fn;
 303         -- it's almost always a variable
 304
 305 ----------------------
 306 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 307 scBind env (Rec [(fn,rhs)])
 308   | not (null val_bndrs)
 309   = scExpr env' body                    `thenUs` \ (usg@(SCU { calls = calls, occs = occs }), body') ->
 310     specialise env fn bndrs body usg    `thenUs` \ (rules, spec_prs) ->
 311     returnUs (extendBndrs env bndrs,
 312               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 313               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 314   where
 315     (bndrs,body) = collectBinders rhs
 316     val_bndrs    = filter isId bndrs
 317     env' = env `extendVarEnvList` ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs])
 318
 319 scBind env (Rec prs)
 320   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 321     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 322   where
 323     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 324                         returnUs (usg, (bndr,rhs'))
 325
 326 scBind env (NonRec bndr rhs)
 327   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 328     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 329
 330 ----------------------
 331 varUsage env v use
 332   | Just RecArg <- lookupVarEnv env v = SCU { calls = emptyVarEnv, occs = unitVarEnv v use }
 333   | otherwise                         = nullUsage
 334 \end{code}
 335
 336
 337 %************************************************************************
 338 %*                                                                      *
 339 \subsection{The specialiser}
 340 %*                                                                      *
 341 %************************************************************************
 342
 343 \begin{code}
 344 specialise :: ScEnv
 345            -> Id                        -- Functionn
 346            -> [CoreBndr] -> CoreExpr    -- Its RHS
 347            -> ScUsage                   -- Info on usage
 348            -> UniqSM ([CoreRule],       -- Rules
 349                       [(Id,CoreExpr)])  -- Bindings
 350
 351 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 352   = getUs               `thenUs` \ us ->
 353     let
 354         all_calls = lookupVarEnv calls fn `orElse` []
 355
 356         good_calls :: [[CoreArg]]
 357         good_calls = [ pats
 358                      | call_args <- all_calls,
 359                        length call_args >= n_bndrs,     -- App is saturated
 360                        let call = (bndrs `zip` call_args),
 361                        any (good_arg occs) call,
 362                        let (_, pats) = argsToPats us call_args
 363                      ]
 364     in
 365     pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
 366     mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 367                   (nubBy same_call good_calls `zip` [1..])
 368   where
 369     n_bndrs  = length bndrs
 370     same_call as1 as2 = and (zipWith eqExpr as1 as2)
 371
 372 ---------------------
 373 good_arg :: IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
 374 good_arg arg_occs (bndr, arg)
 375   = case exprIsConApp_maybe arg of                      -- exprIsConApp_maybe looks
 376         Just (dc,_) -> not (isExistentialDataCon dc)    -- through unfoldings
 377                        && bndr_usg_ok arg_occs bndr arg
 378         other   -> False
 379
 380 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 381 bndr_usg_ok arg_occs bndr arg
 382   = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
 383     case lookupVarEnv arg_occs bndr of
 384         Just CaseScrut -> True                  -- Used only by case scrutiny
 385         Just Both      -> case arg of           -- Used by case and elsewhere
 386                             App _ _ -> True     -- so the arg should be an explicit con app
 387                             other   -> False
 388         other -> False                          -- Not used, or used wonkily
 389
 390
 391 ---------------------
 392 argsToPats :: UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
 393 argsToPats us args = mapAccumL argToPat us args
 394
 395 argToPat   :: UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
 396 --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 397 argToPat us (Type ty)
 398   = (us, Type ty)
 399
 400 argToPat us arg
 401   | Just (dc,args) <- exprIsConApp_maybe arg
 402   = let
 403         (us',args') = argsToPats us args
 404     in
 405     (us', mkConApp dc args')
 406
 407 argToPat us (Var v)     -- Don't uniqify existing vars,
 408   = (us, Var v)         -- so that we can spot when we pass them twice
 409
 410 argToPat us arg
 411   = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
 412   where
 413     (us1,us2) = splitUniqSupply us
 414
 415 ---------------------
 416 spec_one :: ScEnv
 417          -> Id                                  -- Function
 418          -> CoreExpr                            -- Rhs of the original function
 419          -> ([CoreArg], Int)
 420          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 421
 422 {-
 423   Example
 424
 425      In-scope: a, x::a
 426      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
 427           [c is presumably bound by the (...) part]
 428   ==>
 429      f_spec = /\ b c \ v::(a,(b,c)) ->
 430                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
 431
 432      RULE:  forall b c,
 433                    y::[(a,(b,c))],
 434                    v::(a,(b,c)),
 435                    h::(a,(b,c))->[(a,(b,c))] .
 436
 437             f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
 438 -}
 439
 440 spec_one env fn rhs (pats, n)
 441   = getUniqueUs                         `thenUs` \ spec_uniq ->
 442     let
 443         fn_name      = idName fn
 444         fn_loc       = nameSrcLoc fn_name
 445         spec_occ     = mkSpecOcc (nameOccName fn_name)
 446         pat_fvs      = varSetElems (exprsFreeVars pats)
 447         vars_to_bind = filter not_avail pat_fvs
 448         not_avail v  = not (v `elemVarEnv` env)
 449                 -- Put the type variables first just for tidiness
 450         (tvs, ids)   = partition isTyVar vars_to_bind
 451         bndrs        = tvs ++ ids
 452
 453         rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
 454         spec_rhs  = mkLams bndrs (mkApps rhs pats)
 455         spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
 456         rule      = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
 457     in
 458     returnUs (rule, (spec_id, spec_rhs))
 459 \end{code}