SimplEnv.hs 27.2 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3
{-
(c) The AQUA Project, Glasgow University, 1993-1998

4
\section[SimplMonad]{The simplifier Monad}
Austin Seipp's avatar
Austin Seipp committed
5
-}
6

7 8
{-# LANGUAGE CPP #-}

9
module SimplEnv (
10 11
        InId, InBind, InExpr, InAlt, InArg, InType, InBndr, InVar,
        OutId, OutTyVar, OutBind, OutExpr, OutAlt, OutArg, OutType, OutBndr, OutVar,
12
        InCoercion, OutCoercion,
13

14 15
        -- The simplifier mode
        setMode, getMode, updMode,
16

17
        -- Environments
18
        SimplEnv(..), StaticEnv, pprSimplEnv,   -- Temp not abstract
19 20
        mkSimplEnv, extendIdSubst,
        SimplEnv.extendTvSubst, SimplEnv.extendCvSubst,
21 22
        zapSubstEnv, setSubstEnv,
        getInScope, setInScope, setInScopeSet, modifyInScope, addNewInScopeIds,
23
        getSimplRules,
24

25
        SimplSR(..), mkContEx, substId, lookupRecBndr, refineFromInScope,
26

27 28
        simplNonRecBndr, simplRecBndrs,
        simplBinder, simplBinders,
29 30
        substTy, substTyVar, getTCvSubst,
        substCo, substCoVar,
31

32 33
        -- Floats
        Floats, emptyFloats, isEmptyFloats, addNonRec, addFloats, extendFloats,
Peter Wortmann's avatar
Peter Wortmann committed
34
        wrapFloats, setFloats, zapFloats, addRecFloats, mapFloats,
35
        doFloatFromRhs, getFloatBinds
36 37 38 39
    ) where

#include "HsVersions.h"

40
import SimplMonad
41
import CoreMonad        ( SimplifierMode(..) )
42
import CoreSyn
43 44
import CoreUtils
import Var
45
import VarEnv
46
import VarSet
47
import OrdList
48
import Id
49
import MkCore                   ( mkWildValBinder )
50
import TysWiredIn
51
import qualified Type
52
import Type hiding              ( substTy, substTyVar, substTyVarBndr )
53
import qualified Coercion
54
import Coercion hiding          ( substCo, substCoVar, substCoVarBndr )
55
import BasicTypes
56
import MonadUtils
57
import Outputable
58
import Util
59 60

import Data.List
61

Austin Seipp's avatar
Austin Seipp committed
62 63 64
{-
************************************************************************
*                                                                      *
65
\subsection[Simplify-types]{Type declarations}
Austin Seipp's avatar
Austin Seipp committed
66 67 68
*                                                                      *
************************************************************************
-}
69

70
type InBndr     = CoreBndr
71 72 73
type InVar      = Var                   -- Not yet cloned
type InId       = Id                    -- Not yet cloned
type InType     = Type                  -- Ditto
74 75 76 77 78 79
type InBind     = CoreBind
type InExpr     = CoreExpr
type InAlt      = CoreAlt
type InArg      = CoreArg
type InCoercion = Coercion

80
type OutBndr     = CoreBndr
81 82 83 84
type OutVar      = Var                  -- Cloned
type OutId       = Id                   -- Cloned
type OutTyVar    = TyVar                -- Cloned
type OutType     = Type                 -- Cloned
85
type OutCoercion = Coercion
86 87 88 89
type OutBind     = CoreBind
type OutExpr     = CoreExpr
type OutAlt      = CoreAlt
type OutArg      = CoreArg
90

Austin Seipp's avatar
Austin Seipp committed
91 92 93
{-
************************************************************************
*                                                                      *
94
\subsubsection{The @SimplEnv@ type}
Austin Seipp's avatar
Austin Seipp committed
95 96 97
*                                                                      *
************************************************************************
-}
98 99 100

data SimplEnv
  = SimplEnv {
101
     ----------- Static part of the environment -----------
102
     -- Static in the sense of lexically scoped,
103 104
     -- wrt the original expression

105
        seMode      :: SimplifierMode,
106

107 108
        -- The current substitution
        seTvSubst   :: TvSubstEnv,      -- InTyVar |--> OutType
109
        seCvSubst   :: CvSubstEnv,      -- InCoVar |--> OutCoercion
110
        seIdSubst   :: SimplIdSubst,    -- InId    |--> OutExpr
111 112 113 114 115

     ----------- Dynamic part of the environment -----------
     -- Dynamic in the sense of describing the setup where
     -- the expression finally ends up

116 117 118 119 120 121
        -- The current set of in-scope variables
        -- They are all OutVars, and all bound in this module
        seInScope   :: InScopeSet,      -- OutVars only
                -- Includes all variables bound by seFloats
        seFloats    :: Floats
                -- See Note [Simplifier floats]
122 123
    }

124
type StaticEnv = SimplEnv       -- Just the static part is relevant
125

126 127 128
pprSimplEnv :: SimplEnv -> SDoc
-- Used for debugging; selective
pprSimplEnv env
129 130 131 132
  = vcat [text "TvSubst:" <+> ppr (seTvSubst env),
          text "CvSubst:" <+> ppr (seCvSubst env),
          text "IdSubst:" <+> ppr (seIdSubst env),
          text "InScope:" <+> vcat (map ppr_one in_scope_vars)
133 134 135 136 137
    ]
  where
   in_scope_vars = varEnvElts (getInScopeVars (seInScope env))
   ppr_one v | isId v = ppr v <+> ppr (idUnfolding v)
             | otherwise = ppr v
138

139 140
type SimplIdSubst = IdEnv SimplSR       -- IdId |--> OutExpr
        -- See Note [Extending the Subst] in CoreSubst
141 142

data SimplSR
143 144 145
  = DoneEx OutExpr              -- Completed term
  | DoneId OutId                -- Completed term variable
  | ContEx TvSubstEnv           -- A suspended substitution
146
           CvSubstEnv
147 148
           SimplIdSubst
           InExpr
149

150
instance Outputable SimplSR where
151 152 153
  ppr (DoneEx e) = text "DoneEx" <+> ppr e
  ppr (DoneId v) = text "DoneId" <+> ppr v
  ppr (ContEx _tv _cv _id e) = vcat [text "ContEx" <+> ppr e {-,
154 155 156 157 158
                                ppr (filter_env tv), ppr (filter_env id) -}]
        -- where
        -- fvs = exprFreeVars e
        -- filter_env env = filterVarEnv_Directly keep env
        -- keep uniq _ = uniq `elemUFM_Directly` fvs
159

Austin Seipp's avatar
Austin Seipp committed
160
{-
161 162
Note [SimplEnv invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~
163 164 165 166 167 168
seInScope:
        The in-scope part of Subst includes *all* in-scope TyVars and Ids
        The elements of the set may have better IdInfo than the
        occurrences of in-scope Ids, and (more important) they will
        have a correctly-substituted type.  So we use a lookup in this
        set to replace occurrences
169

170 171 172
        The Ids in the InScopeSet are replete with their Rules,
        and as we gather info about the unfolding of an Id, we replace
        it in the in-scope set.
173

174 175
        The in-scope set is actually a mapping OutVar -> OutVar, and
        in case expressions we sometimes bind
176 177

seIdSubst:
178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197
        The substitution is *apply-once* only, because InIds and OutIds
        can overlap.
        For example, we generally omit mappings
                a77 -> a77
        from the substitution, when we decide not to clone a77, but it's quite
        legitimate to put the mapping in the substitution anyway.

        Furthermore, consider
                let x = case k of I# x77 -> ... in
                let y = case k of I# x77 -> ... in ...
        and suppose the body is strict in both x and y.  Then the simplifier
        will pull the first (case k) to the top; so the second (case k) will
        cancel out, mapping x77 to, well, x77!  But one is an in-Id and the
        other is an out-Id.

        Of course, the substitution *must* applied! Things in its domain
        simply aren't necessarily bound in the result.

* substId adds a binding (DoneId new_id) to the substitution if
        the Id's unique has changed
198

199
  Note, though that the substitution isn't necessarily extended
200
  if the type of the Id changes.  Why not?  Because of the next point:
201

202
* We *always, always* finish by looking up in the in-scope set
203
  any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
204 205 206
  Reason: so that we never finish up with a "old" Id in the result.
  An old Id might point to an old unfolding and so on... which gives a space
  leak.
207 208 209 210 211 212 213 214

  [The DoneEx and DoneVar hits map to "new" stuff.]

* It follows that substExpr must not do a no-op if the substitution is empty.
  substType is free to do so, however.

* When we come to a let-binding (say) we generate new IdInfo, including an
  unfolding, attach it to the binder, and add this newly adorned binder to
215 216 217
  the in-scope set.  So all subsequent occurrences of the binder will get
  mapped to the full-adorned binder, which is also the one put in the
  binding site.
218 219 220 221

* The in-scope "set" usually maps x->x; we use it simply for its domain.
  But sometimes we have two in-scope Ids that are synomyms, and should
  map to the same target:  x->x, y->x.  Notably:
222
        case y of x { ... }
223
  That's why the "set" is actually a VarEnv Var
Austin Seipp's avatar
Austin Seipp committed
224
-}
225

226 227
mkSimplEnv :: SimplifierMode -> SimplEnv
mkSimplEnv mode
228
  = SimplEnv { seMode = mode
229 230 231
             , seInScope = init_in_scope
             , seFloats = emptyFloats
             , seTvSubst = emptyVarEnv
232
             , seCvSubst = emptyVarEnv
233
             , seIdSubst = emptyVarEnv }
234
        -- The top level "enclosing CC" is "SUBSUMED".
235

236 237 238 239
init_in_scope :: InScopeSet
init_in_scope = mkInScopeSet (unitVarSet (mkWildValBinder unitTy))
              -- See Note [WildCard binders]

Austin Seipp's avatar
Austin Seipp committed
240
{-
241 242 243 244 245
Note [WildCard binders]
~~~~~~~~~~~~~~~~~~~~~~~
The program to be simplified may have wild binders
    case e of wild { p -> ... }
We want to *rename* them away, so that there are no
246
occurrences of 'wild-id' (with wildCardKey).  The easy
247 248 249
way to do that is to start of with a representative
Id in the in-scope set

250 251 252 253 254 255 256 257
There can be be *occurrences* of wild-id.  For example,
MkCore.mkCoreApp transforms
   e (a /# b)   -->   case (a /# b) of wild { DEFAULT -> e wild }
This is ok provided 'wild' isn't free in 'e', and that's the delicate
thing. Generally, you want to run the simplifier to get rid of the
wild-ids before doing much else.

It's a very dark corner of GHC.  Maybe it should be cleaned up.
Austin Seipp's avatar
Austin Seipp committed
258
-}
259

260 261 262 263 264 265
getMode :: SimplEnv -> SimplifierMode
getMode env = seMode env

setMode :: SimplifierMode -> SimplEnv -> SimplEnv
setMode mode env = env { seMode = mode }

266 267 268
updMode :: (SimplifierMode -> SimplifierMode) -> SimplEnv -> SimplEnv
updMode upd env = env { seMode = upd (seMode env) }

269 270 271
---------------------
extendIdSubst :: SimplEnv -> Id -> SimplSR -> SimplEnv
extendIdSubst env@(SimplEnv {seIdSubst = subst}) var res
272 273
  = ASSERT2( isId var && not (isCoVar var), ppr var )
    env {seIdSubst = extendVarEnv subst var res}
274

275 276 277 278 279 280 281 282 283
extendTvSubst :: SimplEnv -> TyVar -> Type -> SimplEnv
extendTvSubst env@(SimplEnv {seTvSubst = tsubst}) var res
  = ASSERT( isTyVar var )
    env {seTvSubst = extendVarEnv tsubst var res}

extendCvSubst :: SimplEnv -> CoVar -> Coercion -> SimplEnv
extendCvSubst env@(SimplEnv {seCvSubst = csubst}) var co
  = ASSERT( isCoVar var )
    env {seCvSubst = extendVarEnv csubst var co}
284

285 286 287 288 289 290 291 292
---------------------
getInScope :: SimplEnv -> InScopeSet
getInScope env = seInScope env

setInScopeSet :: SimplEnv -> InScopeSet -> SimplEnv
setInScopeSet env in_scope = env {seInScope = in_scope}

setInScope :: SimplEnv -> SimplEnv -> SimplEnv
293 294 295
-- Set the in-scope set, and *zap* the floats
setInScope env env_with_scope
  = env { seInScope = seInScope env_with_scope,
296
          seFloats = emptyFloats }
297 298 299 300 301

setFloats :: SimplEnv -> SimplEnv -> SimplEnv
-- Set the in-scope set *and* the floats
setFloats env env_with_floats
  = env { seInScope = seInScope env_with_floats,
302
          seFloats  = seFloats  env_with_floats }
303 304

addNewInScopeIds :: SimplEnv -> [CoreBndr] -> SimplEnv
305
        -- The new Ids are guaranteed to be freshly allocated
306 307
addNewInScopeIds env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst }) vs
  = env { seInScope = in_scope `extendInScopeSetList` vs,
308 309 310 311 312 313
          seIdSubst = id_subst `delVarEnvList` vs }
        -- Why delete?  Consider
        --      let x = a*b in (x, \x -> x+3)
        -- We add [x |-> a*b] to the substitution, but we must
        -- _delete_ it from the substitution when going inside
        -- the (\x -> ...)!
314

315
modifyInScope :: SimplEnv -> CoreBndr -> SimplEnv
316
-- The variable should already be in scope, but
317 318
-- replace the existing version with this new one
-- which has more information
319
modifyInScope env@(SimplEnv {seInScope = in_scope}) v
320
  = env {seInScope = extendInScopeSet in_scope v}
321 322 323

---------------------
zapSubstEnv :: SimplEnv -> SimplEnv
324
zapSubstEnv env = env {seTvSubst = emptyVarEnv, seCvSubst = emptyVarEnv, seIdSubst = emptyVarEnv}
325

326 327
setSubstEnv :: SimplEnv -> TvSubstEnv -> CvSubstEnv -> SimplIdSubst -> SimplEnv
setSubstEnv env tvs cvs ids = env { seTvSubst = tvs, seCvSubst = cvs, seIdSubst = ids }
328 329

mkContEx :: SimplEnv -> InExpr -> SimplSR
330
mkContEx (SimplEnv { seTvSubst = tvs, seCvSubst = cvs, seIdSubst = ids }) e = ContEx tvs cvs ids e
331

Austin Seipp's avatar
Austin Seipp committed
332 333 334
{-
************************************************************************
*                                                                      *
335
\subsection{Floats}
Austin Seipp's avatar
Austin Seipp committed
336 337
*                                                                      *
************************************************************************
338 339 340 341 342

Note [Simplifier floats]
~~~~~~~~~~~~~~~~~~~~~~~~~
The Floats is a bunch of bindings, classified by a FloatFlag.

343 344 345 346
* All of them satisfy the let/app invariant

Examples

347 348
  NonRec x (y:ys)       FltLifted
  Rec [(x,rhs)]         FltLifted
349

350
  NonRec x* (p:q)       FltOKSpec   -- RHS is WHNF.  Question: why not FltLifted?
351
  NonRec x# (y +# 3)    FltOkSpec   -- Unboxed, but ok-for-spec'n
352

353
  NonRec x* (f y)       FltCareful  -- Strict binding; might fail or diverge
354 355 356 357

Can't happen:
  NonRec x# (a /# b)    -- Might fail; does not satisfy let/app
  NonRec x# (f y)       -- Might diverge; does not satisfy let/app
Austin Seipp's avatar
Austin Seipp committed
358
-}
359 360

data Floats = Floats (OrdList OutBind) FloatFlag
361
        -- See Note [Simplifier floats]
362 363

data FloatFlag
364 365
  = FltLifted   -- All bindings are lifted and lazy
                --  Hence ok to float to top level, or recursive
366

367 368 369 370 371 372 373
  | FltOkSpec   -- All bindings are FltLifted *or*
                --      strict (perhaps because unlifted,
                --      perhaps because of a strict binder),
                --        *and* ok-for-speculation
                --  Hence ok to float out of the RHS
                --  of a lazy non-recursive let binding
                --  (but not to top level, or into a rec group)
374

375 376 377
  | FltCareful  -- At least one binding is strict (or unlifted)
                --      and not guaranteed cheap
                --      Do not float these bindings out of a lazy let
378 379 380 381 382

instance Outputable Floats where
  ppr (Floats binds ff) = ppr ff $$ ppr (fromOL binds)

instance Outputable FloatFlag where
383 384 385
  ppr FltLifted = text "FltLifted"
  ppr FltOkSpec = text "FltOkSpec"
  ppr FltCareful = text "FltCareful"
386

387
andFF :: FloatFlag -> FloatFlag -> FloatFlag
388
andFF FltCareful _          = FltCareful
389
andFF FltOkSpec  FltCareful = FltCareful
390 391
andFF FltOkSpec  _          = FltOkSpec
andFF FltLifted  flt        = flt
392

393
doFloatFromRhs :: TopLevelFlag -> RecFlag -> Bool -> OutExpr -> SimplEnv -> Bool
394
-- If you change this function look also at FloatIn.noFloatFromRhs
395
doFloatFromRhs lvl rec str rhs (SimplEnv {seFloats = Floats fs ff})
396 397
  =  not (isNilOL fs) && want_to_float && can_float
  where
Austin Seipp's avatar
Austin Seipp committed
398
     want_to_float = isTopLevel lvl || exprIsCheap rhs || exprIsExpandable rhs
Simon Peyton Jones's avatar
Simon Peyton Jones committed
399
                     -- See Note [Float when cheap or expandable]
400
     can_float = case ff of
401 402 403
                   FltLifted  -> True
                   FltOkSpec  -> isNotTopLevel lvl && isNonRec rec
                   FltCareful -> isNotTopLevel lvl && isNonRec rec && str
404

Austin Seipp's avatar
Austin Seipp committed
405
{-
406 407 408 409 410
Note [Float when cheap or expandable]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to float a let from a let if the residual RHS is
   a) cheap, such as (\x. blah)
   b) expandable, such as (f b) if f is CONLIKE
