Fix desugaring of record updates on data families

This fixes a long-standing bug in the desugaring of record
updates for data families, when the latter involves a GADT. It's
all explained in Note [Update for GADTs] in GHC.HsToCore.Expr.

Building the correct cast is surprisingly tricky, as that Note
explains.

Fixes #18809.  The test case (in indexed-types/should_compile/T18809)
contains several examples that exercise the dark corners.

(cherry picked from commit bfdccac6acce84e15292a454d12f4e0d87ef6f10)

8 files changed, 156 insertions, 36 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 63fbe2812e..487c8a746a 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -41,6 +41,7 @@ module GHC.Core.Coercion (
         downgradeRole, mkAxiomRuleCo,
         mkGReflRightCo, mkGReflLeftCo, mkCoherenceLeftCo, mkCoherenceRightCo,
         mkKindCo, castCoercionKind, castCoercionKindI,
+        mkFamilyTyConAppCo,
 
         mkHeteroCoercionType,
         mkPrimEqPred, mkReprPrimEqPred, mkPrimEqPredRole,
@@ -1505,6 +1506,27 @@ castCoercionKindI g h1 h2
   = mkCoherenceRightCo r t2 h2 (mkCoherenceLeftCo r t1 h1 g)
   where (Pair t1 t2, r) = coercionKindRole g
 
+mkFamilyTyConAppCo :: TyCon -> [CoercionN] -> CoercionN
+-- ^ Given a family instance 'TyCon' and its arg 'Coercion's, return the
+-- corresponding family 'Coercion'.  E.g:
+--
+-- > data family T a
+-- > data instance T (Maybe b) = MkT b
+--
+-- Where the instance 'TyCon' is :RTL, so:
+--
+-- > mkFamilyTyConAppCo :RTL (co :: a ~# Int) = T (Maybe a) ~# T (Maybe Int)
+--
+-- cf. 'mkFamilyTyConApp'
+mkFamilyTyConAppCo tc cos
+  | Just (fam_tc, fam_tys) <- tyConFamInst_maybe tc
+  , let tvs = tyConTyVars tc
+        fam_cos = ASSERT2( tvs `equalLength` cos, ppr tc <+> ppr cos )
+                  map (liftCoSubstWith Nominal tvs cos) fam_tys
+  = mkTyConAppCo Nominal fam_tc fam_cos
+  | otherwise
+  = mkTyConAppCo Nominal tc cos
+
 -- See note [Newtype coercions] in GHC.Core.TyCon
 
 mkPiCos :: Role -> [Var] -> Coercion -> Coercion
diff --git a/compiler/GHC/Core/DataCon.hs b/compiler/GHC/Core/DataCon.hs
index 12e98a517f..2fd51a2052 100644
--- a/compiler/GHC/Core/DataCon.hs
+++ b/compiler/GHC/Core/DataCon.hs
@@ -580,6 +580,7 @@ variables:
   purposes of TypeApplications, and as a consequence, they do not come equipped
   with visibilities (that is, they are TyVars/TyCoVars instead of
   TyCoVarBinders).
+
 * dcUserTyVarBinders, for the type variables binders in the order in which they
   originally arose in the user-written type signature. Their order *does* matter
   for TypeApplications, so they are full TyVarBinders, complete with
@@ -600,10 +601,10 @@ dcExTyCoVars. That is, the tyvars in dcUserTyVarBinders are a permutation of
 ordering, they in fact share the same type variables (with the same Uniques). We
 sometimes refer to this as "the dcUserTyVarBinders invariant".
 
-dcUserTyVarBinders, as the name suggests, is the one that users will see most of
-the time. It's used when computing the type signature of a data constructor (see
-dataConWrapperType), and as a result, it's what matters from a TypeApplications
-perspective.
+dcUserTyVarBinders, as the name suggests, is the one that users will
+see most of the time. It's used when computing the type signature of a
+data constructor wrapper (see dataConWrapperType), and as a result,
+it's what matters from a TypeApplications perspective.
 
 Note [The dcEqSpec domain invariant]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/GHC/Core/TyCo/Subst.hs b/compiler/GHC/Core/TyCo/Subst.hs
index b3f51739b5..b06add3172 100644
--- a/compiler/GHC/Core/TyCo/Subst.hs
+++ b/compiler/GHC/Core/TyCo/Subst.hs
@@ -435,8 +435,8 @@ mkTvSubstPrs prs =
 zipTyEnv :: HasDebugCallStack => [TyVar] -> [Type] -> TvSubstEnv
 zipTyEnv tyvars tys
   | debugIsOn
-  , not (all isTyVar tyvars)
-  = pprPanic "zipTyEnv" (ppr tyvars <+> ppr tys)
+  , not (all isTyVar tyvars && (tyvars `equalLength` tys))
+  = pprPanic "zipTyEnv" (ppr tyvars $$ ppr tys)
   | otherwise
   = ASSERT( all (not . isCoercionTy) tys )
     mkVarEnv (zipEqual "zipTyEnv" tyvars tys)
diff --git a/compiler/GHC/Hs/Expr.hs b/compiler/GHC/Hs/Expr.hs
index 565065e2c2..f762897f99 100644
--- a/compiler/GHC/Hs/Expr.hs
+++ b/compiler/GHC/Hs/Expr.hs
@@ -574,8 +574,9 @@ data RecordUpdTc = RecordUpdTc
 
       , rupd_in_tys  :: [Type]  -- Argument types of *input* record type
       , rupd_out_tys :: [Type]  --             and  *output* record type
-                                -- The original type can be reconstructed
-                                -- with conLikeResTy
+                -- For a data family, these are the type args of the
+                -- /representation/ type constructor
+
       , rupd_wrap :: HsWrapper  -- See note [Record Update HsWrapper]
       }
 
