Fix substitution in bindAuxiliaryDict

In GHC.Core.Opt.Specialise.bindAuxiliaryDict we were unnecessarily
calling `extendInScope` to bring into scope variables that were
/already/ in scope.  Worse, GHC.Core.Subst.extendInScope strangely
deleted the newly-in-scope variables from the substitution -- and that
was fatal in #21391.

I removed the redundant calls to extendInScope.

More ambitiously, I changed GHC.Core.Subst.extendInScope (and cousins)
to stop deleting variables from the substitution.  I even changed the
names of the function to extendSubstInScope (and cousins) and audited
all the calls to check that deleting from the substitution was wrong.
In fact there are very few such calls, and they are all about
introducing a fresh non-in-scope variable.  These are "OutIds"; it is
utterly wrong to mess with the "InId" substitution.

I have not added a Note, because I'm deleting wrong code, and it'd be
distracting to document a bug.

9 files changed, 74 insertions, 44 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index c78285c6f9..6c0729ec5b 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -45,7 +45,7 @@ import GHC.Core.Multiplicity
 -- We have two sorts of substitution:
 --   GHC.Core.Subst.Subst, and GHC.Core.TyCo.TCvSubst
 -- Both have substTy, substCo  Hence need for qualification
-import GHC.Core.Subst    as Core hiding ( extendInScopeSet )
+import GHC.Core.Subst    as Core
 import GHC.Core.Type     as Type
 import GHC.Core.Coercion as Type
 
diff --git a/compiler/GHC/Core/Opt/SpecConstr.hs b/compiler/GHC/Core/Opt/SpecConstr.hs
index 90f492ffea..d9429053a0 100644
--- a/compiler/GHC/Core/Opt/SpecConstr.hs
+++ b/compiler/GHC/Core/Opt/SpecConstr.hs
@@ -956,7 +956,8 @@ zapScSubst env = env { sc_subst = zapSubstEnv (sc_subst env) }
 
 extendScInScope :: ScEnv -> [Var] -> ScEnv
         -- Bring the quantified variables into scope
-extendScInScope env qvars = env { sc_subst = extendInScopeList (sc_subst env) qvars }
+extendScInScope env qvars
+  = env { sc_subst = extendSubstInScopeList (sc_subst env) qvars }
 
         -- Extend the substitution
 extendScSubst :: ScEnv -> Var -> OutExpr -> ScEnv
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 36ee54b7c0..5c515dd95b 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -1307,7 +1307,7 @@ bringFloatedDictsIntoScope env uds =
   env{se_subst=subst'}
   where
    dx_bndrs = ud_bs_of_binds uds
-   subst'   = se_subst env `Core.extendInScopeSet` dx_bndrs
+   subst'   = se_subst env `Core.extendSubstInScopeSet` dx_bndrs
 
 specBind :: SpecEnv                     -- Use this for RHSs
          -> CoreBind                    -- Binders are already cloned by cloneBindSM,
@@ -1516,6 +1516,7 @@ specCalls spec_imp env existing_rules calls_for_me fn rhs
 --                                        , text "spec_args: " <+> ppr spec_args
 --                                        , text "dx_binds:  " <+> ppr dx_binds
 --                                        , text "rhs_env2:  " <+> ppr (se_subst rhs_env2)
+--                                        , text "leftover_bndrs:" <+> pprIds leftover_bndrs
 --                                        , ppr dx_binds ]) $
 --             return ()
 
@@ -2354,6 +2355,9 @@ specHeader env (bndr : bndrs) (SpecType ty : args)
              qvars    = scopedSort $
                         filterOut (`elemInScopeSet` in_scope) $
                         tyCoVarsOfTypeList ty
