SimplEnv.lhs 30.5 KB
Newer Older
1
%
Simon Peyton Jones's avatar
Simon Peyton Jones committed
2
% (c) The AQUA Project, Glasgow University, 1993-1998
3 4 5 6 7
%
\section[SimplMonad]{The simplifier Monad}

\begin{code}
module SimplEnv (
8 9
        InId, InBind, InExpr, InAlt, InArg, InType, InBndr, InVar,
        OutId, OutTyVar, OutBind, OutExpr, OutAlt, OutArg, OutType, OutBndr, OutVar,
10
        InCoercion, OutCoercion,
11

12 13
        -- The simplifier mode
        setMode, getMode, updMode,
14

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

22
        SimplSR(..), mkContEx, substId, lookupRecBndr,
23

24 25 26 27 28
        simplNonRecBndr, simplRecBndrs, simplLamBndr, simplLamBndrs,
        simplBinder, simplBinders, addBndrRules,
        substExpr, substTy, substTyVar, getTvSubst,
        getCvSubst, substCo, substCoVar,
        mkCoreSubst,
29

30 31 32
        -- Floats
        Floats, emptyFloats, isEmptyFloats, addNonRec, addFloats, extendFloats,
        wrapFloats, floatBinds, setFloats, zapFloats, addRecFloats,
33
        doFloatFromRhs, getFloatBinds, getFloats, mapFloats
34 35 36 37
    ) where

#include "HsVersions.h"

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

import Data.List
62 63 64
\end{code}

%************************************************************************
65
%*                                                                      *
66
\subsection[Simplify-types]{Type declarations}
67
%*                                                                      *
68 69 70
%************************************************************************

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

81
type OutBndr     = CoreBndr
82 83 84 85
type OutVar      = Var                  -- Cloned
type OutId       = Id                   -- Cloned
type OutTyVar    = TyVar                -- Cloned
type OutType     = Type                 -- Cloned
86
type OutCoercion = Coercion
87 88 89 90
type OutBind     = CoreBind
type OutExpr     = CoreExpr
type OutAlt      = CoreAlt
type OutArg      = CoreArg
91 92 93
\end{code}

%************************************************************************
94
%*                                                                      *
95
\subsubsection{The @SimplEnv@ type}
96
%*                                                                      *
97 98 99 100 101 102
%************************************************************************


\begin{code}
data SimplEnv
  = SimplEnv {
103
     ----------- Static part of the environment -----------
104
     -- Static in the sense of lexically scoped,
105 106
     -- wrt the original expression

107
        seMode      :: SimplifierMode,
108

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

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

118 119 120 121 122 123
        -- 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]
124 125
    }

126
type StaticEnv = SimplEnv       -- Just the static part is relevant
127

128 129 130
pprSimplEnv :: SimplEnv -> SDoc
-- Used for debugging; selective
pprSimplEnv env
Ian Lynagh's avatar
Ian Lynagh committed
131
  = vcat [ptext (sLit "TvSubst:") <+> ppr (seTvSubst env),
132
          ptext (sLit "IdSubst:") <+> ppr (seIdSubst env),
133 134 135 136 137 138
          ptext (sLit "InScope:") <+> vcat (map ppr_one in_scope_vars)
    ]
  where
   in_scope_vars = varEnvElts (getInScopeVars (seInScope env))
   ppr_one v | isId v = ppr v <+> ppr (idUnfolding v)
             | otherwise = ppr v
139

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

data SimplSR
144 145 146
  = DoneEx OutExpr              -- Completed term
  | DoneId OutId                -- Completed term variable
  | ContEx TvSubstEnv           -- A suspended substitution
147
           CvSubstEnv
148 149
           SimplIdSubst
           InExpr
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
150

151
instance Outputable SimplSR where
Ian Lynagh's avatar
Ian Lynagh committed
152 153
  ppr (DoneEx e) = ptext (sLit "DoneEx") <+> ppr e
  ppr (DoneId v) = ptext (sLit "DoneId") <+> ppr v
