Add an ambient Id substitution to Substwip/T10527-2

After a struggle, I fixed Trac #5113 (again) on the 7.10 branch,
by adding an ambient substitution to Subst; see CoreSubst,
esp Note [IdSubstEnv].

This allowed me to do the impedence-matching in SimplEnv.substExpr
efficiently (fixing #10370) as well correctly (fixing the latest
problem with #5113).

This cost me more time than I like to say.  Sigh.

4 files changed, 155 insertions, 109 deletions

diff --git a/compiler/coreSyn/CoreSubst.hs b/compiler/coreSyn/CoreSubst.hs
index 35dbb50229..26732a2e48 100644
--- a/compiler/coreSyn/CoreSubst.hs
+++ b/compiler/coreSyn/CoreSubst.hs
@@ -20,7 +20,7 @@ module CoreSubst (
         substTickish, substVarSet,
 
         -- ** Operations on substitutions
-        emptySubst, mkEmptySubst, mkSubst, mkOpenSubst, substInScope, isEmptySubst,
+        emptySubst, mkEmptySubst, mkGblSubst, mkOpenSubst, substInScope, isEmptySubst,
         extendIdSubst, extendIdSubstList, extendTvSubst, extendTvSubstList,
         extendCvSubst, extendCvSubstList,
         extendSubst, extendSubstList, extendSubstWithVar, zapSubstEnv,
@@ -178,24 +178,106 @@ TvSubstEnv and CvSubstEnv?
 
 * For TyVars, only coercion variables can possibly change, and they are
   easy to spot
+
+Note [IdSubstEnv]
+~~~~~~~~~~~~~~~~~
+An IdSubstEnv has a "local environment" of type (IdEnv CoreExpr);
+this is extended when we meet a binder, in the usual way.  But it also
+has a "global environment" of type GblIdSubst.  This global envt is
+never modified during substitution.  Rather:
+
+ * The GblIdSubst is used when initialising the substitution via
+   mkGblSubst, to give an "ambient substitution" for the enclosing
+   context.
+
+ * On lookup, we look first in the local envt and then in the global envt
+   (see lookupIdSubst)
+
+ * The GblIdSubst is just a function; but since we need to delete things
+   from the substitution when passing a binder, we have to record a set
+   of Ids gis_del that must *not* be looked up in the gbl envt.
+
+All this is needed to support SimplEnv.substExpr, which starts off
+with a SimplIdSubst, which provides the ambient subsitution.
 -}
 
 -- | An environment for substituting for 'Id's
-type IdSubstEnv = IdEnv CoreExpr
+-- See Note [IdSubstEnv]
+data IdSubstEnv = ISE { ise_env  :: !(IdEnv CoreExpr)
+                      , ise_gbl  :: !GblIdSubst }
+
+data GblIdSubst = NoGIS
+                | GIS { gis_env :: !(InScopeSet -> Id -> Maybe CoreExpr)
+                      , gis_del :: !IdSet }  -- Deletions from gis_env
+
+instance Outputable IdSubstEnv where
+  ppr (ISE { ise_env = lcl, ise_gbl = gbl })
+    = ppr gbl $$ ppr lcl
+
+instance Outputable GblIdSubst where
+  ppr NoGIS = empty
+  ppr (GIS { gis_del = dels }) = ptext (sLit "GIS") <+> ppr dels
+
+lookupGIS :: GblIdSubst -> InScopeSet -> Id -> Maybe CoreExpr
+lookupGIS NoGIS _ _ = Nothing
+lookupGIS (GIS { gis_env = gbl_fn, gis_del = dels }) in_scope v
+  | v `elemVarSet` dels = Nothing
+  | otherwise           = gbl_fn in_scope v
+
+isEmptyIdSubst :: IdSubstEnv -> Bool
+isEmptyIdSubst (ISE { ise_env = lcl, ise_gbl = NoGIS }) = isEmptyVarEnv lcl
+isEmptyIdSubst _ = False
+
+emptyIdSubst :: IdSubstEnv
+emptyIdSubst = ISE { ise_env = emptyVarEnv, ise_gbl = NoGIS }
+
+extendIdSubstEnv :: IdSubstEnv -> Id -> CoreExpr -> IdSubstEnv
+extendIdSubstEnv ise v e = ise { ise_env = extendVarEnv (ise_env ise) v e }
+
+extendIdSubstEnvList :: IdSubstEnv -> [(Id,CoreExpr)] -> IdSubstEnv
+extendIdSubstEnvList ise prs = ise { ise_env = extendVarEnvList (ise_env ise) prs }
+
+delIdSubst :: IdSubstEnv -> Id -> IdSubstEnv
+delIdSubst (ISE { ise_env = lcl, ise_gbl = gbl }) v
+  = ISE { ise_env = delVarEnv lcl v, ise_gbl = delGIS gbl v }
+
+delIdSubstList :: IdSubstEnv -> [Id] -> IdSubstEnv
+delIdSubstList (ISE { ise_env = lcl, ise_gbl = gbl }) vs
+  = ISE { ise_env = delVarEnvList lcl vs, ise_gbl = delGISList gbl vs }
+
+delGIS :: GblIdSubst -> Id -> GblIdSubst
+delGIS NoGIS _ = NoGIS
+delGIS (GIS { gis_env = gbl, gis_del = dels }) v
+  = GIS { gis_env = gbl, gis_del = if isJust (gbl emptyInScopeSet v)
+                                   then extendVarSet dels v
+                                   else dels }
+
+delGISList :: GblIdSubst -> [Id] -> GblIdSubst
+delGISList NoGIS _ = NoGIS
+delGISList (GIS { gis_env = gbl, gis_del = dels }) vs
+  = GIS { gis_env = gbl, gis_del = extendVarSetList dels del_vs }
+  where
+    del_vs = [ v | v <- vs, isJust (gbl emptyInScopeSet v)]
 
 ----------------------------
 isEmptySubst :: Subst -> Bool
 isEmptySubst (Subst _ id_env tv_env cv_env)
-  = isEmptyVarEnv id_env && isEmptyVarEnv tv_env && isEmptyVarEnv cv_env
+  = isEmptyIdSubst id_env && isEmptyVarEnv tv_env && isEmptyVarEnv cv_env
 
 emptySubst :: Subst
-emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv emptyVarEnv
+emptySubst = Subst emptyInScopeSet emptyIdSubst emptyVarEnv emptyVarEnv
 
 mkEmptySubst :: InScopeSet -> Subst
-mkEmptySubst in_scope = Subst in_scope emptyVarEnv emptyVarEnv emptyVarEnv
+mkEmptySubst in_scope = Subst in_scope emptyIdSubst emptyVarEnv emptyVarEnv
 
-mkSubst :: InScopeSet -> TvSubstEnv -> CvSubstEnv -> IdSubstEnv -> Subst
-mkSubst in_scope tvs cvs ids = Subst in_scope ids tvs cvs
+mkGblSubst :: InScopeSet -> TvSubstEnv -> CvSubstEnv
+           -> (InScopeSet -> Id -> Maybe CoreExpr)
+           -> Subst
+mkGblSubst in_scope tvs cvs lookup_id
+ = Subst in_scope id_subst tvs cvs
+ where
+    id_subst = ISE { ise_env = emptyVarEnv
+                   , ise_gbl = GIS { gis_env = lookup_id, gis_del = emptyVarSet } }
 
 -- | Find the in-scope set: see "CoreSubst#in_scope_invariant"
 substInScope :: Subst -> InScopeSet
@@ -204,17 +286,17 @@ substInScope (Subst in_scope _ _ _) = in_scope
 -- | Remove all substitutions for 'Id's and 'Var's that might have been built up
 -- while preserving the in-scope set
 zapSubstEnv :: Subst -> Subst
-zapSubstEnv (Subst in_scope _ _ _) = Subst in_scope emptyVarEnv emptyVarEnv emptyVarEnv
+zapSubstEnv (Subst in_scope _ _ _) = Subst in_scope emptyIdSubst emptyVarEnv emptyVarEnv
 
 -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the in-scope set is
 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
 extendIdSubst :: Subst -> Id -> CoreExpr -> Subst
 -- ToDo: add an ASSERT that fvs(subst-result) is already in the in-scope set
-extendIdSubst (Subst in_scope ids tvs cvs) v r = Subst in_scope (extendVarEnv ids v r) tvs cvs
+extendIdSubst (Subst in_scope ids tvs cvs) v r = Subst in_scope (extendIdSubstEnv ids v r) tvs cvs
 
 -- | Adds multiple 'Id' substitutions to the 'Subst': see also 'extendIdSubst'
 extendIdSubstList :: Subst -> [(Id, CoreExpr)] -> Subst
-extendIdSubstList (Subst in_scope ids tvs cvs) prs = Subst in_scope (extendVarEnvList ids prs) tvs cvs
+extendIdSubstList (Subst in_scope ids tvs cvs) prs = Subst in_scope (extendIdSubstEnvList ids prs) tvs cvs
 
 -- | Add a substitution for a 'TyVar' to the 'Subst': you must ensure that the in-scope set is
 -- such that the "CoreSubst#in_scope_invariant" is true after extending the substitution like this
@@ -260,9 +342,10 @@ extendSubstList subst ((var,rhs):prs) = extendSubstList (extendSubst subst var r
 
 -- | Find the substitution for an 'Id' in the 'Subst'
 lookupIdSubst :: SDoc -> Subst -> Id -> CoreExpr
-lookupIdSubst doc (Subst in_scope ids _ _) v
+lookupIdSubst doc (Subst in_scope (ISE { ise_env = lcl, ise_gbl = gbl }) _ _) v
   | not (isLocalId v) = Var v
-  | Just e  <- lookupVarEnv ids       v = e
+  | Just e <- lookupVarEnv lcl        v = e
+  | Just e <- lookupGIS gbl in_scope  v = e
   | Just v' <- lookupInScope in_scope v = Var v'
         -- Vital! See Note [Extending the Subst]
   | otherwise = WARN( True, ptext (sLit "CoreSubst.lookupIdSubst") <+> doc <+> ppr v
@@ -278,14 +361,15 @@ lookupCvSubst :: Subst -> CoVar -> Coercion
 lookupCvSubst (Subst _ _ _ cvs) v = ASSERT( isCoVar v ) lookupVarEnv cvs v `orElse` mkCoVarCo v
 
 delBndr :: Subst -> Var -> Subst
+-- Doesn't work for gbl_ids
 delBndr (Subst in_scope ids tvs cvs) v
   | isCoVar v = Subst in_scope ids tvs (delVarEnv cvs v)
   | isTyVar v = Subst in_scope ids (delVarEnv tvs v) cvs
-  | otherwise = Subst in_scope (delVarEnv ids v) tvs cvs
+  | otherwise = Subst in_scope (delIdSubst ids v) tvs cvs
 
 delBndrs :: Subst -> [Var] -> Subst
 delBndrs (Subst in_scope ids tvs cvs) vs
-  = Subst in_scope (delVarEnvList ids vs) (delVarEnvList tvs vs) (delVarEnvList cvs vs)
+  = Subst in_scope (delIdSubstList ids vs) (delVarEnvList tvs vs) (delVarEnvList cvs vs)
       -- Easiest thing is just delete all from all!
 
 -- | Simultaneously substitute for a bunch of variables
@@ -293,10 +377,11 @@ delBndrs (Subst in_scope ids tvs cvs) vs
 --   ie the substitution for   (\x \y. e) a1 a2
 --      so neither x nor y scope over a1 a2
 mkOpenSubst :: InScopeSet -> [(Var,CoreArg)] -> Subst
-mkOpenSubst in_scope pairs = Subst in_scope
-                                   (mkVarEnv [(id,e)  | (id, e) <- pairs, isId id])
-                                   (mkVarEnv [(tv,ty) | (tv, Type ty) <- pairs])
-                                   (mkVarEnv [(v,co)  | (v, Coercion co) <- pairs])
+mkOpenSubst in_scope pairs
+  = Subst in_scope
+          (ISE { ise_env = mkVarEnv [(id,e)  | (id, e) <- pairs, isId id], ise_gbl = NoGIS})
+          (mkVarEnv [(tv,ty) | (tv, Type ty) <- pairs])
+          (mkVarEnv [(v,co)  | (v, Coercion co) <- pairs])
 
 ------------------------------
 isInScope :: Var -> Subst -> Bool
@@ -313,20 +398,20 @@ addInScopeSet (Subst in_scope ids tvs cvs) vs
 extendInScope :: Subst -> Var -> Subst
 extendInScope (Subst in_scope ids tvs cvs) v
   = Subst (in_scope `extendInScopeSet` v)
-          (ids `delVarEnv` v) (tvs `delVarEnv` v) (cvs `delVarEnv` v)
+          (ids `delIdSubst` v) (tvs `delVarEnv` v) (cvs `delVarEnv` v)
 
 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
 extendInScopeList :: Subst -> [Var] -> Subst
 extendInScopeList (Subst in_scope ids tvs cvs) vs
   = Subst (in_scope `extendInScopeSetList` vs)
-          (ids `delVarEnvList` vs) (tvs `delVarEnvList` vs) (cvs `delVarEnvList` vs)
+          (ids `delIdSubstList` vs) (tvs `delVarEnvList` vs) (cvs `delVarEnvList` vs)
 
 -- | Optimized version of 'extendInScopeList' that can be used if you are certain
 -- all the things being added are 'Id's and hence none are 'TyVar's or 'CoVar's
 extendInScopeIds :: Subst -> [Id] -> Subst
 extendInScopeIds (Subst in_scope ids tvs cvs) vs
   = Subst (in_scope `extendInScopeSetList` vs)
-          (ids `delVarEnvList` vs) tvs cvs
+          (ids `delIdSubstList` vs) tvs cvs
 
 setInScope :: Subst -> InScopeSet -> Subst
 setInScope (Subst _ ids tvs cvs) in_scope = Subst in_scope ids tvs cvs
@@ -497,8 +582,8 @@ substIdBndr _doc rec_subst subst@(Subst in_scope env tvs cvs) old_id
 
         -- Extend the substitution if the unique has changed
         -- See the notes with substTyVarBndr for the delVarEnv
-    new_env | no_change = delVarEnv env old_id
-            | otherwise = extendVarEnv env old_id (Var new_id)
+    new_env | no_change = delIdSubst env old_id
+            | otherwise = extendIdSubstEnv env old_id (Var new_id)
 
     no_change = id1 == old_id
         -- See Note [Extending the Subst]
@@ -553,7 +638,7 @@ clone_id rec_subst subst@(Subst in_scope idvs tvs cvs) (old_id, uniq)
     id2     = substIdType subst id1
     new_id  = maybeModifyIdInfo (substIdInfo rec_subst id2 (idInfo old_id)) id2
     (new_idvs, new_cvs) | isCoVar old_id = (idvs, extendVarEnv cvs old_id (mkCoVarCo new_id))
-                        | otherwise      = (extendVarEnv idvs old_id (Var new_id), cvs)
+                        | otherwise      = (extendIdSubstEnv idvs old_id (Var new_id), cvs)
 
 {-
 ************************************************************************
@@ -1066,9 +1151,9 @@ subst_opt_id_bndr subst@(Subst in_scope id_subst tv_subst cv_subst) old_id
         -- or there's some useful occurrence information
         -- See the notes with substTyVarBndr for the delSubstEnv
     new_id_subst | new_id /= old_id
-                 = extendVarEnv id_subst old_id (Var new_id)
+                 = extendIdSubstEnv id_subst old_id (Var new_id)
                  | otherwise
-                 = delVarEnv id_subst old_id
+                 = delIdSubst id_subst old_id
 
 ----------------------
 subst_opt_bndrs :: Subst -> [InVar] -> (Subst, [OutVar])
diff --git a/compiler/simplCore/SimplEnv.hs b/compiler/simplCore/SimplEnv.hs
index a3489b671d..aa2f99f04e 100644
--- a/compiler/simplCore/SimplEnv.hs
+++ b/compiler/simplCore/SimplEnv.hs
@@ -23,7 +23,7 @@ module SimplEnv (
 
         SimplSR(..), mkContEx, substId, lookupRecBndr, refineFromInScope,
 
-        substExprS,
+        substExpr,
         simplNonRecBndr, simplRecBndrs,
         simplBinder, simplBinders,
         substTy, substTyVar, getTvSubst,
@@ -46,6 +46,7 @@ import VarEnv
 import VarSet
 import OrdList
 import Id
+import qualified CoreSubst
 import MkCore                   ( mkWildValBinder )
 import TysWiredIn
 import qualified Type
@@ -538,72 +539,6 @@ lookupRecBndr (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
         Just _ -> pprPanic "lookupRecBndr" (ppr v)
         Nothing -> refineFromInScope in_scope v
 
-
-substExprS :: SimplEnv -> CoreExpr -> CoreExpr
--- This entire substExprS thing is called in just one place
--- but we can't use substExpr because it uses a different shape
--- of substitution  Better solution coming in HEAD.
-substExprS env expr
-  = go expr
-  where
-    go (Var v) = case substId env v of
-                   DoneId v'            -> Var v'
-                   DoneEx e             -> e
-                   ContEx tvs cvs ids e -> substExprS (setSubstEnv env tvs cvs ids) e
-
-    go (Type ty)       = Type (substTy env ty)
-    go (Coercion co)   = Coercion (substCo env co)
-    go (Lit lit)       = Lit lit
-    go (App fun arg)   = App (go fun) (go arg)
-    go (Tick tickish e) = mkTick (substTickishS env tickish) (go e)
-    go (Cast e co)     = Cast (go e) (substCo env co)
-       -- Do not optimise even identity coercions
-       -- Reason: substitution applies to the LHS of RULES, and
-       --         if you "optimise" an identity coercion, you may
-       --         lose a binder. We optimise the LHS of rules at
-       --         construction time
-
-    go (Lam bndr body) = Lam bndr' (substExprS env' body)
-                       where
-                         (env', bndr') = substBndr env bndr
-
-    go (Let bind body) = Let bind' (substExprS env' body)
-                       where
-                         (env', bind') = substBindS env bind
-
-    go (Case scrut bndr ty alts)
-      = Case (go scrut) bndr' (substTy env ty)
-                              (map (go_alt env') alts)
-      where
-       (env', bndr') = substBndr env bndr
-
-    go_alt env (con, bndrs, rhs) = (con, bndrs', substExprS env' rhs)
-                                 where
-                                   (env', bndrs') = substBndrs env bndrs
-
-substTickishS :: SimplEnv -> Tickish Id -> Tickish Id
-substTickishS env (Breakpoint n ids) = Breakpoint n (map do_one ids)
- where
-   do_one = getIdFromTrivialExpr . substExprS env . Var  -- Ugh
-substTickishS _subst other = other
-
--- | Apply a substitution to an entire 'CoreBind', additionally returning an updated 'Subst'
--- that should be used by subsequent substitutions.
-substBindS :: SimplEnv -> CoreBind -> (SimplEnv, CoreBind)
-
-substBindS env (NonRec bndr rhs) = (env', NonRec bndr' (substExprS env rhs))
-                                 where
-                                   (env', bndr') = substBndr env bndr
-
-substBindS env (Rec pairs)
-  = (env', Rec (bndrs' `zip` rhss'))
-  where
-    (bndrs, rhss)  = unzip pairs
-    (env', bndrs') = substBndrs env bndrs
-    rhss'          = map (substExprS env') rhss
-    -- No need for the complexity of CoreSubst.substRecBndrs, because
-    -- we zap all IdInfo that depends on free variables
-
 {-
 ************************************************************************
 *                                                                      *
@@ -612,8 +547,6 @@ substBindS env (Rec pairs)
 ************************************************************************
 
 
-* substBndr, substBndrs: non-monadic version
-
 * sinplBndr, simplBndrs: monadic version, only so that they
                          can be made strict via seq.
 
@@ -647,15 +580,6 @@ simplRecBndrs env@(SimplEnv {}) ids
   = do  { let (env1, ids1) = mapAccumL substIdBndr env ids
         ; seqIds ids1 `seq` return env1 }
 
----------------
-substBndr :: SimplEnv -> InBndr -> (SimplEnv, OutBndr)
-substBndr env bndr
-  | isTyVar bndr  = substTyVarBndr env bndr
-  | otherwise     = substIdBndr env bndr
-
-substBndrs :: SimplEnv -> [InBndr] -> (SimplEnv, [OutBndr])
-substBndrs  env bndrs = mapAccumL substBndr env bndrs
-
 substIdBndr :: SimplEnv -> InBndr -> (SimplEnv, OutBndr)
 -- Might be a coercion variable
 substIdBndr env bndr
@@ -804,3 +728,40 @@ substIdType (SimplEnv { seInScope = in_scope,  seTvSubst = tv_env }) id
                 -- in a Note in the id's type itself
   where
     old_ty = idType id
+
+substExpr :: SimplEnv -> CoreExpr -> CoreExpr
+-- See Note [Substitution in the simplifier]
+substExpr (SimplEnv { seInScope = in_scope
+                    , seTvSubst = tv_env
+                    , seCvSubst = cv_env
+                    , seIdSubst = id_env })
+  = subst_expr in_scope tv_env cv_env id_env
+  where
+    subst_expr :: InScopeSet -> TvSubstEnv -> CvSubstEnv -> SimplIdSubst
+               -> CoreExpr -> CoreExpr
+    subst_expr is tvs cvs id_env
+      = CoreSubst.substExpr (text "SimplEnv.substExpr")
+                            (CoreSubst.mkGblSubst is tvs cvs lookup_id)
+      where
+        lookup_id in_scope v
+          = case lookupVarEnv id_env v of
+              Nothing                  -> Nothing
+              Just (DoneEx e)          -> Just e
+              Just (DoneId v)          -> Just (Var v)
+              Just (ContEx tv cv id e) -> Just (subst_expr in_scope tv cv id e)
+
+{- Note [Substitution in the simplifier]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+In just one place (sigh) we need to lazily substitute over a CoreExpr.
+For that we need CoreSubst.substExpr.  But there is a difficulty: SimplEnv
+has a SimplIdSubst, whose range is SimplSR, not just CoreExpr.
+
+So SimplEnv.substExpr has to perform impedence-matching, via the ambient
+substitution provided by mkGblSubst.  It seems like a lot of work for
+a small thing.  Previously we attempted to construct a (VarEnv CoreExpr)
+from the SimplIdSubst, but that had absolutely terrible performance
+(Trac #10370 comment:12).  Then I tried to write a complete new substExpr
+that used SimplIdSubst insead of (VarEnv CoreExpr), but that got out of
+hand because we need to substitute over rules and unfoldings too
+(Trac #5113, comment:7 and following).
+-}
\ No newline at end of file
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 2e1dcefbdb..40a68d4e6c 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -1179,7 +1179,7 @@ simplCast env body co0 cont0
                 -- But it isn't a common case.
                 --
                 -- Example of use: Trac #995
-         = do { let arg' = substExprS arg_se arg
+         = do { let arg' = substExpr arg_se arg
                            -- It's important that this is lazy, because this argument
                            -- may be disarded if turns out to be the argument of
                            -- (\_ -> e)     This can make a huge difference;
diff --git a/compiler/specialise/Rules.hs b/compiler/specialise/Rules.hs
index 2169dc7e5c..2c68772b91 100644
--- a/compiler/specialise/Rules.hs
+++ b/compiler/specialise/Rules.hs
@@ -591,10 +591,10 @@ data RuleMatchEnv
 rvInScopeEnv :: RuleMatchEnv -> InScopeEnv
 rvInScopeEnv renv = (rnInScopeSet (rv_lcl renv), rv_unf renv)
 
-data RuleSubst = RS { rs_tv_subst :: TvSubstEnv   -- Range is the
-                    , rs_id_subst :: IdSubstEnv   --   template variables
-                    , rs_binds    :: BindWrapper  -- Floated bindings
-                    , rs_bndrs    :: VarSet       -- Variables bound by floated lets
+data RuleSubst = RS { rs_tv_subst :: TvSubstEnv      -- Range is the
+                    , rs_id_subst :: IdEnv CoreExpr  --   template variables
+                    , rs_binds    :: BindWrapper     -- Floated bindings
+                    , rs_bndrs    :: VarSet          -- Variables bound by floated lets
                     }
 
 type BindWrapper = CoreExpr -> CoreExpr