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, mkPiTypes )
  16 import CoreFVs          ( exprsFreeVars )
  17 import CoreTidy         ( tidyRules )
  18 import PprCore          ( pprRules )
  19 import WwLib            ( mkWorkerArgs )
  20 import DataCon          ( dataConRepArity, dataConUnivTyVars )
  21 import Type             ( Type, tyConAppArgs )
  22 import Coercion         ( coercionKind )
  23 import Rules            ( matchN )
  24 import Id               ( Id, idName, idType, isDataConWorkId_maybe,
  25                           mkUserLocal, mkSysLocal, idUnfolding, isLocalId )
  26 import Var              ( Var )
  27 import VarEnv
  28 import VarSet
  29 import Name             ( nameOccName, nameSrcLoc )
  30 import Rules            ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
  31 import OccName          ( mkSpecOcc )
  32 import ErrUtils         ( dumpIfSet_dyn )
  33 import DynFlags         ( DynFlags, DynFlag(..) )
  34 import BasicTypes       ( Activation(..) )
  35 import Maybes           ( orElse, catMaybes, isJust )
  36 import Util             ( zipWithEqual, lengthAtLeast, notNull )
  37 import List             ( nubBy, partition )
  38 import UniqSupply
  39 import Outputable
  40 import FastString
  41 import UniqFM
  42 \end{code}
  43
  44 -----------------------------------------------------
  45                         Game plan
  46 -----------------------------------------------------
  47
  48 Consider
  49         drop n []     = []
  50         drop 0 xs     = []
  51         drop n (x:xs) = drop (n-1) xs
  52
  53 After the first time round, we could pass n unboxed.  This happens in
  54 numerical code too.  Here's what it looks like in Core:
  55
  56         drop n xs = case xs of
  57                       []     -> []
  58                       (y:ys) -> case n of
  59                                   I# n# -> case n# of
  60                                              0 -> []
  61                                              _ -> drop (I# (n# -# 1#)) xs
  62
  63 Notice that the recursive call has an explicit constructor as argument.
  64 Noticing this, we can make a specialised version of drop
  65
  66         RULE: drop (I# n#) xs ==> drop' n# xs
  67
  68         drop' n# xs = let n = I# n# in ...orig RHS...
  69
  70 Now the simplifier will apply the specialisation in the rhs of drop', giving
  71
  72         drop' n# xs = case xs of
  73                       []     -> []
  74                       (y:ys) -> case n# of
  75                                   0 -> []
  76                                   _ -> drop (n# -# 1#) xs
  77
  78 Much better!
  79
  80 We'd also like to catch cases where a parameter is carried along unchanged,
  81 but evaluated each time round the loop:
  82
  83         f i n = if i>0 || i>n then i else f (i*2) n
  84
  85 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  86 In Core, by the time we've w/wd (f is strict in i) we get
  87
  88         f i# n = case i# ># 0 of
  89                    False -> I# i#
  90                    True  -> case n of n' { I# n# ->
  91                             case i# ># n# of
  92                                 False -> I# i#
  93                                 True  -> f (i# *# 2#) n'
  94
  95 At the call to f, we see that the argument, n is know to be (I# n#),
  96 and n is evaluated elsewhere in the body of f, so we can play the same
  97 trick as above.
  98
  99
 100 Note [Reboxing]
 101 ~~~~~~~~~~~~~~~
 102 We must be careful not to allocate the same constructor twice.  Consider
 103         f p = (...(case p of (a,b) -> e)...p...,
 104                ...let t = (r,s) in ...t...(f t)...)
 105 At the recursive call to f, we can see that t is a pair.  But we do NOT want
 106 to make a specialised copy:
 107         f' a b = let p = (a,b) in (..., ...)
 108 because now t is allocated by the caller, then r and s are passed to the
 109 recursive call, which allocates the (r,s) pair again.
 110
 111 This happens if
 112   (a) the argument p is used in other than a case-scrutinsation way.
 113   (b) the argument to the call is not a 'fresh' tuple; you have to
 114         look into its unfolding to see that it's a tuple
 115
 116 Hence the "OR" part of Note [Good arguments] below.
 117
 118 ALTERNATIVE 2: pass both boxed and unboxed versions.  This no longer saves
 119 allocation, but does perhaps save evals. In the RULE we'd have
 120 something like
 121
 122   f (I# x#) = f' (I# x#) x#
 123
 124 If at the call site the (I# x) was an unfolding, then we'd have to
 125 rely on CSE to eliminate the duplicate allocation.... This alternative
 126 doesn't look attractive enough to pursue.
 127
 128 ALTERNATIVE 3: ignore the reboxing problem.  The trouble is that
 129 the conservative reboxing story prevents many useful functions from being
 130 specialised.  Example:
 131         foo :: Maybe Int -> Int -> Int
 132         foo   (Just m) 0 = 0
 133         foo x@(Just m) n = foo x (n-m)
 134 Here the use of 'x' will clearly not require boxing in the specialised function.
 135
 136 The strictness analyser has the same problem, in fact.  Example:
 137         f p@(a,b) = ...
 138 If we pass just 'a' and 'b' to the worker, it might need to rebox the
 139 pair to create (a,b).  A more sophisticated analysis might figure out
 140 precisely the cases in which this could happen, but the strictness
 141 analyser does no such analysis; it just passes 'a' and 'b', and hopes
 142 for the best.
 143
 144 So my current choice is to make SpecConstr similarly aggressive, and
 145 ignore the bad potential of reboxing.
 146
 147
 148 Note [Good arguments]
 149 ~~~~~~~~~~~~~~~~~~~~~
 150 So we look for
 151
 152 * A self-recursive function.  Ignore mutual recursion for now,
 153   because it's less common, and the code is simpler for self-recursion.
 154
 155 * EITHER
 156
 157    a) At a recursive call, one or more parameters is an explicit
 158       constructor application
 159         AND
 160       That same parameter is scrutinised by a case somewhere in
 161       the RHS of the function
 162
 163   OR
 164
 165     b) At a recursive call, one or more parameters has an unfolding
 166        that is an explicit constructor application
 167         AND
 168       That same parameter is scrutinised by a case somewhere in
 169       the RHS of the function
 170         AND
 171       Those are the only uses of the parameter (see Note [Reboxing])
 172
 173
 174 What to abstract over
 175 ~~~~~~~~~~~~~~~~~~~~~
 176 There's a bit of a complication with type arguments.  If the call
 177 site looks like
 178
 179         f p = ...f ((:) [a] x xs)...
 180
 181 then our specialised function look like
 182
 183         f_spec x xs = let p = (:) [a] x xs in ....as before....
 184
 185 This only makes sense if either
 186   a) the type variable 'a' is in scope at the top of f, or
 187   b) the type variable 'a' is an argument to f (and hence fs)
 188
 189 Actually, (a) may hold for value arguments too, in which case
 190 we may not want to pass them.  Supose 'x' is in scope at f's
 191 defn, but xs is not.  Then we'd like
 192
 193         f_spec xs = let p = (:) [a] x xs in ....as before....
 194
 195 Similarly (b) may hold too.  If x is already an argument at the
 196 call, no need to pass it again.
 197
 198 Finally, if 'a' is not in scope at the call site, we could abstract
 199 it as we do the term variables:
 200
 201         f_spec a x xs = let p = (:) [a] x xs in ...as before...
 202
 203 So the grand plan is:
 204
 205         * abstract the call site to a constructor-only pattern
 206           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
 207
 208         * Find the free variables of the abstracted pattern
 209
 210         * Pass these variables, less any that are in scope at
 211           the fn defn.  But see Note [Shadowing] below.
 212
 213
 214 NOTICE that we only abstract over variables that are not in scope,
 215 so we're in no danger of shadowing variables used in "higher up"
 216 in f_spec's RHS.
 217
 218
 219 Note [Shadowing]
 220 ~~~~~~~~~~~~~~~~
 221 In this pass we gather up usage information that may mention variables
 222 that are bound between the usage site and the definition site; or (more
 223 seriously) may be bound to something different at the definition site.
 224 For example:
 225
 226         f x = letrec g y v = let x = ...
 227                              in ...(g (a,b) x)...
 228
 229 Since 'x' is in scope at the call site, we may make a rewrite rule that
 230 looks like
 231         RULE forall a,b. g (a,b) x = ...
 232 But this rule will never match, because it's really a different 'x' at
 233 the call site -- and that difference will be manifest by the time the
 234 simplifier gets to it.  [A worry: the simplifier doesn't *guarantee*
 235 no-shadowing, so perhaps it may not be distinct?]
 236
 237 Anyway, the rule isn't actually wrong, it's just not useful.  One possibility
 238 is to run deShadowBinds before running SpecConstr, but instead we run the
 239 simplifier.  That gives the simplest possible program for SpecConstr to
 240 chew on; and it virtually guarantees no shadowing.
 241
 242 Note [Specialising for constant parameters]
 243 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 244 This one is about specialising on a *constant* (but not necessarily
 245 constructor) argument
 246
 247     foo :: Int -> (Int -> Int) -> Int
 248     foo 0 f = 0
 249     foo m f = foo (f m) (+1)
 250
 251 It produces
 252
 253     lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
 254     lvl_rmV =
 255       \ (ds_dlk :: GHC.Base.Int) ->
 256         case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
 257         GHC.Base.I# (GHC.Prim.+# x_alG 1)
 258
 259     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 260     GHC.Prim.Int#
 261     T.$wfoo =
 262       \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
 263         case ww_sme of ds_Xlw {
 264           __DEFAULT ->
 265         case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
 266         T.$wfoo ww1_Xmz lvl_rmV
 267         };
 268           0 -> 0
 269         }
 270
 271 The recursive call has lvl_rmV as its argument, so we could create a specialised copy
 272 with that argument baked in; that is, not passed at all.   Now it can perhaps be inlined.
 273
 274 When is this worth it?  Call the constant 'lvl'
 275 - If 'lvl' has an unfolding that is a constructor, see if the corresponding
 276   parameter is scrutinised anywhere in the body.
 277
 278 - If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
 279   parameter is applied (...to enough arguments...?)
 280
 281   Also do this is if the function has RULES?
 282
 283 Also
 284
 285 Note [Specialising for lambda parameters]
 286 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 287     foo :: Int -> (Int -> Int) -> Int
 288     foo 0 f = 0
 289     foo m f = foo (f m) (\n -> n-m)
 290
 291 This is subtly different from the previous one in that we get an
 292 explicit lambda as the argument:
 293
 294     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 295     GHC.Prim.Int#
 296     T.$wfoo =
 297       \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
 298         case ww_sm8 of ds_Xlr {
 299           __DEFAULT ->
 300         case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
 301         T.$wfoo
 302           ww1_Xmq
 303           (\ (n_ad3 :: GHC.Base.Int) ->
 304              case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
 305              GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
 306              })
 307         };
 308           0 -> 0
 309         }
 310
 311 I wonder if SpecConstr couldn't be extended to handle this? After all,
 312 lambda is a sort of constructor for functions and perhaps it already
 313 has most of the necessary machinery?
 314
 315 Furthermore, there's an immediate win, because you don't need to allocate the lamda
 316 at the call site; and if perchance it's called in the recursive call, then you
 317 may avoid allocating it altogether.  Just like for constructors.
 318
 319 Looks cool, but probably rare...but it might be easy to implement.
 320
 321
 322 Note [SpecConstr for casts]
 323 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 324 Consider
 325     data family T a :: *
 326     data instance T Int = T Int
 327
 328     foo n = ...
 329        where
 330          go (T 0) = 0
 331          go (T n) = go (T (n-1))
 332
 333 The recursive call ends up looking like
 334         go (T (I# ...) `cast` g)
 335 So we want to spot the construtor application inside the cast.
 336 That's why we have the Cast case in argToPat
 337
 338
 339 -----------------------------------------------------
 340                 Stuff not yet handled
 341 -----------------------------------------------------
 342
 343 Here are notes arising from Roman's work that I don't want to lose.
 344
 345 Example 1
 346 ~~~~~~~~~
 347     data T a = T !a
 348
 349     foo :: Int -> T Int -> Int
 350     foo 0 t = 0
 351     foo x t | even x    = case t of { T n -> foo (x-n) t }
 352             | otherwise = foo (x-1) t
 353
 354 SpecConstr does no specialisation, because the second recursive call
 355 looks like a boxed use of the argument.  A pity.
 356
 357     $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 358     $wfoo_sFw =
 359       \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
 360          case ww_sFo of ds_Xw6 [Just L] {
 361            __DEFAULT ->
 362                 case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
 363                   __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
 364                   0 ->
 365                     case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
 366                     case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
 367                     $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
 368                     } } };
 369            0 -> 0
 370
 371 Example 2
 372 ~~~~~~~~~
 373     data a :*: b = !a :*: !b
 374     data T a = T !a
 375
 376     foo :: (Int :*: T Int) -> Int
 377     foo (0 :*: t) = 0
 378     foo (x :*: t) | even x    = case t of { T n -> foo ((x-n) :*: t) }
 379                   | otherwise = foo ((x-1) :*: t)
 380
 381 Very similar to the previous one, except that the parameters are now in
 382 a strict tuple. Before SpecConstr, we have
 383
 384     $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 385     $wfoo_sG3 =
 386       \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
 387     GHC.Base.Int) ->
 388         case ww_sFU of ds_Xws [Just L] {
 389           __DEFAULT ->
 390         case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
 391           __DEFAULT ->
 392             case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
 393             $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2             -- $wfoo1
 394             };
 395           0 ->
 396             case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
 397             case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
 398             $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB        -- $wfoo2
 399             } } };
 400           0 -> 0 }
 401
 402 We get two specialisations:
 403 "SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
 404                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
 405                   = Foo.$s$wfoo1 a_sFB sc_sGC ;
 406 "SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
 407                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
 408                   = Foo.$s$wfoo y_aFp sc_sGC ;
 409
 410 But perhaps the first one isn't good.  After all, we know that tpl_B2 is
 411 a T (I# x) really, because T is strict and Int has one constructor.  (We can't
 412 unbox the strict fields, becuase T is polymorphic!)
 413
 414
 415
 416 %************************************************************************
 417 %*                                                                      *
 418 \subsection{Top level wrapper stuff}
 419 %*                                                                      *
 420 %************************************************************************
 421
 422 \begin{code}
 423 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
 424 specConstrProgram dflags us binds
 425   = do
 426         showPass dflags "SpecConstr"
 427
 428         let (binds', _) = initUs us (go emptyScEnv binds)
 429
 430         endPass dflags "SpecConstr" Opt_D_dump_spec binds'
 431
 432         dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
 433                   (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
 434
 435         return binds'
 436   where
 437     go env []           = returnUs []
 438     go env (bind:binds) = scBind env bind       `thenUs` \ (env', _, bind') ->
 439                           go env' binds         `thenUs` \ binds' ->
 440                           returnUs (bind' : binds')
 441 \end{code}
 442
 443
 444 %************************************************************************
 445 %*                                                                      *
 446 \subsection{Environment: goes downwards}
 447 %*                                                                      *
 448 %************************************************************************
 449
 450 \begin{code}
 451 data ScEnv = SCE { scope :: InScopeEnv,
 452                         -- Binds all non-top-level variables in scope
 453
 454                    cons  :: ConstrEnv
 455              }
 456
 457 type InScopeEnv = VarEnv HowBound
 458
 459 type ConstrEnv = IdEnv ConValue
 460 data ConValue  = CV AltCon [CoreArg]
 461         -- Variables known to be bound to a constructor
 462         -- in a particular case alternative
 463
 464
 465 instance Outputable ConValue where
 466    ppr (CV con args) = ppr con <+> interpp'SP args
 467
 468 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
 469
 470 data HowBound = RecFun  -- These are the recursive functions for which
 471                         -- we seek interesting call patterns
 472
 473               | RecArg  -- These are those functions' arguments, or their sub-components;
 474                         -- we gather occurrence information for these
 475
 476               | Other   -- We track all others so we know what's in scope
 477                         -- This is used in spec_one to check what needs to be
 478                         -- passed as a parameter and what is in scope at the
 479                         -- function definition site
 480
 481 instance Outputable HowBound where
 482   ppr RecFun = text "RecFun"
 483   ppr RecArg = text "RecArg"
 484   ppr Other = text "Other"
 485
 486 lookupScopeEnv env v = lookupVarEnv (scope env) v
 487
 488 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
 489 extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }
 490
 491     -- When we encounter
 492     --  case scrut of b
 493     --      C x y -> ...
 494     -- we want to bind b, and perhaps scrut too, to (C x y)
 495 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
 496 extendCaseBndrs env case_bndr scrut con alt_bndrs
 497   = case con of
 498         DEFAULT    -> env1
 499         LitAlt lit -> extendCons env1 scrut case_bndr (CV con [])
 500         DataAlt dc -> extend_data_con dc
 501   where
 502     cur_scope = scope env
 503     env1 = env { scope = extendVarEnvList cur_scope
 504                                 [(b,how_bound) | b <- case_bndr:alt_bndrs] }
 505
 506         -- Record RecArg for the components iff the scrutinee is RecArg
 507         -- I think the only reason for this is to keep the usage envt small
 508         -- so is it worth it at all?
 509         --      [This comment looks plain wrong to me, so I'm ignoring it
 510         --           "Also forget if the scrutinee is a RecArg, because we're
 511         --           now in the branch of a case, and we don't want to
 512         --           record a non-scrutinee use of v if we have
 513         --              case v of { (a,b) -> ...(f v)... }" ]
 514     how_bound = get_how scrut
 515         where
 516             get_how (Var v)    = lookupVarEnv cur_scope v `orElse` Other
 517             get_how (Cast e _) = get_how e
 518             get_how (Note _ e) = get_how e
 519             get_how other      = Other
 520
 521     extend_data_con data_con =
 522       extendCons env1 scrut case_bndr (CV con vanilla_args)
 523         where
 524             vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
 525                            varsToCoreExprs alt_bndrs
 526
 527 extendCons :: ScEnv -> CoreExpr -> Id -> ConValue -> ScEnv
 528 extendCons env scrut case_bndr val
 529   = case scrut of
 530         Var v -> env { cons = extendVarEnv cons1 v val }
 531         other -> env { cons = cons1 }
 532   where
 533     cons1 = extendVarEnv (cons env) case_bndr val
 534
 535     -- When we encounter a recursive function binding
 536     --  f = \x y -> ...
 537     -- we want to extend the scope env with bindings
 538     -- that record that f is a RecFn and x,y are RecArgs
 539 extendRecBndr env fn bndrs
 540   =  env { scope = scope env `extendVarEnvList`
 541                    ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
 542 \end{code}
 543
 544
 545 %************************************************************************
 546 %*                                                                      *
 547 \subsection{Usage information: flows upwards}
 548 %*                                                                      *
 549 %************************************************************************
 550
 551 \begin{code}
 552 data ScUsage
 553    = SCU {
 554         calls :: !(IdEnv ([Call])),     -- Calls
 555                                         -- The functions are a subset of the
 556                                         --      RecFuns in the ScEnv
 557
 558         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 559      }                                  -- The variables are a subset of the
 560                                         --      RecArg in the ScEnv
 561
 562 type Call = (ConstrEnv, [CoreArg])
 563         -- The arguments of the call, together with the
 564         -- env giving the constructor bindings at the call site
 565
 566 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 567
 568 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 569                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 570
 571 combineUsages [] = nullUsage
 572 combineUsages us = foldr1 combineUsage us
 573
 574 lookupOcc :: ScUsage -> Var -> (ScUsage, ArgOcc)
 575 lookupOcc (SCU { calls = sc_calls, occs = sc_occs }) bndr
 576   = (SCU {calls = sc_calls, occs = delVarEnv sc_occs bndr},
 577      lookupVarEnv sc_occs bndr `orElse` NoOcc)
 578
 579 lookupOccs :: ScUsage -> [Var] -> (ScUsage, [ArgOcc])
 580 lookupOccs (SCU { calls = sc_calls, occs = sc_occs }) bndrs
 581   = (SCU {calls = sc_calls, occs = delVarEnvList sc_occs bndrs},
 582      [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
 583
 584 data ArgOcc = NoOcc     -- Doesn't occur at all; or a type argument
 585             | UnkOcc    -- Used in some unknown way
 586
 587             | ScrutOcc (UniqFM [ArgOcc])        -- See Note [ScrutOcc]
 588
 589             | BothOcc   -- Definitely taken apart, *and* perhaps used in some other way
 590
 591 {-      Note  [ScrutOcc]
 592
 593 An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
 594 is *only* taken apart or applied.
 595
 596   Functions, literal: ScrutOcc emptyUFM
 597   Data constructors:  ScrutOcc subs,
 598
 599 where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
 600 The domain of the UniqFM is the Unique of the data constructor
 601
 602 The [ArgOcc] is the occurrences of the *pattern-bound* components
 603 of the data structure.  E.g.
 604         data T a = forall b. MkT a b (b->a)
 605 A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
 606
 607 -}
 608
 609 instance Outputable ArgOcc where
 610   ppr (ScrutOcc xs) = ptext SLIT("scrut-occ") <> ppr xs
 611   ppr UnkOcc        = ptext SLIT("unk-occ")
 612   ppr BothOcc       = ptext SLIT("both-occ")
 613   ppr NoOcc         = ptext SLIT("no-occ")
 614
 615 -- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
 616 -- that if the thing is scrutinised anywhere then we get to see that
 617 -- in the overall result, even if it's also used in a boxed way
 618 -- This might be too agressive; see Note [Reboxing] Alternative 3
 619 combineOcc NoOcc         occ           = occ
 620 combineOcc occ           NoOcc         = occ
 621 combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
 622 combineOcc occ           (ScrutOcc ys) = ScrutOcc ys
 623 combineOcc (ScrutOcc xs) occ           = ScrutOcc xs
 624 combineOcc UnkOcc        UnkOcc        = UnkOcc
 625 combineOcc _        _                  = BothOcc
 626
 627 combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
 628 combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
 629
 630 conArgOccs :: ArgOcc -> AltCon -> [ArgOcc]
 631 -- Find usage of components of data con; returns [UnkOcc...] if unknown
 632 -- See Note [ScrutOcc] for the extra UnkOccs in the vanilla datacon case
 633
 634 conArgOccs (ScrutOcc fm) (DataAlt dc)
 635   | Just pat_arg_occs <- lookupUFM fm dc
 636   = [UnkOcc | tv <- dataConUnivTyVars dc] ++ pat_arg_occs
 637
 638 conArgOccs other con = repeat UnkOcc
 639 \end{code}
 640
 641
 642 %************************************************************************
 643 %*                                                                      *
 644 \subsection{The main recursive function}
 645 %*                                                                      *
 646 %************************************************************************
 647
 648 The main recursive function gathers up usage information, and
 649 creates specialised versions of functions.
 650
 651 \begin{code}
 652 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 653         -- The unique supply is needed when we invent
 654         -- a new name for the specialised function and its args
 655
 656 scExpr env e@(Type t) = returnUs (nullUsage, e)
 657 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 658 scExpr env e@(Var v)  = returnUs (varUsage env v UnkOcc, e)
 659 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 660                         returnUs (usg, Note n e')
 661 scExpr env (Cast e co)= scExpr env e    `thenUs` \ (usg,e') ->
 662                         returnUs (usg, Cast e' co)
 663 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 664                         returnUs (usg, Lam b e')
 665
 666 scExpr env (Case scrut b ty alts)
 667   = do  { (alt_usgs, alt_occs, alts') <- mapAndUnzip3Us sc_alt alts
 668         ; let (alt_usg, b_occ) = lookupOcc (combineUsages alt_usgs) b
 669               scrut_occ = foldr combineOcc b_occ alt_occs
 670                 -- The combined usage of the scrutinee is given
 671                 -- by scrut_occ, which is passed to scScrut, which
 672                 -- in turn treats a bare-variable scrutinee specially
 673         ; (scrut_usg, scrut') <- scScrut env scrut scrut_occ
 674         ; return (alt_usg `combineUsage` scrut_usg,
 675                   Case scrut' b ty alts') }
 676   where
 677     sc_alt (con,bs,rhs)
 678       = do { let env1 = extendCaseBndrs env b scrut con bs
 679            ; (usg,rhs') <- scExpr env1 rhs
 680            ; let (usg', arg_occs) = lookupOccs usg bs
 681                  scrut_occ = case con of
 682                                 DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
 683                                 other      -> ScrutOcc emptyUFM
 684            ; return (usg', scrut_occ, (con,bs,rhs')) }
 685
 686 scExpr env (Let bind body)
 687   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 688     scExpr env' body    `thenUs` \ (body_usg, body') ->
 689     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 690
 691 scExpr env e@(App _ _)
 692   = do  { let (fn, args) = collectArgs e
 693         ; (fn_usg, fn') <- scScrut env fn (ScrutOcc emptyUFM)
 694         -- Process the function too.   It's almost always a variable,
 695         -- but not always.  In particular, if this pass follows float-in,
 696         -- which it may, we can get
 697         --      (let f = ...f... in f) arg1 arg2
 698         -- We use scScrut to record the fact that the function is called
 699         -- Perhpas we should check that it has at least one value arg,
 700         -- but currently we don't bother
 701
 702         ; (arg_usgs, args') <- mapAndUnzipUs (scExpr env) args
 703         ; let call_usg = case fn of
 704                            Var f | Just RecFun <- lookupScopeEnv env f
 705                                  -> SCU { calls = unitVarEnv f [(cons env, args)],
 706                                           occs  = emptyVarEnv }
 707                            other -> nullUsage
 708         ; return (combineUsages arg_usgs `combineUsage` fn_usg
 709                                          `combineUsage` call_usg,
 710                   mkApps fn' args') }
 711
 712
 713 ----------------------
 714 scScrut :: ScEnv -> CoreExpr -> ArgOcc -> UniqSM (ScUsage, CoreExpr)
 715 -- Used for the scrutinee of a case,
 716 -- or the function of an application.
 717 -- Remember to look through casts
 718 scScrut env e@(Var v)   occ = returnUs (varUsage env v occ, e)
 719 scScrut env (Cast e co) occ = do { (usg, e') <- scScrut env e occ
 720                                  ; returnUs (usg, Cast e' co) }
 721 scScrut env e           occ = scExpr env e
 722
 723
 724 ----------------------
 725 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 726 scBind env (Rec [(fn,rhs)])
 727   | notNull val_bndrs
 728   = scExpr env_fn_body body             `thenUs` \ (usg, body') ->
 729     specialise env fn bndrs body' usg   `thenUs` \ (rules, spec_prs) ->
 730         -- Note body': the specialised copies should be based on the
 731         --             optimised version of the body, in case there were
 732         --             nested functions inside.
 733     let
 734         SCU { calls = calls, occs = occs } = usg
 735     in
 736     returnUs (extendBndr env fn,        -- For the body of the letrec, just
 737                                         -- extend the env with Other to record
 738                                         -- that it's in scope; no funny RecFun business
 739               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 740               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 741   where
 742     (bndrs,body) = collectBinders rhs
 743     val_bndrs    = filter isId bndrs
 744     env_fn_body  = extendRecBndr env fn bndrs
 745
 746 scBind env (Rec prs)
 747   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 748     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 749   where
 750     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 751                         returnUs (usg, (bndr,rhs'))
 752
 753 scBind env (NonRec bndr rhs)
 754   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 755     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 756
 757 ----------------------
 758 varUsage env v use
 759   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
 760                                                 occs = unitVarEnv v use }
 761   | otherwise                           = nullUsage
 762 \end{code}
 763
 764
 765 %************************************************************************
 766 %*                                                                      *
 767 \subsection{The specialiser}
 768 %*                                                                      *
 769 %************************************************************************
 770
 771 \begin{code}
 772 specialise :: ScEnv
 773            -> Id                        -- Functionn
 774            -> [CoreBndr] -> CoreExpr    -- Its RHS
 775            -> ScUsage                   -- Info on usage
 776            -> UniqSM ([CoreRule],       -- Rules
 777                       [(Id,CoreExpr)])  -- Bindings
 778
 779 specialise env fn bndrs body body_usg
 780   = do  { let (_, bndr_occs) = lookupOccs body_usg bndrs
 781               all_calls = lookupVarEnv (calls body_usg) fn `orElse` []
 782
 783         ; mb_pats <- mapM (callToPats (scope env) bndr_occs) all_calls
 784
 785         ; let good_pats :: [([Var], [CoreArg])]
 786               good_pats = catMaybes mb_pats
 787               in_scope = mkInScopeSet $ unionVarSets $
 788                          [ exprsFreeVars pats `delVarSetList` vs
 789                          | (vs,pats) <- good_pats ]
 790               uniq_pats = nubBy (same_pat in_scope) good_pats
 791         ; -- pprTrace "specialise" (vcat [ppr fn <+> ppr bndrs <+> ppr bndr_occs,
 792           --                            text "calls" <+> ppr all_calls,
 793           --                            text "good pats" <+> ppr good_pats,
 794           --                            text "uniq pats" <+> ppr uniq_pats])  $
 795           mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 796                         (uniq_pats `zip` [1..]) }
 797   where
 798         -- Two pats are the same if they match both ways
 799     same_pat in_scope (vs1,as1)(vs2,as2)
 800          =  isJust (matchN in_scope vs1 as1 as2)
 801          && isJust (matchN in_scope vs2 as2 as1)
 802
 803 callToPats :: InScopeEnv -> [ArgOcc] -> Call
 804            -> UniqSM (Maybe ([Var], [CoreExpr]))
 805         -- The VarSet is the variables to quantify over in the rule
 806         -- The [CoreExpr] are the argument patterns for the rule
 807 callToPats in_scope bndr_occs (con_env, args)
 808   | length args < length bndr_occs      -- Check saturated
 809   = return Nothing
 810   | otherwise
 811   = do  { prs <- argsToPats in_scope con_env (args `zip` bndr_occs)
 812         ; let (good_pats, pats) = unzip prs
 813               pat_fvs = varSetElems (exprsFreeVars pats)
 814               qvars   = filter (not . (`elemVarEnv` in_scope)) pat_fvs
 815                 -- Quantify over variables that are not in sccpe
 816                 -- See Note [Shadowing] at the top
 817
 818         ; -- pprTrace "callToPats"  (ppr args $$ ppr prs $$ ppr bndr_occs) $
 819           if or good_pats
 820           then return (Just (qvars, pats))
 821           else return Nothing }
 822
 823 ---------------------
 824 spec_one :: ScEnv
 825          -> Id                                  -- Function
 826          -> CoreExpr                            -- Rhs of the original function
 827          -> (([Var], [CoreArg]), Int)
 828          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 829
 830 -- spec_one creates a specialised copy of the function, together
 831 -- with a rule for using it.  I'm very proud of how short this
 832 -- function is, considering what it does :-).
 833
 834 {-
 835   Example
 836
 837      In-scope: a, x::a
 838      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
 839           [c::*, v::(b,c) are presumably bound by the (...) part]
 840   ==>
 841      f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
 842                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
 843
 844      RULE:  forall b::* c::*,           -- Note, *not* forall a, x
 845                    v::(b,c),
 846                    hw::[(a,(b,c))] .
 847
 848             f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
 849 -}
 850
 851 spec_one env fn rhs ((vars_to_bind, pats), rule_number)
 852   = getUniqueUs                 `thenUs` \ spec_uniq ->
 853     let
 854         fn_name      = idName fn
 855         fn_loc       = nameSrcLoc fn_name
 856         spec_occ     = mkSpecOcc (nameOccName fn_name)
 857
 858                 -- Put the type variables first; the type of a term
 859                 -- variable may mention a type variable
 860         (tvs, ids)   = partition isTyVar vars_to_bind
 861         bndrs        = tvs ++ ids
 862         spec_body    = mkApps rhs pats
 863         body_ty      = exprType spec_body
 864
 865         (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
 866                 -- Usual w/w hack to avoid generating
 867                 -- a spec_rhs of unlifted type and no args
 868
 869         rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
 870         spec_rhs  = mkLams spec_lam_args spec_body
 871         spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
 872         rule_rhs  = mkVarApps (Var spec_id) spec_call_args
 873         rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
 874     in
 875     returnUs (rule, (spec_id, spec_rhs))
 876
 877 -- In which phase should the specialise-constructor rules be active?
 878 -- Originally I made them always-active, but Manuel found that
 879 -- this defeated some clever user-written rules.  So Plan B
 880 -- is to make them active only in Phase 0; after all, currently,
 881 -- the specConstr transformation is only run after the simplifier
 882 -- has reached Phase 0.  In general one would want it to be
 883 -- flag-controllable, but for now I'm leaving it baked in
 884 --                                      [SLPJ Oct 01]
 885 specConstrActivation :: Activation
 886 specConstrActivation = ActiveAfter 0    -- Baked in; see comments above
 887 \end{code}
 888
 889 %************************************************************************
 890 %*                                                                      *
 891 \subsection{Argument analysis}
 892 %*                                                                      *
 893 %************************************************************************
 894
 895 This code deals with analysing call-site arguments to see whether
 896 they are constructor applications.
 897
 898
 899 \begin{code}
 900     -- argToPat takes an actual argument, and returns an abstracted
 901     -- version, consisting of just the "constructor skeleton" of the
 902     -- argument, with non-constructor sub-expression replaced by new
 903     -- placeholder variables.  For example:
 904     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 905
 906 argToPat :: InScopeEnv                  -- What's in scope at the fn defn site
 907          -> ConstrEnv                   -- ConstrEnv at the call site
 908          -> CoreArg                     -- A call arg (or component thereof)
 909          -> ArgOcc
 910          -> UniqSM (Bool, CoreArg)
 911 -- Returns (interesting, pat),
 912 -- where pat is the pattern derived from the argument
 913 --            intersting=True if the pattern is non-trivial (not a variable or type)
 914 -- E.g.         x:xs         --> (True, x:xs)
 915 --              f xs         --> (False, w)        where w is a fresh wildcard
 916 --              (f xs, 'c')  --> (True, (w, 'c'))  where w is a fresh wildcard
 917 --              \x. x+y      --> (True, \x. x+y)
 918 --              lvl7         --> (True, lvl7)      if lvl7 is bound
 919 --                                                 somewhere further out
 920
 921 argToPat in_scope con_env arg@(Type ty) arg_occ
 922   = return (False, arg)
 923
 924 argToPat in_scope con_env (Let _ arg) arg_occ
 925   = argToPat in_scope con_env arg arg_occ
 926         -- Look through let expressions
 927         -- e.g.         f (let v = rhs in \y -> ...v...)
 928         -- Here we can specialise for f (\y -> ...)
 929         -- because the rule-matcher will look through the let.
 930
 931 argToPat in_scope con_env (Cast arg co) arg_occ
 932   = do  { (interesting, arg') <- argToPat in_scope con_env arg arg_occ
 933         ; if interesting then
 934                 return (interesting, Cast arg' co)
 935           else
 936                 wildCardPat (snd (coercionKind co)) }
 937
 938 argToPat in_scope con_env arg arg_occ
 939   | is_value_lam arg
 940   = return (True, arg)
 941   where
 942     is_value_lam (Lam v e)      -- Spot a value lambda, even if
 943         | isId v = True         -- it is inside a type lambda
 944         | otherwise = is_value_lam e
 945     is_value_lam other = False
 946
 947   -- Check for a constructor application
 948   -- NB: this *precedes* the Var case, so that we catch nullary constrs
 949 argToPat in_scope con_env arg arg_occ
 950   | Just (CV dc args) <- is_con_app_maybe con_env arg
 951   , case arg_occ of
 952         ScrutOcc _ -> True              -- Used only by case scrutinee
 953         BothOcc    -> case arg of       -- Used elsewhere
 954                         App {} -> True  --     see Note [Reboxing]
 955                         other  -> False
 956         other      -> False     -- No point; the arg is not decomposed
 957   = do  { args' <- argsToPats in_scope con_env (args `zip` conArgOccs arg_occ dc)
 958         ; return (True, mk_con_app dc (map snd args')) }
 959
 960   -- Check if the argument is a variable that
 961   -- is in scope at the function definition site
 962   -- It's worth specialising on this if
 963   --    (a) it's used in an interesting way in the body
 964   --    (b) we know what its value is
 965 argToPat in_scope con_env (Var v) arg_occ
 966   | not (isLocalId v) || v `elemVarEnv` in_scope,
 967     case arg_occ of { UnkOcc -> False; other -> True }, -- (a)
 968     isValueUnfolding (idUnfolding v)                    -- (b)
 969   = return (True, Var v)
 970
 971   -- Check for a variable bound inside the function.
 972   -- Don't make a wild-card, because we may usefully share
 973   --    e.g.  f a = let x = ... in f (x,x)
 974   -- NB: this case follows the lambda and con-app cases!!
 975 argToPat in_scope con_env (Var v) arg_occ
 976   = return (False, Var v)
 977
 978   -- The default case: make a wild-card
 979 argToPat in_scope con_env arg arg_occ
 980   = wildCardPat (exprType arg)
 981
 982 wildCardPat :: Type -> UniqSM (Bool, CoreArg)
 983 wildCardPat ty = do { uniq <- getUniqueUs
 984                     ; let id = mkSysLocal FSLIT("sc") uniq ty
 985                     ; return (False, Var id) }
 986
 987 argsToPats :: InScopeEnv -> ConstrEnv
 988            -> [(CoreArg, ArgOcc)]
 989            -> UniqSM [(Bool, CoreArg)]
 990 argsToPats in_scope con_env args
 991   = mapUs do_one args
 992   where
 993     do_one (arg,occ) = argToPat in_scope con_env arg occ
 994 \end{code}
 995
 996
 997 \begin{code}
 998 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
 999 is_con_app_maybe env (Lit lit)
1000   = Just (CV (LitAlt lit) [])
1001
1002 is_con_app_maybe env expr       -- Maybe it's a constructor application
1003   | (Var fun, args) <- collectArgs expr,
1004     Just con <- isDataConWorkId_maybe fun,
1005     args `lengthAtLeast` dataConRepArity con
1006         -- Might be > because the arity excludes type args
1007   = Just (CV (DataAlt con) args)
1008
1009 is_con_app_maybe env (Var v)
1010   | Just stuff <- lookupVarEnv env v
1011   = Just stuff  -- You might think we could look in the idUnfolding here
1012                 -- but that doesn't take account of which branch of a
1013                 -- case we are in, which is the whole point
1014
1015   | isCheapUnfolding unf
1016   = is_con_app_maybe env (unfoldingTemplate unf)
1017   where
1018     unf = idUnfolding v
1019         -- However we do want to consult the unfolding
1020         -- as well, for let-bound constructors!
1021
1022 is_con_app_maybe env expr = Nothing
1023
1024 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
1025 mk_con_app (LitAlt lit)  []   = Lit lit
1026 mk_con_app (DataAlt con) args = mkConApp con args
1027 mk_con_app other args = panic "SpecConstr.mk_con_app"
1028 \end{code}