154
  ppr (ContEx _tv _cv _id e) = vcat [ptext (sLit "ContEx") <+> ppr e {-,
155 156 157 158 159
                                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
160 161
\end{code}

162 163
Note [SimplEnv invariants]
~~~~~~~~~~~~~~~~~~~~~~~~~~
164 165 166 167 168 169
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
170

171 172 173
        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.
174

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

seIdSubst:
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198
        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
199

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

203
* We *always, always* finish by looking up in the in-scope set
204
  any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
205 206 207
  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.
208 209 210 211 212 213 214 215

  [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
216 217 218
  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.
219 220 221 222

* 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:
223
        case y of x { ... }
224 225 226 227
  That's why the "set" is actually a VarEnv Var


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

238 239 240 241 242 243 244 245 246 247
init_in_scope :: InScopeSet
init_in_scope = mkInScopeSet (unitVarSet (mkWildValBinder unitTy))
              -- See Note [WildCard binders]
\end{code}

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
248
occurrences of 'wild-id' (with wildCardKey).  The easy
249 250 251
way to do that is to start of with a representative
Id in the in-scope set

252 253 254 255 256 257 258 259
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.
260 261

\begin{code}
262 263 264 265 266 267
getMode :: SimplEnv -> SimplifierMode
getMode env = seMode env

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

268 269 270
updMode :: (SimplifierMode -> SimplifierMode) -> SimplEnv -> SimplEnv
updMode upd env = env { seMode = upd (seMode env) }

271 272 273
---------------------
extendIdSubst :: SimplEnv -> Id -> SimplSR -> SimplEnv
extendIdSubst env@(SimplEnv {seIdSubst = subst}) var res
274 275
  = ASSERT2( isId var && not (isCoVar var), ppr var )
    env {seIdSubst = extendVarEnv subst var res}
276 277 278 279 280

extendTvSubst :: SimplEnv -> TyVar -> Type -> SimplEnv
extendTvSubst env@(SimplEnv {seTvSubst = subst}) var res
  = env {seTvSubst = extendVarEnv subst var res}

281 282 283 284
extendCvSubst :: SimplEnv -> CoVar -> Coercion -> SimplEnv
extendCvSubst env@(SimplEnv {seCvSubst = subst}) var res
  = env {seCvSubst = extendVarEnv subst var res}

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 332 333
\end{code}


334 335

%************************************************************************
336
%*                                                                      *
337
\subsection{Floats}
338
%*                                                                      *
339 340 341 342 343 344
%************************************************************************

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

345 346
  NonRec x (y:ys)       FltLifted
  Rec [(x,rhs)]         FltLifted
347

348
  NonRec x# (y +# 3)    FltOkSpec   -- Unboxed, but ok-for-spec'n
349

350 351 352 353
  NonRec x# (a /# b)    FltCareful
  NonRec x* (f y)       FltCareful  -- Strict binding; might fail or diverge
  NonRec x# (f y)       FltCareful  -- Unboxed binding: might fail or diverge
                                    --  (where f :: Int -> Int#)
354 355 356

\begin{code}
data Floats = Floats (OrdList OutBind) FloatFlag
357
        -- See Note [Simplifier floats]
358 359

data FloatFlag
360 361
  = FltLifted   -- All bindings are lifted and lazy
                --  Hence ok to float to top level, or recursive
362

363 364 365 366 367 368 369
  | 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)
370

371 372 373
  | FltCareful  -- At least one binding is strict (or unlifted)
                --      and not guaranteed cheap
                --      Do not float these bindings out of a lazy let
374 375 376 377 378

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

instance Outputable FloatFlag where
Ian Lynagh's avatar
Ian Lynagh committed
379 380 381
  ppr FltLifted = ptext (sLit "FltLifted")
  ppr FltOkSpec = ptext (sLit "FltOkSpec")
  ppr FltCareful = ptext (sLit "FltCareful")
382

383
andFF :: FloatFlag -> FloatFlag -> FloatFlag
384
andFF FltCareful _          = FltCareful
385
andFF FltOkSpec  FltCareful = FltCareful
386 387
andFF FltOkSpec  _          = FltOkSpec
andFF FltLifted  flt        = flt
388 389 390

classifyFF :: CoreBind -> FloatFlag
classifyFF (Rec _) = FltLifted
391
classifyFF (NonRec bndr rhs)
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
392
  | not (isStrictId bndr)    = FltLifted
393
  | exprOkForSpeculation rhs = FltOkSpec
394
  | otherwise                = FltCareful
395

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

409 410 411 412 413 414 415 416 417
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
But there are 
  - 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.
418 419 420 421 422 423 424 425 426 427 428 429 430 431 432

\begin{code}
emptyFloats :: Floats
emptyFloats = Floats nilOL FltLifted

unitFloat :: OutBind -> Floats
-- A single-binding float
unitFloat bind = Floats (unitOL bind) (classifyFF bind)

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
433
  = id `seq`   -- This seq forces the Id, and hence its IdInfo,
434
               -- and hence any inner substitutions
435
    env { seFloats = seFloats env `addFlts` unitFloat (NonRec id rhs),
436
          seInScope = extendInScopeSet (seInScope env) id }
437

438 439
mapFloats :: SimplEnv -> ((Id,CoreExpr) -> (Id,CoreExpr)) -> SimplEnv
mapFloats env@SimplEnv { seFloats = Floats fs ff } fun
440 441
   = env { seFloats = Floats (mapOL app fs) ff }
   where
442 443
     app (NonRec b e) = case fun (b,e) of (b',e') -> NonRec b' e'
     app (Rec bs)     = Rec (map fun bs)
444

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

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

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
474
  = ASSERT2( case ff of { FltLifted -> True; _ -> False }, ppr (fromOL bs) )
475 476 477 478 479 480 481 482 483 484 485 486
    env2 {seFloats = seFloats env1 `addFlts` unitFloat (Rec (flattenBinds (fromOL bs)))}

wrapFloats :: SimplEnv -> OutExpr -> OutExpr
wrapFloats env expr = wrapFlts (seFloats env) expr

wrapFlts :: Floats -> OutExpr -> OutExpr
-- Wrap the floats around the expression, using case-binding where necessary
wrapFlts (Floats bs _) body = foldrOL wrap body bs
  where
    wrap (Rec prs)    body = Let (Rec prs) body
    wrap (NonRec b r) body = bindNonRec b r body

487 488 489 490 491
getFloatBinds :: SimplEnv -> [CoreBind]
getFloatBinds env = floatBinds (seFloats env)

getFloats :: SimplEnv -> Floats
getFloats env = seFloats env
492 493 494 495 496

isEmptyFloats :: SimplEnv -> Bool
isEmptyFloats env = isEmptyFlts (seFloats env)

isEmptyFlts :: Floats -> Bool
497
isEmptyFlts (Floats bs _) = isNilOL bs
498 499 500 501 502 503

floatBinds :: Floats -> [OutBind]
floatBinds (Floats bs _) = fromOL bs
\end{code}


504
%************************************************************************
505 506 507
%*                                                                      *
                Substitution of Vars
%*                                                                      *
508 509
%************************************************************************

510 511 512 513
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) -> ...} ... }
Simon Peyton Jones's avatar
Simon Peyton Jones committed
514
The binder-swap in the occurrence analyser will add a binding
515
for a LocalId version of g (with the same unique though):
516 517
   case X.g_34 of b { (a,b) ->  let g_34 = b in
                                ... case X.g_34 of { (p,q) -> ...} ... }
