Get the in-scope set right in FamInstEnv.injectiveBrancheswip/T21896

There was an assert error, as Gergo pointed out in #21896.

I fixed this by adding an InScopeSet argument to tcUnifyTyWithTFs.
And also to GHC.Core.Unify.niFixTCvSubst.

I also took the opportunity to get a couple more InScopeSets right,
and to change some substTyUnchecked into substTy.

This MR touches a lot of other files, but only because I also took the
opportunity to introduce mkInScopeSetList, and use it.

16 files changed, 193 insertions, 130 deletions

diff --git a/compiler/GHC/Core/FamInstEnv.hs b/compiler/GHC/Core/FamInstEnv.hs
index 0e6670f7ec..c4ff9ee1f7 100644
--- a/compiler/GHC/Core/FamInstEnv.hs
+++ b/compiler/GHC/Core/FamInstEnv.hs
@@ -556,30 +556,34 @@ data InjectivityCheckResult
 injectiveBranches :: [Bool] -> CoAxBranch -> CoAxBranch
                   -> InjectivityCheckResult
 injectiveBranches injectivity
-                  ax1@(CoAxBranch { cab_lhs = lhs1, cab_rhs = rhs1 })
-                  ax2@(CoAxBranch { cab_lhs = lhs2, cab_rhs = rhs2 })
+                  ax1@(CoAxBranch { cab_tvs = tvs1, cab_lhs = lhs1, cab_rhs = rhs1 })
+                  ax2@(CoAxBranch { cab_tvs = tvs2, cab_lhs = lhs2, cab_rhs = rhs2 })
   -- See Note [Verifying injectivity annotation], case 1.
   = let getInjArgs  = filterByList injectivity
-    in case tcUnifyTyWithTFs True rhs1 rhs2 of -- True = two-way pre-unification
+        in_scope    = mkInScopeSetList (tvs1 ++ tvs2)
+    in case tcUnifyTyWithTFs True in_scope rhs1 rhs2 of -- True = two-way pre-unification
        Nothing -> InjectivityAccepted
          -- RHS are different, so equations are injective.
          -- This is case 1A from Note [Verifying injectivity annotation]
-       Just subst -> -- RHS unify under a substitution
-        let lhs1Subst = Type.substTys subst (getInjArgs lhs1)
-            lhs2Subst = Type.substTys subst (getInjArgs lhs2)
-        -- If LHSs are equal under the substitution used for RHSs then this pair
-        -- of equations does not violate injectivity annotation. If LHSs are not
-        -- equal under that substitution then this pair of equations violates
-        -- injectivity annotation, but for closed type families it still might
-        -- be the case that one LHS after substitution is unreachable.
-        in if eqTypes lhs1Subst lhs2Subst  -- check case 1B1 from Note.
-           then InjectivityAccepted
-           else InjectivityUnified ( ax1 { cab_lhs = Type.substTys subst lhs1
-                                         , cab_rhs = Type.substTy  subst rhs1 })
-                                   ( ax2 { cab_lhs = Type.substTys subst lhs2
-                                         , cab_rhs = Type.substTy  subst rhs2 })
-                -- payload of InjectivityUnified used only for check 1B2, only
-                -- for closed type families
+
+       Just subst -- RHS unify under a substitution
+          -- If LHSs are equal under the substitution used for RHSs then this pair
+          -- of equations does not violate injectivity annotation. If LHSs are not
+          -- equal under that substitution then this pair of equations violates
+          -- injectivity annotation, but for closed type families it still might
+          -- be the case that one LHS after substitution is unreachable.
+          | eqTypes lhs1Subst lhs2Subst  -- check case 1B1 from Note.
+          -> InjectivityAccepted
+          | otherwise
+          -> InjectivityUnified ( ax1 { cab_lhs = Type.substTys subst lhs1
+                                      , cab_rhs = Type.substTy  subst rhs1 })
+                                ( ax2 { cab_lhs = Type.substTys subst lhs2
+                                      , cab_rhs = Type.substTy  subst rhs2 })
+                  -- Payload of InjectivityUnified used only for check 1B2, only
+                  -- for closed type families
+        where
+          lhs1Subst = Type.substTys subst (getInjArgs lhs1)
+          lhs2Subst = Type.substTys subst (getInjArgs lhs2)
 
 -- takes a CoAxiom with unknown branch incompatibilities and computes
 -- the compatibilities
diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index c6f4fdf42f..cf87106e45 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -2600,7 +2600,7 @@ compatible_branches (CoAxBranch { cab_tvs = tvs1
                                 , cab_rhs = rhs2 })
   = -- we need to freshen ax2 w.r.t. ax1
     -- do this by pretending tvs1 are in scope when processing tvs2
-    let in_scope       = mkInScopeSet (mkVarSet tvs1)
+    let in_scope       = mkInScopeSetList tvs1
         subst0         = mkEmptyTCvSubst in_scope
         (subst, _)     = substTyVarBndrs subst0 tvs2
         lhs2'          = substTys subst lhs2
