Preserve ShadowInfo when rewriting evidence

When the canonicaliser rewrites evidence of a Wanted, it
should preserve the ShadowInfo (ctev_nosh) field.  That is,
a WDerive should rewrite to WDerive, and WOnly to WOnly.

Previously we were unconditionally making a WDeriv, thereby
rewriting WOnly to WDeriv.  This bit Nick Frisby (issue #16735)
in the context of his plugin, but we don't have a compact test
case.

The fix is simple, but does involve a bit more plumbing,
to pass the old ShadowInfo around, to use when building
the new Wanted.

3 files changed, 44 insertions, 21 deletions

diff --git a/compiler/typecheck/TcCanonical.hs b/compiler/typecheck/TcCanonical.hs
index 47c4cc807c..bbb29f5e38 100644
--- a/compiler/typecheck/TcCanonical.hs
+++ b/compiler/typecheck/TcCanonical.hs
@@ -2285,8 +2285,12 @@ rewriteEvidence ev@(CtGiven { ctev_evar = old_evar, ctev_loc = loc }) new_pred c
                                                        (mkTcSymCo co))
 
 rewriteEvidence ev@(CtWanted { ctev_dest = dest
+                             , ctev_nosh = si
                              , ctev_loc = loc }) new_pred co
-  = do { mb_new_ev <- newWanted loc new_pred
+  = do { mb_new_ev <- newWanted_SI si loc new_pred
+               -- The "_SI" varant ensures that we make a new Wanted
+               -- with the same shadow-info as the existing one
+               -- with the same shadow-info as the existing one (#16735)
        ; MASSERT( tcCoercionRole co == ctEvRole ev )
        ; setWantedEvTerm dest
             (mkEvCast (getEvExpr mb_new_ev)
@@ -2334,8 +2338,10 @@ rewriteEqEvidence old_ev swapped nlhs nrhs lhs_co rhs_co
                                   `mkTcTransCo` mkTcSymCo rhs_co)
        ; newGivenEvVar loc' (new_pred, new_tm) }
 
-  | CtWanted { ctev_dest = dest } <- old_ev
-  = do { (new_ev, hole_co) <- newWantedEq loc' (ctEvRole old_ev) nlhs nrhs
+  | CtWanted { ctev_dest = dest, ctev_nosh = si } <- old_ev
+  = do { (new_ev, hole_co) <- newWantedEq_SI si loc' (ctEvRole old_ev) nlhs nrhs
+               -- The "_SI" varant ensures that we make a new Wanted
+               -- with the same shadow-info as the existing one (#16735)
        ; let co = maybeSym swapped $
                   mkSymCo lhs_co
                   `mkTransCo` hole_co
diff --git a/compiler/typecheck/TcSMonad.hs b/compiler/typecheck/TcSMonad.hs
index c52e624d8d..8d98a17149 100644
--- a/compiler/typecheck/TcSMonad.hs
+++ b/compiler/typecheck/TcSMonad.hs
@@ -33,8 +33,10 @@ module TcSMonad (
     MaybeNew(..), freshGoals, isFresh, getEvExpr,
 
     newTcEvBinds, newNoTcEvBinds,
-    newWantedEq, emitNewWantedEq,
-    newWanted, newWantedEvVar, newWantedNC, newWantedEvVarNC, newDerivedNC,
+    newWantedEq, newWantedEq_SI, emitNewWantedEq,
+    newWanted, newWanted_SI, newWantedEvVar,
+    newWantedNC, newWantedEvVarNC,
+    newDerivedNC,
     newBoundEvVarId,
     unifyTyVar, unflattenFmv, reportUnifications,
     setEvBind, setWantedEq,
@@ -3404,12 +3406,18 @@ emitNewWantedEq loc role ty1 ty2
        ; return co }
 
 -- | Make a new equality CtEvidence
-newWantedEq :: CtLoc -> Role -> TcType -> TcType -> TcS (CtEvidence, Coercion)
-newWantedEq loc role ty1 ty2
+newWantedEq :: CtLoc -> Role -> TcType -> TcType
+            -> TcS (CtEvidence, Coercion)
+newWantedEq = newWantedEq_SI WDeriv
+
+newWantedEq_SI :: ShadowInfo -> CtLoc -> Role
+               -> TcType -> TcType
+               -> TcS (CtEvidence, Coercion)
+newWantedEq_SI si loc role ty1 ty2
   = do { hole <- wrapTcS $ TcM.newCoercionHole pty
        ; traceTcS "Emitting new coercion hole" (ppr hole <+> dcolon <+> ppr pty)
        ; return ( CtWanted { ctev_pred = pty, ctev_dest = HoleDest hole
-                           , ctev_nosh = WDeriv
+                           , ctev_nosh = si
                            , ctev_loc = loc}
                 , mkHoleCo hole ) }
   where
@@ -3417,35 +3425,44 @@ newWantedEq loc role ty1 ty2
 
 -- no equalities here. Use newWantedEq instead
 newWantedEvVarNC :: CtLoc -> TcPredType -> TcS CtEvidence
+newWantedEvVarNC = newWantedEvVarNC_SI WDeriv
+
+newWantedEvVarNC_SI :: ShadowInfo -> CtLoc -> TcPredType -> TcS CtEvidence
 -- Don't look up in the solved/inerts; we know it's not there
-newWantedEvVarNC loc pty
+newWantedEvVarNC_SI si loc pty
   = do { new_ev <- newEvVar pty
        ; traceTcS "Emitting new wanted" (ppr new_ev <+> dcolon <+> ppr pty $$
                                          pprCtLoc loc)
        ; return (CtWanted { ctev_pred = pty, ctev_dest = EvVarDest new_ev
-                          , ctev_nosh = WDeriv
+                          , ctev_nosh = si
                           , ctev_loc = loc })}
 
 newWantedEvVar :: CtLoc -> TcPredType -> TcS MaybeNew
+newWantedEvVar = newWantedEvVar_SI WDeriv
+
+newWantedEvVar_SI :: ShadowInfo -> CtLoc -> TcPredType -> TcS MaybeNew
 -- For anything except ClassPred, this is the same as newWantedEvVarNC
-newWantedEvVar loc pty
+newWantedEvVar_SI si loc pty
   = do { mb_ct <- lookupInInerts loc pty
        ; case mb_ct of
             Just ctev
               | not (isDerived ctev)
               -> do { traceTcS "newWantedEvVar/cache hit" $ ppr ctev
                     ; return $ Cached (ctEvExpr ctev) }
-            _ -> do { ctev <- newWantedEvVarNC loc pty
+            _ -> do { ctev <- newWantedEvVarNC_SI si loc pty
                     ; return (Fresh ctev) } }
 
--- deals with both equalities and non equalities. Tries to look
--- up non-equalities in the cache
 newWanted :: CtLoc -> PredType -> TcS MaybeNew
-newWanted loc pty
+-- Deals with both equalities and non equalities. Tries to look
+-- up non-equalities in the cache
+newWanted = newWanted_SI WDeriv
+
+newWanted_SI :: ShadowInfo -> CtLoc -> PredType -> TcS MaybeNew
+newWanted_SI si loc pty
   | Just (role, ty1, ty2) <- getEqPredTys_maybe pty
-  = Fresh . fst <$> newWantedEq loc role ty1 ty2
+  = Fresh . fst <$> newWantedEq_SI si loc role ty1 ty2
   | otherwise
-  = newWantedEvVar loc pty
+  = newWantedEvVar_SI si loc pty
 
 -- deals with both equalities and non equalities. Doesn't do any cache lookups.
 newWantedNC :: CtLoc -> PredType -> TcS CtEvidence
diff --git a/testsuite/tests/polykinds/T14172.stderr b/testsuite/tests/polykinds/T14172.stderr
index 41790105cb..6a0d3927eb 100644
--- a/testsuite/tests/polykinds/T14172.stderr
+++ b/testsuite/tests/polykinds/T14172.stderr
@@ -11,12 +11,12 @@ T14172.hs:6:46: error:
       In the type ‘(a -> f b) -> g a -> f (h _)’
 
 T14172.hs:7:19: error:
-    • Occurs check: cannot construct the infinite type: a ~ g'1 a
+    • Occurs check: cannot construct the infinite type: a ~ g'0 a
       Expected type: (f'0 a -> f (f'0 b))
-                     -> Compose f'0 g'1 a -> f (h a')
-        Actual type: (Unwrapped (Compose f'0 g'1 a)
+                     -> Compose f'0 g'0 a -> f (h a')
+        Actual type: (Unwrapped (Compose f'0 g'0 a)
                       -> f (Unwrapped (h a')))
-                     -> Compose f'0 g'1 a -> f (h a')
+                     -> Compose f'0 g'0 a -> f (h a')
     • In the first argument of ‘(.)’, namely ‘_Wrapping Compose’
       In the expression: _Wrapping Compose . traverse
       In an equation for ‘traverseCompose’: