Apply the right substitution in ty-fam improvement

Trac #13135 showed that we were failing to apply the
correct substitution to the un-substituted tyvars during
type-family improvement using injectivity.  Specifically
in TcInteractlinjImproveEqns we need to use instFlexiX.

An outright bug, easy to fix.

Slight refactoring along the way.  The quantified tyars of the axiom are
readily to hand; we don't need to take the free tyvars of the LHS

5 files changed, 171 insertions, 36 deletions

diff --git a/compiler/typecheck/TcInteract.hs b/compiler/typecheck/TcInteract.hs
index 0ff7a976ab..2b66126fa4 100644
--- a/compiler/typecheck/TcInteract.hs
+++ b/compiler/typecheck/TcInteract.hs
@@ -1367,7 +1367,7 @@ emitFunDepDeriveds fd_eqns
           ; mapM_ (unifyDerived loc Nominal) eqs }
      | otherwise
      = do { traceTcS "emitFunDepDeriveds 2" (ppr (ctl_depth loc) $$ ppr eqs)
-          ; (subst, _) <- instFlexiTcS tvs  -- Takes account of kind substitution
+          ; subst <- instFlexi tvs  -- Takes account of kind substitution
           ; mapM_ (do_one_eq loc subst) eqs }
 
     do_one_eq loc subst (Pair ty1 ty2)
@@ -1500,13 +1500,13 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
     -- family but we only want to derive equalities from one such equation.
     concatMapM (injImproveEqns injective_args) (take 1 $
       buildImprovementData (lookupFamInstEnvByTyCon fam_envs fam_tc)
-                           fi_tys fi_rhs (const Nothing))
+                           fi_tvs fi_tys fi_rhs (const Nothing))
 
   | Just ax <- isClosedSynFamilyTyConWithAxiom_maybe fam_tc
   , Injective injective_args <- familyTyConInjectivityInfo fam_tc
   = concatMapM (injImproveEqns injective_args) $
       buildImprovementData (fromBranches (co_ax_branches ax))
-                           cab_lhs cab_rhs Just
+                           cab_tvs cab_lhs cab_rhs Just
 
   | otherwise
   = return []
@@ -1514,6 +1514,7 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
   where
       buildImprovementData
           :: [a]                     -- axioms for a TF (FamInst or CoAxBranch)
+          -> (a -> [TyVar])          -- get bound tyvars of an axiom
           -> (a -> [Type])           -- get LHS of an axiom
           -> (a -> Type)             -- get RHS of an axiom
           -> (a -> Maybe CoAxBranch) -- Just => apartness check required
@@ -1523,34 +1524,35 @@ improve_top_fun_eqs fam_envs fam_tc args rhs_ty
              -- , RHS-unifying substitution
              -- , axiom variables without substitution
              -- , Maybe matching axiom [Nothing - open TF, Just - closed TF ] )
-      buildImprovementData axioms axiomLHS axiomRHS wrap =
+      buildImprovementData axioms axiomTVs axiomLHS axiomRHS wrap =
           [ (ax_args, subst, unsubstTvs, wrap axiom)
           | axiom <- axioms
           , let ax_args = axiomLHS axiom
-          , let ax_rhs  = axiomRHS axiom
+                ax_rhs  = axiomRHS axiom
+                ax_tvs  = axiomTVs axiom
           , Just subst <- [tcUnifyTyWithTFs False ax_rhs rhs_ty]
-          , let tvs           = tyCoVarsOfTypesList ax_args
-                notInSubst tv = not (tv `elemVarEnv` getTvSubstEnv subst)
-                unsubstTvs    = filter (notInSubst <&&> isTyVar) tvs ]
+          , let notInSubst tv = not (tv `elemVarEnv` getTvSubstEnv subst)
+                unsubstTvs    = filter (notInSubst <&&> isTyVar) ax_tvs ]
+                   -- The order of unsubstTvs is important; it must be
+                   -- in telescope order e.g. (k:*) (a:k)
 
       injImproveEqns :: [Bool]
                      -> ([Type], TCvSubst, [TyCoVar], Maybe CoAxBranch)
                      -> TcS [TypeEqn]
