SpecConstr.lhs 77.1 KB
Newer Older
1
ToDo [Oct 2013]
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
2
~~~~~~~~~~~~~~~
3
1. Nuke ForceSpecConstr for good (it is subsumed by GHC.Types.SPEC in ghc-prim)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
4
5
2. Nuke NoSpecConstr

6
7
8
9
10
11
12
%
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section[SpecConstr]{Specialise over constructors}

\begin{code}
module SpecConstr(
13
        specConstrProgram
14
15
16
#ifdef GHCI
        , SpecConstrAnnotation(..)
#endif
17
18
19
20
21
    ) where

#include "HsVersions.h"

import CoreSyn
22
23
import CoreSubst
import CoreUtils
24
25
import CoreUnfold       ( couldBeSmallEnoughToInline )
import CoreFVs          ( exprsFreeVars )
26
import CoreMonad
27
import Literal          ( litIsLifted )
28
import HscTypes         ( ModGuts(..) )
29
import WwLib            ( mkWorkerArgs )
30
import DataCon
31
import Coercion         hiding( substTy, substCo )
32
import Rules
33
import Type             hiding ( substTy )
34
import TyCon            ( isRecursiveTyCon, tyConName )
35
import Id
36
import MkCore           ( mkImpossibleExpr )
37
import Var
38
39
import VarEnv
import VarSet
40
import Name
41
import BasicTypes
42
43
44
import DynFlags         ( DynFlags(..) )
import StaticFlags      ( opt_PprStyle_Debug )
import Maybes           ( orElse, catMaybes, isJust, isNothing )
45
import Demand
46
import Serialized       ( deserializeWithData )
47
import Util
48
import Pair
49
50
import UniqSupply
import Outputable
51
import FastString
52
import UniqFM
Ian Lynagh's avatar
Ian Lynagh committed
53
import MonadUtils
54
import Control.Monad    ( zipWithM )
Ian Lynagh's avatar
Ian Lynagh committed
55
import Data.List
56
import PrelNames        ( specTyConName )
57

58
-- See Note [Forcing specialisation]
59
60
61
#ifndef GHCI
type SpecConstrAnnotation = ()
#else
Austin Seipp's avatar
Austin Seipp committed
62
import TyCon ( TyCon )
63
64
import GHC.Exts( SpecConstrAnnotation(..) )
#endif
65
66
67
\end{code}

-----------------------------------------------------
68
                        Game plan
69
70
71
-----------------------------------------------------

Consider
72
73
74
        drop n []     = []
        drop 0 xs     = []
        drop n (x:xs) = drop (n-1) xs
75
76
77
78

After the first time round, we could pass n unboxed.  This happens in
numerical code too.  Here's what it looks like in Core:

