Broaden the in-scope sets for liftEnvSubst and composeTCvSubst

This patch fixes two distinct (but closely related) buglets that were uncovered
in #22235:

* `liftEnvSubst` used an empty in-scope set, which was not wide enough to cover
  the variables in the range of the substitution. This patch fixes this by
  populating the in-scope set from the free variables in the range of the
  substitution.
* `composeTCvSubst` applied the first substitution argument to the range of the
  second substitution argument, but the first substitution's in-scope set was
  not wide enough to cover the range of the second substutition. We similarly
  fix this issue in this patch by widening the first substitution's in-scope set
  before applying it.

Fixes #22235.

4 files changed, 37 insertions, 3 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 8a13ecb51b..06bc235913 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -2234,13 +2234,18 @@ liftEnvSubstRight = liftEnvSubst pSnd
 
 liftEnvSubst :: (forall a. Pair a -> a) -> Subst -> LiftCoEnv -> Subst
 liftEnvSubst selector subst lc_env
-  = composeTCvSubst (Subst emptyInScopeSet emptyIdSubstEnv tenv cenv) subst
+  = composeTCvSubst (Subst in_scope emptyIdSubstEnv tenv cenv) subst
   where
     pairs            = nonDetUFMToList lc_env
                        -- It's OK to use nonDetUFMToList here because we
                        -- immediately forget the ordering by creating
                        -- a VarEnv
     (tpairs, cpairs) = partitionWith ty_or_co pairs
+    -- Make sure the in-scope set is wide enough to cover the range of the
+    -- substitution (#22235).
+    in_scope         = mkInScopeSet $
+                       tyCoVarsOfTypes (map snd tpairs) `unionVarSet`
+                       tyCoVarsOfCos (map snd cpairs)
     tenv             = mkVarEnv_Directly tpairs
     cenv             = mkVarEnv_Directly cpairs
 
diff --git a/compiler/GHC/Core/TyCo/Subst.hs b/compiler/GHC/Core/TyCo/Subst.hs
index d9d674bb30..18babd08dc 100644
--- a/compiler/GHC/Core/TyCo/Subst.hs
+++ b/compiler/GHC/Core/TyCo/Subst.hs
@@ -245,8 +245,13 @@ composeTCvSubst subst1@(Subst is1 ids1 tenv1 cenv1) (Subst is2 _ tenv2 cenv2)
   = Subst is3 ids1 tenv3 cenv3
   where
     is3 = is1 `unionInScope` is2
-    tenv3 = tenv1 `plusVarEnv` mapVarEnv (substTy subst1) tenv2
-    cenv3 = cenv1 `plusVarEnv` mapVarEnv (substCo subst1) cenv2
+    tenv3 = tenv1 `plusVarEnv` mapVarEnv (substTy extended_subst1) tenv2
+    cenv3 = cenv1 `plusVarEnv` mapVarEnv (substCo extended_subst1) cenv2
+
+    -- Make sure the in-scope set in the first substitution is wide enough to
+    -- cover the free variables in the range of the second substitution before
+    -- applying it (#22235).
+    extended_subst1 = subst1 `setInScope` is3
 
 emptySubst :: Subst
 emptySubst = Subst emptyInScopeSet emptyVarEnv emptyVarEnv emptyVarEnv
diff --git a/testsuite/tests/gadt/T22235.hs b/testsuite/tests/gadt/T22235.hs
new file mode 100644
index 0000000000..0f4161db72
--- /dev/null
+++ b/testsuite/tests/gadt/T22235.hs
@@ -0,0 +1,23 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+module T22235 (f) where
+
+import Data.Kind (Type)
+import Data.Type.Equality ((:~:)(..))
+
+f :: ST x -> ST y -> x :~: y
+f st1@SMkT st2@SMkT = method st1 st2
+
+type T :: Type -> Type
+data T a where
+  MkT :: T Int
+
+type ST :: T a -> Type
+data ST (t :: T a) where
+  SMkT :: ST MkT
+
+class C f where
+  method :: f a -> f b -> a :~: b
+
+instance C ST where
+  method SMkT SMkT = Refl
diff --git a/testsuite/tests/gadt/all.T b/testsuite/tests/gadt/all.T
index ce30d570f3..1b32762525 100644
--- a/testsuite/tests/gadt/all.T
+++ b/testsuite/tests/gadt/all.T
@@ -125,3 +125,4 @@ test('T20278', normal, compile, [''])
 test('SynDataRec', normal, compile, [''])
 test('T20485', normal, compile, [''])
 test('T20485a', normal, compile, [''])
+test('T22235', normal, compile, [''])