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          ( dataConRepArity )
  18 import Type             ( tyConAppArgs )
  19 import PprCore          ( pprCoreRules )
  20 import Id               ( Id, idName, idType, idSpecialisation,
  21                           isDataConId_maybe,
  22                           mkUserLocal, mkSysLocal )
  23 import Var              ( Var )
  24 import VarEnv
  25 import VarSet
  26 import Name             ( nameOccName, nameSrcLoc )
  27 import Rules            ( addIdSpecialisations )
  28 import OccName          ( mkSpecOcc )
  29 import ErrUtils         ( dumpIfSet_dyn )
  30 import CmdLineOpts      ( DynFlags, DynFlag(..) )
  31 import Outputable
  32
  33 import Maybes           ( orElse )
  34 import Util             ( mapAccumL )
  35 import List             ( nubBy, partition )
  36 import UniqSupply
  37 import Outputable
  38 \end{code}
  39
  40 -----------------------------------------------------
  41                         Game plan
  42 -----------------------------------------------------
  43
  44 Consider
  45         drop n []     = []
  46         drop 0 xs     = []
  47         drop n (x:xs) = drop (n-1) xs
  48
  49 After the first time round, we could pass n unboxed.  This happens in
  50 numerical code too.  Here's what it looks like in Core:
  51
  52         drop n xs = case xs of
  53                       []     -> []
  54                       (y:ys) -> case n of
  55                                   I# n# -> case n# of
  56                                              0 -> []
  57                                              _ -> drop (I# (n# -# 1#)) xs
  58
  59 Notice that the recursive call has an explicit constructor as argument.
  60 Noticing this, we can make a specialised version of drop
  61
  62         RULE: drop (I# n#) xs ==> drop' n# xs
  63
  64         drop' n# xs = let n = I# n# in ...orig RHS...
  65
  66 Now the simplifier will apply the specialisation in the rhs of drop', giving
  67
  68         drop' n# xs = case xs of
  69                       []     -> []
  70                       (y:ys) -> case n# of
  71                                   0 -> []
  72                                   _ -> drop (n# -# 1#) xs
  73
  74 Much better!
  75
  76 We'd also like to catch cases where a parameter is carried along unchanged,
  77 but evaluated each time round the loop:
  78
  79         f i n = if i>0 || i>n then i else f (i*2) n
  80
  81 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  82 In Core, by the time we've w/wd (f is strict in i) we get
  83
  84         f i# n = case i# ># 0 of
  85                    False -> I# i#
  86                    True  -> case n of n' { I# n# ->
  87                             case i# ># n# of
  88                                 False -> I# i#
  89                                 True  -> f (i# *# 2#) n'
  90
  91 At the call to f, we see that the argument, n is know to be (I# n#),
  92 and n is evaluated elsewhere in the body of f, so we can play the same
  93 trick as above.  However we don't want to do that if the boxed version
  94 of n is needed (else we'd avoid the eval but pay more for re-boxing n).
  95 So in this case we want that the *only* uses of n are in case statements.
  96
  97
  98 So we look for
  99
 100 * A self-recursive function.  Ignore mutual recursion for now,
 101   because it's less common, and the code is simpler for self-recursion.
 102
 103 * EITHER
 104
 105    a) At a recursive call, one or more parameters is an explicit
 106       constructor application
 107         AND
 108       That same parameter is scrutinised by a case somewhere in
 109       the RHS of the function
 110
 111   OR
 112
 113     b) At a recursive call, one or more parameters has an unfolding
 114        that is an explicit constructor application
 115         AND
 116       That same parameter is scrutinised by a case somewhere in
 117       the RHS of the function
 118         AND
 119       Those are the only uses of the parameter
 120
 121
 122 There's a bit of a complication with type arguments.  If the call
 123 site looks like
 124
 125         f p = ...f ((:) [a] x xs)...
 126
 127 then our specialised function look like
 128
 129         f_spec x xs = let p = (:) [a] x xs in ....as before....
 130
 131 This only makes sense if either
 132   a) the type variable 'a' is in scope at the top of f, or
 133   b) the type variable 'a' is an argument to f (and hence fs)
 134
 135 Actually, (a) may hold for value arguments too, in which case
 136 we may not want to pass them.  Supose 'x' is in scope at f's
 137 defn, but xs is not.  Then we'd like
 138
 139         f_spec xs = let p = (:) [a] x xs in ....as before....
 140
 141 Similarly (b) may hold too.  If x is already an argument at the
 142 call, no need to pass it again.
 143
 144 Finally, if 'a' is not in scope at the call site, we could abstract
 145 it as we do the term variables:
 146
 147         f_spec a x xs = let p = (:) [a] x xs in ...as before...
 148
 149 So the grand plan is:
 150
 151         * abstract the call site to a constructor-only pattern
 152           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
 153
 154         * Find the free variables of the abstracted pattern
 155
 156         * Pass these variables, less any that are in scope at
 157           the fn defn.
 158
 159
 160 NOTICE that we only abstract over variables that are not in scope,
 161 so we're in no danger of shadowing variables used in "higher up"
 162 in f_spec's RHS.
 163
 164
 165 %************************************************************************
 166 %*                                                                      *
 167 \subsection{Top level wrapper stuff}
 168 %*                                                                      *
 169 %************************************************************************
 170
 171 \begin{code}
 172 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
 173 specConstrProgram dflags us binds
 174   = do
 175         showPass dflags "SpecConstr"
 176
 177         let (binds', _) = initUs us (go emptyScEnv binds)
 178
 179         endPass dflags "SpecConstr" Opt_D_dump_spec binds'
 180
 181         dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
 182                   (vcat (map dump_specs (concat (map bindersOf binds'))))
 183
 184         return binds'
 185   where
 186     go env []           = returnUs []
 187     go env (bind:binds) = scBind env bind       `thenUs` \ (env', _, bind') ->
 188                           go env' binds         `thenUs` \ binds' ->
 189                           returnUs (bind' : binds')
 190
 191 dump_specs var = pprCoreRules var (idSpecialisation var)
 192 \end{code}
 193
 194
 195 %************************************************************************
 196 %*                                                                      *
 197 \subsection{Environment: goes downwards}
 198 %*                                                                      *
 199 %************************************************************************
 200
 201 \begin{code}
 202 data ScEnv = SCE { scope :: VarEnv HowBound,
 203                         -- Binds all non-top-level variables in scope
 204
 205                    cons  :: ConstrEnv
 206              }
 207
 208 type ConstrEnv = IdEnv (AltCon, [CoreArg])
 209         -- Variables known to be bound to a constructor
 210         -- in a particular case alternative
 211
 212 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
 213
 214 data HowBound = RecFun          -- These are the recursive functions for which
 215                                 -- we seek interesting call patterns
 216
 217               | RecArg          -- These are those functions' arguments; we are
 218                                 -- interested to see if those arguments are scrutinised
 219
 220               | Other           -- We track all others so we know what's in scope
 221                                 -- This is used in spec_one to check what needs to be
 222                                 -- passed as a parameter and what is in scope at the
 223                                 -- function definition site
 224
 225 lookupScopeEnv env v = lookupVarEnv (scope env) v
 226
 227 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
 228 extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }
 229
 230     -- When we encounter
 231     --  case scrut of b
 232     --      C x y -> ...
 233     -- we want to bind b, and perhaps scrut too, to (C x y)
 234 extendCaseBndr env case_bndr scrut con alt_bndrs
 235   = case scrut of
 236         Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
 237                    -- Also forget if the scrutinee is a RecArg, because we're
 238                    -- now in the branch of a case, and we don't want to
 239                    -- record a non-scrutinee use of v if we have
 240                    --   case v of { (a,b) -> ...(f v)... }
 241                  SCE { scope = extendVarEnv (scope env1) v Other,
 242                        cons  = extendVarEnv (cons env1)  v (con,args) }
 243         other -> env1
 244
 245   where
 246     env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
 247                  cons  = extendVarEnv     (cons  env) case_bndr (con,args) }
 248
 249     args = map Type (tyConAppArgs (idType case_bndr)) ++
 250            map varToCoreExpr alt_bndrs
 251
 252     -- When we encounter a recursive function binding
 253     --  f = \x y -> ...
 254     -- we want to extend the scope env with bindings
 255     -- that record that f is a RecFn and x,y are RecArgs
 256 extendRecBndr env fn bndrs
 257   =  env { scope = scope env `extendVarEnvList`
 258                    ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
 259 \end{code}
 260
 261
 262 %************************************************************************
 263 %*                                                                      *
 264 \subsection{Usage information: flows upwards}
 265 %*                                                                      *
 266 %************************************************************************
 267
 268 \begin{code}
 269 data ScUsage
 270    = SCU {
 271         calls :: !(IdEnv ([Call])),     -- Calls
 272                                         -- The functions are a subset of the
 273                                         --      RecFuns in the ScEnv
 274
 275         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 276      }                                  -- The variables are a subset of the
 277                                         --      RecArg in the ScEnv
 278
 279 type Call = (ConstrEnv, [CoreArg])
 280         -- The arguments of the call, together with the
 281         -- env giving the constructor bindings at the call site
 282
 283 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 284
 285 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 286                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 287
 288 combineUsages [] = nullUsage
 289 combineUsages us = foldr1 combineUsage us
 290
 291 data ArgOcc = CaseScrut
 292             | OtherOcc
 293             | Both
 294
 295 instance Outputable ArgOcc where
 296   ppr CaseScrut = ptext SLIT("case-scrut")
 297   ppr OtherOcc  = ptext SLIT("other-occ")
 298   ppr Both      = ptext SLIT("case-scrut and other")
 299
 300 combineOcc CaseScrut CaseScrut = CaseScrut
 301 combineOcc OtherOcc  OtherOcc  = OtherOcc
 302 combineOcc _         _         = Both
 303 \end{code}
 304
 305
 306 %************************************************************************
 307 %*                                                                      *
 308 \subsection{The main recursive function}
 309 %*                                                                      *
 310 %************************************************************************
 311
 312 The main recursive function gathers up usage information, and
 313 creates specialised versions of functions.
 314
 315 \begin{code}
 316 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 317         -- The unique supply is needed when we invent
 318         -- a new name for the specialised function and its args
 319
 320 scExpr env e@(Type t) = returnUs (nullUsage, e)
 321 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 322 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
 323 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 324                         returnUs (usg, Note n e')
 325 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 326                         returnUs (usg, Lam b e')
 327
 328 scExpr env (Case scrut b alts)
 329   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
 330     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
 331     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
 332               Case scrut' b alts')
 333   where
 334     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
 335     sc_scrut e         = scExpr env e
 336
 337     sc_alt (con,bs,rhs) = scExpr env1 rhs       `thenUs` \ (usg,rhs') ->
 338                           returnUs (usg, (con,bs,rhs'))
 339                         where
 340                           env1 = extendCaseBndr env b scrut con bs
 341
 342 scExpr env (Let bind body)
 343   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 344     scExpr env' body    `thenUs` \ (body_usg, body') ->
 345     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 346
 347 scExpr env e@(App _ _)
 348   = let
 349         (fn, args) = collectArgs e
 350     in
 351     mapAndUnzipUs (scExpr env) args     `thenUs` \ (usgs, args') ->
 352     let
 353         arg_usg = combineUsages usgs
 354         fn_usg  | Var f <- fn,
 355                   Just RecFun <- lookupScopeEnv env f
 356                 = SCU { calls = unitVarEnv f [(cons env, args)],
 357                         occs  = emptyVarEnv }
 358                 | otherwise
 359                 = nullUsage
 360     in
 361     returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
 362         -- Don't bother to look inside fn;
 363         -- it's almost always a variable
 364
 365 ----------------------
 366 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 367 scBind env (Rec [(fn,rhs)])
 368   | not (null val_bndrs)
 369   = scExpr env' body                    `thenUs` \ (usg, body') ->
 370     let
 371         SCU { calls = calls, occs = occs } = usg
 372     in
 373     specialise env fn bndrs body usg    `thenUs` \ (rules, spec_prs) ->
 374     returnUs (extendBndrs env bndrs,
 375               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 376               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 377   where
 378     (bndrs,body) = collectBinders rhs
 379     val_bndrs    = filter isId bndrs
 380     env'         = extendRecBndr env fn bndrs
 381
 382 scBind env (Rec prs)
 383   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 384     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 385   where
 386     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 387                         returnUs (usg, (bndr,rhs'))
 388
 389 scBind env (NonRec bndr rhs)
 390   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 391     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 392
 393 ----------------------
 394 varUsage env v use
 395   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
 396                                                 occs = unitVarEnv v use }
 397   | otherwise                           = nullUsage
 398 \end{code}
 399
 400
 401 %************************************************************************
 402 %*                                                                      *
 403 \subsection{The specialiser}
 404 %*                                                                      *
 405 %************************************************************************
 406
 407 \begin{code}
 408 specialise :: ScEnv
 409            -> Id                        -- Functionn
 410            -> [CoreBndr] -> CoreExpr    -- Its RHS
 411            -> ScUsage                   -- Info on usage
 412            -> UniqSM ([CoreRule],       -- Rules
 413                       [(Id,CoreExpr)])  -- Bindings
 414
 415 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 416   = getUs               `thenUs` \ us ->
 417     let
 418         all_calls = lookupVarEnv calls fn `orElse` []
 419
 420         good_calls :: [[CoreArg]]
 421         good_calls = [ pats
 422                      | (con_env, call_args) <- all_calls,
 423                        length call_args >= n_bndrs,         -- App is saturated
 424                        let call = (bndrs `zip` call_args),
 425                        any (good_arg con_env occs) call,    -- At least one arg is a constr app
 426                        let (_, pats) = argsToPats con_env us call_args
 427                      ]
 428     in
 429     pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
 430     mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 431                   (nubBy same_call good_calls `zip` [1..])
 432   where
 433     n_bndrs  = length bndrs
 434     same_call as1 as2 = and (zipWith eqExpr as1 as2)
 435
 436 ---------------------
 437 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
 438 good_arg con_env arg_occs (bndr, arg)
 439   = case is_con_app_maybe con_env arg of
 440         Just _ ->  bndr_usg_ok arg_occs bndr arg
 441         other   -> False
 442
 443 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 444 bndr_usg_ok arg_occs bndr arg
 445   = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
 446     case lookupVarEnv arg_occs bndr of
 447         Just CaseScrut -> True                  -- Used only by case scrutiny
 448         Just Both      -> case arg of           -- Used by case and elsewhere
 449                             App _ _ -> True     -- so the arg should be an explicit con app
 450                             other   -> False
 451         other -> False                          -- Not used, or used wonkily
 452
 453
 454 ---------------------
 455 spec_one :: ScEnv
 456          -> Id                                  -- Function
 457          -> CoreExpr                            -- Rhs of the original function
 458          -> ([CoreArg], Int)
 459          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 460
 461 -- spec_one creates a specialised copy of the function, together
 462 -- with a rule for using it.  I'm very proud of how short this
 463 -- function is, considering what it does :-).
 464
 465 {-
 466   Example
 467
 468      In-scope: a, x::a
 469      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
 470           [c is presumably bound by the (...) part]
 471   ==>
 472      f_spec = /\ b c \ v::(a,(b,c)) ->
 473                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
 474
 475      RULE:  forall b c,
 476                    y::[(a,(b,c))],
 477                    v::(a,(b,c)),
 478                    h::(a,(b,c))->[(a,(b,c))] .
 479
 480             f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
 481 -}
 482
 483 spec_one env fn rhs (pats, n)
 484   = getUniqueUs                         `thenUs` \ spec_uniq ->
 485     let
 486         fn_name      = idName fn
 487         fn_loc       = nameSrcLoc fn_name
 488         spec_occ     = mkSpecOcc (nameOccName fn_name)
 489         pat_fvs      = varSetElems (exprsFreeVars pats)
 490         vars_to_bind = filter not_avail pat_fvs
 491         not_avail v  = not (v `elemVarEnv` scope env)
 492                 -- Put the type variables first just for tidiness
 493         (tvs, ids)   = partition isTyVar vars_to_bind
 494         bndrs        = tvs ++ ids
 495
 496         rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
 497         spec_rhs  = mkLams bndrs (mkApps rhs pats)
 498         spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
 499         rule      = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
 500     in
 501     returnUs (rule, (spec_id, spec_rhs))
 502 \end{code}
 503
 504 %************************************************************************
 505 %*                                                                      *
 506 \subsection{Argument analysis}
 507 %*                                                                      *
 508 %************************************************************************
 509
 510 This code deals with analysing call-site arguments to see whether
 511 they are constructor applications.
 512
 513 \begin{code}
 514     -- argToPat takes an actual argument, and returns an abstracted
 515     -- version, consisting of just the "constructor skeleton" of the
 516     -- argument, with non-constructor sub-expression replaced by new
 517     -- placeholder variables.  For example:
 518     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 519
 520 argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
 521 argToPat env us (Type ty)
 522   = (us, Type ty)
 523
 524 argToPat env us arg
 525   | Just (dc,args) <- is_con_app_maybe env arg
 526   = let
 527         (us',args') = argsToPats env us args
 528     in
 529     (us', mk_con_app dc args')
 530
 531 argToPat env us (Var v) -- Don't uniqify existing vars,
 532   = (us, Var v)         -- so that we can spot when we pass them twice
 533
 534 argToPat env us arg
 535   = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
 536   where
 537     (us1,us2) = splitUniqSupply us
 538
 539 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
 540 argsToPats env us args = mapAccumL (argToPat env) us args
 541 \end{code}
 542
 543
 544 \begin{code}
 545 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
 546 is_con_app_maybe env (Var v)
 547   = lookupVarEnv env v
 548         -- You might think we could look in the idUnfolding here
 549         -- but that doesn't take account of which branch of a
 550         -- case we are in, which is the whole point
 551
 552 is_con_app_maybe env (Lit lit)
 553   = Just (LitAlt lit, [])
 554
 555 is_con_app_maybe env expr
 556   = case collectArgs expr of
 557         (Var fun, args) | Just con <- isDataConId_maybe fun,
 558                           length args >= dataConRepArity con
 559                 -- Might be > because the arity excludes type args
 560                         -> Just (DataAlt con,args)
 561
 562         other -> Nothing
 563
 564 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
 565 mk_con_app (LitAlt lit)  []   = Lit lit
 566 mk_con_app (DataAlt con) args = mkConApp con args
 567 \end{code}