Make inert_cycle_breakers into a stack.

Close #20231.

3 files changed, 79 insertions, 19 deletions

diff --git a/compiler/GHC/Data/Maybe.hs b/compiler/GHC/Data/Maybe.hs
index 3163829f75..215a2a0d6f 100644
--- a/compiler/GHC/Data/Maybe.hs
+++ b/compiler/GHC/Data/Maybe.hs
@@ -34,6 +34,7 @@ import Control.Exception (SomeException(..))
 import Data.Maybe
 import Data.Foldable ( foldlM )
 import GHC.Utils.Misc (HasCallStack)
+import Data.List.NonEmpty ( NonEmpty )
 
 infixr 4 `orElse`
 
@@ -50,8 +51,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/InertSet.hs b/compiler/GHC/Tc/Solver/InertSet.hs
index b5aad268b5..b2416e9915 100644
--- a/compiler/GHC/Tc/Solver/InertSet.hs
+++ b/compiler/GHC/Tc/Solver/InertSet.hs
@@ -29,7 +29,13 @@ module GHC.Tc.Solver.InertSet (
     partitionInertEqs, partitionFunEqs,
 
     -- * Kick-out
-    kickOutRewritableLHS
+    kickOutRewritableLHS,
+
+    -- * Cycle breaker vars
+    CycleBreakerVarStack,
+    pushCycleBreakerVarStack,
+    insertCycleBreakerBinding,
+    forAllCycleBreakerBindings_
 
   ) where
 
@@ -56,6 +62,11 @@ import GHC.Utils.Outputable
 import GHC.Utils.Panic
 
 import Data.List          ( partition )
+import Data.List.NonEmpty ( NonEmpty(..), (<|) )
+import qualified Data.List.NonEmpty as NE
+import GHC.Utils.Panic.Plain
+import GHC.Data.Maybe
+import Control.Monad      ( forM_ )
 
 {-
 ************************************************************************
@@ -224,15 +235,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
               -- 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 GHC.Tc.Solver.Canonical
+       , inert_cycle_breakers :: CycleBreakerVarStack
 
        , inert_famapp_cache :: FunEqMap Reduction
               -- Just a hash-cons cache for use when reducing family applications
@@ -273,7 +289,7 @@ emptyInertCans
 emptyInert :: InertSet
 emptyInert
   = IS { inert_cans           = emptyInertCans
-       , inert_cycle_breakers = []
+       , inert_cycle_breakers = [] :| []
        , inert_famapp_cache   = emptyFunEqs
        , inert_solved_dicts   = emptyDictMap }
 
@@ -1553,8 +1569,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
 
@@ -1600,10 +1615,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
@@ -1629,3 +1642,45 @@ prohibitedSuperClassSolve from_loc solve_loc
   = given_size >= wanted_size
   | 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 = ([] <|)
+
+-- | 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
diff --git a/compiler/GHC/Tc/Solver/Monad.hs b/compiler/GHC/Tc/Solver/Monad.hs
index 34f4751f1d..c253770616 100644
--- a/compiler/GHC/Tc/Solver/Monad.hs
+++ b/compiler/GHC/Tc/Solver/Monad.hs
@@ -8,7 +8,7 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ViewPatterns #-}
 
-{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
+{-# OPTIONS_GHC -Wno-incomplete-record-updates -Wno-orphans #-}
 
 -- | Monadic definitions for the constraint solver
 module GHC.Tc.Solver.Monad (
@@ -1098,7 +1098,8 @@ nestImplicTcS ref inner_tclvl (TcS thing_inside)
                    , tcs_abort_on_insoluble = abort_on_insoluble
                    } ->
     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
@@ -1995,8 +1996,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 $ mkReflRedn Nominal new_ty }
                 -- Why reflexive? See Detail (4) of the Note
 
@@ -2016,8 +2017,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