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