518 519
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.)
520 521

\begin{code}
522
substId :: SimplEnv -> InId -> SimplSR
523
-- Returns DoneEx only on a non-Var expression
524 525 526 527 528 529 530 531 532 533
substId (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
  = case lookupVarEnv ids v of  -- Note [Global Ids in the substitution]
        Nothing               -> DoneId (refine in_scope v)
        Just (DoneId v)       -> DoneId (refine in_scope v)
        Just (DoneEx (Var v)) -> DoneId (refine in_scope v)
        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
Ian Lynagh's avatar
Ian Lynagh committed
534
refine :: InScopeSet -> Var -> Var
535
refine in_scope v
536
  | isLocalId v = case lookupInScope in_scope v of
537 538
                  Just v' -> v'
                  Nothing -> WARN( True, ppr v ) v  -- This is an error!
539
  | otherwise = v
540

541
lookupRecBndr :: SimplEnv -> InId -> OutId
542 543 544 545
-- 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
546 547 548
        Just (DoneId v) -> v
        Just _ -> pprPanic "lookupRecBndr" (ppr v)
        Nothing -> refine in_scope v
549 550 551 552
\end{code}


%************************************************************************
553
%*                                                                      *
554
\section{Substituting an Id binder}
555
%*                                                                      *
556 557 558 559 560 561
%************************************************************************


These functions are in the monad only so that they can be made strict via seq.

\begin{code}
562
simplBinders, simplLamBndrs
563
        :: SimplEnv -> [InBndr] -> SimplM (SimplEnv, [OutBndr])
564 565
simplBinders  env bndrs = mapAccumLM simplBinder  env bndrs
simplLamBndrs env bndrs = mapAccumLM simplLamBndr env bndrs
566 567

-------------
568
simplBinder :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
569 570 571 572 573 574
-- 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
575 576 577 578
  | 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) }
