ghc/compiler/coreSyn/Subst.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section[CoreUtils]{Utility functions on @Core@ syntax}
   5
   6 \begin{code}
   7 module Subst (
   8         -- In-scope set
   9         InScopeSet, emptyInScopeSet,
  10         lookupInScope, setInScope, extendInScope, extendInScopes, isInScope, modifyInScope,
  11
  12         -- Substitution stuff
  13         Subst, TyVarSubst, IdSubst,
  14         emptySubst, mkSubst, substEnv, substInScope,
  15         lookupSubst, lookupIdSubst, isEmptySubst, extendSubst, extendSubstList,
  16         zapSubstEnv, setSubstEnv,
  17
  18         bindSubst, unBindSubst, bindSubstList, unBindSubstList,
  19
  20         -- Binders
  21         substBndr, substBndrs, substTyVar, substId, substIds,
  22         substAndCloneId, substAndCloneIds,
  23
  24         -- Type stuff
  25         mkTyVarSubst, mkTopTyVarSubst,
  26         substTy, substTheta,
  27
  28         -- Expression stuff
  29         substExpr, substIdInfo
  30     ) where
  31
  32 #include "HsVersions.h"
  33
  34 import CoreSyn          ( Expr(..), Bind(..), Note(..), CoreExpr, CoreBndr,
  35                           CoreRules(..), CoreRule(..),
  36                           emptyCoreRules, isEmptyCoreRules, seqRules
  37                         )
  38 import CoreFVs          ( exprFreeVars )
  39 import TypeRep          ( Type(..), TyNote(..),
  40                         )  -- friend
  41 import Type             ( ThetaType,
  42                           tyVarsOfType, tyVarsOfTypes, mkAppTy
  43                         )
  44 import VarSet
  45 import VarEnv
  46 import Var              ( setVarUnique, isId )
  47 import Id               ( idType, setIdType, getIdOccInfo, zapFragileIdInfo )
  48 import Name             ( isLocallyDefined )
  49 import IdInfo           ( IdInfo, isFragileOccInfo,
  50                           specInfo, setSpecInfo,
  51                           workerExists, workerInfo, setWorkerInfo, WorkerInfo
  52                         )
  53 import UniqSupply       ( UniqSupply, uniqFromSupply, splitUniqSupply )
  54 import Var              ( Var, IdOrTyVar, Id, TyVar, isTyVar )
  55 import Outputable
  56 import Util             ( mapAccumL, foldl2, seqList, ($!) )
  57 \end{code}
  58
  59 %************************************************************************
  60 %*                                                                      *
  61 \subsection{Substitutions}
  62 %*                                                                      *
  63 %************************************************************************
  64
  65 \begin{code}
  66 type InScopeSet = VarEnv Var
  67
  68 data Subst = Subst InScopeSet           -- In scope
  69                    SubstEnv             -- Substitution itself
  70         -- INVARIANT 1: The (domain of the) in-scope set is a superset
  71         --              of the free vars of the range of the substitution
  72         --              that might possibly clash with locally-bound variables
  73         --              in the thing being substituted in.
  74         -- This is what lets us deal with name capture properly
  75         -- It's a hard invariant to check...
  76         -- There are various ways of causing it to happen:
  77         --      - arrange that the in-scope set really is all the things in scope
  78         --      - arrange that it's the free vars of the range of the substitution
  79         --      - make it empty because all the free vars of the subst are fresh,
  80         --              and hence can't possibly clash.a
  81         --
  82         -- INVARIANT 2: No variable is both in scope and in the domain of the substitution
  83         --              Equivalently, the substitution is idempotent
  84         --
  85
  86 type IdSubst    = Subst
  87 \end{code}
  88
  89 The general plan about the substitution and in-scope set for Ids is as follows
  90
  91 * substId always adds new_id to the in-scope set.
  92   new_id has a correctly-substituted type, but all its fragile IdInfo has been zapped.
  93   That is added back in later.  So new_id is the minimal thing it's
  94   correct to substitute.
  95
  96 * substId adds a binding (DoneVar new_id occ) to the substitution if
  97         EITHER the Id's unique has changed
  98         OR     the Id has interesting occurrence information
  99   Note, though that the substitution isn't necessarily extended
 100   if the type changes.  Why not?  Because of the next point:
 101
 102 * We *always, always* finish by looking up in the in-scope set
 103   any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
 104   Reason: so that we never finish up with a "old" Id in the result.
 105   An old Id might point to an old unfolding and so on... which gives a space leak.
 106
 107   [The DoneEx and DoneVar hits map to "new" stuff.]
 108
 109 * It follows that substExpr must not do a no-op if the substitution is empty.
 110   substType is free to do so, however.
 111
 112 * When we come to a let-binding (say) we generate new IdInfo, including an
 113   unfolding, attach it to the binder, and add this newly adorned binder to
 114   the in-scope set.  So all subsequent occurrences of the binder will get mapped
 115   to the full-adorned binder, which is also the one put in the binding site.
 116
 117 * The in-scope "set" usually maps x->x; we use it simply for its domain.
 118   But sometimes we have two in-scope Ids that are synomyms, and should
 119   map to the same target:  x->x, y->x.  Notably:
 120         case y of x { ... }
 121   That's why the "set" is actually a VarEnv Var
 122
 123 \begin{code}
 124 emptyInScopeSet :: InScopeSet
 125 emptyInScopeSet = emptyVarSet
 126
 127 add_in_scope :: InScopeSet -> Var -> InScopeSet
 128 add_in_scope in_scope v = extendVarEnv in_scope v v
 129 \end{code}
 130
 131
 132
 133 \begin{code}
 134 isEmptySubst :: Subst -> Bool
 135 isEmptySubst (Subst _ env) = isEmptySubstEnv env
 136
 137 emptySubst :: Subst
 138 emptySubst = Subst emptyInScopeSet emptySubstEnv
 139
 140 mkSubst :: InScopeSet -> SubstEnv -> Subst
 141 mkSubst in_scope env = Subst in_scope env
 142
 143 substEnv :: Subst -> SubstEnv
 144 substEnv (Subst _ env) = env
 145
 146 substInScope :: Subst -> InScopeSet
 147 substInScope (Subst in_scope _) = in_scope
 148
 149 zapSubstEnv :: Subst -> Subst
 150 zapSubstEnv (Subst in_scope env) = Subst in_scope emptySubstEnv
 151
 152 extendSubst :: Subst -> Var -> SubstResult -> Subst
 153 extendSubst (Subst in_scope env) v r = Subst in_scope (extendSubstEnv env v r)
 154
 155 extendSubstList :: Subst -> [Var] -> [SubstResult] -> Subst
 156 extendSubstList (Subst in_scope env) v r = Subst in_scope (extendSubstEnvList env v r)
 157
 158 lookupSubst :: Subst -> Var -> Maybe SubstResult
 159 lookupSubst (Subst _ env) v = lookupSubstEnv env v
 160
 161 lookupIdSubst :: Subst -> Id -> SubstResult
 162 -- Does the lookup in the in-scope set too
 163 lookupIdSubst (Subst in_scope env) v
 164   = case lookupSubstEnv env v of
 165         Just (DoneId v' occ) -> case lookupVarEnv in_scope v' of
 166                                   Just v'' -> DoneId v'' occ
 167                                   Nothing  -> DoneId v' occ
 168         Just res             -> res
 169         Nothing              -> DoneId v' (getIdOccInfo v')
 170                              where
 171                                     v' = case lookupVarEnv in_scope v of
 172                                            Just v' -> v'
 173                                            Nothing -> v
 174
 175 lookupInScope :: Subst -> Var -> Maybe Var
 176 lookupInScope (Subst in_scope _) v = lookupVarEnv in_scope v
 177
 178 isInScope :: Var -> Subst -> Bool
 179 isInScope v (Subst in_scope _) = v `elemVarEnv` in_scope
 180
 181 extendInScope :: Subst -> Var -> Subst
 182 extendInScope (Subst in_scope env) v = Subst (in_scope `add_in_scope` v) env
 183
 184 modifyInScope :: Subst -> Var -> Var -> Subst
 185 modifyInScope (Subst in_scope env) old_v new_v = Subst (extendVarEnv in_scope old_v new_v) env
 186         -- make old_v map to new_v
 187
 188 extendInScopes :: Subst -> [Var] -> Subst
 189 extendInScopes (Subst in_scope env) vs = Subst (foldl add_in_scope in_scope vs) env
 190
 191 -------------------------------
 192 bindSubst :: Subst -> Var -> Var -> Subst
 193 -- Extend with a substitution, v1 -> Var v2
 194 -- and extend the in-scopes with v2
 195 bindSubst (Subst in_scope env) old_bndr new_bndr
 196   = Subst (in_scope `add_in_scope` new_bndr)
 197           (extendSubstEnv env old_bndr subst_result)
 198   where
 199     subst_result | isId old_bndr = DoneEx (Var new_bndr)
 200                  | otherwise     = DoneTy (TyVarTy new_bndr)
 201
 202 unBindSubst :: Subst -> Var -> Var -> Subst
 203 -- Reverse the effect of bindSubst
 204 -- If old_bndr was already in the substitution, this doesn't quite work
 205 unBindSubst (Subst in_scope env) old_bndr new_bndr
 206   = Subst (in_scope `delVarEnv` new_bndr) (delSubstEnv env old_bndr)
 207
 208 -- And the "List" forms
 209 bindSubstList :: Subst -> [Var] -> [Var] -> Subst
 210 bindSubstList subst old_bndrs new_bndrs
 211   = foldl2 bindSubst subst old_bndrs new_bndrs
 212
 213 unBindSubstList :: Subst -> [Var] -> [Var] -> Subst
 214 unBindSubstList subst old_bndrs new_bndrs
 215   = foldl2 unBindSubst subst old_bndrs new_bndrs
 216
 217
 218 -------------------------------
 219 setInScope :: Subst     -- Take env part from here
 220            -> InScopeSet
 221            -> Subst
 222 setInScope (Subst in_scope1 env1) in_scope2
 223   = Subst in_scope2 env1
 224
 225 setSubstEnv :: Subst            -- Take in-scope part from here
 226             -> SubstEnv         -- ... and env part from here
 227             -> Subst
 228 setSubstEnv (Subst in_scope1 _) env2 = Subst in_scope1 env2
 229 \end{code}
 230
 231
 232 %************************************************************************
 233 %*                                                                      *
 234 \subsection{Type substitution}
 235 %*                                                                      *
 236 %************************************************************************
 237
 238 \begin{code}
 239 type TyVarSubst    = Subst      -- TyVarSubst are expected to have range elements
 240         -- (We could have a variant of Subst, but it doesn't seem worth it.)
 241
 242 -- mkTyVarSubst generates the in-scope set from
 243 -- the types given; but it's just a thunk so with a bit of luck
 244 -- it'll never be evaluated
 245 mkTyVarSubst :: [TyVar] -> [Type] -> Subst
 246 mkTyVarSubst tyvars tys = Subst (tyVarsOfTypes tys) (zip_ty_env tyvars tys emptySubstEnv)
 247
 248 -- mkTopTyVarSubst is called when doing top-level substitutions.
 249 -- Here we expect that the free vars of the range of the
 250 -- substitution will be empty.
 251 mkTopTyVarSubst :: [TyVar] -> [Type] -> Subst
 252 mkTopTyVarSubst tyvars tys = Subst emptyInScopeSet (zip_ty_env tyvars tys emptySubstEnv)
 253
 254 zip_ty_env []       []       env = env
 255 zip_ty_env (tv:tvs) (ty:tys) env = zip_ty_env tvs tys (extendSubstEnv env tv (DoneTy ty))
 256 \end{code}
 257
 258 substTy works with general Substs, so that it can be called from substExpr too.
 259
 260 \begin{code}
 261 substTy :: Subst -> Type  -> Type
 262 substTy subst ty | isEmptySubst subst = ty
 263                  | otherwise          = subst_ty subst ty
 264
 265 substTheta :: TyVarSubst -> ThetaType -> ThetaType
 266 substTheta subst theta
 267   | isEmptySubst subst = theta
 268   | otherwise          = [(clas, map (subst_ty subst) tys) | (clas, tys) <- theta]
 269
 270 subst_ty subst ty
 271    = go ty
 272   where
 273     go (TyConApp tc tys)          = let args = map go tys
 274                                     in  args `seqList` TyConApp tc args
 275     go (NoteTy (SynNote ty1) ty2) = NoteTy (SynNote $! (go ty1)) $! (go ty2)
 276     go (NoteTy (FTVNote _) ty2)   = go ty2              -- Discard the free tyvar note
 277     go (FunTy arg res)            = (FunTy $! (go arg)) $! (go res)
 278     go (NoteTy (UsgNote usg)  ty2) = (NoteTy $! UsgNote usg) $! go ty2          -- Keep usage annot
 279     go (NoteTy (UsgForAll uv) ty2) = (NoteTy $! UsgForAll uv) $! go ty2         -- Keep uvar bdr
 280     go (AppTy fun arg)            = mkAppTy (go fun) $! (go arg)
 281     go ty@(TyVarTy tv)            = case (lookupSubst subst tv) of
 282                                         Nothing            -> ty
 283                                         Just (DoneTy ty')  -> ty'
 284
 285     go (ForAllTy tv ty)           = case substTyVar subst tv of
 286                                         (subst', tv') -> ForAllTy tv' $! (subst_ty subst' ty)
 287 \end{code}
 288
 289 Here is where we invent a new binder if necessary.
 290
 291 \begin{code}
 292 substTyVar :: Subst -> TyVar -> (Subst, TyVar)
 293 substTyVar subst@(Subst in_scope env) old_var
 294   | old_var == new_var  -- No need to clone
 295                         -- But we *must* zap any current substitution for the variable.
 296                         --  For example:
 297                         --      (\x.e) with id_subst = [x |-> e']
 298                         -- Here we must simply zap the substitution for x
 299                         --
 300                         -- The new_id isn't cloned, but it may have a different type
 301                         -- etc, so we must return it, not the old id
 302   = (Subst (in_scope `add_in_scope` new_var)
 303            (delSubstEnv env old_var),
 304      new_var)
 305
 306   | otherwise   -- The new binder is in scope so
 307                 -- we'd better rename it away from the in-scope variables
 308                 -- Extending the substitution to do this renaming also
 309                 -- has the (correct) effect of discarding any existing
 310                 -- substitution for that variable
 311   = (Subst (in_scope `add_in_scope` new_var)
 312            (extendSubstEnv env old_var (DoneTy (TyVarTy new_var))),
 313      new_var)
 314   where
 315     new_var = uniqAway in_scope old_var
 316         -- The uniqAway part makes sure the new variable is not already in scope
 317 \end{code}
 318
 319
 320 %************************************************************************
 321 %*                                                                      *
 322 \section{Expression substitution}
 323 %*                                                                      *
 324 %************************************************************************
 325
 326 This expression substituter deals correctly with name capture.
 327
 328 BUT NOTE that substExpr silently discards the
 329         unfolding, and
 330         spec env
 331 IdInfo attached to any binders in the expression.  It's quite
 332 tricky to do them 'right' in the case of mutually recursive bindings,
 333 and so far has proved unnecessary.
 334
 335 \begin{code}
 336 substExpr :: Subst -> CoreExpr -> CoreExpr
 337 substExpr subst expr
 338         -- NB: we do not do a no-op when the substitution is empty,
 339         -- because we always want to substitute the variables in the
 340         -- in-scope set for their occurrences.  Why?
 341         --      (a) because they may contain more information
 342         --      (b) because leaving an un-substituted Id might cause
 343         --          a space leak (its unfolding might point to an old version
 344         --          of its right hand side).
 345
 346   = go expr
 347   where
 348     go (Var v) = -- See the notes at the top, with the Subst data type declaration
 349                  case lookupIdSubst subst v of
 350
 351                     ContEx env' e' -> substExpr (setSubstEnv subst env') e'
 352                     DoneId v _     -> Var v
 353                     DoneEx e'      -> e'
 354
 355     go (Type ty)      = Type (go_ty ty)
 356     go (Con con args) = Con con (map go args)
 357     go (App fun arg)  = App (go fun) (go arg)
 358     go (Note note e)  = Note (go_note note) (go e)
 359
 360     go (Lam bndr body) = Lam bndr' (substExpr subst' body)
 361                        where
 362                          (subst', bndr') = substBndr subst bndr
 363
 364     go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (substExpr subst' body)
 365                                     where
 366                                       (subst', bndr') = substBndr subst bndr
 367
 368     go (Let (Rec pairs) body) = Let (Rec pairs') (substExpr subst' body)
 369                               where
 370                                 (subst', bndrs') = substBndrs subst (map fst pairs)
 371                                 pairs'  = bndrs' `zip` rhss'
 372                                 rhss'   = map (substExpr subst' . snd) pairs
 373
 374     go (Case scrut bndr alts) = Case (go scrut) bndr' (map (go_alt subst') alts)
 375                               where
 376                                 (subst', bndr') = substBndr subst bndr
 377
 378     go_alt subst (con, bndrs, rhs) = (con, bndrs', substExpr subst' rhs)
 379                                  where
 380                                    (subst', bndrs') = substBndrs subst bndrs
 381
 382     go_note (Coerce ty1 ty2) = Coerce (go_ty ty1) (go_ty ty2)
 383     go_note note             = note
 384
 385     go_ty ty = substTy subst ty
 386
 387 \end{code}
 388
 389 Substituting in binders is a rather tricky part of the whole compiler.
 390
 391 When we hit a binder we may need to
 392   (a) apply the the type envt (if non-empty) to its type
 393   (c) give it a new unique to avoid name clashes
 394
 395 \begin{code}
 396 substBndr :: Subst -> IdOrTyVar -> (Subst, IdOrTyVar)
 397 substBndr subst bndr
 398   | isTyVar bndr  = substTyVar subst bndr
 399   | otherwise     = substId    subst bndr
 400
 401 substBndrs :: Subst -> [IdOrTyVar] -> (Subst, [IdOrTyVar])
 402 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 403
 404
 405 substIds :: Subst -> [Id] -> (Subst, [Id])
 406 substIds subst bndrs = mapAccumL substId subst bndrs
 407
 408 substId :: Subst -> Id -> (Subst, Id)
 409         -- Returns an Id with empty IdInfo
 410         -- See the notes with the Subst data type decl at the
 411         -- top of this module
 412
 413 substId subst@(Subst in_scope env) old_id
 414   = (Subst (in_scope `add_in_scope` new_id) new_env, new_id)
 415   where
 416     id_ty    = idType old_id
 417     occ_info = getIdOccInfo old_id
 418
 419        -- id1 has its type zapped
 420     id1 |  noTypeSubst env
 421         || isEmptyVarSet (tyVarsOfType id_ty) = old_id
 422                         -- The tyVarsOfType is cheaper than it looks
 423                         -- because we cache the free tyvars of the type
 424                         -- in a Note in the id's type itself
 425         | otherwise  = setIdType old_id (substTy subst id_ty)
 426
 427         -- id2 has its IdInfo zapped
 428     id2 = zapFragileIdInfo id1
 429
 430         -- new_id is cloned if necessary
 431     new_id = uniqAway in_scope id2
 432
 433         -- Extend the substitution if the unique has changed,
 434         -- or there's some useful occurrence information
 435         -- See the notes with substTyVar for the delSubstEnv
 436     new_env | new_id /= old_id || isFragileOccInfo occ_info
 437             = extendSubstEnv env old_id (DoneId new_id occ_info)
 438             | otherwise
 439             = delSubstEnv env old_id
 440 \end{code}
 441
 442 Now a variant that unconditionally allocates a new unique.
 443
 444 \begin{code}
 445 substAndCloneIds :: Subst -> UniqSupply -> [Id] -> (Subst, UniqSupply, [Id])
 446 substAndCloneIds subst us [] = (subst, us, [])
 447 substAndCloneIds subst us (b:bs) = case substAndCloneId  subst  us  b  of { (subst1, us1, b') ->
 448                                    case substAndCloneIds subst1 us1 bs of { (subst2, us2, bs') ->
 449                                    (subst2, us2, (b':bs')) }}
 450
 451 substAndCloneId :: Subst -> UniqSupply -> Id -> (Subst, UniqSupply, Id)
 452 substAndCloneId subst@(Subst in_scope env) us old_id
 453   = (Subst (in_scope `add_in_scope` new_id)
 454            (extendSubstEnv env old_id (DoneEx (Var new_id))),
 455      new_us,
 456      new_id)
 457   where
 458     id_ty    = idType old_id
 459     id1 | noTypeSubst env || isEmptyVarSet (tyVarsOfType id_ty) = old_id
 460         | otherwise                                             = setIdType old_id (substTy subst id_ty)
 461
 462     id2          = zapFragileIdInfo id1
 463     new_id       = setVarUnique id2 (uniqFromSupply us1)
 464     (us1,new_us) = splitUniqSupply us
 465 \end{code}
 466
 467
 468 %************************************************************************
 469 %*                                                                      *
 470 \section{IdInfo substitution}
 471 %*                                                                      *
 472 %************************************************************************
 473
 474 \begin{code}
 475 substIdInfo :: Subst
 476             -> IdInfo           -- Get un-substituted ones from here
 477             -> IdInfo           -- Substitute it and add it to here
 478             -> IdInfo           -- To give this
 479         -- Seq'ing on the returned IdInfo is enough to cause all the
 480         -- substitutions to happen completely
 481
 482 substIdInfo subst old_info new_info
 483   = info2
 484   where
 485     info1 | isEmptyCoreRules old_rules = new_info
 486           | otherwise                  = new_info `setSpecInfo` new_rules
 487                         -- setSpecInfo does a seq
 488           where
 489             new_rules = substRules subst old_rules
 490
 491     info2 | not (workerExists old_wrkr) = info1
 492           | otherwise                   = info1 `setWorkerInfo` new_wrkr
 493                         -- setWorkerInfo does a seq
 494           where
 495             new_wrkr = substWorker subst old_wrkr
 496
 497     old_rules = specInfo   old_info
 498     old_wrkr  = workerInfo old_info
 499
 500 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 501         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 502         -- substitutions to happen completely
 503
 504 substWorker subst Nothing
 505   = Nothing
 506 substWorker subst (Just w)
 507   = case lookupSubst subst w of
 508         Nothing -> Just w
 509         Just (DoneId w1 _)     -> Just w1
 510         Just (DoneEx (Var w1)) -> Just w1
 511         Just (DoneEx other)    -> WARN( True, text "substWorker: DoneEx" <+> ppr w )
 512                                   Nothing       -- Worker has got substituted away altogether
 513         Just (ContEx se1 e)    -> WARN( True, text "substWorker: ContEx" <+> ppr w )
 514                                   Nothing       -- Ditto
 515
 516 substRules :: Subst -> CoreRules -> CoreRules
 517         -- Seq'ing on the returned CoreRules is enough to cause all the
 518         -- substitutions to happen completely
 519
 520 substRules subst rules
 521  | isEmptySubst subst = rules
 522
 523 substRules subst (Rules rules rhs_fvs)
 524   = seqRules new_rules `seq` new_rules
 525   where
 526     new_rules = Rules (map do_subst rules)
 527                       (subst_fvs (substEnv subst) rhs_fvs)
 528
 529     do_subst rule@(BuiltinRule _) = rule
 530     do_subst (Rule name tpl_vars lhs_args rhs)
 531         = Rule name tpl_vars'
 532                (map (substExpr subst') lhs_args)
 533                (substExpr subst' rhs)
 534         where
 535           (subst', tpl_vars') = substBndrs subst tpl_vars
 536
 537     subst_fvs se fvs
 538         = foldVarSet (unionVarSet . subst_fv) emptyVarSet rhs_fvs
 539         where
 540           subst_fv fv = case lookupSubstEnv se fv of
 541                                 Nothing                   -> unitVarSet fv
 542                                 Just (DoneId fv' _)       -> unitVarSet fv'
 543                                 Just (DoneEx expr)        -> exprFreeVars expr
 544                                 Just (DoneTy ty)          -> tyVarsOfType ty
 545                                 Just (ContEx se' expr) -> subst_fvs se' (exprFreeVars expr)
 546 \end{code}