Austin Seipp's avatar
Austin Seipp committed
411
But there are
412 413 414
  - cheap things that are not expandable (eg \x. expensive)
  - expandable things that are not cheap (eg (f b) where b is CONLIKE)
so we must take the 'or' of the two.
Austin Seipp's avatar
Austin Seipp committed
415
-}
416 417 418 419 420

emptyFloats :: Floats
emptyFloats = Floats nilOL FltLifted

unitFloat :: OutBind -> Floats
421 422 423 424 425 426 427
-- This key function constructs a singleton float with the right form
unitFloat bind = Floats (unitOL bind) (flag bind)
  where
    flag (Rec {})                = FltLifted
    flag (NonRec bndr rhs)
      | not (isStrictId bndr)    = FltLifted
      | exprOkForSpeculation rhs = FltOkSpec  -- Unlifted, and lifted but ok-for-spec (eg HNF)
428
      | otherwise                = ASSERT2( not (isUnliftedType (idType bndr)), ppr bndr )
429 430
                                   FltCareful
      -- Unlifted binders can only be let-bound if exprOkForSpeculation holds
431 432 433 434 435 436 437

addNonRec :: SimplEnv -> OutId -> OutExpr -> SimplEnv
-- Add a non-recursive binding and extend the in-scope set
-- The latter is important; the binder may already be in the
-- in-scope set (although it might also have been created with newId)
-- but it may now have more IdInfo
addNonRec env id rhs
438
  = id `seq`   -- This seq forces the Id, and hence its IdInfo,
