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
 489 extendBndrsWith :: HowBound -> ScEnv -> [Var] -> ScEnv
 490 extendBndrsWith how_bound env bndrs
 491   =  env { scope = scope env `extendVarEnvList`
 492                         [(bndr,how_bound) | bndr <- bndrs] }
 493
 494 extendBndrs env bndrs = extendBndrsWith Other env bndrs
 495 extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }
 496
 497     -- When we encounter
 498     --  case scrut of b
 499     --      C x y -> ...
 500     -- we want to bind b, and perhaps scrut too, to (C x y)
 501 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
 502 extendCaseBndrs env case_bndr scrut con alt_bndrs
 503   = case con of
 504         DEFAULT    -> env1
 505         LitAlt lit -> extendCons env1 scrut case_bndr (CV con [])
 506         DataAlt dc -> extendCons env1 scrut case_bndr (CV con vanilla_args)
 507               where
 508                 vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
 509                                varsToCoreExprs alt_bndrs
 510   where
 511     env1 = extendBndrsWith (get_how scrut) env (case_bndr:alt_bndrs)
 512
 513         -- Record RecArg for the components iff the scrutinee is RecArg
 514         -- I think the only reason for this is to keep the usage envt small
 515         -- so is it worth it at all?
 516         --      [This comment looks plain wrong to me, so I'm ignoring it
 517         --           "Also forget if the scrutinee is a RecArg, because we're
 518         --           now in the branch of a case, and we don't want to
 519         --           record a non-scrutinee use of v if we have
 520         --              case v of { (a,b) -> ...(f v)... }" ]
 521     get_how (Var v)    = lookupVarEnv (scope env) v `orElse` Other
 522     get_how (Cast e _) = get_how e
 523     get_how (Note _ e) = get_how e
 524     get_how other      = Other
 525
 526 extendCons :: ScEnv -> CoreExpr -> Id -> ConValue -> ScEnv
 527 extendCons env scrut case_bndr val
 528   = case scrut of
 529         Var v -> env { cons = extendVarEnv cons1 v val }
 530         other -> env { cons = cons1 }
 531   where
 532     cons1 = extendVarEnv (cons env) case_bndr val
 533 \end{code}
 534
 535
 536 %************************************************************************
 537 %*                                                                      *
 538 \subsection{Usage information: flows upwards}
 539 %*                                                                      *
 540 %************************************************************************
 541
 542 \begin{code}
 543 data ScUsage
 544    = SCU {
 545         calls :: !(IdEnv [Call]),       -- Calls
 546                                         -- The functions are a subset of the
 547                                         --      RecFuns in the ScEnv
 548
 549         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 550      }                                  -- The variables are a subset of the
 551                                         --      RecArg in the ScEnv
 552
 553 type Call = (ConstrEnv, [CoreArg])
 554         -- The arguments of the call, together with the
 555         -- env giving the constructor bindings at the call site
 556
 557 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 558
 559 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 560                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 561
 562 combineUsages [] = nullUsage
 563 combineUsages us = foldr1 combineUsage us
 564
 565 lookupOcc :: ScUsage -> Var -> (ScUsage, ArgOcc)
 566 lookupOcc (SCU { calls = sc_calls, occs = sc_occs }) bndr
 567   = (SCU {calls = sc_calls, occs = delVarEnv sc_occs bndr},
 568      lookupVarEnv sc_occs bndr `orElse` NoOcc)
 569
 570 lookupOccs :: ScUsage -> [Var] -> (ScUsage, [ArgOcc])
 571 lookupOccs (SCU { calls = sc_calls, occs = sc_occs }) bndrs
 572   = (SCU {calls = sc_calls, occs = delVarEnvList sc_occs bndrs},
 573      [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
 574
 575 data ArgOcc = NoOcc     -- Doesn't occur at all; or a type argument
 576             | UnkOcc    -- Used in some unknown way
 577
 578             | ScrutOcc (UniqFM [ArgOcc])        -- See Note [ScrutOcc]
 579
 580             | BothOcc   -- Definitely taken apart, *and* perhaps used in some other way
 581
 582 {-      Note  [ScrutOcc]
 583
 584 An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
 585 is *only* taken apart or applied.
 586
 587   Functions, literal: ScrutOcc emptyUFM
 588   Data constructors:  ScrutOcc subs,
 589
 590 where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
 591 The domain of the UniqFM is the Unique of the data constructor
 592
 593 The [ArgOcc] is the occurrences of the *pattern-bound* components
 594 of the data structure.  E.g.
 595         data T a = forall b. MkT a b (b->a)
 596 A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
 597
 598 -}
 599
 600 instance Outputable ArgOcc where
 601   ppr (ScrutOcc xs) = ptext SLIT("scrut-occ") <> ppr xs
 602   ppr UnkOcc        = ptext SLIT("unk-occ")
 603   ppr BothOcc       = ptext SLIT("both-occ")
 604   ppr NoOcc         = ptext SLIT("no-occ")
 605
 606 -- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
 607 -- that if the thing is scrutinised anywhere then we get to see that
 608 -- in the overall result, even if it's also used in a boxed way
 609 -- This might be too agressive; see Note [Reboxing] Alternative 3
 610 combineOcc NoOcc         occ           = occ
 611 combineOcc occ           NoOcc         = occ
 612 combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
 613 combineOcc occ           (ScrutOcc ys) = ScrutOcc ys
 614 combineOcc (ScrutOcc xs) occ           = ScrutOcc xs
 615 combineOcc UnkOcc        UnkOcc        = UnkOcc
 616 combineOcc _        _                  = BothOcc
 617
 618 combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
 619 combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
 620
 621 conArgOccs :: ArgOcc -> AltCon -> [ArgOcc]
 622 -- Find usage of components of data con; returns [UnkOcc...] if unknown
 623 -- See Note [ScrutOcc] for the extra UnkOccs in the vanilla datacon case
 624
 625 conArgOccs (ScrutOcc fm) (DataAlt dc)
 626   | Just pat_arg_occs <- lookupUFM fm dc
 627   = [UnkOcc | tv <- dataConUnivTyVars dc] ++ pat_arg_occs
 628
 629 conArgOccs other con = repeat UnkOcc
 630 \end{code}
 631
 632
 633 %************************************************************************
 634 %*                                                                      *
 635 \subsection{The main recursive function}
 636 %*                                                                      *
 637 %************************************************************************
 638
 639 The main recursive function gathers up usage information, and
 640 creates specialised versions of functions.
 641
 642 \begin{code}
 643 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 644         -- The unique supply is needed when we invent
 645         -- a new name for the specialised function and its args
 646
 647 scExpr env e@(Type t) = returnUs (nullUsage, e)
 648 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 649 scExpr env e@(Var v)  = returnUs (varUsage env v UnkOcc, e)
 650 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 651                         returnUs (usg, Note n e')
 652 scExpr env (Cast e co)= scExpr env e    `thenUs` \ (usg,e') ->
 653                         returnUs (usg, Cast e' co)
 654 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 655                         returnUs (usg, Lam b e')
 656
 657 scExpr env (Case scrut b ty alts)
 658   = do  { (alt_usgs, alt_occs, alts') <- mapAndUnzip3Us sc_alt alts
 659         ; let (alt_usg, b_occ) = lookupOcc (combineUsages alt_usgs) b
 660               scrut_occ = foldr combineOcc b_occ alt_occs
 661                 -- The combined usage of the scrutinee is given
 662                 -- by scrut_occ, which is passed to scScrut, which
 663                 -- in turn treats a bare-variable scrutinee specially
 664         ; (scrut_usg, scrut') <- scScrut env scrut scrut_occ
 665         ; return (alt_usg `combineUsage` scrut_usg,
 666                   Case scrut' b ty alts') }
 667   where
 668     sc_alt (con,bs,rhs)
 669       = do { let env1 = extendCaseBndrs env b scrut con bs
 670            ; (usg,rhs') <- scExpr env1 rhs
 671            ; let (usg', arg_occs) = lookupOccs usg bs
 672                  scrut_occ = case con of
 673                                 DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
 674                                 other      -> ScrutOcc emptyUFM
 675            ; return (usg', scrut_occ, (con,bs,rhs')) }
 676
 677 scExpr env (Let bind body)
 678   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 679     scExpr env' body    `thenUs` \ (body_usg, body') ->
 680     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 681
 682 scExpr env e@(App _ _)
 683   = do  { let (fn, args) = collectArgs e
 684         ; (fn_usg, fn') <- scScrut env fn (ScrutOcc emptyUFM)
 685         -- Process the function too.   It's almost always a variable,
 686         -- but not always.  In particular, if this pass follows float-in,
 687         -- which it may, we can get
 688         --      (let f = ...f... in f) arg1 arg2
 689         -- We use scScrut to record the fact that the function is called
 690         -- Perhpas we should check that it has at least one value arg,
 691         -- but currently we don't bother
 692
 693         ; (arg_usgs, args') <- mapAndUnzipUs (scExpr env) args
 694         ; let call_usg = case fn of
 695                            Var f | Just RecFun <- lookupScopeEnv env f
 696                                  , not (null args)      -- Not a proper call!
 697                                  -> SCU { calls = unitVarEnv f [(cons env, args)],
 698                                           occs  = emptyVarEnv }
 699                            other -> nullUsage
 700         ; return (combineUsages arg_usgs `combineUsage` fn_usg
 701                                          `combineUsage` call_usg,
 702                   mkApps fn' args') }
 703
 704
 705 ----------------------
 706 scScrut :: ScEnv -> CoreExpr -> ArgOcc -> UniqSM (ScUsage, CoreExpr)
 707 -- Used for the scrutinee of a case,
 708 -- or the function of an application.
 709 -- Remember to look through casts
 710 scScrut env e@(Var v)   occ = returnUs (varUsage env v occ, e)
 711 scScrut env (Cast e co) occ = do { (usg, e') <- scScrut env e occ
 712                                  ; returnUs (usg, Cast e' co) }
 713 scScrut env e           occ = scExpr env e
 714
 715
 716 ----------------------
 717 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 718 scBind env (Rec prs)
 719   = do  { let bndrs = map fst prs
 720               rhs_env = extendBndrsWith RecFun env bndrs
 721
 722         ; (rhs_usgs, prs_w_occs) <- mapAndUnzipUs (scRecRhs rhs_env) prs
 723         ; let rhs_usg   = combineUsages rhs_usgs
 724               rhs_calls = calls rhs_usg
 725
 726         ; prs_s <- mapUs (specialise env rhs_calls) prs_w_occs
 727         ; return (extendBndrs env bndrs,
 728                                 -- For the body of the letrec, just
 729                                 -- extend the env with Other to record
 730                                 -- that it's in scope; no funny RecFun business
 731                     rhs_usg { calls = calls rhs_usg `delVarEnvList` bndrs },
 732                     Rec (concat prs_s)) }
 733
 734 scBind env (NonRec bndr rhs)
 735   = do  { (usg, rhs') <- scExpr env rhs
 736         ; return (extendBndr env bndr, usg, NonRec bndr rhs') }
 737
 738 ----------------------
 739 scRecRhs :: ScEnv -> (Id,CoreExpr)
 740          -> UniqSM (ScUsage, (Id, CoreExpr, [ArgOcc]))
 741 -- The returned [ArgOcc] says how the visible,
 742 -- lambda-bound binders of the RHS are used
 743 -- (including the TyVar binders)
 744 scRecRhs env (bndr,rhs)
 745   = do  { let (arg_bndrs,body) = collectBinders rhs
 746               body_env = extendBndrsWith RecArg env arg_bndrs
 747         ; (body_usg, body') <- scExpr body_env body
 748         ; let (rhs_usg, arg_occs) = lookupOccs body_usg arg_bndrs
 749         ; return (rhs_usg, (bndr, mkLams arg_bndrs body', arg_occs)) }
 750
 751 ----------------------
 752 varUsage env v use
 753   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
 754                                                 occs = unitVarEnv v use }
 755   | otherwise                           = nullUsage
 756 \end{code}
 757
 758
 759 %************************************************************************
 760 %*                                                                      *
 761 \subsection{The specialiser}
 762 %*                                                                      *
 763 %************************************************************************
 764
 765 \begin{code}
 766 specialise
 767    :: ScEnv
 768    -> IdEnv [Call]              -- Info on usage
 769    -> (Id, CoreExpr, [ArgOcc])  -- Original binding, plus info on how the rhs's
 770                                 -- lambda-binders are used (includes TyVar bndrs)
 771    -> UniqSM [(Id,CoreExpr)]    -- Original binding (decorated with rules)
 772                                 -- plus specialised bindings
 773
 774 -- Note: the rhs here is the optimised version of the original rhs
 775 -- So when we make a specialised copy of the RHS, we're starting
 776 -- from an RHS whose nested functions have been optimised already.
 777
 778 specialise env calls (fn, rhs, arg_occs)
 779   | notNull arg_occs,   -- Only specialise functions
 780     Just all_calls <- lookupVarEnv calls fn
 781   = do  { mb_pats <- mapM (callToPats (scope env) arg_occs) all_calls
 782
 783         ; let good_pats :: [([Var], [CoreArg])]
 784               good_pats = catMaybes mb_pats
 785               in_scope = mkInScopeSet $ unionVarSets $
 786                          [ exprsFreeVars pats `delVarSetList` vs
 787                          | (vs,pats) <- good_pats ]
 788               uniq_pats = nubBy (same_pat in_scope) good_pats
 789 --      ; pprTrace "specialise" (vcat [ppr fn <+> ppr arg_occs,
 790 --                                      text "calls" <+> ppr all_calls,
 791 --                                      text "good pats" <+> ppr good_pats,
 792 --                              text "uniq pats" <+> ppr uniq_pats])  $
 793 --        return ()
 794
 795         ; (rules, spec_prs) <- mapAndUnzipUs (spec_one fn rhs)
 796                                              (uniq_pats `zip` [1..])
 797
 798         ; return ((fn `addIdSpecialisations` rules, rhs) : spec_prs) }
 799
 800   | otherwise
 801   = return [(fn,rhs)]   -- The boring case
 802   where
 803         -- Two pats are the same if they match both ways
 804     same_pat in_scope (vs1,as1)(vs2,as2)
 805          =  isJust (matchN in_scope vs1 as1 as2)
 806          && isJust (matchN in_scope vs2 as2 as1)
 807
 808 callToPats :: InScopeEnv -> [ArgOcc] -> Call
 809            -> UniqSM (Maybe ([Var], [CoreExpr]))
 810         -- The VarSet is the variables to quantify over in the rule
 811         -- The [CoreExpr] are the argument patterns for the rule
 812 callToPats in_scope bndr_occs (con_env, args)
 813   | length args < length bndr_occs      -- Check saturated
 814   = return Nothing
 815   | otherwise
 816   = do  { prs <- argsToPats in_scope con_env (args `zip` bndr_occs)
 817         ; let (good_pats, pats) = unzip prs
 818               pat_fvs = varSetElems (exprsFreeVars pats)
 819               qvars   = filter (not . (`elemVarEnv` in_scope)) pat_fvs
 820                 -- Quantify over variables that are not in sccpe
 821                 -- See Note [Shadowing] at the top
 822
 823         ; -- pprTrace "callToPats"  (ppr args $$ ppr prs $$ ppr bndr_occs) $
 824           if or good_pats
 825           then return (Just (qvars, pats))
 826           else return Nothing }
 827
 828 ---------------------
 829 spec_one :: Id                                  -- Function
 830          -> CoreExpr                            -- Rhs of the original function
 831          -> (([Var], [CoreArg]), Int)
 832          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 833
 834 -- spec_one creates a specialised copy of the function, together
 835 -- with a rule for using it.  I'm very proud of how short this
 836 -- function is, considering what it does :-).
 837
 838 {-
 839   Example
 840
 841      In-scope: a, x::a
 842      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
 843           [c::*, v::(b,c) are presumably bound by the (...) part]
 844   ==>
 845      f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
 846                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
 847
 848      RULE:  forall b::* c::*,           -- Note, *not* forall a, x
 849                    v::(b,c),
 850                    hw::[(a,(b,c))] .
 851
 852             f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
 853 -}
 854
 855 spec_one fn rhs ((vars_to_bind, pats), rule_number)
 856   = getUniqueUs                 `thenUs` \ spec_uniq ->
 857     let
 858         fn_name      = idName fn
 859         fn_loc       = nameSrcLoc fn_name
 860         spec_occ     = mkSpecOcc (nameOccName fn_name)
 861
 862                 -- Put the type variables first; the type of a term
 863                 -- variable may mention a type variable
 864         (tvs, ids)   = partition isTyVar vars_to_bind
 865         bndrs        = tvs ++ ids
 866         spec_body    = mkApps rhs pats
 867         body_ty      = exprType spec_body
 868
 869         (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
 870                 -- Usual w/w hack to avoid generating
 871                 -- a spec_rhs of unlifted type and no args
 872
 873         rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
 874         spec_rhs  = mkLams spec_lam_args spec_body
 875         spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
 876         rule_rhs  = mkVarApps (Var spec_id) spec_call_args
 877         rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
 878     in
 879     returnUs (rule, (spec_id, spec_rhs))
 880
 881 -- In which phase should the specialise-constructor rules be active?
 882 -- Originally I made them always-active, but Manuel found that
 883 -- this defeated some clever user-written rules.  So Plan B
 884 -- is to make them active only in Phase 0; after all, currently,
 885 -- the specConstr transformation is only run after the simplifier
 886 -- has reached Phase 0.  In general one would want it to be
 887 -- flag-controllable, but for now I'm leaving it baked in
 888 --                                      [SLPJ Oct 01]
 889 specConstrActivation :: Activation
 890 specConstrActivation = ActiveAfter 0    -- Baked in; see comments above
 891 \end{code}
 892
 893 %************************************************************************
 894 %*                                                                      *
 895 \subsection{Argument analysis}
 896 %*                                                                      *
 897 %************************************************************************
 898
 899 This code deals with analysing call-site arguments to see whether
 900 they are constructor applications.
 901
 902
 903 \begin{code}
 904     -- argToPat takes an actual argument, and returns an abstracted
 905     -- version, consisting of just the "constructor skeleton" of the
 906     -- argument, with non-constructor sub-expression replaced by new
 907     -- placeholder variables.  For example:
 908     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 909
 910 argToPat :: InScopeEnv                  -- What's in scope at the fn defn site
 911          -> ConstrEnv                   -- ConstrEnv at the call site
 912          -> CoreArg                     -- A call arg (or component thereof)
 913          -> ArgOcc
 914          -> UniqSM (Bool, CoreArg)
 915 -- Returns (interesting, pat),
 916 -- where pat is the pattern derived from the argument
 917 --            intersting=True if the pattern is non-trivial (not a variable or type)
 918 -- E.g.         x:xs         --> (True, x:xs)
 919 --              f xs         --> (False, w)        where w is a fresh wildcard
 920 --              (f xs, 'c')  --> (True, (w, 'c'))  where w is a fresh wildcard
 921 --              \x. x+y      --> (True, \x. x+y)
 922 --              lvl7         --> (True, lvl7)      if lvl7 is bound
 923 --                                                 somewhere further out
 924
 925 argToPat in_scope con_env arg@(Type ty) arg_occ
 926   = return (False, arg)
 927
 928 argToPat in_scope con_env (Let _ arg) arg_occ
 929   = argToPat in_scope con_env arg arg_occ
 930         -- Look through let expressions
 931         -- e.g.         f (let v = rhs in \y -> ...v...)
 932         -- Here we can specialise for f (\y -> ...)
 933         -- because the rule-matcher will look through the let.
 934
 935 argToPat in_scope con_env (Cast arg co) arg_occ
 936   = do  { (interesting, arg') <- argToPat in_scope con_env arg arg_occ
 937         ; if interesting then
 938                 return (interesting, Cast arg' co)
 939           else
 940                 wildCardPat (snd (coercionKind co)) }
 941
 942 argToPat in_scope con_env arg arg_occ
 943   | is_value_lam arg
 944   = return (True, arg)
 945   where
 946     is_value_lam (Lam v e)      -- Spot a value lambda, even if
 947         | isId v = True         -- it is inside a type lambda
 948         | otherwise = is_value_lam e
 949     is_value_lam other = False
 950
 951   -- Check for a constructor application
 952   -- NB: this *precedes* the Var case, so that we catch nullary constrs
 953 argToPat in_scope con_env arg arg_occ
 954   | Just (CV dc args) <- is_con_app_maybe con_env arg
 955   , case arg_occ of
 956         ScrutOcc _ -> True              -- Used only by case scrutinee
 957         BothOcc    -> case arg of       -- Used elsewhere
 958                         App {} -> True  --     see Note [Reboxing]
 959                         other  -> False
 960         other      -> False     -- No point; the arg is not decomposed
 961   = do  { args' <- argsToPats in_scope con_env (args `zip` conArgOccs arg_occ dc)
 962         ; return (True, mk_con_app dc (map snd args')) }
 963
 964   -- Check if the argument is a variable that
 965   -- is in scope at the function definition site
 966   -- It's worth specialising on this if
 967   --    (a) it's used in an interesting way in the body
 968   --    (b) we know what its value is
 969 argToPat in_scope con_env (Var v) arg_occ
 970   | not (isLocalId v) || v `elemVarEnv` in_scope,
 971     case arg_occ of { UnkOcc -> False; other -> True }, -- (a)
 972     isValueUnfolding (idUnfolding v)                    -- (b)
 973   = return (True, Var v)
 974
 975   -- Check for a variable bound inside the function.
 976   -- Don't make a wild-card, because we may usefully share
 977   --    e.g.  f a = let x = ... in f (x,x)
 978   -- NB: this case follows the lambda and con-app cases!!
 979 argToPat in_scope con_env (Var v) arg_occ
 980   = return (False, Var v)
 981
 982   -- The default case: make a wild-card
 983 argToPat in_scope con_env arg arg_occ
 984   = wildCardPat (exprType arg)
 985
 986 wildCardPat :: Type -> UniqSM (Bool, CoreArg)
 987 wildCardPat ty = do { uniq <- getUniqueUs
 988                     ; let id = mkSysLocal FSLIT("sc") uniq ty
 989                     ; return (False, Var id) }
 990
 991 argsToPats :: InScopeEnv -> ConstrEnv
 992            -> [(CoreArg, ArgOcc)]
 993            -> UniqSM [(Bool, CoreArg)]
 994 argsToPats in_scope con_env args
 995   = mapUs do_one args
 996   where
 997     do_one (arg,occ) = argToPat in_scope con_env arg occ
 998 \end{code}
 999
1000
1001 \begin{code}
1002 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
1003 is_con_app_maybe env (Lit lit)
1004   = Just (CV (LitAlt lit) [])
1005
1006 is_con_app_maybe env expr       -- Maybe it's a constructor application
1007   | (Var fun, args) <- collectArgs expr,
1008     Just con <- isDataConWorkId_maybe fun,
1009     args `lengthAtLeast` dataConRepArity con
1010         -- Might be > because the arity excludes type args
1011   = Just (CV (DataAlt con) args)
1012
1013 is_con_app_maybe env (Var v)
1014   | Just stuff <- lookupVarEnv env v
1015   = Just stuff  -- You might think we could look in the idUnfolding here
1016                 -- but that doesn't take account of which branch of a
1017                 -- case we are in, which is the whole point
1018
1019   | isCheapUnfolding unf
1020   = is_con_app_maybe env (unfoldingTemplate unf)
1021   where
1022     unf = idUnfolding v
1023         -- However we do want to consult the unfolding
1024         -- as well, for let-bound constructors!
1025
1026 is_con_app_maybe env expr = Nothing
1027
1028 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
1029 mk_con_app (LitAlt lit)  []   = Lit lit
1030 mk_con_app (DataAlt con) args = mkConApp con args
1031 mk_con_app other args = panic "SpecConstr.mk_con_app"
1032 \end{code}