diff --git a/compiler/GHC/HsToCore/Expr.hs b/compiler/GHC/HsToCore/Expr.hs
index 2b959006e0..fc07fa412d 100644
--- a/compiler/GHC/HsToCore/Expr.hs
+++ b/compiler/GHC/HsToCore/Expr.hs
@@ -46,6 +46,7 @@ import GHC.Tc.Types.Evidence
 import GHC.Tc.Utils.Monad
 import GHC.Core.Type
 import GHC.Core.Multiplicity
+import GHC.Core.Coercion( Coercion )
 import GHC.Core
 import GHC.Core.Utils
 import GHC.Core.Make
@@ -54,6 +55,7 @@ import GHC.Driver.Session
 import GHC.Types.CostCentre
 import GHC.Types.Id
 import GHC.Types.Id.Make
+import GHC.Types.Var.Env
 import GHC.Unit.Module
 import GHC.Core.ConLike
 import GHC.Core.DataCon
@@ -62,13 +64,11 @@ import GHC.Builtin.Types
 import GHC.Builtin.Names
 import GHC.Types.Basic
 import GHC.Data.Maybe
-import GHC.Types.Var.Env
 import GHC.Types.SrcLoc
 import GHC.Utils.Misc
 import GHC.Data.Bag
 import GHC.Utils.Outputable as Outputable
 import GHC.Core.PatSyn
-
 import Control.Monad
 import Data.List.NonEmpty ( nonEmpty )
 
@@ -614,13 +614,70 @@ Note [Update for GADTs]
 ~~~~~~~~~~~~~~~~~~~~~~~
 Consider
    data T a b where
