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               | RecArg          -- These are those functions' arguments; we are
 206                                 -- interested to see if those arguments are scrutinised
 207               | Other           -- We track all others so we know what's in scope
 208
 209 extendBndrs env bndrs = extendVarEnvList env [(b,Other) | b <- bndrs]
 210 extendBndr  env bndr  = extendVarEnv env bndr Other
 211
 212 data ScUsage
 213    = SCU {
 214         calls :: !(IdEnv ([[CoreArg]])),        -- Calls
 215                                                 -- The functions are a subset of the
 216                                                 --      RecFuns in the ScEnv
 217
 218         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 219      }                                  -- The variables are a subset of the
 220                                         --      RecArg in the ScEnv
 221
 222 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 223
 224 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 225                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 226
 227 combineUsages [] = nullUsage
 228 combineUsages us = foldr1 combineUsage us
 229
 230 data ArgOcc = CaseScrut
 231             | OtherOcc
 232             | Both
 233
 234 instance Outputable ArgOcc where
 235   ppr CaseScrut = ptext SLIT("case-scrut")
 236   ppr OtherOcc  = ptext SLIT("other-occ")
 237   ppr Both      = ptext SLIT("case-scrut and other")
 238
 239 combineOcc CaseScrut CaseScrut = CaseScrut
 240 combineOcc OtherOcc  OtherOcc  = OtherOcc
 241 combineOcc _         _         = Both
 242 \end{code}
 243
 244
 245 %************************************************************************
 246 %*                                                                      *
 247 \subsection{The main recursive function}
 248 %*                                                                      *
 249 %************************************************************************
 250
 251 \begin{code}
 252 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 253         -- The unique supply is needed when we invent
 254         -- a new name for the specialised function and its args
 255
 256 scExpr env e@(Type t) = returnUs (nullUsage, e)
 257 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 258 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
 259 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 260                         returnUs (usg, Note n e')
 261 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 262                         returnUs (usg, Lam b e')
 263
 264 scExpr env (Case scrut b alts)
 265   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
 266     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
 267     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
 268               Case scrut' b alts')
 269   where
 270     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
 271     sc_scrut e         = scExpr env e
 272
 273     sc_alt (con,bs,rhs) = scExpr env rhs        `thenUs` \ (usg,rhs') ->
 274                           returnUs (usg, (con,bs,rhs'))
 275                         where
 276                           env1 = extendBndrs env (b:bs)
 277
 278 scExpr env (Let bind body)
 279   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 280     scExpr env' body    `thenUs` \ (body_usg, body') ->
 281     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 282
 283 scExpr env e@(App _ _)
 284   = let
 285         (fn, args) = collectArgs e
 286     in
 287     mapAndUnzipUs (scExpr env) args     `thenUs` \ (usgs, args') ->
 288     let
 289         arg_usg = combineUsages usgs
 290         fn_usg  | Var f <- fn,
 291                   Just RecFun <- lookupVarEnv env f
 292                 = SCU { calls = unitVarEnv f [args], occs = emptyVarEnv }
 293                 | otherwise
 294                 = nullUsage
 295     in
 296     returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
 297         -- Don't bother to look inside fn;
 298         -- it's almost always a variable
 299
 300 ----------------------
 301 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 302 scBind env (Rec [(fn,rhs)])
 303   | not (null val_bndrs)
 304   = scExpr env' body                    `thenUs` \ (usg@(SCU { calls = calls, occs = occs }), body') ->
 305     specialise env fn bndrs body usg    `thenUs` \ (rules, spec_prs) ->
 306     returnUs (extendBndrs env bndrs,
 307               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 308               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 309   where
 310     (bndrs,body) = collectBinders rhs
 311     val_bndrs    = filter isId bndrs
 312     env' = env `extendVarEnvList` ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs])
 313
 314 scBind env (Rec prs)
 315   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 316     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 317   where
 318     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 319                         returnUs (usg, (bndr,rhs'))
 320
 321 scBind env (NonRec bndr rhs)
 322   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 323     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 324
 325 ----------------------
 326 varUsage env v use
 327   | Just RecArg <- lookupVarEnv env v = SCU { calls = emptyVarEnv, occs = unitVarEnv v use }
 328   | otherwise                         = nullUsage
 329 \end{code}
 330
 331
 332 %************************************************************************
 333 %*                                                                      *
 334 \subsection{The specialiser}
 335 %*                                                                      *
 336 %************************************************************************
 337
 338 \begin{code}
 339 specialise :: ScEnv
 340            -> Id                        -- Functionn
 341            -> [CoreBndr] -> CoreExpr    -- Its RHS
 342            -> ScUsage                   -- Info on usage
 343            -> UniqSM ([CoreRule],       -- Rules
 344                       [(Id,CoreExpr)])  -- Bindings
 345
 346 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 347   = getUs               `thenUs` \ us ->
 348     let
 349         all_calls = lookupVarEnv calls fn `orElse` []
 350
 351         good_calls :: [[CoreArg]]
 352         good_calls = [ pats
 353                      | call_args <- all_calls,
 354                        length call_args >= n_bndrs,     -- App is saturated
 355                        let call = (bndrs `zip` call_args),
 356                        any (good_arg occs) call,
 357                        let (_, pats) = argsToPats us call_args
 358                      ]
 359     in
 360     pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
 361     mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 362                   (nubBy same_call good_calls `zip` [1..])
 363   where
 364     n_bndrs  = length bndrs
 365     same_call as1 as2 = and (zipWith eqExpr as1 as2)
 366
 367 ---------------------
 368 good_arg :: IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
 369 good_arg arg_occs (bndr, arg)
 370   = case exprIsConApp_maybe arg of                      -- exprIsConApp_maybe looks
 371         Just (dc,_) -> not (isExistentialDataCon dc)    -- through unfoldings
 372                        && bndr_usg_ok arg_occs bndr arg
 373         other   -> False
 374
 375 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 376 bndr_usg_ok arg_occs bndr arg
 377   = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
 378     case lookupVarEnv arg_occs bndr of
 379         Just CaseScrut -> True                  -- Used only by case scrutiny
 380         Just Both      -> case arg of           -- Used by case and elsewhere
 381                             App _ _ -> True     -- so the arg should be an explicit con app
 382                             other   -> False
 383         other -> False                          -- Not used, or used wonkily
 384
 385
 386 ---------------------
 387 argsToPats :: UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
 388 argsToPats us args = mapAccumL argToPat us args
 389
 390 argToPat   :: UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
 391 --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 392 argToPat us (Type ty)
 393   = (us, Type ty)
 394
 395 argToPat us arg
 396   | Just (dc,args) <- exprIsConApp_maybe arg
 397   = let
 398         (us',args') = argsToPats us args
 399     in
 400     (us', mkConApp dc args')
 401
 402 argToPat us (Var v)     -- Don't uniqify existing vars,
 403   = (us, Var v)         -- so that we can spot when we pass them twice
 404
 405 argToPat us arg
 406   = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
 407   where
 408     (us1,us2) = splitUniqSupply us
 409
 410 ---------------------
 411 spec_one :: ScEnv
 412          -> Id                                  -- Function
 413          -> CoreExpr                            -- Rhs of the original function
 414          -> ([CoreArg], Int)
 415          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 416
 417 {-
 418   Example
 419
 420      In-scope: a, x::a
 421      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
 422           [c is presumably bound by the (...) part]
 423   ==>
 424      f_spec = /\ b c \ v::(a,(b,c)) ->
 425                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
 426
 427      RULE:  forall b c,
 428                    y::[(a,(b,c))],
 429                    v::(a,(b,c)),
 430                    h::(a,(b,c))->[(a,(b,c))] .
 431
 432             f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
 433 -}
 434
 435 spec_one env fn rhs (pats, n)
 436   = getUniqueUs                         `thenUs` \ spec_uniq ->
 437     let
 438         fn_name      = idName fn
 439         fn_loc       = nameSrcLoc fn_name
 440         spec_occ     = mkSpecOcc (nameOccName fn_name)
 441         pat_fvs      = varSetElems (exprsFreeVars pats)
 442         vars_to_bind = filter not_avail pat_fvs
 443         not_avail v  = not (v `elemVarEnv` env)
 444                 -- Put the type variables first just for tidiness
 445         (tvs, ids)   = partition isTyVar vars_to_bind
 446         bndrs        = tvs ++ ids
 447
 448         rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
 449         spec_rhs  = mkLams bndrs (mkApps rhs pats)
 450         spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
 451         rule      = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
 452     in
 453     returnUs (rule, (spec_id, spec_rhs))
 454 \end{code}