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