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, eqExpr )
  16 import CoreFVs          ( exprsFreeVars )
  17 import DataCon          ( dataConRepArity )
  18 import Type             ( tyConAppArgs )
  19 import PprCore          ( pprCoreRules, pprCoreRule )
  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 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
 235 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
 236   = extendBndrs env (case_bndr : alt_bndrs)
 237
 238 extendCaseBndrs env case_bndr scrut con alt_bndrs
 239   = case scrut of
 240         Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
 241                    -- Also forget if the scrutinee is a RecArg, because we're
 242                    -- now in the branch of a case, and we don't want to
 243                    -- record a non-scrutinee use of v if we have
 244                    --   case v of { (a,b) -> ...(f v)... }
 245                  SCE { scope = extendVarEnv (scope env1) v Other,
 246                        cons  = extendVarEnv (cons env1)  v (con,args) }
 247         other -> env1
 248
 249   where
 250     env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
 251                  cons  = extendVarEnv     (cons  env) case_bndr (con,args) }
 252
 253     args = map Type (tyConAppArgs (idType case_bndr)) ++
 254            map varToCoreExpr alt_bndrs
 255
 256     -- When we encounter a recursive function binding
 257     --  f = \x y -> ...
 258     -- we want to extend the scope env with bindings
 259     -- that record that f is a RecFn and x,y are RecArgs
 260 extendRecBndr env fn bndrs
 261   =  env { scope = scope env `extendVarEnvList`
 262                    ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
 263 \end{code}
 264
 265
 266 %************************************************************************
 267 %*                                                                      *
 268 \subsection{Usage information: flows upwards}
 269 %*                                                                      *
 270 %************************************************************************
 271
 272 \begin{code}
 273 data ScUsage
 274    = SCU {
 275         calls :: !(IdEnv ([Call])),     -- Calls
 276                                         -- The functions are a subset of the
 277                                         --      RecFuns in the ScEnv
 278
 279         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 280      }                                  -- The variables are a subset of the
 281                                         --      RecArg in the ScEnv
 282
 283 type Call = (ConstrEnv, [CoreArg])
 284         -- The arguments of the call, together with the
 285         -- env giving the constructor bindings at the call site
 286
 287 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 288
 289 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 290                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 291
 292 combineUsages [] = nullUsage
 293 combineUsages us = foldr1 combineUsage us
 294
 295 data ArgOcc = CaseScrut
 296             | OtherOcc
 297             | Both
 298
 299 instance Outputable ArgOcc where
 300   ppr CaseScrut = ptext SLIT("case-scrut")
 301   ppr OtherOcc  = ptext SLIT("other-occ")
 302   ppr Both      = ptext SLIT("case-scrut and other")
 303
 304 combineOcc CaseScrut CaseScrut = CaseScrut
 305 combineOcc OtherOcc  OtherOcc  = OtherOcc
 306 combineOcc _         _         = Both
 307 \end{code}
 308
 309
 310 %************************************************************************
 311 %*                                                                      *
 312 \subsection{The main recursive function}
 313 %*                                                                      *
 314 %************************************************************************
 315
 316 The main recursive function gathers up usage information, and
 317 creates specialised versions of functions.
 318
 319 \begin{code}
 320 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 321         -- The unique supply is needed when we invent
 322         -- a new name for the specialised function and its args
 323
 324 scExpr env e@(Type t) = returnUs (nullUsage, e)
 325 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 326 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
 327 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 328                         returnUs (usg, Note n e')
 329 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 330                         returnUs (usg, Lam b e')
 331
 332 scExpr env (Case scrut b alts)
 333   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
 334     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
 335     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
 336               Case scrut' b alts')
 337   where
 338     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
 339     sc_scrut e         = scExpr env e
 340
 341     sc_alt (con,bs,rhs) = scExpr env1 rhs       `thenUs` \ (usg,rhs') ->
 342                           returnUs (usg, (con,bs,rhs'))
 343                         where
 344                           env1 = extendCaseBndrs env b scrut con bs
 345
 346 scExpr env (Let bind body)
 347   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 348     scExpr env' body    `thenUs` \ (body_usg, body') ->
 349     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 350
 351 scExpr env e@(App _ _)
 352   = let
 353         (fn, args) = collectArgs e
 354     in
 355     mapAndUnzipUs (scExpr env) args     `thenUs` \ (usgs, args') ->
 356     let
 357         arg_usg = combineUsages usgs
 358         fn_usg  | Var f <- fn,
 359                   Just RecFun <- lookupScopeEnv env f
 360                 = SCU { calls = unitVarEnv f [(cons env, args)],
 361                         occs  = emptyVarEnv }
 362                 | otherwise
 363                 = nullUsage
 364     in
 365     returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
 366         -- Don't bother to look inside fn;
 367         -- it's almost always a variable
 368
 369 ----------------------
 370 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 371 scBind env (Rec [(fn,rhs)])
 372   | not (null val_bndrs)
 373   = scExpr env' body                    `thenUs` \ (usg, body') ->
 374     let
 375         SCU { calls = calls, occs = occs } = usg
 376     in
 377     specialise env fn bndrs body usg    `thenUs` \ (rules, spec_prs) ->
 378     returnUs (extendBndrs env bndrs,
 379               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 380               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 381   where
 382     (bndrs,body) = collectBinders rhs
 383     val_bndrs    = filter isId bndrs
 384     env'         = extendRecBndr env fn bndrs
 385
 386 scBind env (Rec prs)
 387   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 388     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 389   where
 390     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 391                         returnUs (usg, (bndr,rhs'))
 392
 393 scBind env (NonRec bndr rhs)
 394   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 395     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 396
 397 ----------------------
 398 varUsage env v use
 399   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
 400                                                 occs = unitVarEnv v use }
 401   | otherwise                           = nullUsage
 402 \end{code}
 403
 404
 405 %************************************************************************
 406 %*                                                                      *
 407 \subsection{The specialiser}
 408 %*                                                                      *
 409 %************************************************************************
 410
 411 \begin{code}
 412 specialise :: ScEnv
 413            -> Id                        -- Functionn
 414            -> [CoreBndr] -> CoreExpr    -- Its RHS
 415            -> ScUsage                   -- Info on usage
 416            -> UniqSM ([CoreRule],       -- Rules
 417                       [(Id,CoreExpr)])  -- Bindings
 418
 419 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 420   = getUs               `thenUs` \ us ->
 421     let
 422         all_calls = lookupVarEnv calls fn `orElse` []
 423
 424         good_calls :: [[CoreArg]]
 425         good_calls = [ pats
 426                      | (con_env, call_args) <- all_calls,
 427                        length call_args >= n_bndrs,         -- App is saturated
 428                        let call = (bndrs `zip` call_args),
 429                        any (good_arg con_env occs) call,    -- At least one arg is a constr app
 430                        let (_, pats) = argsToPats con_env us call_args
 431                      ]
 432     in
 433     mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 434                   (nubBy same_call good_calls `zip` [1..])
 435   where
 436     n_bndrs  = length bndrs
 437     same_call as1 as2 = and (zipWith eqExpr as1 as2)
 438
 439 ---------------------
 440 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
 441 good_arg con_env arg_occs (bndr, arg)
 442   = case is_con_app_maybe con_env arg of
 443         Just _ ->  bndr_usg_ok arg_occs bndr arg
 444         other   -> False
 445
 446 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 447 bndr_usg_ok arg_occs bndr arg
 448   = case lookupVarEnv arg_occs bndr of
 449         Just CaseScrut -> True                  -- Used only by case scrutiny
 450         Just Both      -> case arg of           -- Used by case and elsewhere
 451                             App _ _ -> True     -- so the arg should be an explicit con app
 452                             other   -> False
 453         other -> False                          -- Not used, or used wonkily
 454
 455
 456 ---------------------
 457 spec_one :: ScEnv
 458          -> Id                                  -- Function
 459          -> CoreExpr                            -- Rhs of the original function
 460          -> ([CoreArg], Int)
 461          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 462
 463 -- spec_one creates a specialised copy of the function, together
 464 -- with a rule for using it.  I'm very proud of how short this
 465 -- function is, considering what it does :-).
 466
 467 {-
 468   Example
 469
 470      In-scope: a, x::a
 471      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) v (h v))...
 472           [c is presumably bound by the (...) part]
 473   ==>
 474      f_spec = /\ b c \ v::(a,(b,c)) ->
 475                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
 476
 477      RULE:  forall b c,
 478                    y::[(a,(b,c))],
 479                    v::(a,(b,c)),
 480                    h::(a,(b,c))->[(a,(b,c))] .
 481
 482             f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
 483 -}
 484
 485 spec_one env fn rhs (pats, n)
 486   = getUniqueUs                         `thenUs` \ spec_uniq ->
 487     let
 488         fn_name      = idName fn
 489         fn_loc       = nameSrcLoc fn_name
 490         spec_occ     = mkSpecOcc (nameOccName fn_name)
 491         pat_fvs      = varSetElems (exprsFreeVars pats)
 492         vars_to_bind = filter not_avail pat_fvs
 493         not_avail v  = not (v `elemVarEnv` scope env)
 494                 -- Put the type variables first just for tidiness
 495         (tvs, ids)   = partition isTyVar vars_to_bind
 496         bndrs        = tvs ++ ids
 497
 498         rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
 499         spec_rhs  = mkLams bndrs (mkApps rhs pats)
 500         spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
 501         rule      = Rule rule_name pat_fvs pats (mkVarApps (Var spec_id) bndrs)
 502     in
 503     pprTrace "SpecConstr" (pprCoreRule (ppr fn) rule)   $
 504     returnUs (rule, (spec_id, spec_rhs))
 505 \end{code}
 506
 507 %************************************************************************
 508 %*                                                                      *
 509 \subsection{Argument analysis}
 510 %*                                                                      *
 511 %************************************************************************
 512
 513 This code deals with analysing call-site arguments to see whether
 514 they are constructor applications.
 515
 516 \begin{code}
 517     -- argToPat takes an actual argument, and returns an abstracted
 518     -- version, consisting of just the "constructor skeleton" of the
 519     -- argument, with non-constructor sub-expression replaced by new
 520     -- placeholder variables.  For example:
 521     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 522
 523 argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
 524 argToPat env us (Type ty)
 525   = (us, Type ty)
 526
 527 argToPat env us arg
 528   | Just (dc,args) <- is_con_app_maybe env arg
 529   = let
 530         (us',args') = argsToPats env us args
 531     in
 532     (us', mk_con_app dc args')
 533
 534 argToPat env us (Var v) -- Don't uniqify existing vars,
 535   = (us, Var v)         -- so that we can spot when we pass them twice
 536
 537 argToPat env us arg
 538   = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
 539   where
 540     (us1,us2) = splitUniqSupply us
 541
 542 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
 543 argsToPats env us args = mapAccumL (argToPat env) us args
 544 \end{code}
 545
 546
 547 \begin{code}
 548 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
 549 is_con_app_maybe env (Var v)
 550   = lookupVarEnv env v
 551         -- You might think we could look in the idUnfolding here
 552         -- but that doesn't take account of which branch of a
 553         -- case we are in, which is the whole point
 554
 555 is_con_app_maybe env (Lit lit)
 556   = Just (LitAlt lit, [])
 557
 558 is_con_app_maybe env expr
 559   = case collectArgs expr of
 560         (Var fun, args) | Just con <- isDataConId_maybe fun,
 561                           length args >= dataConRepArity con
 562                 -- Might be > because the arity excludes type args
 563                         -> Just (DataAlt con,args)
 564
 565         other -> Nothing
 566
 567 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
 568 mk_con_app (LitAlt lit)  []   = Lit lit
 569 mk_con_app (DataAlt con) args = mkConApp con args
 570 \end{code}