439
               -- and hence any inner substitutions
440
    env { seFloats = seFloats env `addFlts` unitFloat (NonRec id rhs),
441
          seInScope = extendInScopeSet (seInScope env) id }
442

443
extendFloats :: SimplEnv -> OutBind -> SimplEnv
444
-- Add these bindings to the floats, and extend the in-scope env too
445 446
extendFloats env bind
  = env { seFloats  = seFloats env `addFlts` unitFloat bind,
447
          seInScope = extendInScopeSetList (seInScope env) bndrs }
448
  where
449
    bndrs = bindersOf bind
450

451
addFloats :: SimplEnv -> SimplEnv -> SimplEnv
452 453
-- Add the floats for env2 to env1;
-- *plus* the in-scope set for env2, which is bigger
454
-- than that for env1
455
addFloats env1 env2
456
  = env1 {seFloats = seFloats env1 `addFlts` seFloats env2,
457
          seInScope = seInScope env2 }
458 459 460 461 462 463 464 465 466 467 468 469 470 471

addFlts :: Floats -> Floats -> Floats
addFlts (Floats bs1 l1) (Floats bs2 l2)
  = Floats (bs1 `appOL` bs2) (l1 `andFF` l2)

zapFloats :: SimplEnv -> SimplEnv
zapFloats env = env { seFloats = emptyFloats }