79
80
81
82
83
84
        drop n xs = case xs of
                      []     -> []
                      (y:ys) -> case n of
                                  I# n# -> case n# of
                                             0 -> []
                                             _ -> drop (I# (n# -# 1#)) xs
85
86
87
88

Notice that the recursive call has an explicit constructor as argument.
Noticing this, we can make a specialised version of drop

89
90
91
        RULE: drop (I# n#) xs ==> drop' n# xs

        drop' n# xs = let n = I# n# in ...orig RHS...
92
93
94

Now the simplifier will apply the specialisation in the rhs of drop', giving

95
96
97
98
        drop' n# xs = case xs of
                      []     -> []
                      (y:ys) -> case n# of
                                  0 -> []
99
                                  _ -> drop' (n# -# 1#) xs
100

101
Much better!
102
103
104
105

We'd also like to catch cases where a parameter is carried along unchanged,
but evaluated each time round the loop:

106
        f i n = if i>0 || i>n then i else f (i*2) n
107
108
109
110

Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
In Core, by the time we've w/wd (f is strict in i) we get

111
112
        f i# n = case i# ># 0 of
                   False -> I# i#
113
                   True  -> case n of { I# n# ->
114
115
                            case i# ># n# of
                                False -> I# i#
116
                                True  -> f (i# *# 2#) n
117

118
At the call to f, we see that the argument, n is known to be (I# n#),
119
and n is evaluated elsewhere in the body of f, so we can play the same
120
trick as above.
121
122
123
124
125


Note [Reboxing]
~~~~~~~~~~~~~~~
We must be careful not to allocate the same constructor twice.  Consider
126
127
        f p = (...(case p of (a,b) -> e)...p...,
               ...let t = (r,s) in ...t...(f t)...)
128
129
At the recursive call to f, we can see that t is a pair.  But we do NOT want
to make a specialised copy:
130
        f' a b = let p = (a,b) in (..., ...)
131
132
133
134
because now t is allocated by the caller, then r and s are passed to the
recursive call, which allocates the (r,s) pair again.

This happens if
135
  (a) the argument p is used in other than a case-scrutinisation way.
136
  (b) the argument to the call is not a 'fresh' tuple; you have to
137
        look into its unfolding to see that it's a tuple
138
139
140

Hence the "OR" part of Note [Good arguments] below.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
141
ALTERNATIVE 2: pass both boxed and unboxed versions.  This no longer saves
142
143
144
145
146
147
148
149
allocation, but does perhaps save evals. In the RULE we'd have
something like

  f (I# x#) = f' (I# x#) x#

If at the call site the (I# x) was an unfolding, then we'd have to
rely on CSE to eliminate the duplicate allocation.... This alternative
doesn't look attractive enough to pursue.
150

151
ALTERNATIVE 3: ignore the reboxing problem.  The trouble is that
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
152
153
the conservative reboxing story prevents many useful functions from being
specialised.  Example:
154
155
156
        foo :: Maybe Int -> Int -> Int
        foo   (Just m) 0 = 0
        foo x@(Just m) n = foo x (n-m)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
157
158
159
Here the use of 'x' will clearly not require boxing in the specialised function.

The strictness analyser has the same problem, in fact.  Example:
160
        f p@(a,b) = ...
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
161
162
163
164
165
166
167
168
169
If we pass just 'a' and 'b' to the worker, it might need to rebox the
pair to create (a,b).  A more sophisticated analysis might figure out
precisely the cases in which this could happen, but the strictness
analyser does no such analysis; it just passes 'a' and 'b', and hopes
for the best.

So my current choice is to make SpecConstr similarly aggressive, and
ignore the bad potential of reboxing.

170

171
172
Note [Good arguments]
~~~~~~~~~~~~~~~~~~~~~
173
174
So we look for

175
* A self-recursive function.  Ignore mutual recursion for now,
176
177
178
179
  because it's less common, and the code is simpler for self-recursion.

* EITHER

180
   a) At a recursive call, one or more parameters is an explicit
181
      constructor application
182
183
        AND
      That same parameter is scrutinised by a case somewhere in
184
185
186
187
188
189
      the RHS of the function

  OR

    b) At a recursive call, one or more parameters has an unfolding
       that is an explicit constructor application
190
191
        AND
      That same parameter is scrutinised by a case somewhere in
192
      the RHS of the function
193
        AND
194
      Those are the only uses of the parameter (see Note [Reboxing])
195
196


197
198
What to abstract over
~~~~~~~~~~~~~~~~~~~~~
199
200
201
There's a bit of a complication with type arguments.  If the call
site looks like

202
        f p = ...f ((:) [a] x xs)...
203
204
205

then our specialised function look like

206
        f_spec x xs = let p = (:) [a] x xs in ....as before....
207
208
209
210
211
212
213
214
215

This only makes sense if either
  a) the type variable 'a' is in scope at the top of f, or
  b) the type variable 'a' is an argument to f (and hence fs)

Actually, (a) may hold for value arguments too, in which case
we may not want to pass them.  Supose 'x' is in scope at f's
defn, but xs is not.  Then we'd like

216
        f_spec xs = let p = (:) [a] x xs in ....as before....
217
218
219
220
221
222
223

Similarly (b) may hold too.  If x is already an argument at the
call, no need to pass it again.

Finally, if 'a' is not in scope at the call site, we could abstract
it as we do the term variables:

224
        f_spec a x xs = let p = (:) [a] x xs in ...as before...
225
226
227

So the grand plan is:

228
229
        * abstract the call site to a constructor-only pattern
          e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)
230

231
        * Find the free variables of the abstracted pattern
232

233
234
        * Pass these variables, less any that are in scope at
          the fn defn.  But see Note [Shadowing] below.
235
236
237
238
239
240
241


NOTICE that we only abstract over variables that are not in scope,
so we're in no danger of shadowing variables used in "higher up"
in f_spec's RHS.


242
243
244
245
246
247
248
Note [Shadowing]
~~~~~~~~~~~~~~~~
In this pass we gather up usage information that may mention variables
that are bound between the usage site and the definition site; or (more
seriously) may be bound to something different at the definition site.
For example:

249
250
        f x = letrec g y v = let x = ...
                             in ...(g (a,b) x)...
251

252
Since 'x' is in scope at the call site, we may make a rewrite rule that
253
looks like
254
255
        RULE forall a,b. g (a,b) x = ...
But this rule will never match, because it's really a different 'x' at
256
257
258
259
260
261
262
263
264
the call site -- and that difference will be manifest by the time the
simplifier gets to it.  [A worry: the simplifier doesn't *guarantee*
no-shadowing, so perhaps it may not be distinct?]

Anyway, the rule isn't actually wrong, it's just not useful.  One possibility
is to run deShadowBinds before running SpecConstr, but instead we run the
simplifier.  That gives the simplest possible program for SpecConstr to
chew on; and it virtually guarantees no shadowing.

265
266
Note [Specialising for constant parameters]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
This one is about specialising on a *constant* (but not necessarily
constructor) argument

    foo :: Int -> (Int -> Int) -> Int
    foo 0 f = 0
    foo m f = foo (f m) (+1)

It produces

    lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
    lvl_rmV =
      \ (ds_dlk :: GHC.Base.Int) ->
        case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
        GHC.Base.I# (GHC.Prim.+# x_alG 1)

    T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
    GHC.Prim.Int#
    T.$wfoo =
      \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
        case ww_sme of ds_Xlw {
          __DEFAULT ->
288
289
290
        case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
        T.$wfoo ww1_Xmz lvl_rmV
        };
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
          0 -> 0
        }

The recursive call has lvl_rmV as its argument, so we could create a specialised copy
with that argument baked in; that is, not passed at all.   Now it can perhaps be inlined.

When is this worth it?  Call the constant 'lvl'
- If 'lvl' has an unfolding that is a constructor, see if the corresponding
  parameter is scrutinised anywhere in the body.

- If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
  parameter is applied (...to enough arguments...?)

  Also do this is if the function has RULES?

306
Also
307

308
309
Note [Specialising for lambda parameters]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
310
311
312
313
314
315
316
317
318
319
320
321
322
    foo :: Int -> (Int -> Int) -> Int
    foo 0 f = 0
    foo m f = foo (f m) (\n -> n-m)

This is subtly different from the previous one in that we get an
explicit lambda as the argument:

    T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
    GHC.Prim.Int#
    T.$wfoo =
      \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
        case ww_sm8 of ds_Xlr {
          __DEFAULT ->
323
324
325
326
327
328
329
330
        case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
        T.$wfoo
          ww1_Xmq
          (\ (n_ad3 :: GHC.Base.Int) ->
             case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
             GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
             })
        };
331
332
333
334
335
336
337
          0 -> 0
        }

I wonder if SpecConstr couldn't be extended to handle this? After all,
lambda is a sort of constructor for functions and perhaps it already
has most of the necessary machinery?

Gabor Greif's avatar
Gabor Greif committed
338
Furthermore, there's an immediate win, because you don't need to allocate the lambda
339
340
341
342
343
at the call site; and if perchance it's called in the recursive call, then you
may avoid allocating it altogether.  Just like for constructors.

Looks cool, but probably rare...but it might be easy to implement.

344
345
346

Note [SpecConstr for casts]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
347
Consider
348
349
350
351
352
353
354
355
    data family T a :: *
    data instance T Int = T Int

    foo n = ...
       where
         go (T 0) = 0
         go (T n) = go (T (n-1))

