compiler/specialise/SpecConstr.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[SpecConstr]{Specialise over constructors}
   5
   6 \begin{code}
   7 module SpecConstr(
   8         specConstrProgram
   9     ) where
  10
  11 #include "HsVersions.h"
  12
  13 import CoreSyn
  14 import CoreLint         ( showPass, endPass )
  15 import CoreUtils        ( exprType, tcEqExpr, mkPiTypes )
  16 import CoreFVs          ( exprsFreeVars )
  17 import CoreSubst        ( Subst, mkSubst, substExpr )
  18 import CoreTidy         ( tidyRules )
  19 import PprCore          ( pprRules )
  20 import WwLib            ( mkWorkerArgs )
  21 import DataCon          ( dataConRepArity, isVanillaDataCon )
  22 import Type             ( tyConAppArgs, tyVarsOfTypes )
  23 import Unify            ( coreRefineTys )
  24 import Id               ( Id, idName, idType, isDataConWorkId_maybe,
  25                           mkUserLocal, mkSysLocal, idUnfolding )
  26 import Var              ( Var )
  27 import VarEnv
  28 import VarSet
  29 import Name             ( nameOccName, nameSrcLoc )
  30 import Rules            ( addIdSpecialisations, mkLocalRule, rulesOfBinds )
  31 import OccName          ( mkSpecOcc )
  32 import ErrUtils         ( dumpIfSet_dyn )
  33 import DynFlags         ( DynFlags, DynFlag(..) )
  34 import BasicTypes       ( Activation(..) )
  35 import Maybes           ( orElse )
  36 import Util             ( mapAccumL, lengthAtLeast, notNull )
  37 import List             ( nubBy, partition )
  38 import UniqSupply
  39 import Outputable
  40 import FastString
  41 \end{code}
  42
  43 -----------------------------------------------------
  44                         Game plan
  45 -----------------------------------------------------
  46
  47 Consider
  48         drop n []     = []
  49         drop 0 xs     = []
  50         drop n (x:xs) = drop (n-1) xs
  51
  52 After the first time round, we could pass n unboxed.  This happens in
  53 numerical code too.  Here's what it looks like in Core:
  54
  55         drop n xs = case xs of
  56                       []     -> []
  57                       (y:ys) -> case n of
  58                                   I# n# -> case n# of
  59                                              0 -> []
  60                                              _ -> drop (I# (n# -# 1#)) xs
  61
  62 Notice that the recursive call has an explicit constructor as argument.
  63 Noticing this, we can make a specialised version of drop
  64
  65         RULE: drop (I# n#) xs ==> drop' n# xs
  66
  67         drop' n# xs = let n = I# n# in ...orig RHS...
  68
  69 Now the simplifier will apply the specialisation in the rhs of drop', giving
  70
  71         drop' n# xs = case xs of
  72                       []     -> []
  73                       (y:ys) -> case n# of
  74                                   0 -> []
  75                                   _ -> drop (n# -# 1#) xs
  76
  77 Much better!
  78
  79 We'd also like to catch cases where a parameter is carried along unchanged,
  80 but evaluated each time round the loop:
  81
  82         f i n = if i>0 || i>n then i else f (i*2) n
  83
  84 Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  85 In Core, by the time we've w/wd (f is strict in i) we get
  86
  87         f i# n = case i# ># 0 of
  88                    False -> I# i#
  89                    True  -> case n of n' { I# n# ->
  90                             case i# ># n# of
  91                                 False -> I# i#
  92                                 True  -> f (i# *# 2#) n'
  93
  94 At the call to f, we see that the argument, n is know to be (I# n#),
  95 and n is evaluated elsewhere in the body of f, so we can play the same
  96 trick as above.
  97
  98
  99 Note [Reboxing]
 100 ~~~~~~~~~~~~~~~
 101 We must be careful not to allocate the same constructor twice.  Consider
 102         f p = (...(case p of (a,b) -> e)...p...,
 103                ...let t = (r,s) in ...t...(f t)...)
 104 At the recursive call to f, we can see that t is a pair.  But we do NOT want
 105 to make a specialised copy:
 106         f' a b = let p = (a,b) in (..., ...)
 107 because now t is allocated by the caller, then r and s are passed to the
 108 recursive call, which allocates the (r,s) pair again.
 109
 110 This happens if
 111   (a) the argument p is used in other than a case-scrutinsation way.
 112   (b) the argument to the call is not a 'fresh' tuple; you have to
 113         look into its unfolding to see that it's a tuple
 114
 115 Hence the "OR" part of Note [Good arguments] below.
 116
 117 ALTERNATIVE: pass both boxed and unboxed versions.  This no longer saves
 118 allocation, but does perhaps save evals. In the RULE we'd have
 119 something like
 120
 121   f (I# x#) = f' (I# x#) x#
 122
 123 If at the call site the (I# x) was an unfolding, then we'd have to
 124 rely on CSE to eliminate the duplicate allocation.... This alternative
 125 doesn't look attractive enough to pursue.
 126
 127
 128 Note [Good arguments]
 129 ~~~~~~~~~~~~~~~~~~~~~
 130 So we look for
 131
 132 * A self-recursive function.  Ignore mutual recursion for now,
 133   because it's less common, and the code is simpler for self-recursion.
 134
 135 * EITHER
 136
 137    a) At a recursive call, one or more parameters is an explicit
 138       constructor application
 139         AND
 140       That same parameter is scrutinised by a case somewhere in
 141       the RHS of the function
 142
 143   OR
 144
 145     b) At a recursive call, one or more parameters has an unfolding
 146        that is an explicit constructor application
 147         AND
 148       That same parameter is scrutinised by a case somewhere in
 149       the RHS of the function
 150         AND
 151       Those are the only uses of the parameter (see Note [Reboxing])
 152
 153
 154 What to abstract over
 155 ~~~~~~~~~~~~~~~~~~~~~
 156 There's a bit of a complication with type arguments.  If the call
 157 site looks like
 158
 159         f p = ...f ((:) [a] x xs)...
 160
 161 then our specialised function look like
 162
 163         f_spec x xs = let p = (:) [a] x xs in ....as before....
 164
 165 This only makes sense if either
 166   a) the type variable 'a' is in scope at the top of f, or
 167   b) the type variable 'a' is an argument to f (and hence fs)
 168
 169 Actually, (a) may hold for value arguments too, in which case
 170 we may not want to pass them.  Supose 'x' is in scope at f's
 171 defn, but xs is not.  Then we'd like
 172
 173         f_spec xs = let p = (:) [a] x xs in ....as before....
 174
 175 Similarly (b) may hold too.  If x is already an argument at the
 176 call, no need to pass it again.
 177
 178 Finally, if 'a' is not in scope at the call site, we could abstract
 179 it as we do the term variables:
 180
 181         f_spec a x xs = let p = (:) [a] x xs in ...as before...
 182
 183 So the grand plan is:
 184
 185         * abstract the call site to a constructor-only pattern
 186           e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
 187
 188         * Find the free variables of the abstracted pattern
 189
 190         * Pass these variables, less any that are in scope at
 191           the fn defn.  But see Note [Shadowing] below.
 192
 193
 194 NOTICE that we only abstract over variables that are not in scope,
 195 so we're in no danger of shadowing variables used in "higher up"
 196 in f_spec's RHS.
 197
 198
 199 Note [Shadowing]
 200 ~~~~~~~~~~~~~~~~
 201 In this pass we gather up usage information that may mention variables
 202 that are bound between the usage site and the definition site; or (more
 203 seriously) may be bound to something different at the definition site.
 204 For example:
 205
 206         f x = letrec g y v = let x = ...
 207                              in ...(g (a,b) x)...
 208
 209 Since 'x' is in scope at the call site, we may make a rewrite rule that
 210 looks like
 211         RULE forall a,b. g (a,b) x = ...
 212 But this rule will never match, because it's really a different 'x' at
 213 the call site -- and that difference will be manifest by the time the
 214 simplifier gets to it.  [A worry: the simplifier doesn't *guarantee*
 215 no-shadowing, so perhaps it may not be distinct?]
 216
 217 Anyway, the rule isn't actually wrong, it's just not useful.  One possibility
 218 is to run deShadowBinds before running SpecConstr, but instead we run the
 219 simplifier.  That gives the simplest possible program for SpecConstr to
 220 chew on; and it virtually guarantees no shadowing.
 221
 222 -----------------------------------------------------
 223                 Stuff not yet handled
 224 -----------------------------------------------------
 225
 226 Here are notes arising from Roman's work that I don't want to lose.
 227
 228 Specialising for constant parameters
 229 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 230 This one is about specialising on a *constant* (but not necessarily
 231 constructor) argument
 232
 233     foo :: Int -> (Int -> Int) -> Int
 234     foo 0 f = 0
 235     foo m f = foo (f m) (+1)
 236
 237 It produces
 238
 239     lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
 240     lvl_rmV =
 241       \ (ds_dlk :: GHC.Base.Int) ->
 242         case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
 243         GHC.Base.I# (GHC.Prim.+# x_alG 1)
 244
 245     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 246     GHC.Prim.Int#
 247     T.$wfoo =
 248       \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
 249         case ww_sme of ds_Xlw {
 250           __DEFAULT ->
 251         case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
 252         T.$wfoo ww1_Xmz lvl_rmV
 253         };
 254           0 -> 0
 255         }
 256
 257 The recursive call has lvl_rmV as its argument, so we could create a specialised copy
 258 with that argument baked in; that is, not passed at all.   Now it can perhaps be inlined.
 259
 260 When is this worth it?  Call the constant 'lvl'
 261 - If 'lvl' has an unfolding that is a constructor, see if the corresponding
 262   parameter is scrutinised anywhere in the body.
 263
 264 - If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
 265   parameter is applied (...to enough arguments...?)
 266
 267   Also do this is if the function has RULES?
 268
 269 Also
 270
 271 Specialising for lambdas
 272 ~~~~~~~~~~~~~~~~~~~~~~~~
 273     foo :: Int -> (Int -> Int) -> Int
 274     foo 0 f = 0
 275     foo m f = foo (f m) (\n -> n-m)
 276
 277 This is subtly different from the previous one in that we get an
 278 explicit lambda as the argument:
 279
 280     T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
 281     GHC.Prim.Int#
 282     T.$wfoo =
 283       \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
 284         case ww_sm8 of ds_Xlr {
 285           __DEFAULT ->
 286         case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
 287         T.$wfoo
 288           ww1_Xmq
 289           (\ (n_ad3 :: GHC.Base.Int) ->
 290              case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
 291              GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
 292              })
 293         };
 294           0 -> 0
 295         }
 296
 297 I wonder if SpecConstr couldn't be extended to handle this? After all,
 298 lambda is a sort of constructor for functions and perhaps it already
 299 has most of the necessary machinery?
 300
 301 Furthermore, there's an immediate win, because you don't need to allocate the lamda
 302 at the call site; and if perchance it's called in the recursive call, then you
 303 may avoid allocating it altogether.  Just like for constructors.
 304
 305 Looks cool, but probably rare...but it might be easy to implement.
 306
 307 Example 1
 308 ~~~~~~~~~
 309     data T a = T !a
 310
 311     foo :: Int -> T Int -> Int
 312     foo 0 t = 0
 313     foo x t | even x    = case t of { T n -> foo (x-n) t }
 314             | otherwise = foo (x-1) t
 315
 316 SpecConstr does no specialisation, because the second recursive call
 317 looks like a boxed use of the argument.  A pity.
 318
 319     $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 320     $wfoo_sFw =
 321       \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
 322          case ww_sFo of ds_Xw6 [Just L] {
 323            __DEFAULT ->
 324                 case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
 325                   __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
 326                   0 ->
 327                     case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
 328                     case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
 329                     $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
 330                     } } };
 331            0 -> 0
 332
 333 Example 2
 334 ~~~~~~~~~
 335     data a :*: b = !a :*: !b
 336     data T a = T !a
 337
 338     foo :: (Int :*: T Int) -> Int
 339     foo (0 :*: t) = 0
 340     foo (x :*: t) | even x    = case t of { T n -> foo ((x-n) :*: t) }
 341                   | otherwise = foo ((x-1) :*: t)
 342
 343 Very similar to the previous one, except that the parameters are now in
 344 a strict tuple. Before SpecConstr, we have
 345
 346     $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
 347     $wfoo_sG3 =
 348       \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
 349     GHC.Base.Int) ->
 350         case ww_sFU of ds_Xws [Just L] {
 351           __DEFAULT ->
 352         case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
 353           __DEFAULT ->
 354             case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
 355             $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2             -- $wfoo1
 356             };
 357           0 ->
 358             case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
 359             case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
 360             $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB        -- $wfoo2
 361             } } };
 362           0 -> 0 }
 363
 364 We get two specialisations:
 365 "SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
 366                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
 367                   = Foo.$s$wfoo1 a_sFB sc_sGC ;
 368 "SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
 369                   Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
 370                   = Foo.$s$wfoo y_aFp sc_sGC ;
 371
 372 But perhaps the first one isn't good.  After all, we know that tpl_B2 is
 373 a T (I# x) really, because T is strict and Int has one constructor.  (We can't
 374 unbox the strict fields, becuase T is polymorphic!)
 375
 376
 377
 378 %************************************************************************
 379 %*                                                                      *
 380 \subsection{Top level wrapper stuff}
 381 %*                                                                      *
 382 %************************************************************************
 383
 384 \begin{code}
 385 specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
 386 specConstrProgram dflags us binds
 387   = do
 388         showPass dflags "SpecConstr"
 389
 390         let (binds', _) = initUs us (go emptyScEnv binds)
 391
 392         endPass dflags "SpecConstr" Opt_D_dump_spec binds'
 393
 394         dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
 395                   (pprRules (tidyRules emptyTidyEnv (rulesOfBinds binds')))
 396
 397         return binds'
 398   where
 399     go env []           = returnUs []
 400     go env (bind:binds) = scBind env bind       `thenUs` \ (env', _, bind') ->
 401                           go env' binds         `thenUs` \ binds' ->
 402                           returnUs (bind' : binds')
 403 \end{code}
 404
 405
 406 %************************************************************************
 407 %*                                                                      *
 408 \subsection{Environment: goes downwards}
 409 %*                                                                      *
 410 %************************************************************************
 411
 412 \begin{code}
 413 data ScEnv = SCE { scope :: VarEnv HowBound,
 414                         -- Binds all non-top-level variables in scope
 415
 416                    cons  :: ConstrEnv
 417              }
 418
 419 type ConstrEnv = IdEnv ConValue
 420 data ConValue  = CV AltCon [CoreArg]
 421         -- Variables known to be bound to a constructor
 422         -- in a particular case alternative
 423
 424
 425 instance Outputable ConValue where
 426    ppr (CV con args) = ppr con <+> interpp'SP args
 427
 428 refineConstrEnv :: Subst -> ConstrEnv -> ConstrEnv
 429 -- The substitution is a type substitution only
 430 refineConstrEnv subst env = mapVarEnv refine_con_value env
 431   where
 432     refine_con_value (CV con args) = CV con (map (substExpr subst) args)
 433
 434 emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }
 435
 436 data HowBound = RecFun          -- These are the recursive functions for which
 437                                 -- we seek interesting call patterns
 438
 439               | RecArg          -- These are those functions' arguments; we are
 440                                 -- interested to see if those arguments are scrutinised
 441
 442               | Other           -- We track all others so we know what's in scope
 443                                 -- This is used in spec_one to check what needs to be
 444                                 -- passed as a parameter and what is in scope at the
 445                                 -- function definition site
 446
 447 instance Outputable HowBound where
 448   ppr RecFun = text "RecFun"
 449   ppr RecArg = text "RecArg"
 450   ppr Other = text "Other"
 451
 452 lookupScopeEnv env v = lookupVarEnv (scope env) v
 453
 454 extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
 455 extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }
 456
 457     -- When we encounter
 458     --  case scrut of b
 459     --      C x y -> ...
 460     -- we want to bind b, and perhaps scrut too, to (C x y)
 461 extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
 462 extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
 463   = extendBndrs env (case_bndr : alt_bndrs)
 464
 465 extendCaseBndrs env case_bndr scrut con@(LitAlt lit) alt_bndrs
 466   = ASSERT( null alt_bndrs ) extendAlt env case_bndr scrut (CV con []) []
 467
 468 extendCaseBndrs env case_bndr scrut con@(DataAlt data_con) alt_bndrs
 469   | isVanillaDataCon data_con
 470   = extendAlt env case_bndr scrut (CV con vanilla_args) alt_bndrs
 471
 472   | otherwise   -- GADT
 473   = extendAlt env1 case_bndr scrut (CV con gadt_args) alt_bndrs
 474   where
 475     vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
 476                    map varToCoreExpr alt_bndrs
 477
 478     gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
 479         -- This call generates some bogus warnings from substExpr,
 480         -- because it's inconvenient to put all the Ids in scope
 481         -- Will be fixed when we move to FC
 482
 483     (alt_tvs, _) = span isTyVar alt_bndrs
 484     Just (tv_subst, is_local) = coreRefineTys data_con alt_tvs (idType case_bndr)
 485     subst = mkSubst in_scope tv_subst emptyVarEnv       -- No Id substitition
 486     in_scope = mkInScopeSet (tyVarsOfTypes (varEnvElts tv_subst))
 487
 488     env1 | is_local  = env
 489          | otherwise = env { cons = refineConstrEnv subst (cons env) }
 490
 491
 492
 493 extendAlt :: ScEnv -> Id -> CoreExpr -> ConValue -> [Var] -> ScEnv
 494 extendAlt env case_bndr scrut val alt_bndrs
 495   = let
 496        env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
 497                     cons  = extendVarEnv     (cons  env) case_bndr val }
 498     in
 499     case scrut of
 500         Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
 501                    -- Also forget if the scrutinee is a RecArg, because we're
 502                    -- now in the branch of a case, and we don't want to
 503                    -- record a non-scrutinee use of v if we have
 504                    --   case v of { (a,b) -> ...(f v)... }
 505                  SCE { scope = extendVarEnv (scope env1) v Other,
 506                        cons  = extendVarEnv (cons env1)  v val }
 507         other -> env1
 508
 509     -- When we encounter a recursive function binding
 510     --  f = \x y -> ...
 511     -- we want to extend the scope env with bindings
 512     -- that record that f is a RecFn and x,y are RecArgs
 513 extendRecBndr env fn bndrs
 514   =  env { scope = scope env `extendVarEnvList`
 515                    ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
 516 \end{code}
 517
 518
 519 %************************************************************************
 520 %*                                                                      *
 521 \subsection{Usage information: flows upwards}
 522 %*                                                                      *
 523 %************************************************************************
 524
 525 \begin{code}
 526 data ScUsage
 527    = SCU {
 528         calls :: !(IdEnv ([Call])),     -- Calls
 529                                         -- The functions are a subset of the
 530                                         --      RecFuns in the ScEnv
 531
 532         occs :: !(IdEnv ArgOcc)         -- Information on argument occurrences
 533      }                                  -- The variables are a subset of the
 534                                         --      RecArg in the ScEnv
 535
 536 type Call = (ConstrEnv, [CoreArg])
 537         -- The arguments of the call, together with the
 538         -- env giving the constructor bindings at the call site
 539
 540 nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }
 541
 542 combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
 543                            occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }
 544
 545 combineUsages [] = nullUsage
 546 combineUsages us = foldr1 combineUsage us
 547
 548 data ArgOcc = CaseScrut
 549             | OtherOcc
 550             | Both
 551
 552 instance Outputable ArgOcc where
 553   ppr CaseScrut = ptext SLIT("case-scrut")
 554   ppr OtherOcc  = ptext SLIT("other-occ")
 555   ppr Both      = ptext SLIT("case-scrut and other")
 556
 557 combineOcc CaseScrut CaseScrut = CaseScrut
 558 combineOcc OtherOcc  OtherOcc  = OtherOcc
 559 combineOcc _         _         = Both
 560 \end{code}
 561
 562
 563 %************************************************************************
 564 %*                                                                      *
 565 \subsection{The main recursive function}
 566 %*                                                                      *
 567 %************************************************************************
 568
 569 The main recursive function gathers up usage information, and
 570 creates specialised versions of functions.
 571
 572 \begin{code}
 573 scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
 574         -- The unique supply is needed when we invent
 575         -- a new name for the specialised function and its args
 576
 577 scExpr env e@(Type t) = returnUs (nullUsage, e)
 578 scExpr env e@(Lit l)  = returnUs (nullUsage, e)
 579 scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
 580 scExpr env (Note n e) = scExpr env e    `thenUs` \ (usg,e') ->
 581                         returnUs (usg, Note n e')
 582 scExpr env (Lam b e)  = scExpr (extendBndr env b) e     `thenUs` \ (usg,e') ->
 583                         returnUs (usg, Lam b e')
 584
 585 scExpr env (Case scrut b ty alts)
 586   = sc_scrut scrut              `thenUs` \ (scrut_usg, scrut') ->
 587     mapAndUnzipUs sc_alt alts   `thenUs` \ (alts_usgs, alts') ->
 588     returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
 589               Case scrut' b ty alts')
 590   where
 591     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
 592     sc_scrut e         = scExpr env e
 593
 594     sc_alt (con,bs,rhs) = scExpr env1 rhs       `thenUs` \ (usg,rhs') ->
 595                           returnUs (usg, (con,bs,rhs'))
 596                         where
 597                           env1 = extendCaseBndrs env b scrut con bs
 598
 599 scExpr env (Let bind body)
 600   = scBind env bind     `thenUs` \ (env', bind_usg, bind') ->
 601     scExpr env' body    `thenUs` \ (body_usg, body') ->
 602     returnUs (bind_usg `combineUsage` body_usg, Let bind' body')
 603
 604 scExpr env e@(App _ _)
 605   = let
 606         (fn, args) = collectArgs e
 607     in
 608     mapAndUnzipUs (scExpr env) (fn:args)        `thenUs` \ (usgs, (fn':args')) ->
 609         -- Process the function too.   It's almost always a variable,
 610         -- but not always.  In particular, if this pass follows float-in,
 611         -- which it may, we can get
 612         --      (let f = ...f... in f) arg1 arg2
 613     let
 614         call_usg = case fn of
 615                         Var f | Just RecFun <- lookupScopeEnv env f
 616                               -> SCU { calls = unitVarEnv f [(cons env, args)],
 617                                        occs  = emptyVarEnv }
 618                         other -> nullUsage
 619     in
 620     returnUs (combineUsages usgs `combineUsage` call_usg, mkApps fn' args')
 621
 622
 623 ----------------------
 624 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 625 scBind env (Rec [(fn,rhs)])
 626   | notNull val_bndrs
 627   = scExpr env_fn_body body             `thenUs` \ (usg, body') ->
 628     specialise env fn bndrs body' usg   `thenUs` \ (rules, spec_prs) ->
 629         -- Note body': the specialised copies should be based on the
 630         --             optimised version of the body, in case there were
 631         --             nested functions inside.
 632     let
 633         SCU { calls = calls, occs = occs } = usg
 634     in
 635     returnUs (extendBndr env fn,        -- For the body of the letrec, just
 636                                         -- extend the env with Other to record
 637                                         -- that it's in scope; no funny RecFun business
 638               SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
 639               Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
 640   where
 641     (bndrs,body) = collectBinders rhs
 642     val_bndrs    = filter isId bndrs
 643     env_fn_body  = extendRecBndr env fn bndrs
 644
 645 scBind env (Rec prs)
 646   = mapAndUnzipUs do_one prs    `thenUs` \ (usgs, prs') ->
 647     returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
 648   where
 649     do_one (bndr,rhs) = scExpr env rhs  `thenUs` \ (usg, rhs') ->
 650                         returnUs (usg, (bndr,rhs'))
 651
 652 scBind env (NonRec bndr rhs)
 653   = scExpr env rhs      `thenUs` \ (usg, rhs') ->
 654     returnUs (extendBndr env bndr, usg, NonRec bndr rhs')
 655
 656 ----------------------
 657 varUsage env v use
 658   | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv,
 659                                                 occs = unitVarEnv v use }
 660   | otherwise                           = nullUsage
 661 \end{code}
 662
 663
 664 %************************************************************************
 665 %*                                                                      *
 666 \subsection{The specialiser}
 667 %*                                                                      *
 668 %************************************************************************
 669
 670 \begin{code}
 671 specialise :: ScEnv
 672            -> Id                        -- Functionn
 673            -> [CoreBndr] -> CoreExpr    -- Its RHS
 674            -> ScUsage                   -- Info on usage
 675            -> UniqSM ([CoreRule],       -- Rules
 676                       [(Id,CoreExpr)])  -- Bindings
 677
 678 specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 679   = getUs               `thenUs` \ us ->
 680     let
 681         all_calls = lookupVarEnv calls fn `orElse` []
 682
 683         good_calls :: [[CoreArg]]
 684         good_calls = [ pats
 685                      | (con_env, call_args) <- all_calls,
 686                        call_args `lengthAtLeast` n_bndrs,           -- App is saturated
 687                        let call = bndrs `zip` call_args,
 688                        any (good_arg con_env occs) call,    -- At least one arg is a constr app
 689                        let (_, pats) = argsToPats con_env us call_args
 690                      ]
 691     in
 692     mapAndUnzipUs (spec_one env fn (mkLams bndrs body))
 693                   (nubBy same_call good_calls `zip` [1..])
 694   where
 695     n_bndrs  = length bndrs
 696     same_call as1 as2 = and (zipWith tcEqExpr as1 as2)
 697
 698 ---------------------
 699 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
 700 -- See Note [Good arguments] above
 701 good_arg con_env arg_occs (bndr, arg)
 702   = case is_con_app_maybe con_env arg of
 703         Just _ ->  bndr_usg_ok arg_occs bndr arg
 704         other   -> False
 705
 706 bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
 707 bndr_usg_ok arg_occs bndr arg
 708   = case lookupVarEnv arg_occs bndr of
 709         Just CaseScrut -> True                  -- Used only by case scrutiny
 710         Just Both      -> case arg of           -- Used by case and elsewhere
 711                             App _ _ -> True     -- so the arg should be an explicit con app
 712                             other   -> False
 713         other -> False                          -- Not used, or used wonkily
 714
 715
 716 ---------------------
 717 spec_one :: ScEnv
 718          -> Id                                  -- Function
 719          -> CoreExpr                            -- Rhs of the original function
 720          -> ([CoreArg], Int)
 721          -> UniqSM (CoreRule, (Id,CoreExpr))    -- Rule and binding
 722
 723 -- spec_one creates a specialised copy of the function, together
 724 -- with a rule for using it.  I'm very proud of how short this
 725 -- function is, considering what it does :-).
 726
 727 {-
 728   Example
 729
 730      In-scope: a, x::a
 731      f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
 732           [c::*, v::(b,c) are presumably bound by the (...) part]
 733   ==>
 734      f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
 735                   (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
 736
 737      RULE:  forall b::* c::*,           -- Note, *not* forall a, x
 738                    v::(b,c),
 739                    hw::[(a,(b,c))] .
 740
 741             f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
 742 -}
 743
 744 spec_one env fn rhs (pats, rule_number)
 745   = getUniqueUs                 `thenUs` \ spec_uniq ->
 746     let
 747         fn_name      = idName fn
 748         fn_loc       = nameSrcLoc fn_name
 749         spec_occ     = mkSpecOcc (nameOccName fn_name)
 750         pat_fvs      = varSetElems (exprsFreeVars pats)
 751         vars_to_bind = filter not_avail pat_fvs
 752                 -- See Note [Shadowing] at the top
 753
 754         not_avail v  = not (v `elemVarEnv` scope env)
 755                 -- Put the type variables first; the type of a term
 756                 -- variable may mention a type variable
 757         (tvs, ids)   = partition isTyVar vars_to_bind
 758         bndrs        = tvs ++ ids
 759         spec_body    = mkApps rhs pats
 760         body_ty      = exprType spec_body
 761
 762         (spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
 763                 -- Usual w/w hack to avoid generating
 764                 -- a spec_rhs of unlifted type and no args
 765
 766         rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
 767         spec_rhs  = mkLams spec_lam_args spec_body
 768         spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
 769         rule_rhs  = mkVarApps (Var spec_id) spec_call_args
 770         rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
 771     in
 772     returnUs (rule, (spec_id, spec_rhs))
 773
 774 -- In which phase should the specialise-constructor rules be active?
 775 -- Originally I made them always-active, but Manuel found that
 776 -- this defeated some clever user-written rules.  So Plan B
 777 -- is to make them active only in Phase 0; after all, currently,
 778 -- the specConstr transformation is only run after the simplifier
 779 -- has reached Phase 0.  In general one would want it to be
 780 -- flag-controllable, but for now I'm leaving it baked in
 781 --                                      [SLPJ Oct 01]
 782 specConstrActivation :: Activation
 783 specConstrActivation = ActiveAfter 0    -- Baked in; see comments above
 784 \end{code}
 785
 786 %************************************************************************
 787 %*                                                                      *
 788 \subsection{Argument analysis}
 789 %*                                                                      *
 790 %************************************************************************
 791
 792 This code deals with analysing call-site arguments to see whether
 793 they are constructor applications.
 794
 795 \begin{code}
 796     -- argToPat takes an actual argument, and returns an abstracted
 797     -- version, consisting of just the "constructor skeleton" of the
 798     -- argument, with non-constructor sub-expression replaced by new
 799     -- placeholder variables.  For example:
 800     --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
 801
 802 argToPat   :: ConstrEnv -> UniqSupply -> CoreArg -> (UniqSupply, CoreExpr)
 803 argToPat env us (Type ty)
 804   = (us, Type ty)
 805
 806 argToPat env us arg
 807   | Just (CV dc args) <- is_con_app_maybe env arg
 808   = let
 809         (us',args') = argsToPats env us args
 810     in
 811     (us', mk_con_app dc args')
 812
 813 argToPat env us (Var v) -- Don't uniqify existing vars,
 814   = (us, Var v)         -- so that we can spot when we pass them twice
 815
 816 argToPat env us arg
 817   = (us1, Var (mkSysLocal FSLIT("sc") (uniqFromSupply us2) (exprType arg)))
 818   where
 819     (us1,us2) = splitUniqSupply us
 820
 821 argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
 822 argsToPats env us args = mapAccumL (argToPat env) us args
 823 \end{code}
 824
 825
 826 \begin{code}
 827 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
 828 is_con_app_maybe env (Var v)
 829   = case lookupVarEnv env v of
 830         Just stuff -> Just stuff
 831                 -- You might think we could look in the idUnfolding here
 832                 -- but that doesn't take account of which branch of a
 833                 -- case we are in, which is the whole point
 834
 835         Nothing | isCheapUnfolding unf
 836                 -> is_con_app_maybe env (unfoldingTemplate unf)
 837                 where
 838                   unf = idUnfolding v
 839                 -- However we do want to consult the unfolding as well,
 840                 -- for let-bound constructors!
 841
 842         other  -> Nothing
 843
 844 is_con_app_maybe env (Lit lit)
 845   = Just (CV (LitAlt lit) [])
 846
 847 is_con_app_maybe env expr
 848   = case collectArgs expr of
 849         (Var fun, args) | Just con <- isDataConWorkId_maybe fun,
 850                           args `lengthAtLeast` dataConRepArity con
 851                 -- Might be > because the arity excludes type args
 852                         -> Just (CV (DataAlt con) args)
 853
 854         other -> Nothing
 855
 856 mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
 857 mk_con_app (LitAlt lit)  []   = Lit lit
 858 mk_con_app (DataAlt con) args = mkConApp con args
 859 \end{code}