579 580 581 582

-------------
simplLamBndr :: SimplEnv -> Var -> SimplM (SimplEnv, Var)
-- Used for lambda binders.  These sometimes have unfoldings added by
583
-- the worker/wrapper pass that must be preserved, because they can't
584
-- be reconstructed from context.  For example:
585 586
--      f x = case x of (a,b) -> fw a b x
--      fw a b x{=(a,b)} = ...
587 588
-- The "{=(a,b)}" is an unfolding we can't reconstruct otherwise.
simplLamBndr env bndr
589
  | isId bndr && hasSomeUnfolding old_unf = seqId id2 `seq` return (env2, id2)  -- Special case
590
  | otherwise                             = simplBinder env bndr                -- Normal case
591 592
  where
    old_unf = idUnfolding bndr
593
    (env1, id1) = substIdBndr env bndr
Simon Marlow's avatar
Simon Marlow committed
594
    id2  = id1 `setIdUnfolding` substUnfolding env old_unf
595
    env2 = modifyInScope env1 id2
596

597
---------------
598
simplNonRecBndr :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
599
-- A non-recursive let binder
600
simplNonRecBndr env id
601 602
  = do  { let (env1, id1) = substIdBndr env id
        ; seqId id1 `seq` return (env1, id1) }
603 604

---------------
605
simplRecBndrs :: SimplEnv -> [InBndr] -> SimplM SimplEnv
606
-- Recursive let binders
Ian Lynagh's avatar
Ian Lynagh committed
607
simplRecBndrs env@(SimplEnv {}) ids
608 609
  = do  { let (env1, ids1) = mapAccumL substIdBndr env ids
        ; seqIds ids1 `seq` return env1 }
610 611

---------------
612 613 614 615 616 617 618
substIdBndr :: SimplEnv -> InBndr -> (SimplEnv, OutBndr)
-- Might be a coercion variable
substIdBndr env bndr
  | isCoVar bndr  = substCoVarBndr env bndr
  | otherwise     = substNonCoVarIdBndr env bndr

---------------
619 620 621
substNonCoVarIdBndr
   :: SimplEnv
   -> InBndr    -- Env and binder to transform
622
   -> (SimplEnv, OutBndr)