addRecFloats :: SimplEnv -> SimplEnv -> SimplEnv
-- Flattens the floats from env2 into a single Rec group,
-- prepends the floats from env1, and puts the result back in env2
-- This is all very specific to the way recursive bindings are
-- handled; see Simplify.simplRecBind
addRecFloats env1 env2@(SimplEnv {seFloats = Floats bs ff})
Ian Lynagh's avatar
Ian Lynagh committed
472
  = ASSERT2( case ff of { FltLifted -> True; _ -> False }, ppr (fromOL bs) )
473 474 475
    env2 {seFloats = seFloats env1 `addFlts` unitFloat (Rec (flattenBinds (fromOL bs)))}

wrapFloats :: SimplEnv -> OutExpr -> OutExpr
476 477 478 479
-- Wrap the floats around the expression; they should all
-- satisfy the let/app invariant, so mkLets should do the job just fine
wrapFloats (SimplEnv {seFloats = Floats bs _}) body
  = foldrOL Let body bs
480

481
getFloatBinds :: SimplEnv -> [CoreBind]
482 483
getFloatBinds (SimplEnv {seFloats = Floats bs _})
  = fromOL bs
484 485

isEmptyFloats :: SimplEnv -> Bool
486 487
isEmptyFloats (SimplEnv {seFloats = Floats bs _})
  = isNilOL bs
