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