-      injImproveEqns inj_args (ax_args, theta, unsubstTvs, cabr) = do
-        (theta', _) <- instFlexiTcS unsubstTvs
-        -- The use of deterministically ordered list for `unsubstTvs`
-        -- is not strictly necessary here, we only use the substitution
-        -- part of the result of instFlexiTcS. If we used the second
-        -- part of the tuple, which is the range of the substitution then
-        -- the order could be important.
-        let subst = theta `unionTCvSubst` theta'
-        return [ Pair (substTyUnchecked subst ax_arg) arg
-                   -- NB: the ax_arg part is on the left
-                   -- see Note [Improvement orientation]
-               | case cabr of
-                  Just cabr' -> apartnessCheck (substTys subst ax_args) cabr'
-                  _          -> True
-               , (ax_arg, arg, True) <- zip3 ax_args args inj_args ]
+      injImproveEqns inj_args (ax_args, subst, unsubstTvs, cabr)
+        = do { subst <- instFlexiX subst unsubstTvs
+                  -- If the current substitution bind [k -> *], and
+                  -- one of the un-substituted tyvars is (a::k), we'd better
+                  -- be sure to apply the current substitution to a's kind.
+                  -- Hence instFlexiX.   Trac #13135 was an example.
+
+             ; return [ Pair (substTyUnchecked subst ax_arg) arg
+                        -- NB: the ax_arg part is on the left
+                        -- see Note [Improvement orientation]
+                      | case cabr of
+                          Just cabr' -> apartnessCheck (substTys subst ax_args) cabr'
+                          _          -> True
+                      , (ax_arg, arg, True) <- zip3 ax_args args inj_args ] }
 
 
 shortCutReduction :: CtEvidence -> TcTyVar -> TcCoercion
diff --git a/compiler/typecheck/TcSMonad.hs b/compiler/typecheck/TcSMonad.hs
index 1dca0c31ce..5f6a4f16d3 100644
--- a/compiler/typecheck/TcSMonad.hs
+++ b/compiler/typecheck/TcSMonad.hs
@@ -90,7 +90,7 @@ module TcSMonad (
     instDFunType,                              -- Instantiation
 
     -- MetaTyVars
-    newFlexiTcSTy, instFlexiTcS,
+    newFlexiTcSTy, instFlexi, instFlexiX,
     cloneMetaTyVar, demoteUnfilledFmv,
 
     TcLevel, isTouchableMetaTyVarTcS,
@@ -2790,21 +2790,21 @@ demoteUnfilledFmv fmv
                    do { tv_ty <- TcM.newFlexiTyVarTy (tyVarKind fmv)
                       ; TcM.writeMetaTyVar fmv tv_ty } }
 
-instFlexiTcS :: [TKVar] -> TcS (TCvSubst, [TcType])
-instFlexiTcS tvs = wrapTcS (mapAccumLM inst_one emptyTCvSubst tvs)
-  where
-     inst_one subst tv
-         = do { ty' <- instFlexiTcSHelper (tyVarName tv)
-                                          (substTyUnchecked subst (tyVarKind tv))
-              ; return (extendTvSubst subst tv ty', ty') }
+instFlexi :: [TKVar] -> TcS TCvSubst
+instFlexi = instFlexiX emptyTCvSubst
+
+instFlexiX :: TCvSubst -> [TKVar] -> TcS TCvSubst
+instFlexiX subst tvs
+  = wrapTcS (foldlM instFlexiHelper subst tvs)
 
-instFlexiTcSHelper :: Name -> Kind -> TcM TcType
-instFlexiTcSHelper tvname kind
+instFlexiHelper :: TCvSubst -> TKVar -> TcM TCvSubst
+instFlexiHelper subst tv
   = do { uniq <- TcM.newUnique
        ; details <- TcM.newMetaDetails TauTv
-       ; let name = setNameUnique tvname uniq
-       ; return (mkTyVarTy (mkTcTyVar name kind details)) }
-
+       ; let name = setNameUnique (tyVarName tv) uniq
+             kind = substTyUnchecked subst (tyVarKind tv)
+             ty'  = mkTyVarTy (mkTcTyVar name kind details)
+       ; return (extendTvSubst subst tv ty') }
 
 
 -- Creating and setting evidence variables and CtFlavors
diff --git a/testsuite/tests/dependent/should_fail/T13135.hs b/testsuite/tests/dependent/should_fail/T13135.hs
new file mode 100644
index 0000000000..c39b3f5842
--- /dev/null
+++ b/testsuite/tests/dependent/should_fail/T13135.hs
@@ -0,0 +1,122 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TypeInType #-}
+{-# LANGUAGE TypeFamilyDependencies #-}
+
+
+module T12135 where
+
+
+import Data.Kind (Type)
+
+
+class sub :<: sup
+  where
+    -- | Injection from @sub@ to @sup@
+    inj ::  sub sig -> sup sig
+
+instance (sub :<: sup) => (sub :<: AST sup)
+  where
+    inj = Sym . inj
+
+
+data Sig a = Full a | (:->) a (Sig a)
+
+infixr :->
+
+type Sym t = (Sig t -> Type)
+
+data AST :: Sym t -> Sym t
+  where
+    Sym  :: sym sig -> AST sym sig
+    (:$) :: AST sym (a :-> sig) -> AST sym (Full a) -> AST sym sig
+
+-- | Generic N-ary syntactic construct
+--
+-- 'Construct' gives a quick way to introduce a syntactic construct by giving its name and semantic
+-- function.
+data Constr sig = Constr String
+
+smartSym :: (sub :<: sup) => sub sig -> SmartFun sup sig
+smartSym = error "urk"
+
+type family   SmartFun (sym :: Sig t -> Type) sig = f | f -> sym sig  where
+-- In full glory:
+-- type family SmartFun {t} (sym :: Sig t -> Type) (sig :: Sig t) :: *
+--            = f | f -> {t} sym sig  where
+  SmartFun sym (Full a)    = AST sym (Full a)
+  SmartFun sym (a :-> sig) = AST sym (Full a) -> SmartFun sym sig
+
+-- | Get the length of an array
+arrLen ::  AST Constr (Full [a]) -> AST Constr (Full Int)
+arrLen = smartSym sym where
+  sym = Constr "arrLen"
+
+
+
+{- The original bug was a familure to subsitute
+   properly during type-function improvement.
+
+--------------------------------------
+doTopReact
+  [WD] hole{a1y1} {0}:: (SmartFun
+                           t_a1x9[tau:2] sup_a1xb[tau:2] sig_a1xc[tau:2] :: *)
+                        GHC.Prim.~#
+                        (s_a1y0[fuv:0] :: *) (CFunEqCan)
+matchFamTcM
+  Matching: SmartFun t_a1x9[tau:2] sup_a1xb[tau:2] sig_a1xc[tau:2]
+  Match failed
+improveTopFunEqs
+  SmartFun [t_a1x9[tau:2], sup_a1xb[tau:2],
+            sig_a1xc[tau:2]] s_a1y0[fuv:0]
+  [* ~ t_a1x9[tau:2], Constr (Sig *) ~ sup_a1xb[tau:2],
+   (':->) * [a_a1wT[sk:2]] sig_a1y3[tau:2] ~ sig_a1xc[tau:2]]
+
+Emitting new derived equality
+  [D] _ {0}:: (* :: *) GHC.Prim.~# (t_a1x9[tau:2] :: *)
+  [D] _ {0}:: (Constr (Sig *) :: (Sig * -> *))
+                GHC.Prim.~#
+              (sup_a1xb[tau:2] :: (Sig t_a1x9[tau:2] -> *))
+  [D] _ {0}:: ((':->) * [a_a1wT[sk:2]] sig_a1y3[tau:2] :: Sig *)
+              GHC.Prim.~#
+              (sig_a1xc[tau:2] :: Sig t_a1x9[tau:2])
+
+but sig_a1y3 :: Sig t   BAD
+--------------------------------------
+
+Instance
+  forall t (sig :: Sig t) (a :: t) (sym :: Sig t -> *).
+    SmartFun t sym ((:->) t a sig) = AST t sym (Full t a) -> SmartFun t sym sig
+
+Wanted:
+  SmartFun t_a1x9[tau:2] sup_a1xb[tau:2] sig_a1xc[tau:2] ~ s_a1y0[fuv:0]
+  s_a1y0[fuv:0]  ~  AST * (Constr (Sig *))
+                          ('Full * [a_a1wT[sk:2]])
+                    -> AST * (Constr (Sig *)) ('Full * Int)
+
+Substitution after matching RHSs
+  [ t -> *
+  , (sym :: Sig t -> *) -> Constr (Sig *)
+  , a -> a_a1wT ]
+
+add sig -> sig0 :: Sig t   -- Or rather Sig *
+
+Apply to LHS
+
+   SmartFun * (Constr (Sig *)) ((:->) * a_a1wT sig0)
+
+
+improveTopFunEqs
+  SmartFun [t_a1x9[tau:2], sup_a1xb[tau:2],
+            sig_a1xc[tau:2]] s_a1y0[fuv:0]
+  [* ~ t_a1x9[tau:2], Constr (Sig *) ~ sup_a1xb[tau:2],
+   (':->) * [a_a1wT[sk:2]] sig_a1y3[tau:2] ~ sig_a1xc[tau:2]]
+-}
diff --git a/testsuite/tests/dependent/should_fail/T13135.stderr b/testsuite/tests/dependent/should_fail/T13135.stderr
new file mode 100644
index 0000000000..88de56943e
--- /dev/null
+++ b/testsuite/tests/dependent/should_fail/T13135.stderr
@@ -0,0 +1,10 @@
+
+T13135.hs:60:10: error:
+    • No instance for (Constr :<: Constr)
+        arising from a use of ‘smartSym’
+    • In the expression: smartSym sym
+      In an equation for ‘arrLen’:
+          arrLen
+            = smartSym sym
+            where
+                sym = Constr "arrLen"
diff --git a/testsuite/tests/dependent/should_fail/all.T b/testsuite/tests/dependent/should_fail/all.T
index b9c82f131b..c648f9ed1d 100644
--- a/testsuite/tests/dependent/should_fail/all.T
+++ b/testsuite/tests/dependent/should_fail/all.T
@@ -16,3 +16,4 @@ test('T11473', normal, compile_fail, [''])
 test('T11471', normal, compile_fail, [''])
 test('T12174', normal, compile_fail, [''])
 test('T12081', normal, compile_fail, [''])
+test('T13135', normal, compile_fail, [''])