356
357
The recursive call ends up looking like
        go (T (I# ...) `cast` g)
358
So we want to spot the constructor application inside the cast.
359
360
That's why we have the Cast case in argToPat

361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
Note [Local recursive groups]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For a *local* recursive group, we can see all the calls to the
function, so we seed the specialisation loop from the calls in the
body, not from the calls in the RHS.  Consider:

  bar m n = foo n (n,n) (n,n) (n,n) (n,n)
   where
     foo n p q r s
       | n == 0    = m
       | n > 3000  = case p of { (p1,p2) -> foo (n-1) (p2,p1) q r s }
       | n > 2000  = case q of { (q1,q2) -> foo (n-1) p (q2,q1) r s }
       | n > 1000  = case r of { (r1,r2) -> foo (n-1) p q (r2,r1) s }
       | otherwise = case s of { (s1,s2) -> foo (n-1) p q r (s2,s1) }

If we start with the RHSs of 'foo', we get lots and lots of specialisations,
most of which are not needed.  But if we start with the (single) call
in the rhs of 'bar' we get exactly one fully-specialised copy, and all
the recursive calls go to this fully-specialised copy. Indeed, the original
380
function is later collected as dead code.  This is very important in
381
382
383
specialising the loops arising from stream fusion, for example in NDP where
we were getting literally hundreds of (mostly unused) specialisations of
a local function.
384

385
386
387
388
389
390
391
392
In a case like the above we end up never calling the original un-specialised
function.  (Although we still leave its code around just in case.)

However, if we find any boring calls in the body, including *unsaturated*
ones, such as
      letrec foo x y = ....foo...
      in map foo xs
then we will end up calling the un-specialised function, so then we *should*
393
394
use the calls in the un-specialised RHS as seeds.  We call these
"boring call patterns", and callsToPats reports if it finds any of these.
395
396


397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
Note [Top-level recursive groups]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If all the bindings in a top-level recursive group are not exported,
all the calls are in the rest of the top-level bindings.
This means we can specialise with those call patterns instead of with the RHSs
of the recursive group.

To get the call usage information, we work backwards through the top-level bindings
so we see the usage before we get to the binding of the function.
Before we can collect the usage though, we go through all the bindings and add them
to the environment. This is necessary because usage is only tracked for functions
in the environment.

The actual seeding of the specialisation is very similar to Note [Local recursive group].


413
414
415
416
417
418
419
420
Note [Do not specialise diverging functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Specialising a function that just diverges is a waste of code.
Furthermore, it broke GHC (simpl014) thus:
   {-# STR Sb #-}
   f = \x. case x of (a,b) -> f x
If we specialise f we get
   f = \x. case x of (a,b) -> fspec a b
421
But fspec doesn't have decent strictness info.  As it happened,
422
423
424
425
(f x) :: IO t, so the state hack applied and we eta expanded fspec,
and hence f.  But now f's strictness is less than its arity, which
breaks an invariant.

426

427
428
Note [Forcing specialisation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450

With stream fusion and in other similar cases, we want to fully
specialise some (but not necessarily all!) loops regardless of their
size and the number of specialisations.

We allow a library to do this, in one of two ways (one which is
deprecated):

  1) Add a parameter of type GHC.Types.SPEC (from ghc-prim) to the loop body.

  2) (Deprecated) Annotate a type with ForceSpecConstr from GHC.Exts,
     and then add *that* type as a parameter to the loop body

The reason #2 is deprecated is because it requires GHCi, which isn't
available for things like a cross compiler using stage1.

Here's a (simplified) example from the `vector` package. You may bring
the special 'force specialization' type into scope by saying:

  import GHC.Types (SPEC(..))

or by defining your own type (again, deprecated):
451
452
453
454

  data SPEC = SPEC | SPEC2
  {-# ANN type SPEC ForceSpecConstr #-}

455
456
457
458
459
(Note this is the exact same definition of GHC.Types.SPEC, just
without the annotation.)

After that, you say:

460
461
462
463
  foldl :: (a -> b -> a) -> a -> Stream b -> a
  {-# INLINE foldl #-}
  foldl f z (Stream step s _) = foldl_loop SPEC z s
    where
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
464
465
466
      foldl_loop !sPEC z s = case step s of
                              Yield x s' -> foldl_loop sPEC (f z x) s'
                              Skip       -> foldl_loop sPEC z s'
467
468
469
                              Done       -> z

SpecConstr will spot the SPEC parameter and always fully specialise
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
470
471
472
473
474
475
476
477
478
foldl_loop. Note that

  * We have to prevent the SPEC argument from being removed by
    w/w which is why (a) SPEC is a sum type, and (b) we have to seq on
    the SPEC argument.

  * And lastly, the SPEC argument is ultimately eliminated by
    SpecConstr itself so there is no runtime overhead.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
479
This is all quite ugly; we ought to come up with a better design.
480
481

ForceSpecConstr arguments are spotted in scExpr' and scTopBinds which then set
482
sc_force to True when calling specLoop. This flag does four things:
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
483
484
485
486
487
488
  * Ignore specConstrThreshold, to specialise functions of arbitrary size
        (see scTopBind)
  * Ignore specConstrCount, to make arbitrary numbers of specialisations
        (see specialise)
  * Specialise even for arguments that are not scrutinised in the loop
        (see argToPat; Trac #4488)
489
  * Only specialise on recursive types a finite number of times
490
        (see is_too_recursive; Trac #5550; Note [Limit recursive specialisation])
491

492
493
494
495
This flag is inherited for nested non-recursive bindings (which are likely to
be join points and hence should be fully specialised) but reset for nested
recursive bindings.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
496
What alternatives did I consider? Annotating the loop itself doesn't
497
498
work because (a) it is local and (b) it will be w/w'ed and having
w/w propagating annotations somehow doesn't seem like a good idea. The
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
499
500
501
types of the loop arguments really seem to be the most persistent
thing.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
502
Annotating the types that make up the loop state doesn't work,
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
503
504
505
506
507
508
either, because (a) it would prevent us from using types like Either
or tuples here, (b) we don't want to restrict the set of types that
can be used in Stream states and (c) some types are fixed by the user
(e.g., the accumulator here) but we still want to specialise as much
as possible.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
509
510
511
512
513
514
515
516
517
518
519
Alternatives to ForceSpecConstr
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Instead of giving the loop an extra argument of type SPEC, we
also considered *wrapping* arguments in SPEC, thus
  data SPEC a = SPEC a | SPEC2

  loop = \arg -> case arg of
                     SPEC state ->
                        case state of (x,y) -> ... loop (SPEC (x',y')) ...
                        S2 -> error ...
The idea is that a SPEC argument says "specialise this argument
520
regardless of whether the function case-analyses it".  But this
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
521
522
523
524
525
526
527
doesn't work well:
  * SPEC must still be a sum type, else the strictness analyser
    eliminates it
  * But that means that 'loop' won't be strict in its real payload
This loss of strictness in turn screws up specialisation, because
we may end up with calls like
   loop (SPEC (case z of (p,q) -> (q,p)))
528
529
Without the SPEC, if 'loop' were strict, the case would move out
and we'd see loop applied to a pair. But if 'loop' isn't strict
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
530
531
this doesn't look like a specialisable call.

532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
Note [Limit recursive specialisation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It is possible for ForceSpecConstr to cause an infinite loop of specialisation.
Because there is no limit on the number of specialisations, a recursive call with
a recursive constructor as an argument (for example, list cons) will generate
a specialisation for that constructor. If the resulting specialisation also
contains a recursive call with the constructor, this could proceed indefinitely.

For example, if ForceSpecConstr is on:
  loop :: [Int] -> [Int] -> [Int]
  loop z []         = z
  loop z (x:xs)     = loop (x:z) xs
this example will create a specialisation for the pattern
  loop (a:b) c      = loop' a b c

  loop' a b []      = (a:b)
  loop' a b (x:xs)  = loop (x:(a:b)) xs
and a new pattern is found:
  loop (a:(b:c)) d  = loop'' a b c d
which can continue indefinitely.

Roman's suggestion to fix this was to stop after a couple of times on recursive types,
but still specialising on non-recursive types as much as possible.

To implement this, we count the number of recursive constructors in each
function argument. If the maximum is greater than the specConstrRecursive limit,
do not specialise on that pattern.

This is only necessary when ForceSpecConstr is on: otherwise the specConstrCount
will force termination anyway.

See Trac #5550.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
565
566
Note [NoSpecConstr]
~~~~~~~~~~~~~~~~~~~
567
The ignoreDataCon stuff allows you to say
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
568
    {-# ANN type T NoSpecConstr #-}
569
to mean "don't specialise on arguments of this type".  It was added
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
570
571
572
573
574
before we had ForceSpecConstr.  Lacking ForceSpecConstr we specialised
regardless of size; and then we needed a way to turn that *off*.  Now
that we have ForceSpecConstr, this NoSpecConstr is probably redundant.
(Used only for PArray.)

575
-----------------------------------------------------
576
                Stuff not yet handled
577
578
579
580
-----------------------------------------------------

Here are notes arising from Roman's work that I don't want to lose.

581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
Example 1
~~~~~~~~~
    data T a = T !a

    foo :: Int -> T Int -> Int
    foo 0 t = 0
    foo x t | even x    = case t of { T n -> foo (x-n) t }
            | otherwise = foo (x-1) t

SpecConstr does no specialisation, because the second recursive call
looks like a boxed use of the argument.  A pity.

    $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
    $wfoo_sFw =
      \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
596
597
598
599
600
601
602
603
604
605
         case ww_sFo of ds_Xw6 [Just L] {
           __DEFAULT ->
                case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
                  __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
                  0 ->
                    case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
                    case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
                    $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
                    } } };
           0 -> 0
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625

Example 2
~~~~~~~~~
    data a :*: b = !a :*: !b
    data T a = T !a

    foo :: (Int :*: T Int) -> Int
    foo (0 :*: t) = 0
    foo (x :*: t) | even x    = case t of { T n -> foo ((x-n) :*: t) }
                  | otherwise = foo ((x-1) :*: t)

Very similar to the previous one, except that the parameters are now in
a strict tuple. Before SpecConstr, we have

    $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
    $wfoo_sG3 =
      \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
    GHC.Base.Int) ->
        case ww_sFU of ds_Xws [Just L] {
          __DEFAULT ->
626
627
628
629
630
631
632
633
634
635
        case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
          __DEFAULT ->
            case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
            $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2             -- $wfoo1
            };
          0 ->
            case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
            case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
            $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB        -- $wfoo2
            } } };
636
637
638
639
          0 -> 0 }

We get two specialisations:
"SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
640
641
                  Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
                  = Foo.$s$wfoo1 a_sFB sc_sGC ;
642
"SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
643
644
                  Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
                  = Foo.$s$wfoo y_aFp sc_sGC ;
645
646

But perhaps the first one isn't good.  After all, we know that tpl_B2 is
647
a T (I# x) really, because T is strict and Int has one constructor.  (We can't
Gabor Greif's avatar
typos    
Gabor Greif committed
648
unbox the strict fields, because T is polymorphic!)
649

650
%************************************************************************
651
%*                                                                      *
652
\subsection{Top level wrapper stuff}
653
%*                                                                      *
654
655
656
%************************************************************************

\begin{code}
657
658
659
660
661
specConstrProgram :: ModGuts -> CoreM ModGuts
specConstrProgram guts
  = do
      dflags <- getDynFlags
      us     <- getUniqueSupplyM
662
      annos  <- getFirstAnnotations deserializeWithData guts
663
664
665
666
667
      let binds' = reverse $ fst $ initUs us $ do
                    -- Note [Top-level recursive groups]
                    (env, binds) <- goEnv (initScEnv dflags annos) (mg_binds guts)
                    go env nullUsage (reverse binds)

668
      return (guts { mg_binds = binds' })
669
  where
670
671
672
673
674
675
676
677
678
    goEnv env []            = return (env, [])
    goEnv env (bind:binds)  = do (env', bind')   <- scTopBindEnv env bind
                                 (env'', binds') <- goEnv env' binds
                                 return (env'', bind' : binds')

    go _   _   []           = return []
    go env usg (bind:binds) = do (usg', bind') <- scTopBind env usg bind
                                 binds' <- go env usg' binds
                                 return (bind' : binds')
679
680
681
682
\end{code}


%************************************************************************
683
%*                                                                      *
684
\subsection{Environment: goes downwards}
685
%*                                                                      *
686
687
%************************************************************************

688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
Note [Work-free values only in environment]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The sc_vals field keeps track of in-scope value bindings, so 
that if we come across (case x of Just y ->...) we can reduce the
case from knowing that x is bound to a pair.

But only *work-free* values are ok here. For example if the envt had
    x -> Just (expensive v)
then we do NOT want to expand to
     let y = expensive v in ...
because the x-binding still exists and we've now duplicated (expensive v).

This seldom happens because let-bound constructor applications are 
ANF-ised, but it can happen as a result of on-the-fly transformations in
SpecConstr itself.  Here is Trac #7865:

        let {
          a'_shr =
            case xs_af8 of _ {
              [] -> acc_af6;
              : ds_dgt [Dmd=<L,A>] ds_dgu [Dmd=<L,A>] ->
                (expensive x_af7, x_af7
            } } in
        let {
          ds_sht =
            case a'_shr of _ { (p'_afd, q'_afe) ->
            TSpecConstr_DoubleInline.recursive
              (GHC.Types.: @ GHC.Types.Int x_af7 wild_X6) (q'_afe, p'_afd)
            } } in

When processed knowing that xs_af8 was bound to a cons, we simplify to 
   a'_shr = (expensive x_af7, x_af7)
and we do NOT want to inline that at the occurrence of a'_shr in ds_sht.
(There are other occurrences of a'_shr.)  No no no.

It would be possible to do some on-the-fly ANF-ising, so that a'_shr turned
into a work-free value again, thus
   a1 = expensive x_af7
   a'_shr = (a1, x_af7)
but that's more work, so until its shown to be important I'm going to 
leave it for now.

730
\begin{code}
731
732
733
data ScEnv = SCE { sc_dflags    :: DynFlags,
                   sc_size      :: Maybe Int,   -- Size threshold
                   sc_count     :: Maybe Int,   -- Max # of specialisations for any one fn
734
                                                -- See Note [Avoiding exponential blowup]
735
736
737
738
739

                   sc_recursive :: Int,         -- Max # of specialisations over recursive type.
                                                -- Stops ForceSpecConstr from diverging.

                   sc_force     :: Bool,        -- Force specialisation?
740
                                                -- See Note [Forcing specialisation]
741

742
                   sc_subst     :: Subst,       -- Current substitution
743
                                                -- Maps InIds to OutExprs
744

745
746
747
                   sc_how_bound :: HowBoundEnv,
                        -- Binds interesting non-top-level variables
                        -- Domain is OutVars (*after* applying the substitution)
748

749
                   sc_vals      :: ValueEnv,
750
751
                        -- Domain is OutIds (*after* applying the substitution)
                        -- Used even for top-level bindings (but not imported ones)
752
753
754
755
                        -- The range of the ValueEnv is *work-free* values
                        -- such as (\x. blah), or (Just v)
                        -- but NOT (Just (expensive v))
                        -- See Note [Work-free values only in environment]
756

757
                   sc_annotations :: UniqFM SpecConstrAnnotation
758
             }
759

760
761
---------------------
-- As we go, we apply a substitution (sc_subst) to the current term
762
type InExpr = CoreExpr          -- _Before_ applying the subst
763
type InVar  = Var
764

765
type OutExpr = CoreExpr         -- _After_ applying the subst
766
767
768
769
type OutId   = Id
type OutVar  = Var

---------------------
770
type HowBoundEnv = VarEnv HowBound      -- Domain is OutVars
771

772
---------------------
773
774
775
776
type ValueEnv = IdEnv Value             -- Domain is OutIds
data Value    = ConVal AltCon [CoreArg] -- _Saturated_ constructors
                                        --   The AltCon is never DEFAULT
              | LambdaVal               -- Inlinable lambdas or PAPs
777

778
779
instance Outputable Value where
   ppr (ConVal con args) = ppr con <+> interpp'SP args
780
   ppr LambdaVal         = ptext (sLit "<Lambda>")
781

782
---------------------
783
initScEnv :: DynFlags -> UniqFM SpecConstrAnnotation -> ScEnv
784
initScEnv dflags anns
785
786
787
788
789
790
791
792
  = SCE { sc_dflags      = dflags,
          sc_size        = specConstrThreshold dflags,
          sc_count       = specConstrCount     dflags,
          sc_recursive   = specConstrRecursive dflags,
          sc_force       = False,
          sc_subst       = emptySubst,
          sc_how_bound   = emptyVarEnv,
          sc_vals        = emptyVarEnv,
793
          sc_annotations = anns }
794

795
796
data HowBound = RecFun  -- These are the recursive functions for which
                        -- we seek interesting call patterns
797

798
799
              | RecArg  -- These are those functions' arguments, or their sub-components;
                        -- we gather occurrence information for these
800

801
802
803
804
instance Outputable HowBound where
  ppr RecFun = text "RecFun"
  ppr RecArg = text "RecArg"

805
806
807
scForce :: ScEnv -> Bool -> ScEnv
scForce env b = env { sc_force = b }

808
809
810
811
lookupHowBound :: ScEnv -> Id -> Maybe HowBound
lookupHowBound env id = lookupVarEnv (sc_how_bound env) id

scSubstId :: ScEnv -> Id -> CoreExpr
812
scSubstId env v = lookupIdSubst (text "scSubstId") (sc_subst env) v
813
814
815
816

scSubstTy :: ScEnv -> Type -> Type
scSubstTy env ty = substTy (sc_subst env) ty

817
818
819
scSubstCo :: ScEnv -> Coercion -> Coercion
scSubstCo env co = substCo (sc_subst env) co

820
821
zapScSubst :: ScEnv -> ScEnv
zapScSubst env = env { sc_subst = zapSubstEnv (sc_subst env) }
822

823
extendScInScope :: ScEnv -> [Var] -> ScEnv
824
        -- Bring the quantified variables into scope
825
extendScInScope env qvars = env { sc_subst = extendInScopeList (sc_subst env) qvars }
826

827
        -- Extend the substitution
828
829
830
831
832
extendScSubst :: ScEnv -> Var -> OutExpr -> ScEnv
extendScSubst env var expr = env { sc_subst = extendSubst (sc_subst env) var expr }

extendScSubstList :: ScEnv -> [(Var,OutExpr)] -> ScEnv
extendScSubstList env prs = env { sc_subst = extendSubstList (sc_subst env) prs }
833
834
835
836

extendHowBound :: ScEnv -> [Var] -> HowBound -> ScEnv
extendHowBound env bndrs how_bound
  = env { sc_how_bound = extendVarEnvList (sc_how_bound env)
837
                            [(bndr,how_bound) | bndr <- bndrs] }
838
839

extendBndrsWith :: HowBound -> ScEnv -> [Var] -> (ScEnv, [Var])
840
extendBndrsWith how_bound env bndrs
841
  = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndrs')
842
  where
843
    (subst', bndrs') = substBndrs (sc_subst env) bndrs
844
845
    hb_env' = sc_how_bound env `extendVarEnvList`
                    [(bndr,how_bound) | bndr <- bndrs']
846
847

extendBndrWith :: HowBound -> ScEnv -> Var -> (ScEnv, Var)
848
extendBndrWith how_bound env bndr
849
  = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndr')
850
  where
851
852
853
854
855
    (subst', bndr') = substBndr (sc_subst env) bndr
    hb_env' = extendVarEnv (sc_how_bound env) bndr' how_bound

extendRecBndrs :: ScEnv -> [Var] -> (ScEnv, [Var])
extendRecBndrs env bndrs  = (env { sc_subst = subst' }, bndrs')
856
857
                      where
                        (subst', bndrs') = substRecBndrs (sc_subst env) bndrs
858
859
860

extendBndr :: ScEnv -> Var -> (ScEnv, Var)
extendBndr  env bndr  = (env { sc_subst = subst' }, bndr')
861
862
                      where
                        (subst', bndr') = substBndr (sc_subst env) bndr
863

864
extendValEnv :: ScEnv -> Id -> Maybe Value -> ScEnv
865
extendValEnv env _  Nothing   = env
866
867
868
869
extendValEnv env id (Just cv) 
 | valueIsWorkFree cv      -- Don't duplicate work!!  Trac #7865
 = env { sc_vals = extendVarEnv (sc_vals env) id cv }
extendValEnv env _ _ = env
870

871
extendCaseBndrs :: ScEnv -> OutExpr -> OutId -> AltCon -> [Var] -> (ScEnv, [Var])
872
-- When we encounter
873
874
--      case scrut of b
--          C x y -> ...
875
876
877
878
-- we want to bind b, to (C x y)
-- NB1: Extends only the sc_vals part of the envt
-- NB2: Kill the dead-ness info on the pattern binders x,y, since
--      they are potentially made alive by the [b -> C x y] binding
879
880
extendCaseBndrs env scrut case_bndr con alt_bndrs
   = (env2, alt_bndrs')
881
 where
882
883
   live_case_bndr = not (isDeadBinder case_bndr)
   env1 | Var v <- scrut = extendValEnv env v cval
884
        | otherwise      = env  -- See Note [Add scrutinee to ValueEnv too]
885
   env2 | live_case_bndr = extendValEnv env1 case_bndr cval
886
887
888
889
890
891
892
        | otherwise      = env1

   alt_bndrs' | case scrut of { Var {} -> True; _ -> live_case_bndr }
              = map zap alt_bndrs
              | otherwise
              = alt_bndrs

893
   cval = case con of
894
895
896
897
898
899
900
901
                DEFAULT    -> Nothing
                LitAlt {}  -> Just (ConVal con [])
                DataAlt {} -> Just (ConVal con vanilla_args)
                      where
                        vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
                                       varsToCoreExprs alt_bndrs

   zap v | isTyVar v = v                -- See NB2 above
902
903
         | otherwise = zapIdOccInfo v

904

905
906
decreaseSpecCount :: ScEnv -> Int -> ScEnv
-- See Note [Avoiding exponential blowup]
907
decreaseSpecCount env n_specs
908
909
910
  = env { sc_count = case sc_count env of
                       Nothing -> Nothing
                       Just n  -> Just (n `div` (n_specs + 1)) }
911
912
        -- The "+1" takes account of the original function;
        -- See Note [Avoiding exponential blowup]
913
914

---------------------------------------------------
915
-- See Note [Forcing specialisation]
916
ignoreType    :: ScEnv -> Type   -> Bool
917
ignoreDataCon  :: ScEnv -> DataCon -> Bool
918
forceSpecBndr :: ScEnv -> Var    -> Bool
919

920
#ifndef GHCI
921
ignoreType    _ _  = False
922
ignoreDataCon  _ _ = False
923
924
#else /* GHCI */

925
ignoreDataCon env dc = ignoreTyCon env (dataConTyCon dc)
926

927
ignoreType env ty
928
929
930
  = case tyConAppTyCon_maybe ty of
      Just tycon -> ignoreTyCon env tycon
      _          -> False
931

932
933
934
ignoreTyCon :: ScEnv -> TyCon -> Bool
ignoreTyCon env tycon
  = lookupUFM (sc_annotations env) tycon == Just NoSpecConstr
935
#endif /* GHCI */
936

937
forceSpecBndr env var = forceSpecFunTy env . snd . splitForAllTys . varType $ var
938
939
940
941
942
943
944
945
946
947
948

forceSpecFunTy :: ScEnv -> Type -> Bool
forceSpecFunTy env = any (forceSpecArgTy env) . fst . splitFunTys

forceSpecArgTy :: ScEnv -> Type -> Bool
forceSpecArgTy env ty
  | Just ty' <- coreView ty = forceSpecArgTy env ty'

forceSpecArgTy env ty
  | Just (tycon, tys) <- splitTyConApp_maybe ty
  , tycon /= funTyCon
949
950
951
952
      = tyConName tycon == specTyConName
#ifdef GHCI
        || lookupUFM (sc_annotations env) tycon == Just ForceSpecConstr
#endif
953
954
955
        || any (forceSpecArgTy env) tys

forceSpecArgTy _ _ = False
956
957
\end{code}

958
959
960
961
962
963
964
965
966
967
968
969
Note [Add scrutinee to ValueEnv too]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
   case x of y
     (a,b) -> case b of c
                I# v -> ...(f y)...
By the time we get to the call (f y), the ValueEnv
will have a binding for y, and for c
    y -> (a,b)
    c -> I# v
BUT that's not enough!  Looking at the call (f y) we
see that y is pair (a,b), but we also need to know what 'b' is.
970
So in extendCaseBndrs we must *also* add the binding
971
972
973
974
975
976
   b -> I# v
else we lose a useful specialisation for f.  This is necessary even
though the simplifier has systematically replaced uses of 'x' with 'y'
and 'b' with 'c' in the code.  The use of 'b' in the ValueEnv came
from outside the case.  See Trac #4908 for the live example.

977
978
979
980
Note [Avoiding exponential blowup]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The sc_count field of the ScEnv says how many times we are prepared to
duplicate a single function.  But we must take care with recursive
981
specialisations.  Consider
982

983
984
        let $j1 = let $j2 = let $j3 = ...
                            in
985
                            ...$j3...
986
                  in
987
                  ...$j2...
988
        in
989
990
991
992
        ...$j1...

If we specialise $j1 then in each specialisation (as well as the original)
we can specialise $j2, and similarly $j3.  Even if we make just *one*
Gabor Greif's avatar
typos    
Gabor Greif committed
993
specialisation of each, because we also have the original we'll get 2^n
994
995
996
997
998
copies of $j3, which is not good.

So when recursively specialising we divide the sc_count by the number of
copies we are making at this level, including the original.

999
1000

%************************************************************************
1001
%*                                                                      *
1002
\subsection{Usage information: flows upwards}
1003
%*                                                                      *
1004
%************************************************************************
1005

1006
\begin{code}
1007
1008
data ScUsage
   = SCU {
1009
1010
1011
        scu_calls :: CallEnv,           -- Calls
                                        -- The functions are a subset of the
                                        --      RecFuns in the ScEnv
1012

1013
1014
        scu_occs :: !(IdEnv ArgOcc)     -- Information on argument occurrences
     }                                  -- The domain is OutIds
1015

1016
type CallEnv = IdEnv [Call]
1017
type Call = (ValueEnv, [CoreArg])
1018
1019
        -- The arguments of the call, together with the
        -- env giving the constructor bindings at the call site
1020

1021
1022
nullUsage :: ScUsage
nullUsage = SCU { scu_calls = emptyVarEnv, scu_occs = emptyVarEnv }
1023

1024
1025
1026
combineCalls :: CallEnv -> CallEnv -> CallEnv
combineCalls = plusVarEnv_C (++)

1027
1028
combineUsage :: ScUsage -> ScUsage -> ScUsage
combineUsage u1 u2 = SCU { scu_calls = combineCalls (scu_calls u1) (scu_calls u2),
1029
                           scu_occs  = plusVarEnv_C combineOcc (scu_occs u1) (scu_occs u2) }
1030

1031
combineUsages :: [ScUsage] -> ScUsage
1032
1033
1034
combineUsages [] = nullUsage
combineUsages us = foldr1 combineUsage us

1035
1036
1037
lookupOccs :: ScUsage -> [OutVar] -> (ScUsage, [ArgOcc])
lookupOccs (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndrs
  = (SCU {scu_calls = sc_calls, scu_occs = delVarEnvList sc_occs bndrs},
1038
1039
     [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])

1040
1041
data ArgOcc = NoOcc     -- Doesn't occur at all; or a type argument
            | UnkOcc    -- Used in some unknown way
1042

1043
            | ScrutOcc  -- See Note [ScrutOcc]
1044
                 (DataConEnv [ArgOcc])   -- How the sub-components are used
1045

1046
type DataConEnv a = UniqFM a     -- Keyed by DataCon
1047

1048
1049
{- Note  [ScrutOcc]
~~~~~~~~~~~~~~~~~~~
1050
1051
An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
is *only* taken apart or applied.
1052

1053
  Functions, literal: ScrutOcc emptyUFM
1054
1055
1056
1057
1058
  Data constructors:  ScrutOcc subs,

where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
The domain of the UniqFM is the Unique of the data constructor

1059
The [ArgOcc] is the occurrences of the *pattern-bound* components
1060
of the data structure.  E.g.
1061
        data T a = forall b. MkT a b (b->a)
1062
1063
1064
A pattern binds b, x::a, y::b, z::b->a, but not 'a'!

-}
1065
1066

instance Outputable ArgOcc where
Ian Lynagh's avatar
Ian Lynagh committed
1067
  ppr (ScrutOcc xs) = ptext (sLit "scrut-occ") <> ppr xs
1068
1069
  ppr UnkOcc        = ptext (sLit "unk-occ")
  ppr NoOcc         = ptext (sLit "no-occ")
1070

1071
1072
1073
evalScrutOcc :: ArgOcc
evalScrutOcc = ScrutOcc emptyUFM

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
1074
-- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
1075
1076
-- that if the thing is scrutinised anywhere then we get to see that
-- in the overall result, even if it's also used in a boxed way
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
1077
-- This might be too agressive; see Note [Reboxing] Alternative 3
1078
combineOcc :: ArgOcc -> ArgOcc -> ArgOcc
1079
1080
combineOcc NoOcc         occ           = occ
combineOcc occ           NoOcc         = occ
1081
combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
1082
combineOcc UnkOcc        (ScrutOcc ys) = ScrutOcc ys
1083
combineOcc (ScrutOcc xs) UnkOcc        = ScrutOcc xs
1084
1085
1086
1087
1088
combineOcc UnkOcc        UnkOcc        = UnkOcc

combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys

1089
setScrutOcc :: ScEnv -> ScUsage -> OutExpr -> ArgOcc -> ScUsage
Thomas Schilling's avatar
Thomas Schilling committed
1090
-- _Overwrite_ the occurrence info for the scrutinee, if the scrutinee
1091
-- is a variable, and an interesting variable
1092
1093
setScrutOcc env usg (Cast e _) occ      = setScrutOcc env usg e occ
setScrutOcc env usg (Tick _ e) occ      = setScrutOcc env usg e occ
1094
setScrutOcc env usg (Var v)    occ
1095
  | Just RecArg <- lookupHowBound env v = usg { scu_occs = extendVarEnv (scu_occs usg) v occ }
1096
1097
1098
  | otherwise                           = usg
setScrutOcc _env usg _other _occ        -- Catch-all
  = usg
1099
1100
1101
\end{code}

%************************************************************************
1102
%*                                                                      *
1103
\subsection{The main recursive function}
1104
%*                                                                      *
1105
1106
%************************************************************************

1107
1108
1109
The main recursive function gathers up usage information, and
creates specialised versions of functions.

1110
\begin{code}
1111
scExpr, scExpr' :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
1112
1113
        -- The unique supply is needed when we invent
        -- a new name for the specialised function and its args
1114

1115
1116
1117
scExpr env e = scExpr' env e


1118
scExpr' env (Var v)      = case scSubstId env v of
1119
1120
                            Var v' -> return (mkVarUsage env v' [], Var v')
                            e'     -> scExpr (zapScSubst env) e'
1121

1122
scExpr' env (Type t)     = return (nullUsage, Type (scSubstTy env t))
1123
scExpr' env (Coercion c) = return (nullUsage, Coercion (scSubstCo env c))
1124
1125
1126
1127
1128
1129
1130
1131
1132
scExpr' _   e@(Lit {})   = return (nullUsage, e)
scExpr' env (Tick t e)   = do (usg, e') <- scExpr env e
                              return (usg, Tick t e')
scExpr' env (Cast e co)  = do (usg, e') <- scExpr env e
                              return (usg, Cast e' (scSubstCo env co))
scExpr' env e@(App _ _)  = scApp env (collectArgs e)
scExpr' env (Lam b e)    = do let (env', b') = extendBndr env b
                              (usg, e') <- scExpr env' e
                              return (usg, Lam b' e')
1133

1134
1135
1136
1137
1138
1139
scExpr' env (Case scrut b ty alts)
  = do  { (scrut_usg, scrut') <- scExpr env scrut
        ; case isValue (sc_vals env) scrut' of
                Just (ConVal con args) -> sc_con_app con args scrut'
                _other                 -> sc_vanilla scrut_usg scrut'
        }
1140
  where
1141
    sc_con_app con args scrut'  -- Known constructor; simplify
1142
     = do { let (_, bs, rhs) = findAlt con alts
1143
                                  `orElse` (DEFAULT, [], mkImpossibleExpr ty)
1144
1145
                alt_env'     = extendScSubstList env ((b,scrut') : bs `zip` trimConArgs con args)
          ; scExpr alt_env' rhs }
1146
1147

    sc_vanilla scrut_usg scrut' -- Normal case
1148
     = do { let (alt_env,b') = extendBndrWith RecArg env b
1149
                        -- Record RecArg for the components
1150

1151
1152
          ; (alt_usgs, alt_occs, alts')
                <- mapAndUnzip3M (sc_alt alt_env scrut' b') alts
1153

1154
1155
1156
1157
1158
          ; let scrut_occ  = foldr combineOcc NoOcc alt_occs
                scrut_usg' = setScrutOcc env scrut_usg scrut' scrut_occ
                -- The combined usage of the scrutinee is given
                -- by scrut_occ, which is passed to scScrut, which
                -- in turn treats a bare-variable scrutinee specially
1159

1160
1161
          ; return (foldr combineUsage scrut_usg' alt_usgs,
                    Case scrut' b' (scSubstTy env ty) alts') }
1162

1163
    sc_alt env scrut' b' (con,bs,rhs)
1164
1165
1166
1167
1168
1169
1170
1171
     = do { let (env1, bs1) = extendBndrsWith RecArg env bs
                (env2, bs2) = extendCaseBndrs env1 scrut' b' con bs1
          ; (usg, rhs') <- scExpr env2 rhs
          ; let (usg', b_occ:arg_occs) = lookupOccs usg (b':bs2)
                scrut_occ = case con of
                               DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
                               _          -> ScrutOcc emptyUFM
          ; return (usg', b_occ `combineOcc` scrut_occ, (con, bs2, rhs')) }
1172
1173

scExpr' env (Let (NonRec bndr rhs) body)
1174
  | isTyVar bndr        -- Type-lets may be created by doBeta
1175
  = scExpr' (extendScSubst env bndr rhs) body
1176

1177
1178
  | otherwise
  = do  { let (body_env, bndr') = extendBndr env bndr
1179
        ; (rhs_usg, rhs_info)  <- scRecRhs env (bndr',rhs)
1180

1181
        ; let body_env2         = extendHowBound body_env [bndr'] RecFun
1182
                                   -- Note [Local let bindings]
1183
1184
              RI _ rhs' _ _ _   = rhs_info
              body_env3         = extendValEnv body_env2 bndr' (isValue (sc_vals env) rhs')
1185

1186
        ; (body_usg, body') <- scExpr body_env3 body
1187

1188
1189
          -- NB: For non-recursive bindings we inherit sc_force flag from
          -- the parent function (see Note [Forcing specialisation])
1190
1191
1192
        ; (spec_usg, specs) <- specialise env
                                          (scu_calls body_usg)
                                          rhs_info
1193
                                          (SI [] 0 (Just rhs_usg))
1194

1195
1196
1197
1198
        ; return (body_usg { scu_calls = scu_calls body_usg `delVarEnv` bndr' }
                    `combineUsage` rhs_usg `combineUsage` spec_usg,
                  mkLets [NonRec b r | (b,r) <- specInfoBinds rhs_info specs] body')
        }
1199