488

Peter Wortmann's avatar
Peter Wortmann committed
489 490 491 492 493 494
mapFloats :: SimplEnv -> ((Id,CoreExpr) -> (Id,CoreExpr)) -> SimplEnv
mapFloats env@SimplEnv { seFloats = Floats fs ff } fun
   = env { seFloats = Floats (mapOL app fs) ff }
   where
    app (NonRec b e) = case fun (b,e) of (b',e') -> NonRec b' e'
    app (Rec bs)     = Rec (map fun bs)
495

Peter Wortmann's avatar
Peter Wortmann committed
496
{-
Austin Seipp's avatar
Austin Seipp committed
497 498
************************************************************************
*                                                                      *
499
                Substitution of Vars
Austin Seipp's avatar
Austin Seipp committed
500 501
*                                                                      *
************************************************************************
502

503 504 505 506
Note [Global Ids in the substitution]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We look up even a global (eg imported) Id in the substitution. Consider
   case X.g_34 of b { (a,b) ->  ... case X.g_34 of { (p,q) -> ...} ... }
507
The binder-swap in the occurrence analyser will add a binding
508
for a LocalId version of g (with the same unique though):
509 510
   case X.g_34 of b { (a,b) ->  let g_34 = b in
                                ... case X.g_34 of { (p,q) -> ...} ... }
511 512
So we want to look up the inner X.g_34 in the substitution, where we'll
find that it has been substituted by b.  (Or conceivably cloned.)
Austin Seipp's avatar
Austin Seipp committed
513
-}
514

