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