623
-- Clone Id if necessary, substitute its type
624 625 626 627 628 629 630
-- 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
631 632 633 634 635 636
--
-- For the robust info, see Note [Arity robustness]
--
-- Augment the substitution  if the unique changed
-- Extend the in-scope set with the new Id
--
637 638 639
-- Similar to CoreSubst.substIdBndr, except that
--      the type of id_subst differs
--      all fragile info is zapped
640 641 642
substNonCoVarIdBndr env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst })
                    old_id
  = ASSERT2( not (isCoVar old_id), ppr old_id )
643 644
    (env { seInScope = in_scope `extendInScopeSet` new_id,
           seIdSubst = new_subst }, new_id)
645
  where
646
    id1    = uniqAway in_scope old_id
647
    id2    = substIdType env id1
648 649
    new_id = zapFragileIdInfo id2       -- Zaps rules, worker-info, unfolding
                                        -- and fragile OccInfo
650

651 652 653
        -- Extend the substitution if the unique has changed,
        -- or there's some useful occurrence information
        -- See the notes with substTyVarBndr for the delSubstEnv
654
    new_subst | new_id /= old_id
655 656 657
              = extendVarEnv id_subst old_id (DoneId new_id)
              | otherwise
              = delVarEnv id_subst old_id
658 659
\end{code}

660 661 662 663 664 665
\begin{code}
------------------------------------
seqTyVar :: TyVar -> ()
seqTyVar b = b `seq` ()

seqId :: Id -> ()
666 667 668
seqId id = seqType (idType id)  `seq`
           idInfo id            `seq`
           ()
669 670 671 672 673 674 675 676 677 678 679 680

seqIds :: [Id] -> ()
seqIds []       = ()
seqIds (id:ids) = seqId id `seq` seqIds ids
\end{code}


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:
681
        f = \x. ....g (\y. f y)....
Gabor Greif's avatar
typos  
Gabor Greif committed
682
We can eta-reduce the arg to g, because f is a value.  But that
683
needs to be visible.
684 685

This interacts with the 'state hack' too:
686 687 688 689 690
        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
691 692 693
take advantage of the 'state hack' on the result of
(f y) :: State# -> (State#, Int) to expand the arity one more.

694
There is a disadvantage though.  Making the arity visible in the RHS
695
allows us to eta-reduce
696
        f = \x -> f x
697
to
698
        f = f
699
which technically is not sound.   This is very much a corner case, so
700
I'm not worried about it.  Another idea is to ensure that f's arity
701 702
never decreases; its arity started as 1, and we should never eta-reduce
below that.
703 704


705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720
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.


Note [Rules in a letrec]
~~~~~~~~~~~~~~~~~~~~~~~~
After creating fresh binders for the binders of a letrec, we
substitute the RULES and add them back onto the binders; this is done
*before* processing any of the RHSs.  This is important.  Manuel found
cases where he really, really wanted a RULE for a recursive function
to apply in that function's own right-hand side.

See Note [Loop breaking and RULES] in OccAnal.
721 722 723


\begin{code}
724
addBndrRules :: SimplEnv -> InBndr -> OutBndr -> (SimplEnv, OutBndr)
725
-- Rules are added back in to to the bin
726 727
addBndrRules env in_id out_id
  | isEmptySpecInfo old_rules = (env, out_id)
728
  | otherwise = (modifyInScope env final_id, final_id)
729
  where
730
    subst     = mkCoreSubst (text "local rules") env
731
    old_rules = idSpecialisation in_id
732
    new_rules = CoreSubst.substSpec subst out_id old_rules
733
    final_id  = out_id `setIdSpecialisation` new_rules
734 735 736
\end{code}


737
%************************************************************************
738 739 740
%*                                                                      *
                Impedence matching to type substitution
%*                                                                      *
741 742 743
%************************************************************************

\begin{code}
744 745 746 747
getTvSubst :: SimplEnv -> TvSubst
getTvSubst (SimplEnv { seInScope = in_scope, seTvSubst = tv_env })
  = mkTvSubst in_scope tv_env

