Make inert_cycle_breakers into a stack.

Close #20231.

(cherry picked from commit 1617fed3a97cd13b55a180029ab8fb9468d2b797)

2 files changed, 67 insertions, 17 deletions

diff --git a/compiler/GHC/Data/Maybe.hs b/compiler/GHC/Data/Maybe.hs
index ac9c687b62..961641e0fb 100644
--- a/compiler/GHC/Data/Maybe.hs
+++ b/compiler/GHC/Data/Maybe.hs
@@ -33,6 +33,7 @@ import Control.Exception (catch, SomeException(..))
 import Data.Maybe
 import Data.Foldable ( foldlM )
 import GHC.Utils.Misc (HasCallStack)
+import Data.List.NonEmpty ( NonEmpty )
 
 infixr 4 `orElse`
 
@@ -49,8 +50,10 @@ firstJust a b = firstJusts [a, b]
 
 -- | Takes a list of @Maybes@ and returns the first @Just@ if there is one, or
 -- @Nothing@ otherwise.
-firstJusts :: [Maybe a] -> Maybe a
+firstJusts :: Foldable f => f (Maybe a) -> Maybe a
 firstJusts = msum
+{-# SPECIALISE firstJusts :: [Maybe a] -> Maybe a #-}
+{-# SPECIALISE firstJusts :: NonEmpty (Maybe a) -> Maybe a #-}
 
 -- | Takes computations returnings @Maybes@; tries each one in order.
 -- The first one to return a @Just@ wins. Returns @Nothing@ if all computations
diff --git a/compiler/GHC/Tc/Solver/Monad.hs b/compiler/GHC/Tc/Solver/Monad.hs
index 841d2b55bc..91a28389ce 100644
--- a/compiler/GHC/Tc/Solver/Monad.hs
+++ b/compiler/GHC/Tc/Solver/Monad.hs
@@ -383,15 +383,20 @@ instance Outputable WorkList where
 *                                                                      *
 ********************************************************************* -}
 
+type CycleBreakerVarStack = NonEmpty [(TcTyVar, TcType)]
+   -- ^ a stack of (CycleBreakerTv, original family applications) lists
+   -- first element in the stack corresponds to current implication;
+   --   later elements correspond to outer implications
+   -- used to undo the cycle-breaking needed to handle
+   -- Note [Type variable cycles] in GHC.Tc.Solver.Canonical
+   -- Why store the outer implications? For the use in mightEqualLater (only)
+
 data InertSet
   = IS { inert_cans :: InertCans
               -- Canonical Given, Wanted, Derived
               -- Sometimes called "the inert set"
 
-       , inert_cycle_breakers :: [(TcTyVar, TcType)]
-              -- a list of CycleBreakerTv / original family applications
-              -- used to undo the cycle-breaking needed to handle
-              -- Note [Type variable cycles in Givens] in GHC.Tc.Solver.Canonical
+       , inert_cycle_breakers :: CycleBreakerVarStack
 
        , inert_famapp_cache :: FunEqMap (TcCoercion, TcType)
               -- Just a hash-cons cache for use when reducing family applications
@@ -433,7 +438,7 @@ emptyInertCans
 emptyInert :: InertSet
 emptyInert
   = IS { inert_cans           = emptyInertCans
-       , inert_cycle_breakers = []
+       , inert_cycle_breakers = [] :| []
        , inert_famapp_cache   = emptyFunEqs
        , inert_solved_dicts   = emptyDictMap }
 
@@ -2534,8 +2539,7 @@ matchableGivens loc_w pred_w inerts@(IS { inert_cans = inert_cans })
 mightEqualLater :: InertSet -> TcPredType -> CtLoc -> TcPredType -> CtLoc -> Bool
 -- See Note [What might equal later?]
 -- Used to implement logic in Note [Instance and Given overlap] in GHC.Tc.Solver.Interact
-mightEqualLater (IS { inert_cycle_breakers = cbvs })
-                given_pred given_loc wanted_pred wanted_loc
+mightEqualLater inert_set given_pred given_loc wanted_pred wanted_loc
   | prohibitedSuperClassSolve given_loc wanted_loc
   = False
 
@@ -2581,10 +2585,8 @@ mightEqualLater (IS { inert_cycle_breakers = cbvs })
       = case metaTyVarInfo tv of
           -- See Examples 8 and 9 in the Note
           CycleBreakerTv
-            | Just tyfam_app <- lookup tv cbvs
-            -> anyFreeVarsOfType mentions_meta_ty_var tyfam_app
-            | otherwise
-            -> pprPanic "mightEqualLater finds an unbound cbv" (ppr tv $$ ppr cbvs)
+            -> anyFreeVarsOfType mentions_meta_ty_var
+                 (lookupCycleBreakerVar tv inert_set)
           _ -> True
       | otherwise
       = False
@@ -2611,6 +2613,49 @@ prohibitedSuperClassSolve from_loc solve_loc
   | otherwise
   = False
 
+{- *********************************************************************
+*                                                                      *
+    Cycle breakers
+*                                                                      *
+********************************************************************* -}
+
+-- | Return the type family application a CycleBreakerTv maps to.
+lookupCycleBreakerVar :: TcTyVar    -- ^ cbv, must be a CycleBreakerTv
+                      -> InertSet
+                      -> TcType     -- ^ type family application the cbv maps to
+lookupCycleBreakerVar cbv (IS { inert_cycle_breakers = cbvs_stack })
+-- This function looks at every environment in the stack. This is necessary
+-- to avoid #20231. This function (and its one usage site) is the only reason
+-- that we store a stack instead of just the top environment.
+  | Just tyfam_app <- ASSERT( (isCycleBreakerTyVar cbv) )
+                      firstJusts (NE.map (lookup cbv) cbvs_stack)
+  = tyfam_app
+  | otherwise
+  = pprPanic "lookupCycleBreakerVar found an unbound cycle breaker" (ppr cbv $$ ppr cbvs_stack)
+
+-- | Push a fresh environment onto the cycle-breaker var stack. Useful
+-- when entering a nested implication.
+pushCycleBreakerVarStack :: CycleBreakerVarStack -> CycleBreakerVarStack
+pushCycleBreakerVarStack = ([] NE.<|)
+
+-- | Add a new cycle-breaker binding to the top environment on the stack.
+insertCycleBreakerBinding :: TcTyVar   -- ^ cbv, must be a CycleBreakerTv
+                          -> TcType    -- ^ cbv's expansion
+                          -> CycleBreakerVarStack -> CycleBreakerVarStack
+insertCycleBreakerBinding cbv expansion (top_env :| rest_envs)
+  = ASSERT( (isCycleBreakerTyVar cbv) )
+    ((cbv, expansion) : top_env) :| rest_envs
+
+-- | Perform a monadic operation on all pairs in the top environment
+-- in the stack.
+forAllCycleBreakerBindings_ :: Monad m
+                            => CycleBreakerVarStack
+                            -> (TcTyVar -> TcType -> m ()) -> m ()
+forAllCycleBreakerBindings_ (top_env :| _rest_envs) action
+  = forM_ top_env (uncurry action)
+{-# INLINABLE forAllCycleBreakerBindings_ #-}  -- to allow SPECIALISE later
+
+
 {- Note [Unsolved Derived equalities]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 In getUnsolvedInerts, we return a derived equality from the inert_eqs
@@ -3302,7 +3347,8 @@ nestImplicTcS ref inner_tclvl (TcS thing_inside)
                    , tcs_unif_lvl      = unif_lvl
                    } ->
     do { inerts <- TcM.readTcRef old_inert_var
-       ; let nest_inert = inerts { inert_cycle_breakers = []
+       ; let nest_inert = inerts { inert_cycle_breakers = pushCycleBreakerVarStack
+                                                            (inert_cycle_breakers inerts)
                                  , inert_cans = (inert_cans inerts)
                                                    { inert_given_eqs = False } }
                  -- All other InertSet fields are inherited
@@ -4193,8 +4239,8 @@ breakTyVarCycle_maybe ev cte_result (TyVarLHS lhs_tv) rhs
            ; traceTcS "breakTyVarCycle replacing type family in Given" (ppr new_given)
            ; emitWorkNC [new_given]
            ; updInertTcS $ \is ->
-               is { inert_cycle_breakers = (new_tv, fun_app) :
-                                           inert_cycle_breakers is }
+               is { inert_cycle_breakers = insertCycleBreakerBinding new_tv fun_app
+                                             (inert_cycle_breakers is) }
            ; return (mkNomReflCo new_ty, new_ty) }
                 -- Why reflexive? See Detail (4) of the Note
 
@@ -4214,8 +4260,9 @@ breakTyVarCycle_maybe _ _ _ _ = return Nothing
 -- See Note [Type variable cycles] in GHC.Tc.Solver.Canonical.
 restoreTyVarCycles :: InertSet -> TcM ()
 restoreTyVarCycles is
-  = forM_ (inert_cycle_breakers is) $ \ (cycle_breaker_tv, orig_ty) ->
-    TcM.writeMetaTyVar cycle_breaker_tv orig_ty
+  = forAllCycleBreakerBindings_ (inert_cycle_breakers is) TcM.writeMetaTyVar
+{-# SPECIALISE forAllCycleBreakerBindings_ ::
+      CycleBreakerVarStack -> (TcTyVar -> TcType -> TcM ()) -> TcM () #-}
 
 -- Unwrap a type synonym only when either:
 --   The type synonym is forgetful, or