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