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