@@ -2858,7 +2858,7 @@ initL cfg m
   where
     (tcvs, ids) = partition isTyCoVar $ l_vars cfg
     env = LE { le_flags = l_flags cfg
-             , le_subst = mkEmptyTCvSubst (mkInScopeSet (mkVarSet tcvs))
+             , le_subst = mkEmptyTCvSubst (mkInScopeSetList tcvs)
              , le_ids   = mkVarEnv [(id, (id,idType id)) | id <- ids]
              , le_joins = emptyVarSet
              , le_loc = []
diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 74a903fbc8..961d5f4b5b 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -606,7 +606,7 @@ specProgram guts@(ModGuts { mg_module = this_mod
               -- accidentally re-use a unique that's already in use
               -- Easiest thing is to do it all at once, as if all the top-level
               -- decls were mutually recursive
-       ; let top_env = SE { se_subst = Core.mkEmptySubst $ mkInScopeSet $ mkVarSet $
+       ; let top_env = SE { se_subst = Core.mkEmptySubst $ mkInScopeSetList $
                                        bindersOfBinds binds
                           , se_module = this_mod
                           , se_dflags = dflags }
diff --git a/compiler/GHC/Core/TyCo/Subst.hs b/compiler/GHC/Core/TyCo/Subst.hs
index 7322ff5e03..069270a1a5 100644
--- a/compiler/GHC/Core/TyCo/Subst.hs
+++ b/compiler/GHC/Core/TyCo/Subst.hs
@@ -350,13 +350,13 @@ extendTvSubstBinderAndInScope subst (Anon {}) _
   = subst
 
 extendTvSubstWithClone :: TCvSubst -> TyVar -> TyVar -> TCvSubst
--- Adds a new tv -> tv mapping, /and/ extends the in-scope set
+-- Adds a new tv -> tv mapping, /and/ extends the in-scope set with the clone
+-- Does not look in the kind of the new variable;
+--   those variables should be in scope already
 extendTvSubstWithClone (TCvSubst in_scope tenv cenv) tv tv'
-  = TCvSubst (extendInScopeSetSet in_scope new_in_scope)
+  = TCvSubst (extendInScopeSet in_scope tv')
              (extendVarEnv tenv tv (mkTyVarTy tv'))
              cenv
-  where
-    new_in_scope = tyCoVarsOfType (tyVarKind tv') `extendVarSet` tv'
 
 extendCvSubst :: TCvSubst -> CoVar -> Coercion -> TCvSubst
 extendCvSubst (TCvSubst in_scope tenv cenv) v co
diff --git a/compiler/GHC/Core/Unify.hs b/compiler/GHC/Core/Unify.hs
index 27f745d980..0c3e28f0e1 100644
--- a/compiler/GHC/Core/Unify.hs
+++ b/compiler/GHC/Core/Unify.hs
@@ -41,7 +41,7 @@ import GHC.Core.TyCo.FVs   ( tyCoVarsOfCoList, tyCoFVsOfTypes )
 import GHC.Core.TyCo.Subst ( mkTvSubst )
 import GHC.Core.RoughMap
 import GHC.Core.Map.Type
-import GHC.Utils.FV( FV, fvVarSet, fvVarList )
+import GHC.Utils.FV( FV, fvVarList )
 import GHC.Utils.Misc
 import GHC.Data.Pair
 import GHC.Utils.Outputable
@@ -474,25 +474,26 @@ tcUnifyTyKi t1 t2 = tcUnifyTyKis alwaysBindFun [t1] [t2]
 tcUnifyTyWithTFs :: Bool  -- ^ True <=> do two-way unification;
                           --   False <=> do one-way matching.
                           --   See end of sec 5.2 from the paper
-                 -> Type -> Type -> Maybe TCvSubst
+                 -> InScopeSet     -- Should include the free tyvars of both Type args
+                 -> Type -> Type   -- Types to unify
+                 -> Maybe TCvSubst
 -- This algorithm is an implementation of the "Algorithm U" presented in
 -- the paper "Injective type families for Haskell", Figures 2 and 3.
 -- The code is incorporated with the standard unifier for convenience, but
 -- its operation should match the specification in the paper.
-tcUnifyTyWithTFs twoWay t1 t2
+tcUnifyTyWithTFs twoWay in_scope t1 t2
   = case tc_unify_tys alwaysBindFun twoWay True False
                        rn_env emptyTvSubstEnv emptyCvSubstEnv
                        [t1] [t2] of
-      Unifiable          (subst, _) -> Just $ maybe_fix subst
-      MaybeApart _reason (subst, _) -> Just $ maybe_fix subst
+      Unifiable          (tv_subst, _cv_subst) -> Just $ maybe_fix tv_subst
+      MaybeApart _reason (tv_subst, _cv_subst) -> Just $ maybe_fix tv_subst
       -- we want to *succeed* in questionable cases. This is a
       -- pre-unification algorithm.
       SurelyApart      -> Nothing
   where
-    in_scope = mkInScopeSet $ tyCoVarsOfTypes [t1, t2]
     rn_env   = mkRnEnv2 in_scope
 
-    maybe_fix | twoWay    = niFixTCvSubst
+    maybe_fix | twoWay    = niFixTCvSubst in_scope
               | otherwise = mkTvSubst in_scope -- when matching, don't confuse
                                                -- domain with range
 
@@ -587,13 +588,13 @@ tc_unify_tys_fg :: Bool
                 -> [Type] -> [Type]
                 -> UnifyResult
 tc_unify_tys_fg match_kis bind_fn tys1 tys2
-  = do { (env, _) <- tc_unify_tys bind_fn True False match_kis env
+  = do { (env, _) <- tc_unify_tys bind_fn True False match_kis rn_env
                                   emptyTvSubstEnv emptyCvSubstEnv
                                   tys1 tys2
-       ; return $ niFixTCvSubst env }
+       ; return $ niFixTCvSubst in_scope env }
   where
-    vars = tyCoVarsOfTypes tys1 `unionVarSet` tyCoVarsOfTypes tys2
-    env  = mkRnEnv2 $ mkInScopeSet vars
+    in_scope = mkInScopeSet $ tyCoVarsOfTypes tys1 `unionVarSet` tyCoVarsOfTypes tys2
+    rn_env   = mkRnEnv2 in_scope
 
 -- | This function is actually the one to call the unifier -- a little
 -- too general for outside clients, though.
@@ -726,13 +727,13 @@ variables in the in-scope set; it is used only to ensure no
 shadowing.
 -}
 
-niFixTCvSubst :: TvSubstEnv -> TCvSubst
+niFixTCvSubst :: InScopeSet -> TvSubstEnv -> TCvSubst
 -- Find the idempotent fixed point of the non-idempotent substitution
 -- This is surprisingly tricky:
 --   see Note [Finding the substitution fixpoint]
 -- ToDo: use laziness instead of iteration?
-niFixTCvSubst tenv
-  | not_fixpoint = niFixTCvSubst (mapVarEnv (substTy subst) tenv)
+niFixTCvSubst in_scope tenv
+  | not_fixpoint = niFixTCvSubst in_scope (mapVarEnv (substTy subst) tenv)
   | otherwise    = subst
   where
     range_fvs :: FV
@@ -749,9 +750,8 @@ niFixTCvSubst tenv
     free_tvs = scopedSort (filterOut in_domain range_tvs)
 
     -- See Note [Finding the substitution fixpoint], Step 6
-    init_in_scope = mkInScopeSet (fvVarSet range_fvs)
     subst = foldl' add_free_tv
-                  (mkTvSubst init_in_scope tenv)
+                  (mkTvSubst in_scope tenv)
                   free_tvs
 
     add_free_tv :: TCvSubst -> TyVar -> TCvSubst
diff --git a/compiler/GHC/Tc/Deriv.hs b/compiler/GHC/Tc/Deriv.hs
index 272d7b6d4e..f191f74d46 100644
--- a/compiler/GHC/Tc/Deriv.hs
+++ b/compiler/GHC/Tc/Deriv.hs
@@ -1920,7 +1920,7 @@ genFamInsts spec@(DS { ds_tvs = tyvars, ds_mechanism = mechanism
       -- See Note [DeriveAnyClass and default family instances]
       DerivSpecAnyClass -> do
         let mini_env   = mkVarEnv (classTyVars clas `zip` inst_tys)
-            mini_subst = mkTvSubst (mkInScopeSet (mkVarSet tyvars)) mini_env
+            mini_subst = mkTvSubst (mkInScopeSetList tyvars) mini_env
         dflags <- getDynFlags
         tyfam_insts <-
           -- canDeriveAnyClass should ensure that this code can't be reached
diff --git a/compiler/GHC/Tc/Deriv/Generate.hs b/compiler/GHC/Tc/Deriv/Generate.hs
index a8536971bd..810c4c7a32 100644
--- a/compiler/GHC/Tc/Deriv/Generate.hs
+++ b/compiler/GHC/Tc/Deriv/Generate.hs
@@ -73,7 +73,6 @@ import GHC.Builtin.Types
 import GHC.Core.Type
 import GHC.Core.Class
 import GHC.Types.Unique.FM ( lookupUFM, listToUFM )
-import GHC.Types.Var.Set
 import GHC.Types.Var.Env
 import GHC.Utils.Misc
 import GHC.Types.Var
@@ -2079,7 +2078,7 @@ gen_Newtype_fam_insts loc' cls inst_tvs inst_tys rhs_ty
     ats       = classATs cls
     locn      = noAnnSrcSpan loc'
     cls_tvs   = classTyVars cls
-    in_scope  = mkInScopeSet $ mkVarSet inst_tvs
+    in_scope  = mkInScopeSetList inst_tvs
     lhs_env   = zipTyEnv cls_tvs inst_tys
     lhs_subst = mkTvSubst in_scope lhs_env
     rhs_env   = zipTyEnv cls_tvs underlying_inst_tys
@@ -2129,7 +2128,7 @@ mkCoerceClassMethEqn cls inst_tvs inst_tys rhs_ty id
          (substTy lhs_subst user_meth_ty)
   where
     cls_tvs = classTyVars cls
-    in_scope = mkInScopeSet $ mkVarSet inst_tvs
+    in_scope = mkInScopeSetList inst_tvs
     lhs_subst = mkTvSubst in_scope (zipTyEnv cls_tvs inst_tys)
     rhs_subst = mkTvSubst in_scope (zipTyEnv cls_tvs (changeLast inst_tys rhs_ty))
     (_class_tvs, _class_constraint, user_meth_ty)
diff --git a/compiler/GHC/Tc/Gen/HsType.hs b/compiler/GHC/Tc/Gen/HsType.hs
index 8d666fb7b0..d045984024 100644
--- a/compiler/GHC/Tc/Gen/HsType.hs
+++ b/compiler/GHC/Tc/Gen/HsType.hs
@@ -1058,7 +1058,7 @@ tc_infer_hs_type _ (XHsType ty)
            subst_prs = [ (getUnique nm, tv)
                        | ATyVar nm tv <- nonDetNameEnvElts (tcl_env env) ]
            subst = mkTvSubst
-                     (mkInScopeSet $ mkVarSet $ map snd subst_prs)
+                     (mkInScopeSetList $ map snd subst_prs)
                      (listToUFM_Directly $ map (liftSnd mkTyVarTy) subst_prs)
            ty' = substTy subst ty
        return (ty', tcTypeKind ty')
@@ -3663,7 +3663,7 @@ etaExpandAlgTyCon flav skol_info tcbs res_kind
   = return ([], res_kind)
   where
     tyvars     = binderVars tcbs
-    in_scope   = mkInScopeSet (mkVarSet tyvars)
+    in_scope   = mkInScopeSetList tyvars
     avoid_occs = map getOccName tyvars
 
 needsEtaExpansion :: TyConFlavour -> Bool
diff --git a/compiler/GHC/Tc/Instance/FunDeps.hs b/compiler/GHC/Tc/Instance/FunDeps.hs
index 5b215490af..067c87a50f 100644
--- a/compiler/GHC/Tc/Instance/FunDeps.hs
+++ b/compiler/GHC/Tc/Instance/FunDeps.hs
@@ -57,7 +57,6 @@ import Data.Foldable    ( fold )
 *                                                                      *
 ************************************************************************
 
-
 Each functional dependency with one variable in the RHS is responsible
 for generating a single equality. For instance:
      class C a b | a -> b
@@ -83,31 +82,56 @@ Will generate:
 INVARIANT: Corresponding types aren't already equal
 That is, there exists at least one non-identity equality in FDEqs.
 
+Note [Improving against instances]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Assume:
-       class C a b c | a -> b c
-       instance C Int x x
-And:   [Wanted] C Int Bool alpha
-We will /match/ the LHS of fundep equations, producing a matching substitution
-and create equations for the RHS sides. In our last example we'd have generated:
-      ({x}, [fd1,fd2])
-where
-       fd1 = FDEq 1 Bool x
-       fd2 = FDEq 2 alpha x
-To ``execute'' the equation, make fresh type variable for each tyvar in the set,
-instantiate the two types with these fresh variables, and then unify or generate
-a new constraint. In the above example we would generate a new unification
-variable 'beta' for x and produce the following constraints:
-     [Wanted] (Bool ~ beta)
-     [Wanted] (alpha ~ beta)
-
-Notice the subtle difference between the above class declaration and:
-       class C a b c | a -> b, a -> c
-where we would generate:
-      ({x},[fd1]),({x},[fd2])
-This means that the template variable would be instantiated to different
-unification variables when producing the FD constraints.
-
-Finally, the position parameters will help us rewrite the wanted constraint ``on the spot''
+   class C a b | a -> b
+   instance C Int Bool
+   [W] C Int ty
+
+Then `improveFromInstEnv` should return a FDEqn with
+   FDEqn { fd_qtvs = [], fd_eqs = [Pair Bool ty] }
+
+describing an equality (Int ~ ty).  To do this we /match/ the instance head
+against the [W], using just the LHS of the fundep; if we match, we return
+an equality for the RHS.
+
+Wrinkles:
+
+(1) meta_tvs: sometimes the instance mentions variables in the RHS that
+    are not bound in the LHS.  For example
+
+     class C a b | a -> b
+     instance C Int (Maybe x)
+     [W] C Int ty
+
+    Note that although the `Int` parts match, that does not fix what `x` is.
+    So we just make up a fresh unification variable (a meta_tv), to give the
+    "shape" of the RHS.  So we emit the FDEqun
+       FDEqn { fd_qtvs = [x], fd_eqs = [Pair (Maybe x) ty] }
+
+    Note that the fd_qtvs can be free in the /first/ component of the Pair,
+
+    but not in the seconde (which comes from the [W] constraint.
+
+(2) Multi-range fundeps. When these meta_tvs are involved, there is a subtle
+    difference between the fundep (a -> b c) and the two fundeps (a->b, a->c).
+    Consider
+       class D a b c | a -> b c
+       instance D Int x (Maybe x)
+       [W] D Int Bool ty
+
+    Then we'll generate
+       FDEqn { fd_qtvs = [x], fd_eqs = [Pair x Bool, Pair (Maybe x) ty] }
+
+    But if the fundeps had been (a->b, a->c) we'd generate two FDEqns
+       FDEqn { fd_qtvs = [x], fd_eqs = [Pair x Bool] }
+       FDEqn { fd_qtvs = [x], fd_eqs = [Pair (Maybe x) ty] }
+    with two FDEqns, generating two separate unification variables.
+
+(3) improveFromInstEnv doesn't return any equations that already hold.
+    Reason: then we know if any actual improvement has happened, in
+    which case we need to iterate the solver
 -}
 
 data FunDepEqn loc
@@ -116,11 +140,24 @@ data FunDepEqn loc
                                  -- Non-empty only for FunDepEqns arising from instance decls
 
           , fd_eqs   :: [TypeEqn]  -- Make these pairs of types equal
+                                   -- Invariant: In each (Pair ty1 ty2), the fd_qtvs may be
+                                   -- free in ty1 but not in ty2.  See Wrinkle (1) of
+                                   -- Note [Improving against instances]
+
           , fd_pred1 :: PredType   -- The FunDepEqn arose from
           , fd_pred2 :: PredType   --  combining these two constraints
           , fd_loc   :: loc  }
     deriving Functor
 
+instance Outputable (FunDepEqn a) where
+  ppr = pprEquation
+
+pprEquation :: FunDepEqn a -> SDoc
+pprEquation (FDEqn { fd_qtvs = qtvs, fd_eqs = pairs })
+  = vcat [text "forall" <+> braces (pprWithCommas ppr qtvs),
+          nest 2 (vcat [ ppr t1 <+> text "~" <+> ppr t2
+                       | Pair t1 t2 <- pairs])]
+
 {-
 Given a bunch of predicates that must hold, such as
 
@@ -195,20 +232,12 @@ improveFromAnother _ _ _ = []
 -- Improve a class constraint from instance declarations
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-instance Outputable (FunDepEqn a) where
-  ppr = pprEquation
-
-pprEquation :: FunDepEqn a -> SDoc
-pprEquation (FDEqn { fd_qtvs = qtvs, fd_eqs = pairs })
-  = vcat [text "forall" <+> braces (pprWithCommas ppr qtvs),
-          nest 2 (vcat [ ppr t1 <+> text "~" <+> ppr t2
-                       | Pair t1 t2 <- pairs])]
-
 improveFromInstEnv :: InstEnvs
                    -> (PredType -> SrcSpan -> loc)
                    -> Class -> [Type]
                    -> [FunDepEqn loc] -- Needs to be a FunDepEqn because
                                       -- of quantified variables
+-- See Note [Improving against instances]
 -- Post: Equations oriented from the template (matching instance) to the workitem!
 improveFromInstEnv inst_env mk_loc cls tys
   = [ FDEqn { fd_qtvs = meta_tvs, fd_eqs = eqs
@@ -232,37 +261,16 @@ improveFromInstEnv inst_env mk_loc cls tys
     rough_tcs          = RM_KnownTc (className cls) : roughMatchTcs tys
     pred               = mkClassPred cls tys
 
-
-
-
 improveClsFD :: [TyVar] -> FunDep TyVar    -- One functional dependency from the class
              -> ClsInst                    -- An instance template
              -> [Type] -> [RoughMatchTc]   -- Arguments of this (C tys) predicate
              -> [([TyCoVar], [TypeEqn])]   -- Empty or singleton
+-- See Note [Improving against instances]
 
 improveClsFD clas_tvs fd
              (ClsInst { is_tvs = qtvs, is_tys = tys_inst, is_tcs = rough_tcs_inst })
              tys_actual rough_tcs_actual
 
--- Compare instance      {a,b}    C sx sp sy sq
---         with wanted     [W] C tx tp ty tq
---         for fundep (x,y -> p,q)  from class  (C x p y q)
--- If (sx,sy) unifies with (tx,ty), take the subst S
-
--- 'qtvs' are the quantified type variables, the ones which can be instantiated
--- to make the types match.  For example, given
---      class C a b | a->b where ...
---      instance C (Maybe x) (Tree x) where ..
---
--- and a wanted constraint of form (C (Maybe t1) t2),
--- then we will call checkClsFD with
---
---      is_qtvs = {x}, is_tys = [Maybe x,  Tree x]
---                     tys_actual = [Maybe t1, t2]
---
--- We can instantiate x to t1, and then we want to force
---      (Tree x) [t1/x]  ~   t2
-
   | instanceCantMatch rough_tcs_inst rough_tcs_actual
   = []          -- Filter out ones that can't possibly match,
 
@@ -271,12 +279,11 @@ improveClsFD clas_tvs fd
                equalLength tys_inst clas_tvs)
               (ppr tys_inst <+> ppr tys_actual) $
 
-    case tcMatchTyKis ltys1 ltys2 of
+    case tcMatchTyKisX init_subst ltys1 ltys2 of
         Nothing  -> []
         Just subst | isJust (tcMatchTyKisX subst rtys1 rtys2)
                         -- Don't include any equations that already hold.
-                        -- Reason: then we know if any actual improvement has happened,
-                        --         in which case we need to iterate the solver
+                        -- See Note [Improving against instances] wrinkle (3)
                         -- In making this check we must taking account of the fact that any
                         -- qtvs that aren't already instantiated can be instantiated to anything
                         -- at all
@@ -307,7 +314,7 @@ improveClsFD clas_tvs fd
                         -- executed.  What we're doing instead is recording the partial
                         -- work of the ls1/ls2 unification leaving a smaller unification problem
                   where
-                    rtys1' = map (substTyUnchecked subst) rtys1
+                    rtys1' = map (substTy subst) rtys1
 
                     fdeqs = zipAndComputeFDEqs (\_ _ -> False) rtys1' rtys2
                         -- Don't discard anything!
@@ -315,8 +322,10 @@ improveClsFD clas_tvs fd
                         -- eqType again, since we know for sure that /at least one/
                         -- equation in there is useful)
 
-                    meta_tvs = [ setVarType tv (substTyUnchecked subst (varType tv))
-                               | tv <- qtvs, tv `notElemTCvSubst` subst ]
+                    meta_tvs = [ setVarType tv (substTy subst (varType tv))
+                               | tv <- qtvs
+                               , tv `notElemTCvSubst` subst
+                               , tv `elemVarSet` rtys1_tvs ]
                         -- meta_tvs are the quantified type variables
                         -- that have not been substituted out
                         --
@@ -331,11 +340,14 @@ improveClsFD clas_tvs fd
                         --              type variables' kinds
                         --          (b) we must apply 'subst' to the kinds, in case we have
                         --              matched out a kind variable, but not a type variable
-                        --              whose kind mentions that kind variable!
-                        --          #6015, #6068
+                        --              whose kind mentions that kind variable! #6015, #6068
+                        --          (c) no need to include tyvars not in rtys1
   where
+    init_subst     = mkEmptyTCvSubst $ mkInScopeSet $
+                     mkVarSet qtvs `unionVarSet` tyCoVarsOfTypes ltys2
     (ltys1, rtys1) = instFD fd clas_tvs tys_inst
     (ltys2, rtys2) = instFD fd clas_tvs tys_actual
+    rtys1_tvs      = tyCoVarsOfTypes rtys1
 
 {-
 %************************************************************************
diff --git a/compiler/GHC/Tc/Solver/Interact.hs b/compiler/GHC/Tc/Solver/Interact.hs
index 7323585dc2..e60e6993cc 100644
--- a/compiler/GHC/Tc/Solver/Interact.hs
+++ b/compiler/GHC/Tc/Solver/Interact.hs
@@ -1845,7 +1845,7 @@ doTopFundepImprovement work_item@(CDictCan { cc_ev = ev, cc_class = cls
 doTopFundepImprovement work_item = pprPanic "doTopFundepImprovement" (ppr work_item)
 
 emitFunDepWanteds :: RewriterSet  -- from the work item
-                   -> [FunDepEqn (CtLoc, RewriterSet)] -> TcS ()
+                  -> [FunDepEqn (CtLoc, RewriterSet)] -> TcS ()
 -- See Note [FunDep and implicit parameter reactions]
 emitFunDepWanteds work_rewriters fd_eqns
   = mapM_ do_one_FDEqn fd_eqns
@@ -1859,15 +1859,24 @@ emitFunDepWanteds work_rewriters fd_eqns
 
      | otherwise
      = do { traceTcS "emitFunDepWanteds 2" (ppr (ctl_depth loc) $$ ppr tvs $$ ppr eqs)
-          ; subst <- instFlexi tvs  -- Takes account of kind substitution
+          ; subst <- instFlexiX emptyTCvSubst tvs  -- Takes account of kind substitution
           ; mapM_ (do_one_eq loc all_rewriters subst) (reverse eqs) }
                -- See Note [Reverse order of fundep equations]
      where
        all_rewriters = work_rewriters S.<> rewriters
 
     do_one_eq loc rewriters subst (Pair ty1 ty2)
-       = unifyWanted rewriters loc Nominal
-                     (Type.substTyUnchecked subst ty1) (Type.substTyUnchecked subst ty2)
+       = unifyWanted rewriters loc Nominal (Type.substTy subst' ty1) ty2
+         -- ty2 does not mention fd_qtvs, so no need to subst it.
+         -- See GHC.Tc.Instance.Fundeps Note [Improving against instances]
+         --     Wrinkle (1)
+      where
+         subst' = extendTCvInScopeSet subst (tyCoVarsOfType ty1)
+         -- The free vars of ty1 aren't just fd_qtvs: ty1 is the result
+         -- of matching with the [W] constraint. So we add its free
+         -- vars to InScopeSet, to satisfy substTy's invariants, even
+         -- though ty1 will never (currently) be a poytype, so this
+         -- InScopeSet will never be looked at.
 
 {-
 **********************************************************************
@@ -2065,6 +2074,8 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
   = return []
 
   where
+      in_scope = mkInScopeSet (tyCoVarsOfType rhs_ty)
+
       buildImprovementData
           :: [a]                     -- axioms for a TF (FamInst or CoAxBranch)
           -> (a -> [TyVar])          -- get bound tyvars of an axiom
@@ -2083,7 +2094,8 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
           , let ax_args = axiomLHS axiom
                 ax_rhs  = axiomRHS axiom
                 ax_tvs  = axiomTVs axiom
-          , Just subst <- [tcUnifyTyWithTFs False ax_rhs rhs_ty]
+                in_scope1 = in_scope `extendInScopeSetList` ax_tvs
+          , Just subst <- [tcUnifyTyWithTFs False in_scope1 ax_rhs rhs_ty]
           , let notInSubst tv = not (tv `elemVarEnv` getTvSubstEnv subst)
                 unsubstTvs    = filter (notInSubst <&&> isTyVar) ax_tvs ]
                    -- The order of unsubstTvs is important; it must be
@@ -2099,7 +2111,7 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
                   -- be sure to apply the current substitution to a's kind.
                   -- Hence instFlexiX.   #13135 was an example.
 
-             ; return [ Pair (substTyUnchecked subst ax_arg) arg
+             ; return [ Pair (substTy subst ax_arg) arg
                         -- NB: the ax_arg part is on the left
                         -- see Note [Improvement orientation]
                       | case cabr of
diff --git a/compiler/GHC/Tc/Solver/Monad.hs b/compiler/GHC/Tc/Solver/Monad.hs
index 99c35f826d..6621f54317 100644
--- a/compiler/GHC/Tc/Solver/Monad.hs
+++ b/compiler/GHC/Tc/Solver/Monad.hs
@@ -89,7 +89,7 @@ module GHC.Tc.Solver.Monad (
     instDFunType,                              -- Instantiation
 
     -- MetaTyVars
-    newFlexiTcSTy, instFlexi, instFlexiX,
+    newFlexiTcSTy, instFlexiX,
     cloneMetaTyVar,
     tcInstSkolTyVarsX,
 
@@ -1614,22 +1614,21 @@ newFlexiTcSTy knd = wrapTcS (TcM.newFlexiTyVarTy knd)
 cloneMetaTyVar :: TcTyVar -> TcS TcTyVar
 cloneMetaTyVar tv = wrapTcS (TcM.cloneMetaTyVar tv)
 
-instFlexi :: [TKVar] -> TcS TCvSubst
-instFlexi = instFlexiX emptyTCvSubst
-
 instFlexiX :: TCvSubst -> [TKVar] -> TcS TCvSubst
 instFlexiX subst tvs
   = wrapTcS (foldlM instFlexiHelper subst tvs)
 
 instFlexiHelper :: TCvSubst -> TKVar -> TcM TCvSubst
+-- Makes fresh tyvar, extends the substitution, and the in-scope set
 instFlexiHelper subst tv
   = do { uniq <- TcM.newUnique
        ; details <- TcM.newMetaDetails TauTv
-       ; let name = setNameUnique (tyVarName tv) uniq
-             kind = substTyUnchecked subst (tyVarKind tv)
-             ty'  = mkTyVarTy (mkTcTyVar name kind details)
-       ; TcM.traceTc "instFlexi" (ppr ty')
-       ; return (extendTvSubst subst tv ty') }
+       ; let name   = setNameUnique (tyVarName tv) uniq
+             kind   = substTyUnchecked subst (tyVarKind tv)
+             tv'    = mkTcTyVar name kind details
+             subst' = extendTvSubstWithClone subst tv tv'
+       ; TcM.traceTc "instFlexi" (ppr tv')
+       ; return (extendTvSubst subst' tv (mkTyVarTy tv')) }
 
 matchGlobalInst :: DynFlags
                 -> Bool      -- True <=> caller is the short-cut solver
diff --git a/compiler/GHC/Tc/TyCl/PatSyn.hs b/compiler/GHC/Tc/TyCl/PatSyn.hs
index 7fd1f3677f..dad820674e 100644
--- a/compiler/GHC/Tc/TyCl/PatSyn.hs
+++ b/compiler/GHC/Tc/TyCl/PatSyn.hs
@@ -42,7 +42,7 @@ import GHC.Utils.Panic
 import GHC.Utils.Outputable
 import GHC.Data.FastString
 import GHC.Types.Var
-import GHC.Types.Var.Env( emptyTidyEnv, mkInScopeSet )
+import GHC.Types.Var.Env( emptyTidyEnv, mkInScopeSetList )
 import GHC.Types.Id
 import GHC.Types.Id.Info( RecSelParent(..) )
 import GHC.Tc.Gen.Bind
@@ -456,7 +456,7 @@ tcCheckPatSynDecl psb@PSB{ psb_id = lname@(L _ name), psb_args = details
            pushLevelAndCaptureConstraints   $
            tcExtendNameTyVarEnv univ_tv_prs $
            tcCheckPat PatSyn lpat (unrestricted skol_pat_ty)   $
-           do { let in_scope    = mkInScopeSet (mkVarSet skol_univ_tvs)
+           do { let in_scope    = mkInScopeSetList skol_univ_tvs
                     empty_subst = mkEmptyTCvSubst in_scope
               ; (inst_subst, ex_tvs') <- mapAccumLM newMetaTyVarX empty_subst skol_ex_tvs
                     -- newMetaTyVarX: see the "Existential type variables"
diff --git a/compiler/GHC/Types/Var/Env.hs b/compiler/GHC/Types/Var/Env.hs
index 02d3fa2ad5..88337b9e9d 100644
--- a/compiler/GHC/Types/Var/Env.hs
+++ b/compiler/GHC/Types/Var/Env.hs
@@ -50,7 +50,7 @@ module GHC.Types.Var.Env (
         InScopeSet,
 
         -- ** Operations on InScopeSets
-        emptyInScopeSet, mkInScopeSet, delInScopeSet,
+        emptyInScopeSet, mkInScopeSet, mkInScopeSetList, delInScopeSet,
         extendInScopeSet, extendInScopeSetList, extendInScopeSetSet,
         getInScopeVars, lookupInScope, lookupInScope_Directly,
         unionInScope, elemInScopeSet, uniqAway,
@@ -148,6 +148,9 @@ getInScopeVars (InScope vs) = vs
 mkInScopeSet :: VarSet -> InScopeSet
 mkInScopeSet in_scope = InScope in_scope
 
+mkInScopeSetList :: [Var] -> InScopeSet
+mkInScopeSetList vs = InScope (mkVarSet vs)
+
 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
 extendInScopeSet (InScope in_scope) v
    = InScope (extendVarSet in_scope v)
diff --git a/testsuite/tests/indexed-types/should_fail/T21896.hs b/testsuite/tests/indexed-types/should_fail/T21896.hs
new file mode 100644
index 0000000000..64ce18d023
--- /dev/null
+++ b/testsuite/tests/indexed-types/should_fail/T21896.hs
@@ -0,0 +1,9 @@
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE DataKinds, TypeFamilies, TypeFamilyDependencies, PolyKinds #-}
+module Bug where
+
+data T = Foo | Bar
+
+type family F (ns :: T) (ret :: k) = (r :: k) | r -> ret where
+    F Foo r = r
+    F Bar r = r
diff --git a/testsuite/tests/indexed-types/should_fail/T21896.stderr b/testsuite/tests/indexed-types/should_fail/T21896.stderr
new file mode 100644
index 0000000000..971c220171
--- /dev/null
+++ b/testsuite/tests/indexed-types/should_fail/T21896.stderr
@@ -0,0 +1,24 @@
+
+T21896.hs:8:5: error:
+    • Type family equation right-hand sides overlap; this violates
+      the family's injectivity annotation:
+        forall {k} {r :: k}. F 'Foo r = r -- Defined at T21896.hs:8:5
+        forall {k} {r :: k}. F 'Bar r = r -- Defined at T21896.hs:9:5
+    • In the equations for closed type family ‘F’
+      In the type family declaration for ‘F’
+
+T21896.hs:8:5: error:
+    • Type family equation violates the family's injectivity annotation.
+      RHS of injective type family equation is a bare type variable
+      but these LHS type and kind patterns are not bare variables: ‘'Foo’
+        forall {k} {r :: k}. F 'Foo r = r -- Defined at T21896.hs:8:5
+    • In the equations for closed type family ‘F’
+      In the type family declaration for ‘F’
+
+T21896.hs:9:5: error:
+    • Type family equation violates the family's injectivity annotation.
+      RHS of injective type family equation is a bare type variable
+      but these LHS type and kind patterns are not bare variables: ‘'Bar’
+        forall {k} {r :: k}. F 'Bar r = r -- Defined at T21896.hs:9:5
+    • In the equations for closed type family ‘F’
+      In the type family declaration for ‘F’
diff --git a/testsuite/tests/indexed-types/should_fail/all.T b/testsuite/tests/indexed-types/should_fail/all.T
index 5338a8780f..5a8d0cd88d 100644
--- a/testsuite/tests/indexed-types/should_fail/all.T
+++ b/testsuite/tests/indexed-types/should_fail/all.T
@@ -168,3 +168,4 @@ test('T20466', normal, compile_fail, [''])
 test('ExpandTFs', normal, compile_fail, [''])
 test('T20465', normal, compile_fail, [''])
 test('T20521', normal, compile_fail, [''])
+test('T21896', normal, compile_fail, [''])