+             -- qvars are the type variables free in the call that
+             -- are not already in scope.  Quantify over these.
+             -- See Note [Specialising polymorphic dictionaries]
              (env1, qvars') = substBndrs env qvars
              ty'            = substTy env1 ty
              env2           = extendTvSubstList env1 [(bndr, ty')]
@@ -2393,12 +2397,12 @@ specHeader env (bndr : bndrs) (UnspecType : args)
 -- a wildcard binder to match the dictionary (See Note [Specialising Calls] for
 -- the nitty-gritty), as a LHS rule and unfolding details.
 specHeader env (bndr : bndrs) (SpecDict d : args)
-  = do { bndr' <- newDictBndr env bndr -- See Note [Zap occ info in rule binders]
-       ; let (env', dx_bind, spec_dict) = bindAuxiliaryDict env bndr bndr' d
-       ; (_, env'', leftover_bndrs, rule_bs, rule_es, bs', dx, spec_args)
-             <- specHeader env' bndrs args
+  = do { (env1, bndr') <- newDictBndr env bndr -- See Note [Zap occ info in rule binders]
+       ; let (env2, dx_bind, spec_dict) = bindAuxiliaryDict env1 bndr bndr' d
+       ; (_, env3, leftover_bndrs, rule_bs, rule_es, bs', dx, spec_args)
+             <- specHeader env2 bndrs args
        ; pure ( True      -- Ha!  A useful specialisation!
-              , env''
+              , env3
               , leftover_bndrs
               -- See Note [Evidence foralls]
               , exprFreeIdsList (varToCoreExpr bndr') ++ rule_bs
@@ -2464,9 +2468,9 @@ specHeader env bndrs []
 bindAuxiliaryDict
   :: SpecEnv
   -> InId -> OutId -> OutExpr -- Original dict binder, and the witnessing expression
-  -> ( SpecEnv        -- Substitute for orig_dict_id
+  -> ( SpecEnv        -- Substitutes for orig_dict_id
      , Maybe DictBind -- Auxiliary dict binding, if any
-     , OutExpr)        -- Witnessing expression (always trivial)
+     , OutExpr)       -- Witnessing expression (always trivial)
 bindAuxiliaryDict env@(SE { se_subst = subst, se_interesting = interesting })
                   orig_dict_id fresh_dict_id dict_expr
 
@@ -2474,7 +2478,6 @@ bindAuxiliaryDict env@(SE { se_subst = subst, se_interesting = interesting })
   -- don’t bother creating a new dict binding; just substitute
   | Just dict_id <- getIdFromTrivialExpr_maybe dict_expr
   = let env' = env { se_subst = Core.extendSubst subst orig_dict_id dict_expr
-                                `Core.extendInScope` dict_id
                           -- See Note [Keep the old dictionaries interesting]
                    , se_interesting = interesting `extendVarSet` dict_id }
     in -- pprTrace "bindAuxiliaryDict:trivial" (ppr orig_dict_id <+> ppr dict_id) $
@@ -2486,8 +2489,8 @@ bindAuxiliaryDict env@(SE { se_subst = subst, se_interesting = interesting })
           -- See Note [Specialisation modulo dictionary selectors] for the unfolding
         dict_bind = mkDB (NonRec fresh_dict_id' dict_expr)
         env' = env { se_subst = Core.extendSubst subst orig_dict_id (Var fresh_dict_id')
-                                `Core.extendInScope`     fresh_dict_id'
-                                `Core.extendInScopeList` exprFreeVarsList dict_expr
+                                `Core.extendSubstInScope` fresh_dict_id'
+                                -- Ensure the new unfolding is in the in-scope set
                       -- See Note [Make the new dictionaries interesting]
                    , se_interesting = interesting `extendVarSet` fresh_dict_id' }
     in -- pprTrace "bindAuxiliaryDict:non-trivial" (ppr orig_dict_id <+> ppr fresh_dict_id' $$ ppr dict_expr $$ ppr (exprFreeVarsList dict_expr)) $
@@ -3184,13 +3187,16 @@ cloneBindSM env@(SE { se_subst = subst, se_interesting = interesting }) (Rec pai
                                            [ v | (v,r) <- pairs, typeDeterminesValue (idType v), interestingDict env r ] }
        ; return (env', env', Rec (bndrs' `zip` map snd pairs)) }
 
-newDictBndr :: SpecEnv -> CoreBndr -> SpecM CoreBndr
+newDictBndr :: SpecEnv -> CoreBndr -> SpecM (SpecEnv, CoreBndr)
 -- Make up completely fresh binders for the dictionaries
 -- Their bindings are going to float outwards
-newDictBndr env b = do { uniq <- getUniqueM
-                        ; let n   = idName b
-                              ty' = substTy env (idType b)
-                        ; return (mkUserLocal (nameOccName n) uniq Many ty' (getSrcSpan n)) }
+newDictBndr env@(SE { se_subst = subst }) b
+  = do { uniq <- getUniqueM
+       ; let n    = idName b
+             ty'  = Core.substTy subst (idType b)
+             b'   = mkUserLocal (nameOccName n) uniq Many ty' (getSrcSpan n)
+             env' = env { se_subst = subst `Core.extendSubstInScope` b' }
+       ; pure (env', b') }
 
 newSpecIdSM :: Id -> Type -> Maybe JoinArity -> SpecM Id
     -- Give the new Id a similar occurrence name to the old one
diff --git a/compiler/GHC/Core/Ppr.hs b/compiler/GHC/Core/Ppr.hs
index b319abec08..816af025fc 100644
--- a/compiler/GHC/Core/Ppr.hs
+++ b/compiler/GHC/Core/Ppr.hs
@@ -22,7 +22,7 @@ module GHC.Core.Ppr (
         pprCoreExpr, pprParendExpr,
         pprCoreBinding, pprCoreBindings, pprCoreAlt,
         pprCoreBindingWithSize, pprCoreBindingsWithSize,
-        pprCoreBinder, pprCoreBinders,
+        pprCoreBinder, pprCoreBinders, pprId, pprIds,
         pprRule, pprRules, pprOptCo,
         pprOcc, pprOccWithTick
     ) where
@@ -462,6 +462,13 @@ pprKindedTyVarBndr :: TyVar -> SDoc
 pprKindedTyVarBndr tyvar
   = text "@" <> pprTyVar tyvar
 
+-- pprId x prints x :: ty
+pprId :: Id -> SDoc
+pprId x = ppr x <+> dcolon <+> ppr (idType x)
+
+pprIds :: [Id] -> SDoc
+pprIds xs = sep (map pprId xs)
+
 -- pprIdBndr does *not* print the type
 -- When printing any Id binder in debug mode, we print its inline pragma and one-shot-ness
 pprIdBndr :: Id -> SDoc
diff --git a/compiler/GHC/Core/SimpleOpt.hs b/compiler/GHC/Core/SimpleOpt.hs
index 360c868738..8818f51384 100644
--- a/compiler/GHC/Core/SimpleOpt.hs
+++ b/compiler/GHC/Core/SimpleOpt.hs
@@ -22,7 +22,7 @@ import GHC.Prelude
 
 import GHC.Core
 import GHC.Core.Opt.Arity
-import GHC.Core.Subst hiding ( extendInScopeSet )
+import GHC.Core.Subst
 import GHC.Core.Utils
 import GHC.Core.FVs
 import GHC.Core.Unfold
@@ -1260,7 +1260,7 @@ exprIsConApp_maybe (in_scope, id_unf) expr
     subst_in_scope (Right s) = substInScope s
 
     subst_extend_in_scope (Left in_scope) v = Left (in_scope `extendInScopeSet` v)
-    subst_extend_in_scope (Right s) v = Right (s `extendInScope` v)
+    subst_extend_in_scope (Right s) v = Right (s `extendSubstInScope` v)
 
     subst_co (Left {}) co = co
     subst_co (Right s) co = GHC.Core.Subst.substCo s co
diff --git a/compiler/GHC/Core/Subst.hs b/compiler/GHC/Core/Subst.hs
index 2c470c5dcb..e6f0237f32 100644
--- a/compiler/GHC/Core/Subst.hs
+++ b/compiler/GHC/Core/Subst.hs
@@ -24,7 +24,7 @@ module GHC.Core.Subst (
         emptySubst, mkEmptySubst, mkSubst, mkOpenSubst, substInScope, isEmptySubst,
         extendIdSubst, extendIdSubstList, extendTCvSubst, extendTvSubstList,
         extendSubst, extendSubstList, extendSubstWithVar, zapSubstEnv,
-        extendInScope, extendInScopeList, extendInScopeIds, GHC.Core.Subst.extendInScopeSet,
+        extendSubstInScope, extendSubstInScopeList, extendSubstInScopeSet,
         isInScope, setInScope, getTCvSubst, extendTvSubst, extendCvSubst,
         delBndr, delBndrs,
 
@@ -57,7 +57,6 @@ import GHC.Types.Var
 import GHC.Types.Tickish
 import GHC.Types.Id.Info
 import GHC.Types.Unique.Supply
-import GHC.Types.Unique.Set
 
 import GHC.Builtin.Names
 import GHC.Data.Maybe
@@ -285,33 +284,23 @@ mkOpenSubst in_scope pairs = Subst in_scope
 isInScope :: Var -> Subst -> Bool
 isInScope v (Subst in_scope _ _ _) = v `elemInScopeSet` in_scope
 
--- | Add the 'Var' to the in-scope set: as a side effect,
--- and remove any existing substitutions for it
-extendInScope :: Subst -> Var -> Subst
-extendInScope (Subst in_scope ids tvs cvs) v
+-- | Add the 'Var' to the in-scope set
+extendSubstInScope :: Subst -> Var -> Subst
+extendSubstInScope (Subst in_scope ids tvs cvs) v
   = Subst (in_scope `InScopeSet.extendInScopeSet` v)
-          (ids `delVarEnv` v) (tvs `delVarEnv` v) (cvs `delVarEnv` v)
+          ids tvs cvs
 
 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
-extendInScopeList :: Subst -> [Var] -> Subst
-extendInScopeList (Subst in_scope ids tvs cvs) vs
+extendSubstInScopeList :: Subst -> [Var] -> Subst
+extendSubstInScopeList (Subst in_scope ids tvs cvs) vs
   = Subst (in_scope `extendInScopeSetList` vs)
-          (ids `delVarEnvList` 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 tvs cvs
 
 -- | Add the 'Var's to the in-scope set: see also 'extendInScope'
-extendInScopeSet :: Subst -> VarSet -> Subst
-extendInScopeSet (Subst in_scope ids tvs cvs) vs
+extendSubstInScopeSet :: Subst -> VarSet -> Subst
+extendSubstInScopeSet (Subst in_scope ids tvs cvs) vs
   = Subst (in_scope `extendInScopeSetSet` vs)
-          (ids `minus` vs) (tvs `minus` vs) (cvs `minus` vs)
-  where
-    minus env set = minusVarEnv env (getUniqSet set)
+          ids tvs cvs
 
 setInScope :: Subst -> InScopeSet -> Subst
 setInScope (Subst _ ids tvs cvs) in_scope = Subst in_scope ids tvs cvs
diff --git a/compiler/GHC/Plugins.hs b/compiler/GHC/Plugins.hs
index e79362b9d1..9c384317d2 100644
--- a/compiler/GHC/Plugins.hs
+++ b/compiler/GHC/Plugins.hs
@@ -88,7 +88,7 @@ import GHC.Core.DataCon
 import GHC.Core.Utils
 import GHC.Core.Make
 import GHC.Core.FVs
-import GHC.Core.Subst hiding( substTyVarBndr, substCoVarBndr, extendCvSubst, extendInScopeSet )
+import GHC.Core.Subst hiding( substTyVarBndr, substCoVarBndr, extendCvSubst, extendSubstInScopeSet )
        -- These names are also exported by Type
 
 import GHC.Core.Rules
diff --git a/testsuite/tests/simplCore/should_compile/T21391.hs b/testsuite/tests/simplCore/should_compile/T21391.hs
new file mode 100644
index 0000000000..3a974eddb7
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T21391.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE RankNTypes #-}
+module Web.Routing.SafeRouting where
+
+import Control.DeepSeq (NFData (..))
+import Data.Kind (Constraint, Type)
+import Data.Typeable (Typeable)
+
+class FromHttpApiData a where
+
+data PolyMap (c :: Type -> Constraint) (f :: Type -> Type) (a :: Type) where
+  PMNil :: PolyMap c f a
+  PMCons :: (Typeable p, c p) => f (p -> a) -> PolyMap c f a -> PolyMap c f a
+
+rnfHelper :: (forall p. c p => f (p -> a) -> ()) -> PolyMap c f a -> ()
+rnfHelper _ PMNil = ()
+rnfHelper h (PMCons v pm) = h v `seq` rnfHelper h pm
+
+data PathMap x =
+  PathMap [x] (PolyMap FromHttpApiData PathMap x)
+
+instance NFData x => NFData (PathMap x) where
+  rnf (PathMap a b) = rnf a `seq` rnfHelper rnf b
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 02a5de56c6..cbe344df93 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -403,3 +403,5 @@ test('T21261', [ grep_errmsg(r'^(yes|no)') ], compile, ['-O -ddump-simpl -dno-ty
 test('T17966',  [ grep_errmsg(r'SPEC') ], compile, ['-O -ddump-spec'])
 # We expect to see a SPEC rule for $cm
 test('T19644',  [ grep_errmsg(r'SPEC') ], compile, ['-O -ddump-spec'])
+
+test('T21391', normal, compile, ['-O -dcore-lint'])