ghc/compiler/specialise/Specialise.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1996
   3 %
   4 \section[Specialise]{Stamping out overloading, and (optionally) polymorphism}
   5
   6 \begin{code}
   7 module Specialise (
   8         specProgram,
   9         idSpecVars,
  10         substSpecEnvRhs
  11     ) where
  12
  13 #include "HsVersions.h"
  14
  15 import Id               ( Id, DictVar, idType, mkUserLocal,
  16
  17                           getIdSpecialisation, setIdSpecialisation,
  18
  19                           IdSet, mkIdSet, addOneToIdSet, intersectIdSets, isEmptyIdSet,
  20                                  emptyIdSet, unionIdSets, minusIdSet, unitIdSet, elementOfIdSet,
  21
  22                           IdEnv, mkIdEnv, lookupIdEnv, addOneToIdEnv, delOneFromIdEnv
  23                         )
  24
  25 import Type             ( Type, mkTyVarTy, splitSigmaTy, instantiateTy, isDictTy,
  26                           tyVarsOfType, tyVarsOfTypes, applyTys, mkForAllTys
  27                         )
  28 import TyCon            ( TyCon )
  29 import TyVar            ( TyVar,
  30                           TyVarSet, mkTyVarSet, isEmptyTyVarSet, intersectTyVarSets,
  31                                     elementOfTyVarSet, unionTyVarSets, emptyTyVarSet,
  32                           TyVarEnv, mkTyVarEnv, delFromTyVarEnv
  33                         )
  34 import CoreSyn
  35 import PprCore          ()      -- Instances
  36 import Name             ( NamedThing(..), getSrcLoc )
  37 import SpecEnv          ( addToSpecEnv, lookupSpecEnv, specEnvValues )
  38
  39 import UniqSupply       ( UniqSupply,
  40                           UniqSM, initUs, thenUs, returnUs, getUnique, mapUs
  41                         )
  42
  43 import FiniteMap
  44 import Maybes           ( MaybeErr(..), maybeToBool )
  45 import Bag
  46 import List             ( partition )
  47 import Util             ( zipEqual )
  48 import Outputable
  49
  50
  51 infixr 9 `thenSM`
  52 \end{code}
  53
  54 %************************************************************************
  55 %*                                                                      *
  56 \subsection[notes-Specialise]{Implementation notes [SLPJ, Aug 18 1993]}
  57 %*                                                                      *
  58 %************************************************************************
  59
  60 These notes describe how we implement specialisation to eliminate
  61 overloading, and optionally to eliminate unboxed polymorphism, and
  62 full polymorphism.
  63
  64 The specialisation pass is a partial evaluator which works on Core
  65 syntax, complete with all the explicit dictionary application,
  66 abstraction and construction as added by the type checker.  The
  67 existing type checker remains largely as it is.
  68
  69 One important thought: the {\em types} passed to an overloaded
  70 function, and the {\em dictionaries} passed are mutually redundant.
  71 If the same function is applied to the same type(s) then it is sure to
  72 be applied to the same dictionary(s)---or rather to the same {\em
  73 values}.  (The arguments might look different but they will evaluate
  74 to the same value.)
  75
  76 Second important thought: we know that we can make progress by
  77 treating dictionary arguments as static and worth specialising on.  So
  78 we can do without binding-time analysis, and instead specialise on
  79 dictionary arguments and no others.
  80
  81 The basic idea
  82 ~~~~~~~~~~~~~~
  83 Suppose we have
  84
  85         let f = <f_rhs>
  86         in <body>
  87
  88 and suppose f is overloaded.
  89
  90 STEP 1: CALL-INSTANCE COLLECTION
  91
  92 We traverse <body>, accumulating all applications of f to types and
  93 dictionaries.
  94
  95 (Might there be partial applications, to just some of its types and
  96 dictionaries?  In principle yes, but in practice the type checker only
  97 builds applications of f to all its types and dictionaries, so partial
  98 applications could only arise as a result of transformation, and even
  99 then I think it's unlikely.  In any case, we simply don't accumulate such
 100 partial applications.)
 101
 102 There's a choice of whether to collect details of all *polymorphic* functions
 103 or simply all *overloaded* ones.  How to sort this out?
 104   Pass in a predicate on the function to say if it is "interesting"?
 105   This is dependent on the user flags: SpecialiseOverloaded
 106                                        SpecialiseUnboxed
 107                                        SpecialiseAll
 108
 109 STEP 2: EQUIVALENCES
 110
 111 So now we have a collection of calls to f:
 112         f t1 t2 d1 d2
 113         f t3 t4 d3 d4
 114         ...
 115 Notice that f may take several type arguments.  To avoid ambiguity, we
 116 say that f is called at type t1/t2 and t3/t4.
 117
 118 We take equivalence classes using equality of the *types* (ignoring
 119 the dictionary args, which as mentioned previously are redundant).
 120
 121 STEP 3: SPECIALISATION
 122
 123 For each equivalence class, choose a representative (f t1 t2 d1 d2),
 124 and create a local instance of f, defined thus:
 125
 126         f@t1/t2 = <f_rhs> t1 t2 d1 d2
 127
 128 (f_rhs presumably has some big lambdas and dictionary lambdas, so lots
 129 of simplification will now result.)  Then we should recursively do
 130 everything again.
 131
 132 The new id has its own unique, but its print-name (if exported) has
 133 an explicit representation of the instance types t1/t2.
 134
 135 Add this new id to f's IdInfo, to record that f has a specialised version.
 136
 137 Before doing any of this, check that f's IdInfo doesn't already
 138 tell us about an existing instance of f at the required type/s.
 139 (This might happen if specialisation was applied more than once, or
 140 it might arise from user SPECIALIZE pragmas.)
 141
 142 Recursion
 143 ~~~~~~~~~
 144 Wait a minute!  What if f is recursive?  Then we can't just plug in
 145 its right-hand side, can we?
 146
 147 But it's ok.  The type checker *always* creates non-recursive definitions
 148 for overloaded recursive functions.  For example:
 149
 150         f x = f (x+x)           -- Yes I know its silly
 151
 152 becomes
 153
 154         f a (d::Num a) = let p = +.sel a d
 155                          in
 156                          letrec fl (y::a) = fl (p y y)
 157                          in
 158                          fl
 159
 160 We still have recusion for non-overloadd functions which we
 161 speciailise, but the recursive call should get speciailised to the
 162 same recursive version.
 163
 164
 165 Polymorphism 1
 166 ~~~~~~~~~~~~~~
 167
 168 All this is crystal clear when the function is applied to *constant
 169 types*; that is, types which have no type variables inside.  But what if
 170 it is applied to non-constant types?  Suppose we find a call of f at type
 171 t1/t2.  There are two possibilities:
 172
 173 (a) The free type variables of t1, t2 are in scope at the definition point
 174 of f.  In this case there's no problem, we proceed just as before.  A common
 175 example is as follows.  Here's the Haskell:
 176
 177         g y = let f x = x+x
 178               in f y + f y
 179
 180 After typechecking we have
 181
 182         g a (d::Num a) (y::a) = let f b (d'::Num b) (x::b) = +.sel b d' x x
 183                                 in +.sel a d (f a d y) (f a d y)
 184
 185 Notice that the call to f is at type type "a"; a non-constant type.
 186 Both calls to f are at the same type, so we can specialise to give:
 187
 188         g a (d::Num a) (y::a) = let f@a (x::a) = +.sel a d x x
 189                                 in +.sel a d (f@a y) (f@a y)
 190
 191
 192 (b) The other case is when the type variables in the instance types
 193 are *not* in scope at the definition point of f.  The example we are
 194 working with above is a good case.  There are two instances of (+.sel a d),
 195 but "a" is not in scope at the definition of +.sel.  Can we do anything?
 196 Yes, we can "common them up", a sort of limited common sub-expression deal.
 197 This would give:
 198
 199         g a (d::Num a) (y::a) = let +.sel@a = +.sel a d
 200                                     f@a (x::a) = +.sel@a x x
 201                                 in +.sel@a (f@a y) (f@a y)
 202
 203 This can save work, and can't be spotted by the type checker, because
 204 the two instances of +.sel weren't originally at the same type.
 205
 206 Further notes on (b)
 207
 208 * There are quite a few variations here.  For example, the defn of
 209   +.sel could be floated ouside the \y, to attempt to gain laziness.
 210   It certainly mustn't be floated outside the \d because the d has to
 211   be in scope too.
 212
 213 * We don't want to inline f_rhs in this case, because
 214 that will duplicate code.  Just commoning up the call is the point.
 215
 216 * Nothing gets added to +.sel's IdInfo.
 217
 218 * Don't bother unless the equivalence class has more than one item!
 219
 220 Not clear whether this is all worth it.  It is of course OK to
 221 simply discard call-instances when passing a big lambda.
 222
 223 Polymorphism 2 -- Overloading
 224 ~~~~~~~~~~~~~~
 225 Consider a function whose most general type is
 226
 227         f :: forall a b. Ord a => [a] -> b -> b
 228
 229 There is really no point in making a version of g at Int/Int and another
 230 at Int/Bool, because it's only instancing the type variable "a" which
 231 buys us any efficiency. Since g is completely polymorphic in b there
 232 ain't much point in making separate versions of g for the different
 233 b types.
 234
 235 That suggests that we should identify which of g's type variables
 236 are constrained (like "a") and which are unconstrained (like "b").
 237 Then when taking equivalence classes in STEP 2, we ignore the type args
 238 corresponding to unconstrained type variable.  In STEP 3 we make
 239 polymorphic versions.  Thus:
 240
 241         f@t1/ = /\b -> <f_rhs> t1 b d1 d2
 242
 243 This seems pretty simple, and a Good Thing.
 244
 245 Polymorphism 3 -- Unboxed
 246 ~~~~~~~~~~~~~~
 247
 248 If we are speciailising at unboxed types we must speciailise
 249 regardless of the overloading constraint.  In the exaple above it is
 250 worth speciailising at types Int/Int#, Int/Bool# and a/Int#, Int#/Int#
 251 etc.
 252
 253 Note that specialising an overloaded type at an uboxed type requires
 254 an unboxed instance -- we cannot default to an unspecialised version!
 255
 256
 257 Dictionary floating
 258 ~~~~~~~~~~~~~~~~~~~
 259 Consider
 260
 261         f x = let g p q = p==q
 262                   h r s = (r+s, g r s)
 263               in
 264               h x x
 265
 266
 267 Before specialisation, leaving out type abstractions we have
 268
 269         f df x = let g :: Eq a => a -> a -> Bool
 270                      g dg p q = == dg p q
 271                      h :: Num a => a -> a -> (a, Bool)
 272                      h dh r s = let deq = eqFromNum dh
 273                                 in (+ dh r s, g deq r s)
 274               in
 275               h df x x
 276
 277 After specialising h we get a specialised version of h, like this:
 278
 279                     h' r s = let deq = eqFromNum df
 280                              in (+ df r s, g deq r s)
 281
 282 But we can't naively make an instance for g from this, because deq is not in scope
 283 at the defn of g.  Instead, we have to float out the (new) defn of deq
 284 to widen its scope.  Notice that this floating can't be done in advance -- it only
 285 shows up when specialisation is done.
 286
 287 DELICATE MATTER: the way we tell a dictionary binding is by looking to
 288 see if it has a Dict type.  If the type has been "undictify'd", so that
 289 it looks like a tuple, then the dictionary binding won't be floated, and
 290 an opportunity to specialise might be lost.
 291
 292 User SPECIALIZE pragmas
 293 ~~~~~~~~~~~~~~~~~~~~~~~
 294 Specialisation pragmas can be digested by the type checker, and implemented
 295 by adding extra definitions along with that of f, in the same way as before
 296
 297         f@t1/t2 = <f_rhs> t1 t2 d1 d2
 298
 299 Indeed the pragmas *have* to be dealt with by the type checker, because
 300 only it knows how to build the dictionaries d1 and d2!  For example
 301
 302         g :: Ord a => [a] -> [a]
 303         {-# SPECIALIZE f :: [Tree Int] -> [Tree Int] #-}
 304
 305 Here, the specialised version of g is an application of g's rhs to the
 306 Ord dictionary for (Tree Int), which only the type checker can conjure
 307 up.  There might not even *be* one, if (Tree Int) is not an instance of
 308 Ord!  (All the other specialision has suitable dictionaries to hand
 309 from actual calls.)
 310
 311 Problem.  The type checker doesn't have to hand a convenient <f_rhs>, because
 312 it is buried in a complex (as-yet-un-desugared) binding group.
 313 Maybe we should say
 314
 315         f@t1/t2 = f* t1 t2 d1 d2
 316
 317 where f* is the Id f with an IdInfo which says "inline me regardless!".
 318 Indeed all the specialisation could be done in this way.
 319 That in turn means that the simplifier has to be prepared to inline absolutely
 320 any in-scope let-bound thing.
 321
 322
 323 Again, the pragma should permit polymorphism in unconstrained variables:
 324
 325         h :: Ord a => [a] -> b -> b
 326         {-# SPECIALIZE h :: [Int] -> b -> b #-}
 327
 328 We *insist* that all overloaded type variables are specialised to ground types,
 329 (and hence there can be no context inside a SPECIALIZE pragma).
 330 We *permit* unconstrained type variables to be specialised to
 331         - a ground type
 332         - or left as a polymorphic type variable
 333 but nothing in between.  So
 334
 335         {-# SPECIALIZE h :: [Int] -> [c] -> [c] #-}
 336
 337 is *illegal*.  (It can be handled, but it adds complication, and gains the
 338 programmer nothing.)
 339
 340
 341 SPECIALISING INSTANCE DECLARATIONS
 342 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 343 Consider
 344
 345         instance Foo a => Foo [a] where
 346                 ...
 347         {-# SPECIALIZE instance Foo [Int] #-}
 348
 349 The original instance decl creates a dictionary-function
 350 definition:
 351
 352         dfun.Foo.List :: forall a. Foo a -> Foo [a]
 353
 354 The SPECIALIZE pragma just makes a specialised copy, just as for
 355 ordinary function definitions:
 356
 357         dfun.Foo.List@Int :: Foo [Int]
 358         dfun.Foo.List@Int = dfun.Foo.List Int dFooInt
 359
 360 The information about what instance of the dfun exist gets added to
 361 the dfun's IdInfo in the same way as a user-defined function too.
 362
 363 In fact, matters are a little bit more complicated than this.
 364 When we make one of these specialised instances, we are defining
 365 a constant dictionary, and so we want immediate access to its constant
 366 methods and superclasses.  Indeed, these constant methods and superclasses
 367 must be in the IdInfo for the class selectors!  We need help from the
 368 typechecker to sort this out, perhaps by generating a separate IdInfo
 369 for each.
 370
 371 Automatic instance decl specialisation?
 372 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 373 Can instance decls be specialised automatically?  It's tricky.
 374 We could collect call-instance information for each dfun, but
 375 then when we specialised their bodies we'd get new call-instances
 376 for ordinary functions; and when we specialised their bodies, we might get
 377 new call-instances of the dfuns, and so on.  This all arises because of
 378 the unrestricted mutual recursion between instance decls and value decls.
 379
 380 Furthermore, instance decls are usually exported and used non-locally,
 381 so we'll want to compile enough to get those specialisations done.
 382
 383 Lastly, there's no such thing as a local instance decl, so we can
 384 survive solely by spitting out *usage* information, and then reading that
 385 back in as a pragma when next compiling the file.  So for now,
 386 we only specialise instance decls in response to pragmas.
 387
 388 That means that even if an instance decl ain't otherwise exported it
 389 needs to be spat out as with a SPECIALIZE pragma.  Furthermore, it needs
 390 something to say which module defined the instance, so the usage info
 391 can be fed into the right reqts info file.  Blegh.
 392
 393
 394 SPECIAILISING DATA DECLARATIONS
 395 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 396
 397 With unboxed specialisation (or full specialisation) we also require
 398 data types (and their constructors) to be speciailised on unboxed
 399 type arguments.
 400
 401 In addition to normal call instances we gather TyCon call instances at
 402 unboxed types, determine equivalence classes for the locally defined
 403 TyCons and build speciailised data constructor Ids for each TyCon and
 404 substitute these in the Con calls.
 405
 406 We need the list of local TyCons to partition the TyCon instance info.
 407 We pass out a FiniteMap from local TyCons to Specialised Instances to
 408 give to the interface and code genertors.
 409
 410 N.B. The specialised data constructors reference the original data
 411 constructor and type constructor which do not have the updated
 412 specialisation info attached.  Any specialisation info must be
 413 extracted from the TyCon map returned.
 414
 415
 416 SPITTING OUT USAGE INFORMATION
 417 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 418
 419 To spit out usage information we need to traverse the code collecting
 420 call-instance information for all imported (non-prelude?) functions
 421 and data types. Then we equivalence-class it and spit it out.
 422
 423 This is done at the top-level when all the call instances which escape
 424 must be for imported functions and data types.
 425
 426
 427 Partial specialisation by pragmas
 428 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 429 What about partial specialisation:
 430
 431         k :: (Ord a, Eq b) => [a] -> b -> b -> [a]
 432         {-# SPECIALIZE k :: Eq b => [Int] -> b -> b -> [a] #-}
 433
 434 or even
 435
 436         {-# SPECIALIZE k :: Eq b => [Int] -> [b] -> [b] -> [a] #-}
 437
 438 Seems quite reasonable.  Similar things could be done with instance decls:
 439
 440         instance (Foo a, Foo b) => Foo (a,b) where
 441                 ...
 442         {-# SPECIALIZE instance Foo a => Foo (a,Int) #-}
 443         {-# SPECIALIZE instance Foo b => Foo (Int,b) #-}
 444
 445 Ho hum.  Things are complex enough without this.  I pass.
 446
 447
 448 Requirements for the simplifer
 449 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 450 The simplifier has to be able to take advantage of the specialisation.
 451
 452 * When the simplifier finds an application of a polymorphic f, it looks in
 453 f's IdInfo in case there is a suitable instance to call instead.  This converts
 454
 455         f t1 t2 d1 d2   ===>   f_t1_t2
 456
 457 Note that the dictionaries get eaten up too!
 458
 459 * Dictionary selection operations on constant dictionaries must be
 460   short-circuited:
 461
 462         +.sel Int d     ===>  +Int
 463
 464 The obvious way to do this is in the same way as other specialised
 465 calls: +.sel has inside it some IdInfo which tells that if it's applied
 466 to the type Int then it should eat a dictionary and transform to +Int.
 467
 468 In short, dictionary selectors need IdInfo inside them for constant
 469 methods.
 470
 471 * Exactly the same applies if a superclass dictionary is being
 472   extracted:
 473
 474         Eq.sel Int d   ===>   dEqInt
 475
 476 * Something similar applies to dictionary construction too.  Suppose
 477 dfun.Eq.List is the function taking a dictionary for (Eq a) to
 478 one for (Eq [a]).  Then we want
 479
 480         dfun.Eq.List Int d      ===> dEq.List_Int
 481
 482 Where does the Eq [Int] dictionary come from?  It is built in
 483 response to a SPECIALIZE pragma on the Eq [a] instance decl.
 484
 485 In short, dfun Ids need IdInfo with a specialisation for each
 486 constant instance of their instance declaration.
 487
 488
 489 What does the specialisation IdInfo look like?
 490 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 491
 492         SpecInfo
 493                 [Maybe Type] -- Instance types
 494                 Int             -- No of dicts to eat
 495                 Id              -- Specialised version
 496
 497 For example, if f has this SpecInfo:
 498
 499         SpecInfo [Just t1, Nothing, Just t3] 2 f'
 500
 501 then
 502
 503         f t1 t2 t3 d1 d2  ===>  f t2
 504
 505 The "Nothings" identify type arguments in which the specialised
 506 version is polymorphic.
 507
 508 What can't be done this way?
 509 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 510 There is no way, post-typechecker, to get a dictionary for (say)
 511 Eq a from a dictionary for Eq [a].  So if we find
 512
 513         ==.sel [t] d
 514
 515 we can't transform to
 516
 517         eqList (==.sel t d')
 518
 519 where
 520         eqList :: (a->a->Bool) -> [a] -> [a] -> Bool
 521
 522 Of course, we currently have no way to automatically derive
 523 eqList, nor to connect it to the Eq [a] instance decl, but you
 524 can imagine that it might somehow be possible.  Taking advantage
 525 of this is permanently ruled out.
 526
 527 Still, this is no great hardship, because we intend to eliminate
 528 overloading altogether anyway!
 529
 530
 531 Mutter mutter
 532 ~~~~~~~~~~~~~
 533 What about types/classes mentioned in SPECIALIZE pragmas spat out,
 534 but not otherwise exported.  Even if they are exported, what about
 535 their original names.
 536
 537 Suggestion: use qualified names in pragmas, omitting module for
 538 prelude and "this module".
 539
 540
 541 Mutter mutter 2
 542 ~~~~~~~~~~~~~~~
 543 Consider this
 544
 545         f a (d::Num a) = let g = ...
 546                          in
 547                          ...(let d1::Ord a = Num.Ord.sel a d in g a d1)...
 548
 549 Here, g is only called at one type, but the dictionary isn't in scope at the
 550 definition point for g.  Usually the type checker would build a
 551 definition for d1 which enclosed g, but the transformation system
 552 might have moved d1's defn inward.
 553
 554
 555 Unboxed bindings
 556 ~~~~~~~~~~~~~~~~
 557
 558 What should we do when a value is specialised to a *strict* unboxed value?
 559
 560         map_*_* f (x:xs) = let h = f x
 561                                t = map f xs
 562                            in h:t
 563
 564 Could convert let to case:
 565
 566         map_*_Int# f (x:xs) = case f x of h# ->
 567                               let t = map f xs
 568                               in h#:t
 569
 570 This may be undesirable since it forces evaluation here, but the value
 571 may not be used in all branches of the body. In the general case this
 572 transformation is impossible since the mutual recursion in a letrec
 573 cannot be expressed as a case.
 574
 575 There is also a problem with top-level unboxed values, since our
 576 implementation cannot handle unboxed values at the top level.
 577
 578 Solution: Lift the binding of the unboxed value and extract it when it
 579 is used:
 580
 581         map_*_Int# f (x:xs) = let h = case (f x) of h# -> _Lift h#
 582                                   t = map f xs
 583                               in case h of
 584                                  _Lift h# -> h#:t
 585
 586 Now give it to the simplifier and the _Lifting will be optimised away.
 587
 588 The benfit is that we have given the specialised "unboxed" values a
 589 very simplep lifted semantics and then leave it up to the simplifier to
 590 optimise it --- knowing that the overheads will be removed in nearly
 591 all cases.
 592
 593 In particular, the value will only be evaluted in the branches of the
 594 program which use it, rather than being forced at the point where the
 595 value is bound. For example:
 596
 597         filtermap_*_* p f (x:xs)
 598           = let h = f x
 599                 t = ...
 600             in case p x of
 601                 True  -> h:t
 602                 False -> t
 603    ==>
 604         filtermap_*_Int# p f (x:xs)
 605           = let h = case (f x) of h# -> _Lift h#
 606                 t = ...
 607             in case p x of
 608                 True  -> case h of _Lift h#
 609                            -> h#:t
 610                 False -> t
 611
 612 The binding for h can still be inlined in the one branch and the
 613 _Lifting eliminated.
 614
 615
 616 Question: When won't the _Lifting be eliminated?
 617
 618 Answer: When they at the top-level (where it is necessary) or when
 619 inlining would duplicate work (or possibly code depending on
 620 options). However, the _Lifting will still be eliminated if the
 621 strictness analyser deems the lifted binding strict.
 622
 623
 624 A note about non-tyvar dictionaries
 625 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 626 Some Ids have types like
 627
 628         forall a,b,c. Eq a -> Ord [a] -> tau
 629
 630 This seems curious at first, because we usually only have dictionary
 631 args whose types are of the form (C a) where a is a type variable.
 632 But this doesn't hold for the functions arising from instance decls,
 633 which sometimes get arguements with types of form (C (T a)) for some
 634 type constructor T.
 635
 636 Should we specialise wrt this compound-type dictionary?  We used to say
 637 "no", saying:
 638         "This is a heuristic judgement, as indeed is the fact that we
 639         specialise wrt only dictionaries.  We choose *not* to specialise
 640         wrt compound dictionaries because at the moment the only place
 641         they show up is in instance decls, where they are simply plugged
 642         into a returned dictionary.  So nothing is gained by specialising
 643         wrt them."
 644
 645 But it is simpler and more uniform to specialise wrt these dicts too;
 646 and in future GHC is likely to support full fledged type signatures
 647 like
 648         f ;: Eq [(a,b)] => ...
 649
 650
 651 %************************************************************************
 652 %*                                                                      *
 653 \subsubsection{The new specialiser}
 654 %*                                                                      *
 655 %************************************************************************
 656
 657 Our basic game plan is this.  For let(rec) bound function
 658         f :: (C a, D c) => (a,b,c,d) -> Bool
 659
 660 * Find any specialised calls of f, (f ts ds), where
 661   ts are the type arguments t1 .. t4, and
 662   ds are the dictionary arguments d1 .. d2.
 663
 664 * Add a new definition for f1 (say):
 665
 666         f1 = /\ b d -> (..body of f..) t1 b t3 d d1 d2
 667
 668   Note that we abstract over the unconstrained type arguments.
 669
 670 * Add the mapping
 671
 672         [t1,b,t3,d]  |->  \d1 d2 -> f1 b d
 673
 674   to the specialisations of f.  This will be used by the
 675   simplifier to replace calls
 676                 (f t1 t2 t3 t4) da db
 677   by
 678                 (\d1 d1 -> f1 t2 t4) da db
 679
 680   All the stuff about how many dictionaries to discard, and what types
 681   to apply the specialised function to, are handled by the fact that the
 682   SpecEnv contains a template for the result of the specialisation.
 683
 684 We don't build *partial* specialisations for f.  For example:
 685
 686   f :: Eq a => a -> a -> Bool
 687   {-# SPECIALISE f :: (Eq b, Eq c) => (b,c) -> (b,c) -> Bool #-}
 688
 689 Here, little is gained by making a specialised copy of f.
 690 There's a distinct danger that the specialised version would
 691 first build a dictionary for (Eq b, Eq c), and then select the (==)
 692 method from it!  Even if it didn't, not a great deal is saved.
 693
 694 We do, however, generate polymorphic, but not overloaded, specialisations:
 695
 696   f :: Eq a => [a] -> b -> b -> b
 697   {#- SPECIALISE f :: [Int] -> b -> b -> b #-}
 698
 699 Hence, the invariant is this:
 700
 701         *** no specialised version is overloaded ***
 702
 703
 704 %************************************************************************
 705 %*                                                                      *
 706 \subsubsection{The exported function}
 707 %*                                                                      *
 708 %************************************************************************
 709
 710 \begin{code}
 711 specProgram :: UniqSupply -> [CoreBinding] -> [CoreBinding]
 712 specProgram us binds
 713   = initSM us (go binds         `thenSM` \ (binds', _) ->
 714                returnSM binds'
 715               )
 716   where
 717     go []           = returnSM ([], emptyUDs)
 718     go (bind:binds) = go binds          `thenSM` \ (binds', uds) ->
 719                       specBind bind uds `thenSM` \ (bind', uds') ->
 720                       returnSM (bind' ++ binds', uds')
 721 \end{code}
 722
 723 %************************************************************************
 724 %*                                                                      *
 725 \subsubsection{@specExpr@: the main function}
 726 %*                                                                      *
 727 %************************************************************************
 728
 729 \begin{code}
 730 specExpr :: CoreExpr -> SpecM (CoreExpr, UsageDetails)
 731
 732 ---------------- First the easy cases --------------------
 733 specExpr e@(Var _)    = returnSM (e, emptyUDs)
 734 specExpr e@(Lit _)    = returnSM (e, emptyUDs)
 735 specExpr e@(Con _ _)  = returnSM (e, emptyUDs)
 736 specExpr e@(Prim _ _) = returnSM (e, emptyUDs)
 737
 738 specExpr (Coerce co ty body)
 739   = specExpr body       `thenSM` \ (body', uds) ->
 740     returnSM (Coerce co ty body', uds)
 741
 742 specExpr (SCC cc body)
 743   = specExpr body       `thenSM` \ (body', uds) ->
 744     returnSM (SCC cc body', uds)
 745
 746
 747 ---------------- Applications might generate a call instance --------------------
 748 specExpr e@(App fun arg)
 749   = go fun [arg]
 750   where
 751     go (App fun arg) args = go fun (arg:args)
 752     go (Var f)       args = returnSM (e, mkCallUDs f args)
 753     go other         args = specExpr other      `thenSM` \ (e', uds) ->
 754                             returnSM (foldl App e' args, uds)
 755
 756 ---------------- Lambda/case require dumping of usage details --------------------
 757 specExpr e@(Lam _ _)
 758   = specExpr body       `thenSM` \ (body', uds) ->
 759     let
 760         (filtered_uds, body'') = dumpUDs bndrs uds body'
 761     in
 762     returnSM (foldr Lam body'' bndrs, filtered_uds)
 763   where
 764     (bndrs, body) = go [] e
 765
 766         -- More efficient to collect a group of binders together all at once
 767     go bndrs (Lam bndr e) = go (bndr:bndrs) e
 768     go bndrs e            = (reverse bndrs, e)
 769
 770
 771 specExpr (Case scrut alts)
 772   = specExpr scrut      `thenSM` \ (scrut', uds_scrut) ->
 773     spec_alts alts      `thenSM` \ (alts', uds_alts) ->
 774     returnSM (Case scrut' alts', uds_scrut `plusUDs` uds_alts)
 775   where
 776     spec_alts (AlgAlts alts deflt)
 777         = mapAndCombineSM spec_alg_alt alts     `thenSM` \ (alts', uds1) ->
 778           spec_deflt deflt                      `thenSM` \ (deflt', uds2) ->
 779           returnSM (AlgAlts alts' deflt', uds1 `plusUDs` uds2)
 780
 781     spec_alts (PrimAlts alts deflt)
 782         = mapAndCombineSM spec_prim_alt alts    `thenSM` \ (alts', uds1) ->
 783           spec_deflt deflt                      `thenSM` \ (deflt', uds2) ->
 784           returnSM (PrimAlts alts' deflt', uds1 `plusUDs` uds2)
 785
 786     spec_alg_alt (con, args, rhs)
 787         = specExpr rhs          `thenSM` \ (rhs', uds) ->
 788           let
 789              (uds', rhs'') = dumpUDs (map ValBinder args) uds rhs'
 790           in
 791           returnSM ((con, args, rhs''), uds')
 792
 793     spec_prim_alt (lit, rhs)
 794         = specExpr rhs          `thenSM` \ (rhs', uds) ->
 795           returnSM ((lit, rhs'), uds)
 796
 797     spec_deflt NoDefault = returnSM (NoDefault, emptyUDs)
 798     spec_deflt (BindDefault arg rhs)
 799         = specExpr rhs          `thenSM` \ (rhs', uds) ->
 800           let
 801              (uds', rhs'') = dumpUDs [ValBinder arg] uds rhs'
 802           in
 803           returnSM (BindDefault arg rhs'', uds')
 804
 805 ---------------- Finally, let is the interesting case --------------------
 806 specExpr (Let bind body)
 807   =     -- Deal with the body
 808     specExpr body                               `thenSM` \ (body', body_uds) ->
 809
 810         -- Deal with the bindings
 811     specBind bind body_uds                      `thenSM` \ (binds', uds) ->
 812
 813         -- All done
 814     returnSM (foldr Let body' binds', uds)
 815 \end{code}
 816
 817 %************************************************************************
 818 %*                                                                      *
 819 \subsubsection{Dealing with a binding}
 820 %*                                                                      *
 821 %************************************************************************
 822
 823 \begin{code}
 824 specBind :: CoreBinding
 825          -> UsageDetails                -- Info on how the scope of the binding
 826          -> SpecM ([CoreBinding],       -- New bindings
 827                    UsageDetails)        -- And info to pass upstream
 828
 829 specBind (NonRec bndr rhs) body_uds
 830   | isDictTy (idType bndr)
 831   =     -- It's a dictionary binding
 832         -- Pick it up and float it outwards.
 833     specExpr rhs                                `thenSM` \ (rhs', rhs_uds) ->
 834     let
 835         all_uds = rhs_uds `plusUDs` addDictBind body_uds bndr rhs'
 836     in
 837     returnSM ([], all_uds)
 838
 839   | otherwise
 840   =   -- Deal with the RHS, specialising it according
 841       -- to the calls found in the body
 842     specDefn (calls body_uds) (bndr,rhs)        `thenSM` \ ((bndr',rhs'), spec_defns, spec_uds) ->
 843     let
 844         (all_uds, (dict_binds, dump_calls))
 845                 = splitUDs [ValBinder bndr] (spec_uds `plusUDs` body_uds)
 846
 847         -- If we make specialisations then we Rec the whole lot together
 848         -- If not, leave it as a NonRec
 849         new_bind | null spec_defns = NonRec bndr' rhs'
 850                  | otherwise       = Rec ((bndr',rhs'):spec_defns)
 851     in
 852     returnSM ( new_bind : dict_binds, all_uds )
 853
 854 specBind (Rec pairs) body_uds
 855   = mapSM (specDefn (calls body_uds)) pairs     `thenSM` \ stuff ->
 856     let
 857         (pairs', spec_defns_s, spec_uds_s) = unzip3 stuff
 858         spec_defns = concat spec_defns_s
 859         spec_uds   = plusUDList spec_uds_s
 860         (all_uds, (dict_binds, dump_calls))
 861                 = splitUDs (map (ValBinder . fst) pairs) (spec_uds `plusUDs` body_uds)
 862         new_bind = Rec (spec_defns ++ pairs')
 863     in
 864     returnSM (  new_bind : dict_binds, all_uds )
 865
 866 specDefn :: CallDetails                 -- Info on how it is used in its scope
 867          -> (Id, CoreExpr)              -- The thing being bound and its un-processed RHS
 868          -> SpecM ((Id, CoreExpr),      -- The thing and its processed RHS
 869                                         --      the Id may now have specialisations attached
 870                    [(Id,CoreExpr)],     -- Extra, specialised bindings
 871                    UsageDetails         -- Stuff to fling upwards from the RHS and its
 872             )                           --      specialised versions
 873
 874 specDefn calls (fn, rhs)
 875         -- The first case is the interesting one
 876   |  n_tyvars == length rhs_tyvars      -- Rhs of fn's defn has right number of big lambdas
 877   && n_dicts  <= length rhs_bndrs       -- and enough dict args
 878   && not (null calls_for_me)            -- And there are some calls to specialise
 879   =   -- Specialise the body of the function
 880     specExpr body                                       `thenSM` \ (body', body_uds) ->
 881     let
 882         (float_uds, bound_uds@(dict_binds,_)) = splitUDs rhs_bndrs body_uds
 883     in
 884
 885       -- Make a specialised version for each call in calls_for_me
 886     mapSM (spec_call bound_uds) calls_for_me            `thenSM` \ stuff ->
 887     let
 888         (spec_defns, spec_uds, spec_env_stuff) = unzip3 stuff
 889
 890         fn'  = addIdSpecialisations fn spec_env_stuff
 891         rhs' = foldr Lam (foldr Let body' dict_binds) rhs_bndrs
 892     in
 893     returnSM ((fn',rhs'),
 894               spec_defns,
 895               float_uds `plusUDs` plusUDList spec_uds)
 896
 897   | otherwise   -- No calls or RHS doesn't fit our preconceptions
 898   = specExpr rhs                        `thenSM` \ (rhs', rhs_uds) ->
 899     returnSM ((fn, rhs'), [], rhs_uds)
 900
 901   where
 902     fn_type              = idType fn
 903     (tyvars, theta, tau) = splitSigmaTy fn_type
 904     n_tyvars             = length tyvars
 905     n_dicts              = length theta
 906     mk_spec_tys call_ts  = zipWith mk_spec_ty call_ts tyvars
 907                          where
 908                            mk_spec_ty (Just ty) _     = ty
 909                            mk_spec_ty Nothing   tyvar = mkTyVarTy tyvar
 910
 911     (rhs_tyvars, rhs_ids, rhs_body) = collectBinders rhs
 912     rhs_dicts = take n_dicts rhs_ids
 913     rhs_bndrs = map TyBinder rhs_tyvars ++ map ValBinder rhs_dicts
 914     body      = mkValLam (drop n_dicts rhs_ids) rhs_body
 915                 -- Glue back on the non-dict lambdas
 916
 917     calls_for_me = case lookupFM calls fn of
 918                         Nothing -> []
 919                         Just cs -> fmToList cs
 920
 921     -- Filter out calls for which we already have a specialisation
 922     calls_to_spec        = filter spec_me calls_for_me
 923     spec_me (call_ts, _) = not (maybeToBool (lookupSpecEnv id_spec_env (mk_spec_tys call_ts)))
 924     id_spec_env          = getIdSpecialisation fn
 925
 926     ----------------------------------------------------------
 927         -- Specialise to one particular call pattern
 928     spec_call :: ProtoUsageDetails          -- From the original body, captured by
 929                                             -- the dictionary lambdas
 930               -> ([Maybe Type], [DictVar])  -- Call instance
 931               -> SpecM ((Id,CoreExpr),            -- Specialised definition
 932                         UsageDetails,             -- Usage details from specialised body
 933                         ([TyVar], [Type], CoreExpr))       -- Info for the Id's SpecEnv
 934     spec_call bound_uds (call_ts, call_ds)
 935       = ASSERT( length call_ts == n_tyvars && length call_ds == n_dicts )
 936                 -- Calls are only recorded for properly-saturated applications
 937
 938         -- Supppose the call is for f [Just t1, Nothing, Just t3, Nothing] [d1, d2]
 939
 940                 -- Construct the new binding
 941                 --      f1 = /\ b d -> (..rhs of f..) t1 b t3 d d1 d2
 942                 -- and the type of this binder
 943         let
 944            spec_tyvars = [tyvar | (tyvar, Nothing) <- tyvars `zip` call_ts]
 945            spec_tys    = mk_spec_tys call_ts
 946            spec_rhs    = mkTyLam spec_tyvars $
 947                          mkGenApp rhs (map TyArg spec_tys ++ map VarArg call_ds)
 948            spec_id_ty  = mkForAllTys spec_tyvars (instantiateTy ty_env tau)
 949            ty_env      = mkTyVarEnv (zipEqual "spec_call" tyvars spec_tys)
 950         in
 951         newIdSM fn spec_id_ty           `thenSM` \ spec_f ->
 952
 953
 954                 -- Construct the stuff for f's spec env
 955                 --      [b,d] [t1,b,t3,d]  |->  \d1 d2 -> f1 b d
 956         let
 957            spec_env_rhs  = mkValLam call_ds $
 958                            mkTyApp (Var spec_f) $
 959                            map mkTyVarTy spec_tyvars
 960            spec_env_info = (spec_tyvars, spec_tys, spec_env_rhs)
 961         in
 962
 963                 -- Specialise the UDs from f's RHS
 964         let
 965                 -- Only the overloaded tyvars should be free in the uds
 966            ty_env   = [ (rhs_tyvar,ty)
 967                       | (rhs_tyvar, Just ty) <- zipEqual "specUDs1" rhs_tyvars call_ts
 968                       ]
 969            dict_env = zipEqual "specUDs2" rhs_dicts call_ds
 970         in
 971         specUDs ty_env dict_env bound_uds                       `thenSM` \ spec_uds ->
 972
 973         returnSM ((spec_f, spec_rhs),
 974                   spec_uds,
 975                   spec_env_info
 976         )
 977 \end{code}
 978
 979 %************************************************************************
 980 %*                                                                      *
 981 \subsubsection{UsageDetails and suchlike}
 982 %*                                                                      *
 983 %************************************************************************
 984
 985 \begin{code}
 986 type FreeDicts = IdSet
 987
 988 data UsageDetails
 989   = MkUD {
 990         dict_binds :: !(Bag (DictVar, CoreExpr, TyVarSet, FreeDicts)),
 991                         -- Floated dictionary bindings
 992                         -- The order is important;
 993                         -- in ds1 `union` ds2, bindings in ds2 can depend on those in ds1
 994                         -- (Remember, Bags preserve order in GHC.)
 995                         -- The FreeDicts is the free vars of the RHS
 996
 997         calls     :: !CallDetails
 998     }
 999
1000 emptyUDs = MkUD { dict_binds = emptyBag, calls = emptyFM }
1001
1002 type ProtoUsageDetails = ([CoreBinding],                -- Dict bindings
1003                           [(Id, [Maybe Type], [DictVar])]
1004                          )
1005
1006 ------------------------------------------------------------
1007 type CallDetails  = FiniteMap Id CallInfo
1008 type CallInfo     = FiniteMap [Maybe Type]      -- Nothing => unconstrained type argument
1009                               [DictVar]         -- Dict args
1010         -- The finite maps eliminate duplicates
1011         -- The list of types and dictionaries is guaranteed to
1012         -- match the type of f
1013
1014 callDetailsToList calls = [ (id,tys,dicts)
1015                           | (id,fm) <- fmToList calls,
1016                             (tys,dicts) <- fmToList fm
1017                           ]
1018
1019 listToCallDetails calls  = foldr (unionCalls . singleCall) emptyFM calls
1020
1021 unionCalls :: CallDetails -> CallDetails -> CallDetails
1022 unionCalls c1 c2 = plusFM_C plusFM c1 c2
1023
1024 singleCall (id, tys, dicts) = unitFM id (unitFM tys dicts)
1025
1026 mkCallUDs f args
1027   | null theta
1028   || length spec_tys /= n_tyvars
1029   || length dicts    /= n_dicts
1030   = emptyUDs    -- Not overloaded
1031
1032   | otherwise
1033   = MkUD {dict_binds = emptyBag,
1034           calls = singleCall (f, spec_tys, dicts)
1035     }
1036   where
1037     (tyvars, theta, tau) = splitSigmaTy (idType f)
1038     constrained_tyvars   = foldr (unionTyVarSets . tyVarsOfTypes . snd) emptyTyVarSet theta
1039     n_tyvars             = length tyvars
1040     n_dicts              = length theta
1041
1042     spec_tys = [mk_spec_ty tv ty | (tv, TyArg ty) <- tyvars `zip` args]
1043     dicts    = [d | (_, VarArg d) <- theta `zip` (drop n_tyvars args)]
1044
1045     mk_spec_ty tyvar ty | tyvar `elementOfTyVarSet` constrained_tyvars
1046                         = Just ty
1047                         | otherwise
1048                         = Nothing
1049
1050 ------------------------------------------------------------
1051 plusUDs :: UsageDetails -> UsageDetails -> UsageDetails
1052 plusUDs (MkUD {dict_binds = db1, calls = calls1})
1053         (MkUD {dict_binds = db2, calls = calls2})
1054   = MkUD {dict_binds, calls}
1055   where
1056     dict_binds = db1    `unionBags`   db2
1057     calls      = calls1 `unionCalls`  calls2
1058
1059 plusUDList = foldr plusUDs emptyUDs
1060
1061 mkDB dict rhs = (dict, rhs, db_ftvs, db_fvs)
1062               where
1063                 db_ftvs = tyVarsOfType (idType dict)    -- Superset of RHS fvs
1064                 db_fvs  = dictRhsFVs rhs
1065
1066 addDictBind uds dict rhs = uds { dict_binds = mkDB dict rhs `consBag` dict_binds uds }
1067
1068 dumpUDs :: [CoreBinder]
1069         -> UsageDetails -> CoreExpr
1070         -> (UsageDetails, CoreExpr)
1071 dumpUDs bndrs uds body
1072   = (free_uds, foldr Let body dict_binds)
1073   where
1074     (free_uds, (dict_binds, _)) = splitUDs bndrs uds
1075
1076 splitUDs :: [CoreBinder]
1077          -> UsageDetails
1078          -> (UsageDetails,              -- These don't mention the binders
1079              ProtoUsageDetails)         -- These do
1080
1081 splitUDs bndrs uds@(MkUD {dict_binds = orig_dbs,
1082                           calls      = orig_calls})
1083
1084   = if isEmptyBag dump_dbs && null dump_calls then
1085         -- Common case: binder doesn't affect floats
1086         (uds, ([],[]))
1087
1088     else
1089         -- Binders bind some of the fvs of the floats
1090         (MkUD {dict_binds = free_dbs,
1091                calls      = listToCallDetails free_calls},
1092          (bagToList dump_dbs, dump_calls)
1093         )
1094
1095   where
1096     tyvar_set    = mkTyVarSet [tv | TyBinder tv <- bndrs]
1097     id_set       = mkIdSet    [id | ValBinder id <- bndrs]
1098
1099     (free_dbs, dump_dbs, dump_idset)
1100           = foldlBag dump_db (emptyBag, emptyBag, id_set) orig_dbs
1101                 -- Important that it's foldl not foldr;
1102                 -- we're accumulating the set of dumped ids in dump_set
1103
1104         -- Filter out any calls that mention things that are being dumped
1105         -- Don't need to worry about the tyvars because the dicts will
1106         -- spot the captured ones; any fully polymorphic arguments will
1107         -- be Nothings in the call details
1108     orig_call_list = callDetailsToList orig_calls
1109     (dump_calls, free_calls) = partition captured orig_call_list
1110     captured (id,tys,dicts)  = any (`elementOfIdSet` dump_idset) (id:dicts)
1111
1112     dump_db (free_dbs, dump_dbs, dump_idset) db@(dict, rhs, ftvs, fvs)
1113         |  isEmptyIdSet    (dump_idset `intersectIdSets`    fvs)
1114         && isEmptyTyVarSet (tyvar_set  `intersectTyVarSets` ftvs)
1115         = (free_dbs `snocBag` db, dump_dbs, dump_idset)
1116
1117         | otherwise     -- Dump it
1118         = (free_dbs, dump_dbs `snocBag` NonRec dict rhs,
1119            dump_idset `addOneToIdSet` dict)
1120 \end{code}
1121
1122 Given a type and value substitution, specUDs creates a specialised copy of
1123 the given UDs
1124
1125 \begin{code}
1126 specUDs :: [(TyVar,Type)] -> [(DictVar,DictVar)] -> ProtoUsageDetails -> SpecM UsageDetails
1127 specUDs tv_env_list dict_env_list (dbs, calls)
1128   = specDBs dict_env dbs                `thenSM` \ (dict_env', dbs') ->
1129     returnSM (MkUD { dict_binds = dbs',
1130                      calls      = listToCallDetails (map (inst_call dict_env') calls)
1131     })
1132   where
1133     tv_env   = mkTyVarEnv tv_env_list
1134     dict_env = mkIdEnv dict_env_list
1135
1136     inst_call dict_env (id, tys, dicts) = (id, map inst_maybe_ty tys,
1137                                                map (lookupId dict_env) dicts)
1138
1139     inst_maybe_ty Nothing   = Nothing
1140     inst_maybe_ty (Just ty) = Just (instantiateTy tv_env ty)
1141
1142     specDBs dict_env []
1143         = returnSM (dict_env, emptyBag)
1144     specDBs dict_env (NonRec dict rhs : dbs)
1145         = newIdSM dict (instantiateTy tv_env (idType dict))     `thenSM` \ dict' ->
1146           let
1147             dict_env' = addOneToIdEnv dict_env dict dict'
1148             rhs'      = instantiateDictRhs tv_env dict_env rhs
1149           in
1150           specDBs dict_env' dbs         `thenSM` \ (dict_env'', dbs') ->
1151           returnSM ( dict_env'', mkDB dict' rhs' `consBag` dbs' )
1152 \end{code}
1153
1154 %************************************************************************
1155 %*                                                                      *
1156 \subsubsection{Boring helper functions}
1157 %*                                                                      *
1158 %************************************************************************
1159
1160 \begin{code}
1161 lookupId:: IdEnv Id -> Id -> Id
1162 lookupId env id = case lookupIdEnv env id of
1163                         Nothing  -> id
1164                         Just id' -> id'
1165
1166 instantiateDictRhs :: TyVarEnv Type -> IdEnv Id -> CoreExpr -> CoreExpr
1167         -- Cheapo function for simple RHSs
1168 instantiateDictRhs ty_env id_env rhs
1169   = go rhs
1170   where
1171     go (App e1 (VarArg a)) = App (go e1) (VarArg (lookupId id_env a))
1172     go (App e1 (TyArg t))  = App (go e1) (TyArg (instantiateTy ty_env t))
1173     go (Var v)             = Var (lookupId id_env v)
1174     go (Lit l)             = Lit l
1175
1176 dictRhsFVs :: CoreExpr -> IdSet
1177         -- Cheapo function for simple RHSs
1178 dictRhsFVs (App e1 (VarArg a)) = dictRhsFVs e1 `addOneToIdSet` a
1179 dictRhsFVs (App e1 (TyArg t))  = dictRhsFVs e1
1180 dictRhsFVs (Var v)             = unitIdSet v
1181 dictRhsFVs (Lit l)             = emptyIdSet
1182
1183
1184 addIdSpecialisations id spec_stuff
1185   = (if not (null errs) then
1186         pprTrace "Duplicate specialisations" (vcat (map ppr errs))
1187      else \x -> x
1188     )
1189     setIdSpecialisation id new_spec_env
1190   where
1191     (new_spec_env, errs) = foldr add (getIdSpecialisation id, []) spec_stuff
1192
1193     add (tyvars, tys, template) (spec_env, errs)
1194         = case addToSpecEnv True spec_env tyvars tys template of
1195                 Succeeded spec_env' -> (spec_env', errs)
1196                 Failed err          -> (spec_env, err:errs)
1197
1198 -- Given an Id, isSpecVars returns all its specialisations.
1199 -- We extract these from its SpecEnv.
1200 -- This is used by the occurrence analyser and free-var finder;
1201 -- we regard an Id's specialisations as free in the Id's definition.
1202
1203 idSpecVars :: Id -> [Id]
1204 idSpecVars id
1205   = map get_spec (specEnvValues (getIdSpecialisation id))
1206   where
1207     -- get_spec is another cheapo function like dictRhsFVs
1208     -- It knows what these specialisation temlates look like,
1209     -- and just goes for the jugular
1210     get_spec (App f _) = get_spec f
1211     get_spec (Lam _ b) = get_spec b
1212     get_spec (Var v)   = v
1213
1214 -- substSpecEnvRhs applies a substitution to the RHS's of a SpecEnv
1215 -- It's placed here because Specialise.lhs built that RHS, so
1216 -- it knows its structure.  (Fully general subst
1217
1218 substSpecEnvRhs te ve rhs
1219   = go te ve rhs
1220   where
1221     go te ve (App f (TyArg ty)) = App (go te ve f) (TyArg (instantiateTy te ty))
1222     go te ve (App f (VarArg v)) = App (go te ve f) (case lookupIdEnv ve v of
1223                                                         Just arg' -> arg'
1224                                                         Nothing   -> VarArg v)
1225     go te ve (Var v)              = case lookupIdEnv ve v of
1226                                                 Just (VarArg v') -> Var v'
1227                                                 Just (LitArg l)  -> Lit l
1228                                                 Nothing          -> Var v
1229
1230         -- These equations are a bit half baked, because
1231         -- they don't deal properly wih capture.
1232         -- But I'm sure it'll never matter... sigh.
1233     go te ve (Lam b@(TyBinder tyvar) e) = Lam b (go te' ve e)
1234                                         where
1235                                           te' = delFromTyVarEnv te tyvar
1236
1237     go te ve (Lam b@(ValBinder v) e) = Lam b (go te ve' e)
1238                                      where
1239                                        ve' = delOneFromIdEnv ve v
1240
1241 ----------------------------------------
1242 type SpecM a = UniqSM a
1243
1244 thenSM    = thenUs
1245 returnSM  = returnUs
1246 getUniqSM = getUnique
1247 mapSM     = mapUs
1248 initSM    = initUs
1249
1250 mapAndCombineSM f []     = returnSM ([], emptyUDs)
1251 mapAndCombineSM f (x:xs) = f x  `thenSM` \ (y, uds1) ->
1252                            mapAndCombineSM f xs `thenSM` \ (ys, uds2) ->
1253                            returnSM (y:ys, uds1 `plusUDs` uds2)
1254
1255 newIdSM old_id new_ty
1256   = getUnique           `thenSM` \ uniq ->
1257     returnSM (mkUserLocal (getOccName old_id)
1258                           uniq
1259                           new_ty
1260                           (getSrcLoc old_id)
1261     )
1262 \end{code}
1263
1264