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 -- The above warning supression flag is a temporary kludge.
   8 -- While working on this module you are encouraged to remove it and fix
   9 -- any warnings in the module. See
  10 --     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
  11 -- for details
  12
  13 module SpecConstr(
  14         specConstrProgram
  15     ) where
  16
  17 #include "HsVersions.h"
  18
  19 import CoreSyn
  20 import CoreSubst
  21 import CoreUtils
  22 import CoreUnfold       ( couldBeSmallEnoughToInline )
  23 import CoreFVs          ( exprsFreeVars )
  24 import WwLib            ( mkWorkerArgs )
  25 import DataCon          ( dataConRepArity, dataConUnivTyVars )
  26 import Coercion
  27 import Rules
  28 import Type             hiding( substTy )
  29 import Id
  30 import MkId             ( mkImpossibleExpr )
  31 import Var
  32 import VarEnv
  33 import VarSet
  34 import Name
  35 import OccName          ( mkSpecOcc )
  36 import DynFlags         ( DynFlags(..) )
  37 import StaticFlags      ( opt_PprStyle_Debug )
  38 import StaticFlags      ( opt_SpecInlineJoinPoints )
  39 import BasicTypes       ( Activation(..) )
  40 import Maybes           ( orElse, catMaybes, isJust, isNothing )
  41 import Util
  42 import List             ( nubBy, partition )
  43 import UniqSupply
  44 import Outputable
  45 import FastString
  46 import UniqFM
  47 import MonadUtils
  48 import Control.Monad    ( zipWithM )
  49 \end{code}
  50
  51 -----------------------------------------------------
  52                         Game plan
  53 -----------------------------------------------------
  54
  55 Consider
  56         drop n []     = []
  57         drop 0 xs     = []
  58         drop n (x:xs) = drop (n-1) xs
  59
  60 After the first time round, we could pass n unboxed.  This happens in
  61 numerical code too.  Here's what it looks like in Core:
  62
  63         drop n xs = case xs of
  64                       []     -> []
  65                       (y:ys) -> case n of
  66                                   I# n# -> case n# of
  67                                              0 -> []
  68                                              _ -> drop (I# (n# -# 1#)) xs
  69
  70 Notice that the recursive call has an explicit constructor as argument.
  71 Noticing this, we can make a specialised version of drop
  72
  73         RULE: drop (I# n#) xs ==> drop' n# xs
  74
  75         drop' n# xs = let n = I# n# in ...orig RHS...
  76
  77 Now the simplifier will apply the specialisation in the rhs of drop', giving
  78
  79         drop' n# xs = case xs of
  80                       []     -> []
  81                       (y:ys) -> case n# of
  82                                   0 -> []
  83                                   _ -> drop (n# -# 1#) xs
  84
  85 Much better!
  86
  87 We'd also like to catch cases where a parameter is carried along unchanged,
  88 but evaluated each time round the loop:
  89
  90         f i n = if i>0 || i>n then i else f (i*2) n
  91
  92 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  93 In Core, by the time we've w/wd (f is strict in i) we get
  94
  95         f i# n = case i# ># 0 of
  96                    False -> I# i#
  97                    True  -> case n of n' { I# n# ->
  98                             case i# ># n# of
  99                                 False -> I# i#
 100                                 True  -> f (i# *# 2#) n'
 101
 102 At the call to f, we see that the argument, n is know to be (I# n#),
 103 and n is evaluated elsewhere in the body of f, so we can play the same
 104 trick as above.
 105
 106
 107 Note [Reboxing]
 108 ~~~~~~~~~~~~~~~
 109 We must be careful not to allocate the same constructor twice.  Consider
 110         f p = (...(case p of (a,b) -> e)...p...,
 111                ...let t = (r,s) in ...t...(f t)...)
 112 At the recursive call to f, we can see that t is a pair.  But we do NOT want
 113 to make a specialised copy:
 114         f' a b = let p = (a,b) in (..., ...)
 115 because now t is allocated by the caller, then r and s are passed to the
 116 recursive call, which allocates the (r,s) pair again.
 117
 118 This happens if
 119   (a) the argument p is used in other than a case-scrutinsation way.
 120   (b) the argument to the call is not a 'fresh' tuple; you have to
 121         look into its unfolding to see that it's a tuple
 122
 123 Hence the "OR" part of Note [Good arguments] below.
 124
 125 ALTERNATIVE 2: pass both boxed and unboxed versions.  This no longer saves
 126 allocation, but does perhaps save evals. In the RULE we'd have
 127 something like
 128
 129   f (I# x#) = f' (I# x#) x#
 130
 131 If at the call site the (I# x) was an unfolding, then we'd have to
 132 rely on CSE to eliminate the duplicate allocation.... This alternative
 133 doesn't look attractive enough to pursue.
 134
 135 ALTERNATIVE 3: ignore the reboxing problem.  The trouble is that
 136 the conservative reboxing story prevents many useful functions from being
 137 specialised.  Example:
 138         foo :: Maybe Int -> Int -> Int
 139         foo   (Just m) 0 = 0
 140         foo x@(Just m) n = foo x (n-m)
 141 Here the use of 'x' will clearly not require boxing in the specialised function.
 142
 143 The strictness analyser has the same problem, in fact.  Example:
 144         f p@(a,b) = ...
 145 If we pass just 'a' and 'b' to the worker, it might need to rebox the
 146 pair to create (a,b).  A more sophisticated analysis might figure out
 147 precisely the cases in which this could happen, but the strictness
 148 analyser does no such analysis; it just passes 'a' and 'b', and hopes
 149 for the best.
 150
 151 So my current choice is to make SpecConstr similarly aggressive, and
 152 ignore the bad potential of reboxing.
 153
 154
 155 Note [Good arguments]
 156 ~~~~~~~~~~~~~~~~~~~~~
 157 So we look for
 158
 159 * A self-recursive function.  Ignore mutual recursion for now,
 160   because it's less common, and the code is simpler for self-recursion.
 161
 162 * EITHER
 163
 164    a) At a recursive call, one or more parameters is an explicit
 165       constructor application
 166         AND
 167       That same parameter is scrutinised by a case somewhere in
 168       the RHS of the function
 169
 170   OR
 171
 172     b) At a recursive call, one or more parameters has an unfolding
 173        that is an explicit constructor application
 174         AND
 175       That same parameter is scrutinised by a case somewhere in
 176       the RHS of the function
 177         AND
 178       Those are the only uses of the parameter (see Note [Reboxing])
 179
 180
 181 What to abstract over
 182 ~~~~~~~~~~~~~~~~~~~~~
 183 There's a bit of a complication with type arguments.  If the call
 184 site looks like
 185
 186         f p = ...f ((:) [a] x xs)...
 187
 188 then our specialised function look like
 189
 190         f_spec x xs = let p = (:) [a] x xs in ....as before....
 191
 192 This only makes sense if either
 193   a) the type variable 'a' is in scope at the top of f, or
 194   b) the type variable 'a' is an argument to f (and hence fs)
 195
 196 Actually, (a) may hold for value arguments too, in which case
 197 we may not want to pass them.  Supose 'x' is in scope at f's
 198 defn, but xs is not.  Then we'd like
 199
 200         f_spec xs = let p = (:) [a] x xs in ....as before....
 201
 202 Similarly (b) may hold too.  If x is already an argument at the
 203 call, no need to pass it again.
 204
 205 Finally, if 'a' is not in scope at the call site, we could abstract
 206 it as we do the term variables:
 207
 208         f_spec a x xs = let p = (:) [a] x xs in ...as before...
 209
 210 So the grand plan is:
 211
 212         * abstract the call site to a constructor-only pattern
 213           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
 214
 215         * Find the free variables of the abstracted pattern
 216
 217         * Pass these variables, less any that are in scope at
 218           the fn defn.  But see Note [Shadowing] below.
 219
 220
 221 NOTICE that we only abstract over variables that are not in scope,
 222 so we're in no danger of shadowing variables used in "higher up"
 223 in f_spec's RHS.
 224
 225
 226 Note [Shadowing]
 227 ~~~~~~~~~~~~~~~~
 228 In this pass we gather up usage information that may mention variables
 229 that are bound between the usage site and the definition site; or (more
 230 seriously) may be bound to something different at the definition site.
 231 For example:
 232
 233         f x = letrec g y v = let x = ...
 234                              in ...(g (a,b) x)...
 235
 236 Since 'x' is in scope at the call site, we may make a rewrite rule that
 237 looks like
 238         RULE forall a,b. g (a,b) x = ...
 239 But this rule will never match, because it's really a different 'x' at
 240 the call site -- and that difference will be manifest by the time the
 241 simplifier gets to it.  [A worry: the simplifier doesn't *guarantee*
 242 no-shadowing, so perhaps it may not be distinct?]
 243
 244 Anyway, the rule isn't actually wrong, it's just not useful.  One possibility
 245 is to run deShadowBinds before running SpecConstr, but instead we run the
 246 simplifier.  That gives the simplest possible program for SpecConstr to
 247 chew on; and it virtually guarantees no shadowing.
 248
 249 Note [Specialising for constant parameters]
 250 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 251 This one is about specialising on a *constant* (but not necessarily
 252 constructor) argument
 253
 254     foo :: Int -> (Int -> Int) -> Int
 255     foo 0 f = 0
 256     foo m f = foo (f m) (+1)
 257
 258 It produces
 259
 260     lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
 261     lvl_rmV =
 262       \ (ds_dlk :: GHC.Base.Int) ->
 263         case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
 264         GHC.Base.I# (GHC.Prim.+# x_alG 1)
 265
 266     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 267     GHC.Prim.Int#
 268     T.$wfoo =
 269       \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
 270         case ww_sme of ds_Xlw {
 271           __DEFAULT ->
 272         case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
 273         T.$wfoo ww1_Xmz lvl_rmV
 274         };
 275           0 -> 0
 276         }
 277
 278 The recursive call has lvl_rmV as its argument, so we could create a specialised copy
 279 with that argument baked in; that is, not passed at all.   Now it can perhaps be inlined.
 280
 281 When is this worth it?  Call the constant 'lvl'
 282 - If 'lvl' has an unfolding that is a constructor, see if the corresponding
 283   parameter is scrutinised anywhere in the body.
 284
 285 - If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
 286   parameter is applied (...to enough arguments...?)
 287
 288   Also do this is if the function has RULES?
 289
 290 Also
 291
 292 Note [Specialising for lambda parameters]
 293 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 294     foo :: Int -> (Int -> Int) -> Int
 295     foo 0 f = 0
 296     foo m f = foo (f m) (\n -> n-m)
 297
 298 This is subtly different from the previous one in that we get an
 299 explicit lambda as the argument:
 300
 301     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 302     GHC.Prim.Int#
 303     T.$wfoo =
 304       \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
 305         case ww_sm8 of ds_Xlr {
 306           __DEFAULT ->
 307         case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
 308         T.$wfoo
 309           ww1_Xmq
 310           (\ (n_ad3 :: GHC.Base.Int) ->
 311              case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
 312              GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
 313              })
 314         };
 315           0 -> 0
 316         }
 317
 318 I wonder if SpecConstr couldn't be extended to handle this? After all,
 319 lambda is a sort of constructor for functions and perhaps it already
 320 has most of the necessary machinery?
 321
 322 Furthermore, there's an immediate win, because you don't need to allocate the lamda
 323 at the call site; and if perchance it's called in the recursive call, then you
 324 may avoid allocating it altogether.  Just like for constructors.
 325
 326 Looks cool, but probably rare...but it might be easy to implement.
 327
 328
 329 Note [SpecConstr for casts]
 330 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 331 Consider
 332     data family T a :: *
 333     data instance T Int = T Int
 334
 335     foo n = ...
 336        where
 337          go (T 0) = 0
 338          go (T n) = go (T (n-1))
 339
 340 The recursive call ends up looking like
 341         go (T (I# ...) `cast` g)
 342 So we want to spot the construtor application inside the cast.
 343 That's why we have the Cast case in argToPat
 344
 345 Note [Local recursive groups]
 346 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 347 For a *local* recursive group, we can see all the calls to the
 348 function, so we seed the specialisation loop from the calls in the
 349 body, not from the calls in the RHS.  Consider:
 350
 351   bar m n = foo n (n,n) (n,n) (n,n) (n,n)
 352    where
 353      foo n p q r s
 354        | n == 0    = m
 355        | n > 3000  = case p of { (p1,p2) -> foo (n-1) (p2,p1) q r s }
 356        | n > 2000  = case q of { (q1,q2) -> foo (n-1) p (q2,q1) r s }
 357        | n > 1000  = case r of { (r1,r2) -> foo (n-1) p q (r2,r1) s }
 358        | otherwise = case s of { (s1,s2) -> foo (n-1) p q r (s2,s1) }
 359
 360 If we start with the RHSs of 'foo', we get lots and lots of specialisations,
 361 most of which are not needed.  But if we start with the (single) call
 362 in the rhs of 'bar' we get exactly one fully-specialised copy, and all
 363 the recursive calls go to this fully-specialised copy. Indeed, the original
 364 function is later collected as dead code.  This is very important in
 365 specialising the loops arising from stream fusion, for example in NDP where
 366 we were getting literally hundreds of (mostly unused) specialisations of
 367 a local function.
 368
 369 Note [Do not specialise diverging functions]
 370 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 371 Specialising a function that just diverges is a waste of code.
 372 Furthermore, it broke GHC (simpl014) thus:
 373    {-# STR Sb #-}
 374    f = \x. case x of (a,b) -> f x
 375 If we specialise f we get
 376    f = \x. case x of (a,b) -> fspec a b
 377 But fspec doesn't have decent strictnes info.  As it happened,
 378 (f x) :: IO t, so the state hack applied and we eta expanded fspec,
 379 and hence f.  But now f's strictness is less than its arity, which
 380 breaks an invariant.
 381
 382 -----------------------------------------------------
 383                 Stuff not yet handled
 384 -----------------------------------------------------
 385
 386 Here are notes arising from Roman's work that I don't want to lose.
 387
 388 Example 1
 389 ~~~~~~~~~
 390     data T a = T !a
 391
 392     foo :: Int -> T Int -> Int
 393     foo 0 t = 0
 394     foo x t | even x    = case t of { T n -> foo (x-n) t }
 395             | otherwise = foo (x-1) t
 396
 397 SpecConstr does no specialisation, because the second recursive call
 398 looks like a boxed use of the argument.  A pity.
 399
 400     $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 401     $wfoo_sFw =
 402       \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
 403          case ww_sFo of ds_Xw6 [Just L] {
 404            __DEFAULT ->
 405                 case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
 406                   __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
 407                   0 ->
 408                     case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
 409                     case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
 410                     $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
 411                     } } };
 412            0 -> 0
 413
 414 Example 2
 415 ~~~~~~~~~
 416     data a :*: b = !a :*: !b
 417     data T a = T !a
 418
 419     foo :: (Int :*: T Int) -> Int
 420     foo (0 :*: t) = 0
 421     foo (x :*: t) | even x    = case t of { T n -> foo ((x-n) :*: t) }
 422                   | otherwise = foo ((x-1) :*: t)
 423
 424 Very similar to the previous one, except that the parameters are now in
 425 a strict tuple. Before SpecConstr, we have
 426
 427     $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 428     $wfoo_sG3 =
 429       \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
 430     GHC.Base.Int) ->
 431         case ww_sFU of ds_Xws [Just L] {
 432           __DEFAULT ->
 433         case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
 434           __DEFAULT ->
 435             case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
 436             $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2             -- $wfoo1
 437             };
 438           0 ->
 439             case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
 440             case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
 441             $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB        -- $wfoo2
 442             } } };
 443           0 -> 0 }
 444
 445 We get two specialisations:
 446 "SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
 447                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
 448                   = Foo.$s$wfoo1 a_sFB sc_sGC ;
 449 "SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
 450                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
 451                   = Foo.$s$wfoo y_aFp sc_sGC ;
 452
 453 But perhaps the first one isn't good.  After all, we know that tpl_B2 is
 454 a T (I# x) really, because T is strict and Int has one constructor.  (We can't
 455 unbox the strict fields, becuase T is polymorphic!)
 456
 457
 458
 459 %************************************************************************
 460 %*                                                                      *
 461 \subsection{Top level wrapper stuff}
 462 %*                                                                      *
 463 %************************************************************************
 464
 465 \begin{code}
 466 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> [CoreBind]
 467 specConstrProgram dflags us binds = fst $ initUs us (go (initScEnv dflags) binds)
 468   where
 469     go _   []           = return []
 470     go env (bind:binds) = do (env', bind') <- scTopBind env bind
 471                              binds' <- go env' binds
 472                              return (bind' : binds')
 473 \end{code}
 474
 475
 476 %************************************************************************
 477 %*                                                                      *
 478 \subsection{Environment: goes downwards}
 479 %*                                                                      *
 480 %************************************************************************
 481
 482 \begin{code}
 483 data ScEnv = SCE { sc_size  :: Maybe Int,       -- Size threshold
 484                    sc_count :: Maybe Int,       -- Max # of specialisations for any one fn
 485
 486                    sc_subst :: Subst,           -- Current substitution
 487                                                 -- Maps InIds to OutExprs
 488
 489                    sc_how_bound :: HowBoundEnv,
 490                         -- Binds interesting non-top-level variables
 491                         -- Domain is OutVars (*after* applying the substitution)
 492
 493                    sc_vals  :: ValueEnv
 494                         -- Domain is OutIds (*after* applying the substitution)
 495                         -- Used even for top-level bindings (but not imported ones)
 496              }
 497
 498 ---------------------
 499 -- As we go, we apply a substitution (sc_subst) to the current term
 500 type InExpr = CoreExpr          -- _Before_ applying the subst
 501
 502 type OutExpr = CoreExpr         -- _After_ applying the subst
 503 type OutId   = Id
 504 type OutVar  = Var
 505
 506 ---------------------
 507 type HowBoundEnv = VarEnv HowBound      -- Domain is OutVars
 508
 509 ---------------------
 510 type ValueEnv = IdEnv Value             -- Domain is OutIds
 511 data Value    = ConVal AltCon [CoreArg] -- _Saturated_ constructors
 512               | LambdaVal               -- Inlinable lambdas or PAPs
 513
 514 instance Outputable Value where
 515    ppr (ConVal con args) = ppr con <+> interpp'SP args
 516    ppr LambdaVal         = ptext (sLit "<Lambda>")
 517
 518 ---------------------
 519 initScEnv :: DynFlags -> ScEnv
 520 initScEnv dflags
 521   = SCE { sc_size = specConstrThreshold dflags,
 522           sc_count = specConstrCount dflags,
 523           sc_subst = emptySubst,
 524           sc_how_bound = emptyVarEnv,
 525           sc_vals = emptyVarEnv }
 526
 527 data HowBound = RecFun  -- These are the recursive functions for which
 528                         -- we seek interesting call patterns
 529
 530               | RecArg  -- These are those functions' arguments, or their sub-components;
 531                         -- we gather occurrence information for these
 532
 533 instance Outputable HowBound where
 534   ppr RecFun = text "RecFun"
 535   ppr RecArg = text "RecArg"
 536
 537 lookupHowBound :: ScEnv -> Id -> Maybe HowBound
 538 lookupHowBound env id = lookupVarEnv (sc_how_bound env) id
 539
 540 scSubstId :: ScEnv -> Id -> CoreExpr
 541 scSubstId env v = lookupIdSubst (sc_subst env) v
 542
 543 scSubstTy :: ScEnv -> Type -> Type
 544 scSubstTy env ty = substTy (sc_subst env) ty
 545
 546 zapScSubst :: ScEnv -> ScEnv
 547 zapScSubst env = env { sc_subst = zapSubstEnv (sc_subst env) }
 548
 549 extendScInScope :: ScEnv -> [Var] -> ScEnv
 550         -- Bring the quantified variables into scope
 551 extendScInScope env qvars = env { sc_subst = extendInScopeList (sc_subst env) qvars }
 552
 553         -- Extend the substitution
 554 extendScSubst :: ScEnv -> Var -> OutExpr -> ScEnv
 555 extendScSubst env var expr = env { sc_subst = extendSubst (sc_subst env) var expr }
 556
 557 extendScSubstList :: ScEnv -> [(Var,OutExpr)] -> ScEnv
 558 extendScSubstList env prs = env { sc_subst = extendSubstList (sc_subst env) prs }
 559
 560 extendHowBound :: ScEnv -> [Var] -> HowBound -> ScEnv
 561 extendHowBound env bndrs how_bound
 562   = env { sc_how_bound = extendVarEnvList (sc_how_bound env)
 563                             [(bndr,how_bound) | bndr <- bndrs] }
 564
 565 extendBndrsWith :: HowBound -> ScEnv -> [Var] -> (ScEnv, [Var])
 566 extendBndrsWith how_bound env bndrs
 567   = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndrs')
 568   where
 569     (subst', bndrs') = substBndrs (sc_subst env) bndrs
 570     hb_env' = sc_how_bound env `extendVarEnvList`
 571                     [(bndr,how_bound) | bndr <- bndrs']
 572
 573 extendBndrWith :: HowBound -> ScEnv -> Var -> (ScEnv, Var)
 574 extendBndrWith how_bound env bndr
 575   = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndr')
 576   where
 577     (subst', bndr') = substBndr (sc_subst env) bndr
 578     hb_env' = extendVarEnv (sc_how_bound env) bndr' how_bound
 579
 580 extendRecBndrs :: ScEnv -> [Var] -> (ScEnv, [Var])
 581 extendRecBndrs env bndrs  = (env { sc_subst = subst' }, bndrs')
 582                       where
 583                         (subst', bndrs') = substRecBndrs (sc_subst env) bndrs
 584
 585 extendBndr :: ScEnv -> Var -> (ScEnv, Var)
 586 extendBndr  env bndr  = (env { sc_subst = subst' }, bndr')
 587                       where
 588                         (subst', bndr') = substBndr (sc_subst env) bndr
 589
 590 extendValEnv :: ScEnv -> Id -> Maybe Value -> ScEnv
 591 extendValEnv env _  Nothing   = env
 592 extendValEnv env id (Just cv) = env { sc_vals = extendVarEnv (sc_vals env) id cv }
 593
 594 extendCaseBndrs :: ScEnv -> Id -> AltCon -> [Var] -> (ScEnv, [Var])
 595 -- When we encounter
 596 --      case scrut of b
 597 --          C x y -> ...
 598 -- we want to bind b, to (C x y)
 599 -- NB1: Extends only the sc_vals part of the envt
 600 -- NB2: Kill the dead-ness info on the pattern binders x,y, since
 601 --      they are potentially made alive by the [b -> C x y] binding
 602 extendCaseBndrs env case_bndr con alt_bndrs
 603   | isDeadBinder case_bndr
 604   = (env, alt_bndrs)
 605   | otherwise
 606   = (env1, map zap alt_bndrs)
 607         -- NB: We used to bind v too, if scrut = (Var v); but
 608         --     the simplifer has already done this so it seems
 609         --     redundant to do so here
 610         -- case scrut of
 611         --      Var v  -> extendValEnv env1 v cval
 612         --      _other -> env1
 613  where
 614    zap v | isTyVar v = v                -- See NB2 above
 615          | otherwise = zapIdOccInfo v
 616    env1 = extendValEnv env case_bndr cval
 617    cval = case con of
 618                 DEFAULT    -> Nothing
 619                 LitAlt {}  -> Just (ConVal con [])
 620                 DataAlt {} -> Just (ConVal con vanilla_args)
 621                       where
 622                         vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
 623                                        varsToCoreExprs alt_bndrs
 624 \end{code}
 625
 626
 627 %************************************************************************
 628 %*                                                                      *
 629 \subsection{Usage information: flows upwards}
 630 %*                                                                      *
 631 %************************************************************************
 632
 633 \begin{code}
 634 data ScUsage
 635    = SCU {
 636         scu_calls :: CallEnv,           -- Calls
 637                                         -- The functions are a subset of the
 638                                         --      RecFuns in the ScEnv
 639
 640         scu_occs :: !(IdEnv ArgOcc)     -- Information on argument occurrences
 641      }                                  -- The domain is OutIds
 642
 643 type CallEnv = IdEnv [Call]
 644 type Call = (ValueEnv, [CoreArg])
 645         -- The arguments of the call, together with the
 646         -- env giving the constructor bindings at the call site
 647
 648 nullUsage :: ScUsage
 649 nullUsage = SCU { scu_calls = emptyVarEnv, scu_occs = emptyVarEnv }
 650
 651 combineCalls :: CallEnv -> CallEnv -> CallEnv
 652 combineCalls = plusVarEnv_C (++)
 653
 654 combineUsage :: ScUsage -> ScUsage -> ScUsage
 655 combineUsage u1 u2 = SCU { scu_calls = combineCalls (scu_calls u1) (scu_calls u2),
 656                            scu_occs  = plusVarEnv_C combineOcc (scu_occs u1) (scu_occs u2) }
 657
 658 combineUsages :: [ScUsage] -> ScUsage
 659 combineUsages [] = nullUsage
 660 combineUsages us = foldr1 combineUsage us
 661
 662 lookupOcc :: ScUsage -> OutVar -> (ScUsage, ArgOcc)
 663 lookupOcc (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndr
 664   = (SCU {scu_calls = sc_calls, scu_occs = delVarEnv sc_occs bndr},
 665      lookupVarEnv sc_occs bndr `orElse` NoOcc)
 666
 667 lookupOccs :: ScUsage -> [OutVar] -> (ScUsage, [ArgOcc])
 668 lookupOccs (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndrs
 669   = (SCU {scu_calls = sc_calls, scu_occs = delVarEnvList sc_occs bndrs},
 670      [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
 671
 672 data ArgOcc = NoOcc     -- Doesn't occur at all; or a type argument
 673             | UnkOcc    -- Used in some unknown way
 674
 675             | ScrutOcc (UniqFM [ArgOcc])        -- See Note [ScrutOcc]
 676
 677             | BothOcc   -- Definitely taken apart, *and* perhaps used in some other way
 678
 679 {-      Note  [ScrutOcc]
 680
 681 An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
 682 is *only* taken apart or applied.
 683
 684   Functions, literal: ScrutOcc emptyUFM
 685   Data constructors:  ScrutOcc subs,
 686
 687 where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
 688 The domain of the UniqFM is the Unique of the data constructor
 689
 690 The [ArgOcc] is the occurrences of the *pattern-bound* components
 691 of the data structure.  E.g.
 692         data T a = forall b. MkT a b (b->a)
 693 A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
 694
 695 -}
 696
 697 instance Outputable ArgOcc where
 698   ppr (ScrutOcc xs) = ptext (sLit "scrut-occ") <> ppr xs
 699   ppr UnkOcc        = ptext (sLit "unk-occ")
 700   ppr BothOcc       = ptext (sLit "both-occ")
 701   ppr NoOcc         = ptext (sLit "no-occ")
 702
 703 -- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
 704 -- that if the thing is scrutinised anywhere then we get to see that
 705 -- in the overall result, even if it's also used in a boxed way
 706 -- This might be too agressive; see Note [Reboxing] Alternative 3
 707 combineOcc :: ArgOcc -> ArgOcc -> ArgOcc
 708 combineOcc NoOcc         occ           = occ
 709 combineOcc occ           NoOcc         = occ
 710 combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
 711 combineOcc _occ          (ScrutOcc ys) = ScrutOcc ys
 712 combineOcc (ScrutOcc xs) _occ          = ScrutOcc xs
 713 combineOcc UnkOcc        UnkOcc        = UnkOcc
 714 combineOcc _        _                  = BothOcc
 715
 716 combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
 717 combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
 718
 719 setScrutOcc :: ScEnv -> ScUsage -> OutExpr -> ArgOcc -> ScUsage
 720 -- _Overwrite_ the occurrence info for the scrutinee, if the scrutinee
 721 -- is a variable, and an interesting variable
 722 setScrutOcc env usg (Cast e _) occ = setScrutOcc env usg e occ
 723 setScrutOcc env usg (Note _ e) occ = setScrutOcc env usg e occ
 724 setScrutOcc env usg (Var v)    occ
 725   | Just RecArg <- lookupHowBound env v = usg { scu_occs = extendVarEnv (scu_occs usg) v occ }
 726   | otherwise                           = usg
 727 setScrutOcc _env usg _other _occ        -- Catch-all
 728   = usg
 729
 730 conArgOccs :: ArgOcc -> AltCon -> [ArgOcc]
 731 -- Find usage of components of data con; returns [UnkOcc...] if unknown
 732 -- See Note [ScrutOcc] for the extra UnkOccs in the vanilla datacon case
 733
 734 conArgOccs (ScrutOcc fm) (DataAlt dc)
 735   | Just pat_arg_occs <- lookupUFM fm dc
 736   = [UnkOcc | _ <- dataConUnivTyVars dc] ++ pat_arg_occs
 737
 738 conArgOccs _other _con = repeat UnkOcc
 739 \end{code}
 740
 741 %************************************************************************
 742 %*                                                                      *
 743 \subsection{The main recursive function}
 744 %*                                                                      *
 745 %************************************************************************
 746
 747 The main recursive function gathers up usage information, and
 748 creates specialised versions of functions.
 749
 750 \begin{code}
 751 scExpr, scExpr' :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 752         -- The unique supply is needed when we invent
 753         -- a new name for the specialised function and its args
 754
 755 scExpr env e = scExpr' env e
 756
 757
 758 scExpr' env (Var v)     = case scSubstId env v of
 759                             Var v' -> return (varUsage env v' UnkOcc, Var v')
 760                             e'     -> scExpr (zapScSubst env) e'
 761
 762 scExpr' env (Type t)    = return (nullUsage, Type (scSubstTy env t))
 763 scExpr' _   e@(Lit {})  = return (nullUsage, e)
 764 scExpr' env (Note n e)  = do (usg,e') <- scExpr env e
 765                              return (usg, Note n e')
 766 scExpr' env (Cast e co) = do (usg, e') <- scExpr env e
 767                              return (usg, Cast e' (scSubstTy env co))
 768 scExpr' env e@(App _ _) = scApp env (collectArgs e)
 769 scExpr' env (Lam b e)   = do let (env', b') = extendBndr env b
 770                              (usg, e') <- scExpr env' e
 771                              return (usg, Lam b' e')
 772
 773 scExpr' env (Case scrut b ty alts)
 774   = do  { (scrut_usg, scrut') <- scExpr env scrut
 775         ; case isValue (sc_vals env) scrut' of
 776                 Just (ConVal con args) -> sc_con_app con args scrut'
 777                 _other                 -> sc_vanilla scrut_usg scrut'
 778         }
 779   where
 780     sc_con_app con args scrut'  -- Known constructor; simplify
 781         = do { let (_, bs, rhs) = findAlt con alts
 782                                   `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
 783                    alt_env'  = extendScSubstList env ((b,scrut') : bs `zip` trimConArgs con args)
 784              ; scExpr alt_env' rhs }
 785
 786     sc_vanilla scrut_usg scrut' -- Normal case
 787      = do { let (alt_env,b') = extendBndrWith RecArg env b
 788                         -- Record RecArg for the components
 789
 790           ; (alt_usgs, alt_occs, alts')
 791                 <- mapAndUnzip3M (sc_alt alt_env scrut' b') alts
 792
 793           ; let (alt_usg, b_occ) = lookupOcc (combineUsages alt_usgs) b'
 794                 scrut_occ        = foldr combineOcc b_occ alt_occs
 795                 scrut_usg'       = setScrutOcc env scrut_usg scrut' scrut_occ
 796                 -- The combined usage of the scrutinee is given
 797                 -- by scrut_occ, which is passed to scScrut, which
 798                 -- in turn treats a bare-variable scrutinee specially
 799
 800           ; return (alt_usg `combineUsage` scrut_usg',
 801                     Case scrut' b' (scSubstTy env ty) alts') }
 802
 803     sc_alt env _scrut' b' (con,bs,rhs)
 804       = do { let (env1, bs1)  = extendBndrsWith RecArg env bs
 805                  (env2, bs2) = extendCaseBndrs env1 b' con bs1
 806            ; (usg,rhs') <- scExpr env2 rhs
 807            ; let (usg', arg_occs) = lookupOccs usg bs2
 808                  scrut_occ = case con of
 809                                 DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
 810                                 _          -> ScrutOcc emptyUFM
 811            ; return (usg', scrut_occ, (con, bs2, rhs')) }
 812
 813 scExpr' env (Let (NonRec bndr rhs) body)
 814   | isTyVar bndr        -- Type-lets may be created by doBeta
 815   = scExpr' (extendScSubst env bndr rhs) body
 816   | otherwise
 817   = do  { let (body_env, bndr') = extendBndr env bndr
 818         ; (rhs_usg, (_, args', rhs_body', _)) <- scRecRhs env (bndr',rhs)
 819         ; let rhs' = mkLams args' rhs_body'
 820
 821         ; if not opt_SpecInlineJoinPoints || null args' || isEmptyVarEnv (scu_calls rhs_usg) then do
 822             do  {       -- Vanilla case
 823                   let body_env2 = extendValEnv body_env bndr' (isValue (sc_vals env) rhs')
 824                         -- Record if the RHS is a value
 825                 ; (body_usg, body') <- scExpr body_env2 body
 826                 ; return (body_usg `combineUsage` rhs_usg, Let (NonRec bndr' rhs') body') }
 827           else  -- For now, just brutally inline the join point
 828             do { let body_env2 = extendScSubst env bndr rhs'
 829                ; scExpr body_env2 body } }
 830
 831
 832 {-  Old code
 833             do  {       -- Join-point case
 834                   let body_env2 = extendHowBound body_env [bndr'] RecFun
 835                         -- If the RHS of this 'let' contains calls
 836                         -- to recursive functions that we're trying
 837                         -- to specialise, then treat this let too
 838                         -- as one to specialise
 839                 ; (body_usg, body') <- scExpr body_env2 body
 840
 841                 ; (spec_usg, _, specs) <- specialise env (scu_calls body_usg) ([], rhs_info)
 842
 843                 ; return (body_usg { scu_calls = scu_calls body_usg `delVarEnv` bndr' }
 844                           `combineUsage` rhs_usg `combineUsage` spec_usg,
 845                           mkLets [NonRec b r | (b,r) <- specInfoBinds rhs_info specs] body')
 846         }
 847 -}
 848
 849 -- A *local* recursive group: see Note [Local recursive groups]
 850 scExpr' env (Let (Rec prs) body)
 851   = do  { let (bndrs,rhss) = unzip prs
 852               (rhs_env1,bndrs') = extendRecBndrs env bndrs
 853               rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
 854
 855         ; (rhs_usgs, rhs_infos) <- mapAndUnzipM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
 856         ; (body_usg, body')     <- scExpr rhs_env2 body
 857
 858         -- NB: start specLoop from body_usg
 859         ; (spec_usg, specs) <- specLoop rhs_env2 (scu_calls body_usg) rhs_infos nullUsage
 860                                         [SI [] 0 (Just usg) | usg <- rhs_usgs]
 861
 862         ; let all_usg = spec_usg `combineUsage` body_usg
 863               bind'   = Rec (concat (zipWith specInfoBinds rhs_infos specs))
 864
 865         ; return (all_usg { scu_calls = scu_calls all_usg `delVarEnvList` bndrs' },
 866                   Let bind' body') }
 867
 868 -----------------------------------
 869 scApp :: ScEnv -> (InExpr, [InExpr]) -> UniqSM (ScUsage, CoreExpr)
 870
 871 scApp env (Var fn, args)        -- Function is a variable
 872   = ASSERT( not (null args) )
 873     do  { args_w_usgs <- mapM (scExpr env) args
 874         ; let (arg_usgs, args') = unzip args_w_usgs
 875               arg_usg = combineUsages arg_usgs
 876         ; case scSubstId env fn of
 877             fn'@(Lam {}) -> scExpr (zapScSubst env) (doBeta fn' args')
 878                         -- Do beta-reduction and try again
 879
 880             Var fn' -> return (arg_usg `combineUsage` fn_usg, mkApps (Var fn') args')
 881                 where
 882                   fn_usg = case lookupHowBound env fn' of
 883                                 Just RecFun -> SCU { scu_calls = unitVarEnv fn' [(sc_vals env, args')],
 884                                                      scu_occs  = emptyVarEnv }
 885                                 Just RecArg -> SCU { scu_calls = emptyVarEnv,
 886                                                      scu_occs  = unitVarEnv fn' (ScrutOcc emptyUFM) }
 887                                 Nothing     -> nullUsage
 888
 889
 890             other_fn' -> return (arg_usg, mkApps other_fn' args') }
 891                 -- NB: doing this ignores any usage info from the substituted
 892                 --     function, but I don't think that matters.  If it does
 893                 --     we can fix it.
 894   where
 895     doBeta :: OutExpr -> [OutExpr] -> OutExpr
 896     -- ToDo: adjust for System IF
 897     doBeta (Lam bndr body) (arg : args) = Let (NonRec bndr arg) (doBeta body args)
 898     doBeta fn              args         = mkApps fn args
 899
 900 -- The function is almost always a variable, but not always.
 901 -- In particular, if this pass follows float-in,
 902 -- which it may, we can get
 903 --      (let f = ...f... in f) arg1 arg2
 904 scApp env (other_fn, args)
 905   = do  { (fn_usg,   fn')   <- scExpr env other_fn
 906         ; (arg_usgs, args') <- mapAndUnzipM (scExpr env) args
 907         ; return (combineUsages arg_usgs `combineUsage` fn_usg, mkApps fn' args') }
 908
 909 ----------------------
 910 scTopBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, CoreBind)
 911 scTopBind env (Rec prs)
 912   | Just threshold <- sc_size env
 913   , not (all (couldBeSmallEnoughToInline threshold) rhss)
 914                 -- No specialisation
 915   = do  { let (rhs_env,bndrs') = extendRecBndrs env bndrs
 916         ; (_, rhss') <- mapAndUnzipM (scExpr rhs_env) rhss
 917         ; return (rhs_env, Rec (bndrs' `zip` rhss')) }
 918   | otherwise   -- Do specialisation
 919   = do  { let (rhs_env1,bndrs') = extendRecBndrs env bndrs
 920               rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
 921
 922         ; (rhs_usgs, rhs_infos) <- mapAndUnzipM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
 923         ; let rhs_usg = combineUsages rhs_usgs
 924
 925         ; (_, specs) <- specLoop rhs_env2 (scu_calls rhs_usg) rhs_infos nullUsage
 926                                  [SI [] 0 Nothing | _ <- bndrs]
 927
 928         ; return (rhs_env1,  -- For the body of the letrec, delete the RecFun business
 929                   Rec (concat (zipWith specInfoBinds rhs_infos specs))) }
 930   where
 931     (bndrs,rhss) = unzip prs
 932
 933 scTopBind env (NonRec bndr rhs)
 934   = do  { (_, rhs') <- scExpr env rhs
 935         ; let (env1, bndr') = extendBndr env bndr
 936               env2 = extendValEnv env1 bndr' (isValue (sc_vals env) rhs')
 937         ; return (env2, NonRec bndr' rhs') }
 938
 939 ----------------------
 940 scRecRhs :: ScEnv -> (OutId, InExpr) -> UniqSM (ScUsage, RhsInfo)
 941 scRecRhs env (bndr,rhs)
 942   = do  { let (arg_bndrs,body) = collectBinders rhs
 943               (body_env, arg_bndrs') = extendBndrsWith RecArg env arg_bndrs
 944         ; (body_usg, body') <- scExpr body_env body
 945         ; let (rhs_usg, arg_occs) = lookupOccs body_usg arg_bndrs'
 946         ; return (rhs_usg, (bndr, arg_bndrs', body', arg_occs)) }
 947
 948                 -- The arg_occs says how the visible,
 949                 -- lambda-bound binders of the RHS are used
 950                 -- (including the TyVar binders)
 951                 -- Two pats are the same if they match both ways
 952
 953 ----------------------
 954 specInfoBinds :: RhsInfo -> SpecInfo -> [(Id,CoreExpr)]
 955 specInfoBinds (fn, args, body, _) (SI specs _ _)
 956   = [(id,rhs) | OS _ _ id rhs <- specs] ++
 957     [(fn `addIdSpecialisations` rules, mkLams args body)]
 958   where
 959     rules = [r | OS _ r _ _ <- specs]
 960
 961 ----------------------
 962 varUsage :: ScEnv -> OutVar -> ArgOcc -> ScUsage
 963 varUsage env v use
 964   | Just RecArg <- lookupHowBound env v = SCU { scu_calls = emptyVarEnv
 965                                               , scu_occs = unitVarEnv v use }
 966   | otherwise                           = nullUsage
 967 \end{code}
 968
 969
 970 %************************************************************************
 971 %*                                                                      *
 972                 The specialiser itself
 973 %*                                                                      *
 974 %************************************************************************
 975
 976 \begin{code}
 977 type RhsInfo = (OutId, [OutVar], OutExpr, [ArgOcc])
 978         -- Info about the *original* RHS of a binding we are specialising
 979         -- Original binding f = \xs.body
 980         -- Plus info about usage of arguments
 981
 982 data SpecInfo = SI [OneSpec]            -- The specialisations we have generated
 983                    Int                  -- Length of specs; used for numbering them
 984                    (Maybe ScUsage)      -- Nothing => we have generated specialisations
 985                                         --            from calls in the *original* RHS
 986                                         -- Just cs => we haven't, and this is the usage
 987                                         --            of the original RHS
 988
 989         -- One specialisation: Rule plus definition
 990 data OneSpec  = OS CallPat              -- Call pattern that generated this specialisation
 991                    CoreRule             -- Rule connecting original id with the specialisation
 992                    OutId OutExpr        -- Spec id + its rhs
 993
 994
 995 specLoop :: ScEnv
 996          -> CallEnv
 997          -> [RhsInfo]
 998          -> ScUsage -> [SpecInfo]               -- One per binder; acccumulating parameter
 999          -> UniqSM (ScUsage, [SpecInfo])        -- ...ditto...
1000 specLoop env all_calls rhs_infos usg_so_far specs_so_far
1001   = do  { specs_w_usg <- zipWithM (specialise env all_calls) rhs_infos specs_so_far
1002         ; let (new_usg_s, all_specs) = unzip specs_w_usg
1003               new_usg   = combineUsages new_usg_s
1004               new_calls = scu_calls new_usg
1005               all_usg   = usg_so_far `combineUsage` new_usg
1006         ; if isEmptyVarEnv new_calls then
1007                 return (all_usg, all_specs)
1008           else
1009                 specLoop env new_calls rhs_infos all_usg all_specs }
1010
1011 specialise
1012    :: ScEnv
1013    -> CallEnv                           -- Info on calls
1014    -> RhsInfo
1015    -> SpecInfo                          -- Original RHS plus patterns dealt with
1016    -> UniqSM (ScUsage, SpecInfo)        -- New specialised versions and their usage
1017
1018 -- Note: the rhs here is the optimised version of the original rhs
1019 -- So when we make a specialised copy of the RHS, we're starting
1020 -- from an RHS whose nested functions have been optimised already.
1021
1022 specialise env bind_calls (fn, arg_bndrs, body, arg_occs)
1023                           spec_info@(SI specs spec_count mb_unspec)
1024   | not (isBottomingId fn)      -- Note [Do not specialise diverging functions]
1025   , notNull arg_bndrs           -- Only specialise functions
1026   , Just all_calls <- lookupVarEnv bind_calls fn
1027   = do  { (boring_call, pats) <- callsToPats env specs arg_occs all_calls
1028 --      ; pprTrace "specialise" (vcat [ppr fn <+> ppr arg_occs,
1029 --                                      text "calls" <+> ppr all_calls,
1030 --                                      text "good pats" <+> ppr pats])  $
1031 --        return ()
1032
1033                 -- Bale out if too many specialisations
1034                 -- Rather a hacky way to do so, but it'll do for now
1035         ; let spec_count' = length pats + spec_count
1036         ; case sc_count env of
1037             Just max | spec_count' > max
1038                 -> WARN( True, msg ) return (nullUsage, spec_info)
1039                 where
1040                    msg = vcat [ sep [ ptext (sLit "SpecConstr: specialisation of") <+> quotes (ppr fn)
1041                                     , nest 2 (ptext (sLit "limited by bound of")) <+> int max ]
1042                               , ptext (sLit "Use -fspec-constr-count=n to set the bound")
1043                               , extra ]
1044                    extra | not opt_PprStyle_Debug = ptext (sLit "Use -dppr-debug to see specialisations")
1045                          | otherwise = ptext (sLit "Specialisations:") <+> ppr (pats ++ [p | OS p _ _ _ <- specs])
1046
1047             _normal_case -> do {
1048
1049           (spec_usgs, new_specs) <- mapAndUnzipM (spec_one env fn arg_bndrs body)
1050                                                  (pats `zip` [spec_count..])
1051
1052         ; let spec_usg = combineUsages spec_usgs
1053               (new_usg, mb_unspec')
1054                   = case mb_unspec of
1055                       Just rhs_usg | boring_call -> (spec_usg `combineUsage` rhs_usg, Nothing)
1056                       _                          -> (spec_usg,                      mb_unspec)
1057
1058         ; return (new_usg, SI (new_specs ++ specs) spec_count' mb_unspec') } }
1059   | otherwise
1060   = return (nullUsage, spec_info)               -- The boring case
1061
1062
1063 ---------------------
1064 spec_one :: ScEnv
1065          -> OutId       -- Function
1066          -> [Var]       -- Lambda-binders of RHS; should match patterns
1067          -> CoreExpr    -- Body of the original function
1068          -> (CallPat, Int)
1069          -> UniqSM (ScUsage, OneSpec)   -- Rule and binding
1070
1071 -- spec_one creates a specialised copy of the function, together
1072 -- with a rule for using it.  I'm very proud of how short this
1073 -- function is, considering what it does :-).
1074
1075 {-
1076   Example
1077
1078      In-scope: a, x::a
1079      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
1080           [c::*, v::(b,c) are presumably bound by the (...) part]
1081   ==>
1082      f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
1083                   (...entire body of f...) [b -> (b,c),
1084                                             y -> ((:) (a,(b,c)) (x,v) hw)]
1085
1086      RULE:  forall b::* c::*,           -- Note, *not* forall a, x
1087                    v::(b,c),
1088                    hw::[(a,(b,c))] .
1089
1090             f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
1091 -}
1092
1093 spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
1094   = do  {       -- Specialise the body
1095           let spec_env = extendScSubstList (extendScInScope env qvars)
1096                                            (arg_bndrs `zip` pats)
1097         ; (spec_usg, spec_body) <- scExpr spec_env body
1098
1099 --      ; pprTrace "spec_one" (ppr fn <+> vcat [text "pats" <+> ppr pats,
1100 --                      text "calls" <+> (ppr (scu_calls spec_usg))])
1101 --        (return ())
1102
1103                 -- And build the results
1104         ; spec_uniq <- getUniqueUs
1105         ; let (spec_lam_args, spec_call_args) = mkWorkerArgs qvars body_ty
1106                 -- Usual w/w hack to avoid generating
1107                 -- a spec_rhs of unlifted type and no args
1108
1109               fn_name   = idName fn
1110               fn_loc    = nameSrcSpan fn_name
1111               spec_occ  = mkSpecOcc (nameOccName fn_name)
1112               rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
1113               spec_rhs  = mkLams spec_lam_args spec_body
1114               spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
1115               body_ty   = exprType spec_body
1116               rule_rhs  = mkVarApps (Var spec_id) spec_call_args
1117               rule      = mkLocalRule rule_name specConstrActivation fn_name qvars pats rule_rhs
1118         ; return (spec_usg, OS call_pat rule spec_id spec_rhs) }
1119
1120 -- In which phase should the specialise-constructor rules be active?
1121 -- Originally I made them always-active, but Manuel found that
1122 -- this defeated some clever user-written rules.  So Plan B
1123 -- is to make them active only in Phase 0; after all, currently,
1124 -- the specConstr transformation is only run after the simplifier
1125 -- has reached Phase 0.  In general one would want it to be
1126 -- flag-controllable, but for now I'm leaving it baked in
1127 --                                      [SLPJ Oct 01]
1128 specConstrActivation :: Activation
1129 specConstrActivation = ActiveAfter 0    -- Baked in; see comments above
1130 \end{code}
1131
1132 %************************************************************************
1133 %*                                                                      *
1134 \subsection{Argument analysis}
1135 %*                                                                      *
1136 %************************************************************************
1137
1138 This code deals with analysing call-site arguments to see whether
1139 they are constructor applications.
1140
1141
1142 \begin{code}
1143 type CallPat = ([Var], [CoreExpr])      -- Quantified variables and arguments
1144
1145
1146 callsToPats :: ScEnv -> [OneSpec] -> [ArgOcc] -> [Call] -> UniqSM (Bool, [CallPat])
1147         -- Result has no duplicate patterns,
1148         -- nor ones mentioned in done_pats
1149         -- Bool indicates that there was at least one boring pattern
1150 callsToPats env done_specs bndr_occs calls
1151   = do  { mb_pats <- mapM (callToPats env bndr_occs) calls
1152
1153         ; let good_pats :: [([Var], [CoreArg])]
1154               good_pats = catMaybes mb_pats
1155               done_pats = [p | OS p _ _ _ <- done_specs]
1156               is_done p = any (samePat p) done_pats
1157
1158         ; return (any isNothing mb_pats,
1159                   filterOut is_done (nubBy samePat good_pats)) }
1160
1161 callToPats :: ScEnv -> [ArgOcc] -> Call -> UniqSM (Maybe CallPat)
1162         -- The [Var] is the variables to quantify over in the rule
1163         --      Type variables come first, since they may scope
1164         --      over the following term variables
1165         -- The [CoreExpr] are the argument patterns for the rule
1166 callToPats env bndr_occs (con_env, args)
1167   | length args < length bndr_occs      -- Check saturated
1168   = return Nothing
1169   | otherwise
1170   = do  { let in_scope = substInScope (sc_subst env)
1171         ; prs <- argsToPats in_scope con_env (args `zip` bndr_occs)
1172         ; let (interesting_s, pats) = unzip prs
1173               pat_fvs = varSetElems (exprsFreeVars pats)
1174               qvars   = filterOut (`elemInScopeSet` in_scope) pat_fvs
1175                 -- Quantify over variables that are not in sccpe
1176                 -- at the call site
1177                 -- See Note [Shadowing] at the top
1178
1179               (tvs, ids) = partition isTyVar qvars
1180               qvars'     = tvs ++ ids
1181                 -- Put the type variables first; the type of a term
1182                 -- variable may mention a type variable
1183
1184         ; -- pprTrace "callToPats"  (ppr args $$ ppr prs $$ ppr bndr_occs) $
1185           if or interesting_s
1186           then return (Just (qvars', pats))
1187           else return Nothing }
1188
1189     -- argToPat takes an actual argument, and returns an abstracted
1190     -- version, consisting of just the "constructor skeleton" of the
1191     -- argument, with non-constructor sub-expression replaced by new
1192     -- placeholder variables.  For example:
1193     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
1194
1195 argToPat :: InScopeSet                  -- What's in scope at the fn defn site
1196          -> ValueEnv                    -- ValueEnv at the call site
1197          -> CoreArg                     -- A call arg (or component thereof)
1198          -> ArgOcc
1199          -> UniqSM (Bool, CoreArg)
1200 -- Returns (interesting, pat),
1201 -- where pat is the pattern derived from the argument
1202 --            intersting=True if the pattern is non-trivial (not a variable or type)
1203 -- E.g.         x:xs         --> (True, x:xs)
1204 --              f xs         --> (False, w)        where w is a fresh wildcard
1205 --              (f xs, 'c')  --> (True, (w, 'c'))  where w is a fresh wildcard
1206 --              \x. x+y      --> (True, \x. x+y)
1207 --              lvl7         --> (True, lvl7)      if lvl7 is bound
1208 --                                                 somewhere further out
1209
1210 argToPat _in_scope _val_env arg@(Type {}) _arg_occ
1211   = return (False, arg)
1212
1213 argToPat in_scope val_env (Note _ arg) arg_occ
1214   = argToPat in_scope val_env arg arg_occ
1215         -- Note [Notes in call patterns]
1216         -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1217         -- Ignore Notes.  In particular, we want to ignore any InlineMe notes
1218         -- Perhaps we should not ignore profiling notes, but I'm going to
1219         -- ride roughshod over them all for now.
1220         --- See Note [Notes in RULE matching] in Rules
1221
1222 argToPat in_scope val_env (Let _ arg) arg_occ
1223   = argToPat in_scope val_env arg arg_occ
1224         -- Look through let expressions
1225         -- e.g.         f (let v = rhs in \y -> ...v...)
1226         -- Here we can specialise for f (\y -> ...)
1227         -- because the rule-matcher will look through the let.
1228
1229 argToPat in_scope val_env (Cast arg co) arg_occ
1230   = do  { (interesting, arg') <- argToPat in_scope val_env arg arg_occ
1231         ; let (ty1,ty2) = coercionKind co
1232         ; if not interesting then
1233                 wildCardPat ty2
1234           else do
1235         { -- Make a wild-card pattern for the coercion
1236           uniq <- getUniqueUs
1237         ; let co_name = mkSysTvName uniq (fsLit "sg")
1238               co_var = mkCoVar co_name (mkCoKind ty1 ty2)
1239         ; return (interesting, Cast arg' (mkTyVarTy co_var)) } }
1240
1241 {-      Disabling lambda specialisation for now
1242         It's fragile, and the spec_loop can be infinite
1243 argToPat in_scope val_env arg arg_occ
1244   | is_value_lam arg
1245   = return (True, arg)
1246   where
1247     is_value_lam (Lam v e)      -- Spot a value lambda, even if
1248         | isId v = True         -- it is inside a type lambda
1249         | otherwise = is_value_lam e
1250     is_value_lam other = False
1251 -}
1252
1253   -- Check for a constructor application
1254   -- NB: this *precedes* the Var case, so that we catch nullary constrs
1255 argToPat in_scope val_env arg arg_occ
1256   | Just (ConVal dc args) <- isValue val_env arg
1257   , case arg_occ of
1258         ScrutOcc _ -> True              -- Used only by case scrutinee
1259         BothOcc    -> case arg of       -- Used elsewhere
1260                         App {} -> True  --     see Note [Reboxing]
1261                         _other -> False
1262         _other     -> False     -- No point; the arg is not decomposed
1263   = do  { args' <- argsToPats in_scope val_env (args `zip` conArgOccs arg_occ dc)
1264         ; return (True, mk_con_app dc (map snd args')) }
1265
1266   -- Check if the argument is a variable that
1267   -- is in scope at the function definition site
1268   -- It's worth specialising on this if
1269   --    (a) it's used in an interesting way in the body
1270   --    (b) we know what its value is
1271 argToPat in_scope val_env (Var v) arg_occ
1272   | case arg_occ of { UnkOcc -> False; _other -> True },        -- (a)
1273     is_value                                                    -- (b)
1274   = return (True, Var v)
1275   where
1276     is_value
1277         | isLocalId v = v `elemInScopeSet` in_scope
1278                         && isJust (lookupVarEnv val_env v)
1279                 -- Local variables have values in val_env
1280         | otherwise   = isValueUnfolding (idUnfolding v)
1281                 -- Imports have unfoldings
1282
1283 --      I'm really not sure what this comment means
1284 --      And by not wild-carding we tend to get forall'd
1285 --      variables that are in soope, which in turn can
1286 --      expose the weakness in let-matching
1287 --      See Note [Matching lets] in Rules
1288
1289   -- Check for a variable bound inside the function.
1290   -- Don't make a wild-card, because we may usefully share
1291   --    e.g.  f a = let x = ... in f (x,x)
1292   -- NB: this case follows the lambda and con-app cases!!
1293 -- argToPat _in_scope _val_env (Var v) _arg_occ
1294 --   = return (False, Var v)
1295         -- SLPJ : disabling this to avoid proliferation of versions
1296         -- also works badly when thinking about seeding the loop
1297         -- from the body of the let
1298         --       f x y = letrec g z = ... in g (x,y)
1299         -- We don't want to specialise for that *particular* x,y
1300
1301   -- The default case: make a wild-card
1302 argToPat _in_scope _val_env arg _arg_occ
1303   = wildCardPat (exprType arg)
1304
1305 wildCardPat :: Type -> UniqSM (Bool, CoreArg)
1306 wildCardPat ty = do { uniq <- getUniqueUs
1307                     ; let id = mkSysLocal (fsLit "sc") uniq ty
1308                     ; return (False, Var id) }
1309
1310 argsToPats :: InScopeSet -> ValueEnv
1311            -> [(CoreArg, ArgOcc)]
1312            -> UniqSM [(Bool, CoreArg)]
1313 argsToPats in_scope val_env args
1314   = mapM do_one args
1315   where
1316     do_one (arg,occ) = argToPat in_scope val_env arg occ
1317 \end{code}
1318
1319
1320 \begin{code}
1321 isValue :: ValueEnv -> CoreExpr -> Maybe Value
1322 isValue _env (Lit lit)
1323   = Just (ConVal (LitAlt lit) [])
1324
1325 isValue env (Var v)
1326   | Just stuff <- lookupVarEnv env v
1327   = Just stuff  -- You might think we could look in the idUnfolding here
1328                 -- but that doesn't take account of which branch of a
1329                 -- case we are in, which is the whole point
1330
1331   | not (isLocalId v) && isCheapUnfolding unf
1332   = isValue env (unfoldingTemplate unf)
1333   where
1334     unf = idUnfolding v
1335         -- However we do want to consult the unfolding
1336         -- as well, for let-bound constructors!
1337
1338 isValue env (Lam b e)
1339   | isTyVar b = case isValue env e of
1340                   Just _  -> Just LambdaVal
1341                   Nothing -> Nothing
1342   | otherwise = Just LambdaVal
1343
1344 isValue _env expr       -- Maybe it's a constructor application
1345   | (Var fun, args) <- collectArgs expr
1346   = case isDataConWorkId_maybe fun of
1347
1348         Just con | args `lengthAtLeast` dataConRepArity con
1349                 -- Check saturated; might be > because the
1350                 --                  arity excludes type args
1351                 -> Just (ConVal (DataAlt con) args)
1352
1353         _other | valArgCount args < idArity fun
1354                 -- Under-applied function
1355                -> Just LambdaVal        -- Partial application
1356
1357         _other -> Nothing
1358
1359 isValue _env _expr = Nothing
1360
1361 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
1362 mk_con_app (LitAlt lit)  []   = Lit lit
1363 mk_con_app (DataAlt con) args = mkConApp con args
1364 mk_con_app _other _args = panic "SpecConstr.mk_con_app"
1365
1366 samePat :: CallPat -> CallPat -> Bool
1367 samePat (vs1, as1) (vs2, as2)
1368   = all2 same as1 as2
1369   where
1370     same (Var v1) (Var v2)
1371         | v1 `elem` vs1 = v2 `elem` vs2
1372         | v2 `elem` vs2 = False
1373         | otherwise     = v1 == v2
1374
1375     same (Lit l1)    (Lit l2)    = l1==l2
1376     same (App f1 a1) (App f2 a2) = same f1 f2 && same a1 a2
1377
1378     same (Type {}) (Type {}) = True     -- Note [Ignore type differences]
1379     same (Note _ e1) e2 = same e1 e2    -- Ignore casts and notes
1380     same (Cast e1 _) e2 = same e1 e2
1381     same e1 (Note _ e2) = same e1 e2
1382     same e1 (Cast e2 _) = same e1 e2
1383
1384     same e1 e2 = WARN( bad e1 || bad e2, ppr e1 $$ ppr e2)
1385                  False  -- Let, lambda, case should not occur
1386     bad (Case {}) = True
1387     bad (Let {})  = True
1388     bad (Lam {})  = True
1389     bad _other    = False
1390 \end{code}
1391
1392 Note [Ignore type differences]
1393 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1394 We do not want to generate specialisations where the call patterns
1395 differ only in their type arguments!  Not only is it utterly useless,
1396 but it also means that (with polymorphic recursion) we can generate
1397 an infinite number of specialisations. Example is Data.Sequence.adjustTree,
1398 I think.
1399