748 749 750 751
getCvSubst :: SimplEnv -> CvSubst
getCvSubst (SimplEnv { seInScope = in_scope, seTvSubst = tv_env, seCvSubst = cv_env })
  = CvSubst in_scope tv_env cv_env

752
substTy :: SimplEnv -> Type -> Type
753
substTy env ty = Type.substTy (getTvSubst env) ty
754

755
substTyVar :: SimplEnv -> TyVar -> Type
756 757
substTyVar env tv = Type.substTyVar (getTvSubst env) tv

758
substTyVarBndr :: SimplEnv -> TyVar -> (SimplEnv, TyVar)
759 760
substTyVarBndr env tv
  = case Type.substTyVarBndr (getTvSubst env) tv of
761 762
        (TvSubst in_scope' tv_env', tv')
           -> (env { seInScope = in_scope', seTvSubst = tv_env' }, tv')
763 764 765 766 767 768 769

substCoVar :: SimplEnv -> CoVar -> Coercion
substCoVar env tv = Coercion.substCoVar (getCvSubst env) tv

substCoVarBndr :: SimplEnv -> CoVar -> (SimplEnv, CoVar)
substCoVarBndr env cv
  = case Coercion.substCoVarBndr (getCvSubst env) cv of
770 771
        (CvSubst in_scope' tv_env' cv_env', cv')
           -> (env { seInScope = in_scope', seTvSubst = tv_env', seCvSubst = cv_env' }, cv')
772 773 774

substCo :: SimplEnv -> Coercion -> Coercion
substCo env co = Coercion.substCo (getCvSubst env) co
775 776 777 778 779 780

-- When substituting in rules etc we can get CoreSubst to do the work
-- But CoreSubst uses a simpler form of IdSubstEnv, so we must impedence-match
-- here.  I think the this will not usually result in a lot of work;
-- the substitutions are typically small, and laziness will avoid work in many cases.

781
mkCoreSubst  :: SDoc -> SimplEnv -> CoreSubst.Subst
782 783
mkCoreSubst doc (SimplEnv { seInScope = in_scope, seTvSubst = tv_env, seCvSubst = cv_env, seIdSubst = id_env })
  = mk_subst tv_env cv_env id_env
784
  where
785
    mk_subst tv_env cv_env id_env = CoreSubst.mkSubst in_scope tv_env cv_env (mapVarEnv fiddle id_env)
786

787 788 789
    fiddle (DoneEx e)          = e
    fiddle (DoneId v)          = Var v
    fiddle (ContEx tv cv id e) = CoreSubst.substExpr (text "mkCoreSubst" <+> doc) (mk_subst tv cv id) e
790
                                                -- Don't shortcut here
791

792 793
------------------
substIdType :: SimplEnv -> Id -> Id
794
substIdType (SimplEnv { seInScope = in_scope,  seTvSubst = tv_env }) id
795
  | isEmptyVarEnv tv_env || isEmptyVarSet (tyVarsOfType old_ty) = id
796 797 798 799
  | otherwise = Id.setIdType id (Type.substTy (TvSubst in_scope tv_env) old_ty)
                -- The tyVarsOfType is cheaper than it looks
                -- because we cache the free tyvars of the type
                -- in a Note in the id's type itself
800 801 802 803
  where
    old_ty = idType id

------------------
804 805
substExpr :: SDoc -> SimplEnv -> CoreExpr -> CoreExpr
substExpr doc env
806 807
  = CoreSubst.substExpr (text "SimplEnv.substExpr1" <+> doc)
                        (mkCoreSubst (text "SimplEnv.substExpr2" <+> doc) env)
808
  -- Do *not* short-cut in the case of an empty substitution
809
  -- See Note [SimplEnv invariants]
810 811

substUnfolding :: SimplEnv -> Unfolding -> Unfolding
812 813 814
substUnfolding env unf = CoreSubst.substUnfolding (mkCoreSubst (text "subst-unfolding") env) unf
  -- Do *not* short-cut in the case of an empty substitution
  -- See Note [SimplEnv invariants]
815 816
\end{code}