-     T1 :: { f1 :: a } -> T a Int
+     MkT :: { foo :: a } -> T a Int
+
+   upd :: T s t -> s -> T s t
+   upd z y = z { foo = y}
+
+We need to get this:
+   $WMkT :: a -> T a Int
+   MkT   :: (b ~# Int) => a -> T a b
+
+   upd = /\s t. \(z::T s t) (y::s) ->
+         case z of
+            MkT (co :: t ~# Int) _ -> $WMkT @s y |> T (Refl s) (Sym co)
 
-Then the wrapper function for T1 has type
-   $WT1 :: a -> T a Int
-But if x::T a b, then
-   x { f1 = v } :: T a b   (not T a Int!)
-So we need to cast (T a Int) to (T a b).  Sigh.
+Note the final cast
+   T (Refl s) (Sym co) :: T s Int ~ T s t
+which uses co, bound by the GADT match.  This is the wrap_co coercion
+in wrapped_rhs. How do we produce it?
+
+* Start with raw materials
+    tc, the tycon:                                       T
+    univ_tvs, the universally quantified tyvars of MkT:  a,b
+  NB: these are in 1-1 correspondence with the tyvars of tc
+
+* Form univ_cos, a coercion for each of tc's args: (Refl s) (Sym co)
+  We replaced
+     a  by  (Refl s)    since 's' instantiates 'a'
+     b  by  (Sym co)   since 'b' is in the data-con's EqSpec
+
+* Then form the coercion T (Refl s) (Sym co)
+
+It gets more complicated when data families are involved (#18809).
+Consider
+    data family F x
+    data instance F (a,b) where
+      MkF :: { foo :: Int } -> F (Int,b)
+
+    bar :: F (s,t) -> Int -> F (s,t)
+    bar z y = z { foo = y}
+
+We have
+    data R:FPair a b where
+      MkF :: { foo :: Int } -> R:FPair Int b
+
+    $WMkF :: Int -> F (Int,b)
+    MkF :: forall a b. (a ~# Int) => Int -> R:FPair a b
+
+    bar :: F (s,t) -> Int -> F (s,t)
+    bar = /\s t. \(z::F (s,t)) \(y::Int) ->
+         case z |> co1 of
+            MkF (co2::s ~# Int) _ -> $WMkF @t y |> co3
+
+(Side note: here (z |> co1) is built by typechecking the scrutinee, so
+we ignore it here.  In general the scrutinee is an aribtrary expression.)
+
+The question is: what is co3, the cast for the RHS?
+      co3 :: F (Int,t) ~ F (s,t)
+Again, we can construct it using co2, bound by the GADT match.
+We do /exactly/ the same as the non-family case up to building
+univ_cos.  But that gives us
+     rep_tc:   R:FPair
+     univ_cos: (Sym co2)   (Refl t)
+But then we use mkTcFamilyTyConAppCo to "lift" this to the coercion
+we want, namely
+     F (Sym co2, Refl t) :: F (Int,t) ~ F (s,t)
 
 -}
 
@@ -707,8 +764,7 @@ dsExpr expr@(RecordUpd { rupd_expr = record_expr, rupd_flds = fields
                         dict_req_wrap                                           <.>
                         mkWpTyApps    [ lookupTyVar out_subst tv
                                           `orElse` mkTyVarTy tv
-                                      | tv <- user_tvs
-                                      , not (tv `elemVarEnv` wrap_subst) ]
+                                      | tv <- user_tvs ]
                           -- Be sure to use user_tvs (which may be ordered
                           -- differently than `univ_tvs ++ ex_tvs) above.
                           -- See Note [DataCon user type variable binders]
@@ -719,27 +775,30 @@ dsExpr expr@(RecordUpd { rupd_expr = record_expr, rupd_flds = fields
                         -- Note [Update for GADTs]
                  wrapped_rhs =
                   case con of
-                    RealDataCon data_con ->
-                      let
-                        wrap_co =
-                          mkTcTyConAppCo Nominal
-                            (dataConTyCon data_con)
-                            [ lookup tv ty
-                              | (tv,ty) <- univ_tvs `zip` out_inst_tys ]
-                        lookup univ_tv ty =
-                          case lookupVarEnv wrap_subst univ_tv of
-                            Just co' -> co'
-                            Nothing  -> mkTcReflCo Nominal ty
-                        in if null eq_spec
-                             then rhs
-                             else mkLHsWrap (mkWpCastN wrap_co) rhs
+                    RealDataCon data_con
+                      | null eq_spec -> rhs
+                      | otherwise    -> mkLHsWrap (mkWpCastN wrap_co) rhs
+                                     -- This wrap is the punchline: Note [Update for GADTs]
+                      where
+                        rep_tc   = dataConTyCon data_con
+                        wrap_co  = mkTcFamilyTyConAppCo rep_tc univ_cos
+                        univ_cos = zipWithEqual "dsExpr:upd" mk_univ_co univ_tvs out_inst_tys
+
+                        mk_univ_co :: TyVar   -- Universal tyvar from the DataCon
+                                   -> Type    -- Corresponding instantiating type
+                                   -> Coercion
+                        mk_univ_co univ_tv inst_ty
+                          = case lookupVarEnv eq_spec_env univ_tv of
+                               Just co -> co
+                               Nothing -> mkTcNomReflCo inst_ty
+
+                        eq_spec_env :: VarEnv Coercion
+                        eq_spec_env = mkVarEnv [ (eqSpecTyVar spec, mkTcSymCo (mkTcCoVarCo eqs_var))
+                                               | (spec,eqs_var) <- zipEqual "dsExpr:upd2" eq_spec eqs_vars ]
+
                     -- eq_spec is always null for a PatSynCon
                     PatSynCon _ -> rhs
 
-                 wrap_subst =
-                  mkVarEnv [ (tv, mkTcSymCo (mkTcCoVarCo eq_var))
-                           | (spec, eq_var) <- eq_spec `zip` eqs_vars
-                           , let tv = eqSpecTyVar spec ]
 
                  req_wrap = dict_req_wrap <.> mkWpTyApps in_inst_tys
 
diff --git a/compiler/GHC/Tc/Types/Evidence.hs b/compiler/GHC/Tc/Types/Evidence.hs
index b9d460b8ec..194e9fa4f8 100644
--- a/compiler/GHC/Tc/Types/Evidence.hs
+++ b/compiler/GHC/Tc/Types/Evidence.hs
@@ -48,6 +48,7 @@ module GHC.Tc.Types.Evidence (
   mkTcKindCo,
   tcCoercionKind,
   mkTcCoVarCo,
+  mkTcFamilyTyConAppCo,
   isTcReflCo, isTcReflexiveCo,
   tcCoercionRole,
   unwrapIP, wrapIP,
@@ -139,6 +140,7 @@ mkTcCoherenceRightCo   :: Role -> TcType -> TcCoercionN
 mkTcPhantomCo          :: TcCoercionN -> TcType -> TcType -> TcCoercionP
 mkTcKindCo             :: TcCoercion -> TcCoercionN
 mkTcCoVarCo            :: CoVar -> TcCoercion
+mkTcFamilyTyConAppCo   :: TyCon -> [TcCoercionN] -> TcCoercionN
 
 tcCoercionKind         :: TcCoercion -> Pair TcType
 tcCoercionRole         :: TcCoercion -> Role
@@ -172,6 +174,7 @@ mkTcCoherenceRightCo   = mkCoherenceRightCo
 mkTcPhantomCo          = mkPhantomCo
 mkTcKindCo             = mkKindCo
 mkTcCoVarCo            = mkCoVarCo
+mkTcFamilyTyConAppCo   = mkFamilyTyConAppCo
 
 tcCoercionKind         = coercionKind
 tcCoercionRole         = coercionRole
diff --git a/testsuite/tests/indexed-types/should_compile/T18809.hs b/testsuite/tests/indexed-types/should_compile/T18809.hs
new file mode 100644
index 0000000000..1e56d980f6
--- /dev/null
+++ b/testsuite/tests/indexed-types/should_compile/T18809.hs
@@ -0,0 +1,33 @@
+{-# LANGUAGE GADTs, TypeFamilies #-}
+
+module T18809 where
+
+-- Ordinary
+data F2 s where
+  MkF2 :: {  foo2 :: Int } -> F2 s
+
+bar2 :: F2 s -> Int -> F2 s
+bar2 z y = z { foo2 = y }
+
+-- GADT
+data F1 s where
+  MkF1 :: {  foo1 :: Int } -> F1 Int
+
+bar1 :: F1 s -> Int -> F1 s
+bar1 z y = z { foo1 = y }
+
+-- Orinary data family
+data family F3 a
+data instance F3 (s,t) where
+  MkF2b :: { foo3 :: Int } -> F3 (s,t)
+
+bar3 :: F3 (s,t) -> Int -> F3 (s,t)
+bar3 z y = z {foo3 = y}
+
+-- GADT + data family
+data family F4 a
+data instance F4 (s,t) where
+  MkF2a :: { foo4 :: Int } -> F4 (Int,t)
+
+bar4 :: F4 (s,t) -> Int -> F4 (s,t)
+bar4 z y = z { foo4 = y}
diff --git a/testsuite/tests/indexed-types/should_compile/all.T b/testsuite/tests/indexed-types/should_compile/all.T
index 21703374b9..8b3dd5e866 100644
--- a/testsuite/tests/indexed-types/should_compile/all.T
+++ b/testsuite/tests/indexed-types/should_compile/all.T
@@ -296,3 +296,4 @@ test('T17056', normal, compile, [''])
 test('T17405', normal, multimod_compile, ['T17405c', '-v0'])
 test('T17923', normal, compile, [''])
 test('T18065', normal, compile, ['-O'])
+test('T18809', normal, compile, ['-O'])