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,
  11
  12         -- Substitution stuff
  13         Subst, TyVarSubst, IdSubst,
  14         emptySubst, mkSubst, substEnv, substInScope,
  15         lookupSubst, 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 )
  48 import IdInfo           ( IdInfo, zapFragileIdInfo,
  49                           specInfo, setSpecInfo,
  50                           workerExists, workerInfo, setWorkerInfo, WorkerInfo
  51                         )
  52 import UniqSupply       ( UniqSupply, uniqFromSupply, splitUniqSupply )
  53 import Var              ( Var, IdOrTyVar, Id, TyVar, isTyVar, maybeModifyIdInfo )
  54 import Outputable
  55 import Util             ( mapAccumL, foldl2, seqList, ($!) )
  56 \end{code}
  57
  58 %************************************************************************
  59 %*                                                                      *
  60 \subsection{Substitutions}
  61 %*                                                                      *
  62 %************************************************************************
  63
  64 \begin{code}
  65 type InScopeSet = VarSet
  66
  67 data Subst = Subst InScopeSet           -- In scope
  68                    SubstEnv             -- Substitution itself
  69         -- INVARIANT 1: The in-scope set is a superset
  70         --              of the free vars of the range of the substitution
  71         --              that might possibly clash with locally-bound variables
  72         --              in the thing being substituted in.
  73         -- This is what lets us deal with name capture properly
  74         -- It's a hard invariant to check...
  75         -- There are various ways of causing it to happen:
  76         --      - arrange that the in-scope set really is all the things in scope
  77         --      - arrange that it's the free vars of the range of the substitution
  78         --      - make it empty because all the free vars of the subst are fresh,
  79         --              and hence can't possibly clash.a
  80         --
  81         -- INVARIANT 2: No variable is both in scope and in the domain of the substitution
  82         --              Equivalently, the substitution is idempotent
  83         --
  84
  85 type IdSubst    = Subst
  86 \end{code}
  87
  88 \begin{code}
  89 emptyInScopeSet :: InScopeSet
  90 emptyInScopeSet = emptyVarSet
  91 \end{code}
  92
  93
  94
  95 \begin{code}
  96 isEmptySubst :: Subst -> Bool
  97 isEmptySubst (Subst _ env) = isEmptySubstEnv env
  98
  99 emptySubst :: Subst
 100 emptySubst = Subst emptyVarSet emptySubstEnv
 101
 102 mkSubst :: InScopeSet -> SubstEnv -> Subst
 103 mkSubst in_scope env = Subst in_scope env
 104
 105 substEnv :: Subst -> SubstEnv
 106 substEnv (Subst _ env) = env
 107
 108 substInScope :: Subst -> InScopeSet
 109 substInScope (Subst in_scope _) = in_scope
 110
 111 zapSubstEnv :: Subst -> Subst
 112 zapSubstEnv (Subst in_scope env) = Subst in_scope emptySubstEnv
 113
 114 extendSubst :: Subst -> Var -> SubstResult -> Subst
 115 extendSubst (Subst in_scope env) v r = Subst in_scope (extendSubstEnv env v r)
 116
 117 extendSubstList :: Subst -> [Var] -> [SubstResult] -> Subst
 118 extendSubstList (Subst in_scope env) v r = Subst in_scope (extendSubstEnvList env v r)
 119
 120 lookupSubst :: Subst -> Var -> Maybe SubstResult
 121 lookupSubst (Subst _ env) v = lookupSubstEnv env v
 122
 123 lookupInScope :: Subst -> Var -> Maybe Var
 124 lookupInScope (Subst in_scope _) v = lookupVarSet in_scope v
 125
 126 isInScope :: Var -> Subst -> Bool
 127 isInScope v (Subst in_scope _) = v `elemVarSet` in_scope
 128
 129 extendInScope :: Subst -> Var -> Subst
 130 extendInScope (Subst in_scope env) v = Subst (extendVarSet in_scope v) env
 131
 132 extendInScopes :: Subst -> [Var] -> Subst
 133 extendInScopes (Subst in_scope env) vs = Subst (foldl extendVarSet in_scope vs) env
 134
 135 -------------------------------
 136 bindSubst :: Subst -> Var -> Var -> Subst
 137 -- Extend with a substitution, v1 -> Var v2
 138 -- and extend the in-scopes with v2
 139 bindSubst (Subst in_scope env) old_bndr new_bndr
 140   = Subst (in_scope `extendVarSet` new_bndr)
 141           (extendSubstEnv env old_bndr subst_result)
 142   where
 143     subst_result | isId old_bndr = DoneEx (Var new_bndr)
 144                  | otherwise     = DoneTy (TyVarTy new_bndr)
 145
 146 unBindSubst :: Subst -> Var -> Var -> Subst
 147 -- Reverse the effect of bindSubst
 148 -- If old_bndr was already in the substitution, this doesn't quite work
 149 unBindSubst (Subst in_scope env) old_bndr new_bndr
 150   = Subst (in_scope `delVarSet` new_bndr) (delSubstEnv env old_bndr)
 151
 152 -- And the "List" forms
 153 bindSubstList :: Subst -> [Var] -> [Var] -> Subst
 154 bindSubstList subst old_bndrs new_bndrs
 155   = foldl2 bindSubst subst old_bndrs new_bndrs
 156
 157 unBindSubstList :: Subst -> [Var] -> [Var] -> Subst
 158 unBindSubstList subst old_bndrs new_bndrs
 159   = foldl2 unBindSubst subst old_bndrs new_bndrs
 160
 161
 162 -------------------------------
 163 setInScope :: Subst     -- Take env part from here
 164            -> InScopeSet
 165            -> Subst
 166 setInScope (Subst in_scope1 env1) in_scope2
 167   = ASSERT( in_scope1 `subVarSet` in_scope1 )
 168     Subst in_scope2 env1
 169
 170 setSubstEnv :: Subst            -- Take in-scope part from here
 171             -> SubstEnv         -- ... and env part from here
 172             -> Subst
 173 setSubstEnv (Subst in_scope1 _) env2 = Subst in_scope1 env2
 174 \end{code}
 175
 176
 177 %************************************************************************
 178 %*                                                                      *
 179 \subsection{Type substitution}
 180 %*                                                                      *
 181 %************************************************************************
 182
 183 \begin{code}
 184 type TyVarSubst    = Subst      -- TyVarSubst are expected to have range elements
 185         -- (We could have a variant of Subst, but it doesn't seem worth it.)
 186
 187 -- mkTyVarSubst generates the in-scope set from
 188 -- the types given; but it's just a thunk so with a bit of luck
 189 -- it'll never be evaluated
 190 mkTyVarSubst :: [TyVar] -> [Type] -> Subst
 191 mkTyVarSubst tyvars tys = Subst (tyVarsOfTypes tys) (zip_ty_env tyvars tys emptySubstEnv)
 192
 193 -- mkTopTyVarSubst is called when doing top-level substitutions.
 194 -- Here we expect that the free vars of the range of the
 195 -- substitution will be empty.
 196 mkTopTyVarSubst :: [TyVar] -> [Type] -> Subst
 197 mkTopTyVarSubst tyvars tys = Subst emptyVarSet (zip_ty_env tyvars tys emptySubstEnv)
 198
 199 zip_ty_env []       []       env = env
 200 zip_ty_env (tv:tvs) (ty:tys) env = zip_ty_env tvs tys (extendSubstEnv env tv (DoneTy ty))
 201 \end{code}
 202
 203 substTy works with general Substs, so that it can be called from substExpr too.
 204
 205 \begin{code}
 206 substTy :: Subst -> Type  -> Type
 207 substTy subst ty | isEmptySubst subst = ty
 208                  | otherwise          = subst_ty subst ty
 209
 210 substTheta :: TyVarSubst -> ThetaType -> ThetaType
 211 substTheta subst theta
 212   | isEmptySubst subst = theta
 213   | otherwise          = [(clas, map (subst_ty subst) tys) | (clas, tys) <- theta]
 214
 215 subst_ty subst ty
 216    = go ty
 217   where
 218     go (TyConApp tc tys)          = let args = map go tys
 219                                     in  args `seqList` TyConApp tc args
 220     go (NoteTy (SynNote ty1) ty2) = NoteTy (SynNote $! (go ty1)) $! (go ty2)
 221     go (NoteTy (FTVNote _) ty2)   = go ty2              -- Discard the free tyvar note
 222     go (FunTy arg res)            = (FunTy $! (go arg)) $! (go res)
 223     go (NoteTy (UsgNote usg)  ty2) = (NoteTy $! UsgNote usg) $! go ty2          -- Keep usage annot
 224     go (NoteTy (UsgForAll uv) ty2) = (NoteTy $! UsgForAll uv) $! go ty2         -- Keep uvar bdr
 225     go (AppTy fun arg)            = mkAppTy (go fun) $! (go arg)
 226     go ty@(TyVarTy tv)            = case (lookupSubst subst tv) of
 227                                         Nothing            -> ty
 228                                         Just (DoneTy ty')  -> ty'
 229
 230     go (ForAllTy tv ty)           = case substTyVar subst tv of
 231                                         (subst', tv') -> ForAllTy tv' $! (subst_ty subst' ty)
 232 \end{code}
 233
 234 Here is where we invent a new binder if necessary.
 235
 236 \begin{code}
 237 substTyVar :: Subst -> TyVar -> (Subst, TyVar)
 238 substTyVar subst@(Subst in_scope env) old_var
 239   | old_var == new_var  -- No need to clone
 240                         -- But we *must* zap any current substitution for the variable.
 241                         --  For example:
 242                         --      (\x.e) with id_subst = [x |-> e']
 243                         -- Here we must simply zap the substitution for x
 244                         --
 245                         -- The new_id isn't cloned, but it may have a different type
 246                         -- etc, so we must return it, not the old id
 247   = (Subst (in_scope `extendVarSet` new_var)
 248            (delSubstEnv env old_var),
 249      new_var)
 250
 251   | otherwise   -- The new binder is in scope so
 252                 -- we'd better rename it away from the in-scope variables
 253                 -- Extending the substitution to do this renaming also
 254                 -- has the (correct) effect of discarding any existing
 255                 -- substitution for that variable
 256   = (Subst (in_scope `extendVarSet` new_var)
 257            (extendSubstEnv env old_var (DoneTy (TyVarTy new_var))),
 258      new_var)
 259   where
 260     new_var = uniqAway in_scope old_var
 261         -- The uniqAway part makes sure the new variable is not already in scope
 262 \end{code}
 263
 264
 265 %************************************************************************
 266 %*                                                                      *
 267 \section{Expression substitution}
 268 %*                                                                      *
 269 %************************************************************************
 270
 271 This expression substituter deals correctly with name capture.
 272
 273 BUT NOTE that substExpr silently discards the
 274         unfolding, and
 275         spec env
 276 IdInfo attached to any binders in the expression.  It's quite
 277 tricky to do them 'right' in the case of mutually recursive bindings,
 278 and so far has proved unnecessary.
 279
 280 \begin{code}
 281 substExpr :: Subst -> CoreExpr -> CoreExpr
 282 substExpr subst expr | isEmptySubst subst = expr
 283                      | otherwise          = subst_expr subst expr
 284
 285 subst_expr subst expr
 286   = go expr
 287   where
 288     go (Var v) = case lookupSubst subst v of
 289                     Just (DoneEx e')      -> e'
 290                     Just (ContEx env' e') -> subst_expr (setSubstEnv subst env') e'
 291 --      NO!  NO!  SLPJ 14 July 99
 292                     Nothing               -> case lookupInScope subst v of
 293                                                 Just v' -> Var v'
 294                                                 Nothing -> Var v
 295                         -- NB: we look up in the in_scope set because the variable
 296                         -- there may have more info. In particular, when substExpr
 297                         -- is called from the simplifier, the type inside the *occurrences*
 298                         -- of a variable may not be right; we should replace it with the
 299                         -- binder, from the in_scope set.
 300
 301 --                  Nothing -> Var v
 302
 303     go (Type ty)      = Type (go_ty ty)
 304     go (Con con args) = Con con (map go args)
 305     go (App fun arg)  = App (go fun) (go arg)
 306     go (Note note e)  = Note (go_note note) (go e)
 307
 308     go (Lam bndr body) = Lam bndr' (subst_expr subst' body)
 309                        where
 310                          (subst', bndr') = substBndr subst bndr
 311
 312     go (Let (NonRec bndr rhs) body) = Let (NonRec bndr' (go rhs)) (subst_expr subst' body)
 313                                     where
 314                                       (subst', bndr') = substBndr subst bndr
 315
 316     go (Let (Rec pairs) body) = Let (Rec pairs') (subst_expr subst' body)
 317                               where
 318                                 (subst', bndrs') = substBndrs subst (map fst pairs)
 319                                 pairs'  = bndrs' `zip` rhss'
 320                                 rhss'   = map (subst_expr subst' . snd) pairs
 321
 322     go (Case scrut bndr alts) = Case (go scrut) bndr' (map (go_alt subst') alts)
 323                               where
 324                                 (subst', bndr') = substBndr subst bndr
 325
 326     go_alt subst (con, bndrs, rhs) = (con, bndrs', subst_expr subst' rhs)
 327                                  where
 328                                    (subst', bndrs') = substBndrs subst bndrs
 329
 330     go_note (Coerce ty1 ty2) = Coerce (go_ty ty1) (go_ty ty2)
 331     go_note note             = note
 332
 333     go_ty ty = substTy subst ty
 334
 335 \end{code}
 336
 337 Substituting in binders is a rather tricky part of the whole compiler.
 338
 339 When we hit a binder we may need to
 340   (a) apply the the type envt (if non-empty) to its type
 341   (b) apply the type envt and id envt to its SpecEnv (if it has one)
 342   (c) give it a new unique to avoid name clashes
 343
 344 \begin{code}
 345 substBndr :: Subst -> IdOrTyVar -> (Subst, IdOrTyVar)
 346 substBndr subst bndr
 347   | isTyVar bndr  = substTyVar subst bndr
 348   | otherwise     = substId    subst bndr
 349
 350 substBndrs :: Subst -> [IdOrTyVar] -> (Subst, [IdOrTyVar])
 351 substBndrs subst bndrs = mapAccumL substBndr subst bndrs
 352
 353
 354 substIds :: Subst -> [Id] -> (Subst, [Id])
 355 substIds subst bndrs = mapAccumL substId subst bndrs
 356
 357 substId :: Subst -> Id -> (Subst, Id)
 358
 359 -- Returns an Id with empty unfolding and spec-env.
 360 -- It's up to the caller to sort these out.
 361
 362 substId subst@(Subst in_scope env) old_id
 363   = (Subst (in_scope `extendVarSet` new_id)
 364            (extendSubstEnv env old_id (DoneEx (Var new_id))),
 365      new_id)
 366   where
 367     id_ty    = idType old_id
 368
 369        -- id1 has its type zapped
 370     id1 |  noTypeSubst env
 371         || isEmptyVarSet (tyVarsOfType id_ty) = old_id
 372                         -- The tyVarsOfType is cheaper than it looks
 373                         -- because we cache the free tyvars of the type
 374                         -- in a Note in the id's type itself
 375         | otherwise  = setIdType old_id (substTy subst id_ty)
 376
 377         -- id2 has its fragile IdInfo zapped
 378     id2 = maybeModifyIdInfo zapFragileIdInfo id1
 379
 380         -- new_id is cloned if necessary
 381     new_id = uniqAway in_scope id2
 382 \end{code}
 383
 384 Now a variant that unconditionally allocates a new unique.
 385
 386 \begin{code}
 387 substAndCloneIds :: Subst -> UniqSupply -> [Id] -> (Subst, UniqSupply, [Id])
 388 substAndCloneIds subst us [] = (subst, us, [])
 389 substAndCloneIds subst us (b:bs) = case substAndCloneId  subst  us  b  of { (subst1, us1, b') ->
 390                                    case substAndCloneIds subst1 us1 bs of { (subst2, us2, bs') ->
 391                                    (subst2, us2, (b':bs')) }}
 392
 393 substAndCloneId :: Subst -> UniqSupply -> Id -> (Subst, UniqSupply, Id)
 394 substAndCloneId subst@(Subst in_scope env) us old_id
 395   = (Subst (in_scope `extendVarSet` new_id)
 396            (extendSubstEnv env old_id (DoneEx (Var new_id))),
 397      new_us,
 398      new_id)
 399   where
 400     id_ty    = idType old_id
 401     id1 | noTypeSubst env || isEmptyVarSet (tyVarsOfType id_ty) = old_id
 402         | otherwise                                             = setIdType old_id (substTy subst id_ty)
 403
 404     id2          = maybeModifyIdInfo zapFragileIdInfo id1
 405     new_id       = setVarUnique id2 (uniqFromSupply us1)
 406     (us1,new_us) = splitUniqSupply us
 407 \end{code}
 408
 409
 410 %************************************************************************
 411 %*                                                                      *
 412 \section{IdInfo substitution}
 413 %*                                                                      *
 414 %************************************************************************
 415
 416 \begin{code}
 417 substIdInfo :: Subst
 418             -> IdInfo           -- Get un-substituted ones from here
 419             -> IdInfo           -- Substitute it and add it to here
 420             -> IdInfo           -- To give this
 421         -- Seq'ing on the returned IdInfo is enough to cause all the
 422         -- substitutions to happen completely
 423
 424 substIdInfo subst old_info new_info
 425   = info2
 426   where
 427     info1 | isEmptyCoreRules old_rules = new_info
 428           | otherwise                  = new_info `setSpecInfo` new_rules
 429                         -- setSpecInfo does a seq
 430           where
 431             new_rules = substRules subst old_rules
 432
 433     info2 | not (workerExists old_wrkr) = info1
 434           | otherwise                   = info1 `setWorkerInfo` new_wrkr
 435                         -- setWorkerInfo does a seq
 436           where
 437             new_wrkr = substWorker subst old_wrkr
 438
 439     old_rules = specInfo   old_info
 440     old_wrkr  = workerInfo old_info
 441
 442 substWorker :: Subst -> WorkerInfo -> WorkerInfo
 443         -- Seq'ing on the returned WorkerInfo is enough to cause all the
 444         -- substitutions to happen completely
 445
 446 substWorker subst Nothing
 447   = Nothing
 448 substWorker subst (Just w)
 449   = case lookupSubst subst w of
 450         Nothing -> Just w
 451         Just (DoneEx (Var w1)) -> Just w1
 452         Just (DoneEx other)    -> WARN( True, text "substWorker: DoneEx" <+> ppr w )
 453                                   Nothing       -- Worker has got substituted away altogether
 454         Just (ContEx se1 e)    -> WARN( True, text "substWorker: ContEx" <+> ppr w )
 455                                   Nothing       -- Ditto
 456
 457 substRules :: Subst -> CoreRules -> CoreRules
 458         -- Seq'ing on the returned CoreRules is enough to cause all the
 459         -- substitutions to happen completely
 460
 461 substRules subst rules
 462  | isEmptySubst subst = rules
 463
 464 substRules subst (Rules rules rhs_fvs)
 465   = seqRules new_rules `seq` new_rules
 466   where
 467     new_rules = Rules (map do_subst rules)
 468                       (subst_fvs (substEnv subst) rhs_fvs)
 469
 470     do_subst (Rule name tpl_vars lhs_args rhs)
 471         = Rule name tpl_vars'
 472                (map (substExpr subst') lhs_args)
 473                (substExpr subst' rhs)
 474         where
 475           (subst', tpl_vars') = substBndrs subst tpl_vars
 476
 477     subst_fvs se fvs
 478         = foldVarSet (unionVarSet . subst_fv) emptyVarSet rhs_fvs
 479         where
 480           subst_fv fv = case lookupSubstEnv se fv of
 481                                 Nothing                   -> unitVarSet fv
 482                                 Just (DoneEx expr)        -> exprFreeVars expr
 483                                 Just (DoneTy ty)          -> tyVarsOfType ty
 484                                 Just (ContEx se' expr) -> subst_fvs se' (exprFreeVars expr)
 485 \end{code}