515
substId :: SimplEnv -> InId -> SimplSR
516
-- Returns DoneEx only on a non-Var expression
517 518
substId (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
  = case lookupVarEnv ids v of  -- Note [Global Ids in the substitution]
519 520 521
        Nothing               -> DoneId (refineFromInScope in_scope v)
        Just (DoneId v)       -> DoneId (refineFromInScope in_scope v)
        Just (DoneEx (Var v)) -> DoneId (refineFromInScope in_scope v)
522 523 524 525 526
        Just res              -> res    -- DoneEx non-var, or ContEx

        -- Get the most up-to-date thing from the in-scope set
        -- Even though it isn't in the substitution, it may be in
        -- the in-scope set with better IdInfo
527 528
refineFromInScope :: InScopeSet -> Var -> Var
refineFromInScope in_scope v
529
  | isLocalId v = case lookupInScope in_scope v of
530 531
                  Just v' -> v'
                  Nothing -> WARN( True, ppr v ) v  -- This is an error!
532
  | otherwise = v
533

534
lookupRecBndr :: SimplEnv -> InId -> OutId
535 536 537 538
-- Look up an Id which has been put into the envt by simplRecBndrs,
-- but where we have not yet done its RHS
lookupRecBndr (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
  = case lookupVarEnv ids v of
539 540
        Just (DoneId v) -> v
        Just _ -> pprPanic "lookupRecBndr" (ppr v)
541
        Nothing -> refineFromInScope in_scope v
542

Austin Seipp's avatar
Austin Seipp committed
543 544 545
{-
************************************************************************
*                                                                      *
546
\section{Substituting an Id binder}
Austin Seipp's avatar
Austin Seipp committed
547 548
*                                                                      *
************************************************************************
549 550 551


These functions are in the monad only so that they can be made strict via seq.
Austin Seipp's avatar
Austin Seipp committed
552
-}
553

554
simplBinders :: SimplEnv -> [InBndr] -> SimplM (SimplEnv, [OutBndr])
555
simplBinders  env bndrs = mapAccumLM simplBinder  env bndrs
556 557

-------------
558
simplBinder :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
559 560 561 562 563 564
-- Used for lambda and case-bound variables
-- Clone Id if necessary, substitute type
-- Return with IdInfo already substituted, but (fragile) occurrence info zapped
-- The substitution is extended only if the variable is cloned, because
-- we *don't* need to use it to track occurrence info.
simplBinder env bndr
565 566 567 568
  | isTyVar bndr  = do  { let (env', tv) = substTyVarBndr env bndr
                        ; seqTyVar tv `seq` return (env', tv) }
  | otherwise     = do  { let (env', id) = substIdBndr env bndr
                        ; seqId id `seq` return (env', id) }
569

570
---------------
571
simplNonRecBndr :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
572
-- A non-recursive let binder
573
simplNonRecBndr env id
574 575
  = do  { let (env1, id1) = substIdBndr env id
        ; seqId id1 `seq` return (env1, id1) }
576 577

---------------
578
simplRecBndrs :: SimplEnv -> [InBndr] -> SimplM SimplEnv
579
-- Recursive let binders
Ian Lynagh's avatar
Ian Lynagh committed
580
simplRecBndrs env@(SimplEnv {}) ids
581 582
  = do  { let (env1, ids1) = mapAccumL substIdBndr env ids
        ; seqIds ids1 `seq` return env1 }
583 584

---------------
585 586 587 588 589 590 591
substIdBndr :: SimplEnv -> InBndr -> (SimplEnv, OutBndr)
-- Might be a coercion variable
substIdBndr env bndr
  | isCoVar bndr  = substCoVarBndr env bndr
  | otherwise     = substNonCoVarIdBndr env bndr

---------------
592 593 594
substNonCoVarIdBndr
   :: SimplEnv
   -> InBndr    -- Env and binder to transform
595
   -> (SimplEnv, OutBndr)
596
-- Clone Id if necessary, substitute its type
597 598 599 600 601 602 603
-- Return an Id with its
--      * Type substituted
--      * UnfoldingInfo, Rules, WorkerInfo zapped
--      * Fragile OccInfo (only) zapped: Note [Robust OccInfo]
--      * Robust info, retained especially arity and demand info,
--         so that they are available to occurrences that occur in an
--         earlier binding of a letrec
604 605 606 607 608 609
--
-- For the robust info, see Note [Arity robustness]
--
-- Augment the substitution  if the unique changed
-- Extend the in-scope set with the new Id
--
610 611 612
-- Similar to CoreSubst.substIdBndr, except that
--      the type of id_subst differs
--      all fragile info is zapped
613 614 615
substNonCoVarIdBndr env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst })
                    old_id
  = ASSERT2( not (isCoVar old_id), ppr old_id )
616 617
    (env { seInScope = in_scope `extendInScopeSet` new_id,
           seIdSubst = new_subst }, new_id)
618
  where
619
    id1    = uniqAway in_scope old_id
620
    id2    = substIdType env id1
621 622
    new_id = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
                                        -- and fragile OccInfo
623

624 625 626
        -- Extend the substitution if the unique has changed,
        -- or there's some useful occurrence information
        -- See the notes with substTyVarBndr for the delSubstEnv
627
    new_subst | new_id /= old_id
628 629 630
              = extendVarEnv id_subst old_id (DoneId new_id)
              | otherwise
              = delVarEnv id_subst old_id
631

632 633 634 635 636
------------------------------------
seqTyVar :: TyVar -> ()
seqTyVar b = b `seq` ()

seqId :: Id -> ()
637 638 639
seqId id = seqType (idType id)  `seq`
           idInfo id            `seq`
           ()
640 641 642 643 644

seqIds :: [Id] -> ()
seqIds []       = ()
seqIds (id:ids) = seqId id `seq` seqIds ids

Austin Seipp's avatar
Austin Seipp committed
645
{-
646 647 648 649 650
Note [Arity robustness]
~~~~~~~~~~~~~~~~~~~~~~~
We *do* transfer the arity from from the in_id of a let binding to the
out_id.  This is important, so that the arity of an Id is visible in
its own RHS.  For example:
651
        f = \x. ....g (\y. f y)....
Gabor Greif's avatar
typos  
Gabor Greif committed
652
We can eta-reduce the arg to g, because f is a value.  But that
653
needs to be visible.
654 655

This interacts with the 'state hack' too:
656 657 658 659 660
        f :: Bool -> IO Int
        f = \x. case x of
                  True  -> f y
                  False -> \s -> ...
Can we eta-expand f?  Only if we see that f has arity 1, and then we
661 662 663
take advantage of the 'state hack' on the result of
(f y) :: State# -> (State#, Int) to expand the arity one more.

664
There is a disadvantage though.  Making the arity visible in the RHS
665
allows us to eta-reduce
666
        f = \x -> f x
667
to
668
        f = f
669
which technically is not sound.   This is very much a corner case, so
670
I'm not worried about it.  Another idea is to ensure that f's arity
671 672
never decreases; its arity started as 1, and we should never eta-reduce
below that.
673 674


675 676 677 678 679
Note [Robust OccInfo]
~~~~~~~~~~~~~~~~~~~~~
It's important that we *do* retain the loop-breaker OccInfo, because
that's what stops the Id getting inlined infinitely, in the body of
the letrec.
Austin Seipp's avatar
Austin Seipp committed
680
-}
681

682

Austin Seipp's avatar
Austin Seipp committed
683 684 685
{-
************************************************************************
*                                                                      *
686
                Impedence matching to type substitution
Austin Seipp's avatar
Austin Seipp committed
687 688 689
*                                                                      *
************************************************************************
-}
690

691 692 693
getTCvSubst :: SimplEnv -> TCvSubst
getTCvSubst (SimplEnv { seInScope = in_scope, seTvSubst = tv_env, seCvSubst = cv_env })
  = mkTCvSubst in_scope (tv_env, cv_env)
694

695
substTy :: SimplEnv -> Type -> Type
696
substTy env ty = Type.substTy (getTCvSubst env) ty
697

698
substTyVar :: SimplEnv -> TyVar -> Type
699
substTyVar env tv = Type.substTyVar (getTCvSubst env) tv
700

701
substTyVarBndr :: SimplEnv -> TyVar -> (SimplEnv, TyVar)
702
substTyVarBndr env tv
703 704 705
  = case Type.substTyVarBndr (getTCvSubst env) tv of
        (TCvSubst in_scope' tv_env' cv_env', tv')
           -> (env { seInScope = in_scope', seTvSubst = tv_env', seCvSubst = cv_env' }, tv')
706 707

substCoVar :: SimplEnv -> CoVar -> Coercion
708
substCoVar env tv = Coercion.substCoVar (getTCvSubst env) tv
709 710 711

substCoVarBndr :: SimplEnv -> CoVar -> (SimplEnv, CoVar)
substCoVarBndr env cv
712 713
  = case Coercion.substCoVarBndr (getTCvSubst env) cv of
        (TCvSubst in_scope' tv_env' cv_env', cv')
714
           -> (env { seInScope = in_scope', seTvSubst = tv_env', seCvSubst = cv_env' }, cv')
715 716

substCo :: SimplEnv -> Coercion -> Coercion
717
substCo env co = Coercion.substCo (getTCvSubst env) co
718

719 720
------------------
substIdType :: SimplEnv -> Id -> Id
721 722 723 724 725 726
substIdType (SimplEnv { seInScope = in_scope, seTvSubst = tv_env, seCvSubst = cv_env }) id
  |  (isEmptyVarEnv tv_env && isEmptyVarEnv cv_env)
  || isEmptyVarSet (tyCoVarsOfType old_ty)
  = id
  | otherwise = Id.setIdType id (Type.substTy (TCvSubst in_scope tv_env cv_env) old_ty)
                -- The tyCoVarsOfType is cheaper than it looks
727 728
                -- because we cache the free tyvars of the type
                -- in a Note in the id's type itself
729 730
  where
    old_ty = idType id