Fix treatment of coercion variables

Author: sheaf

changelog.md

@@ -1,4 +1,9 @@
 
+# Version 0.6.0.0 (2025-03-29)
+
+- Fixed the treatment of constraints containing coercion variables; the plugin
+  will no longer crash upon encountering such constraints.
+
 # Version 0.5.2.0 (2024-11-29)
 
 - Update to support GHC 9.10 and 9.12.  

if-instance.cabal

@@ -1,6 +1,6 @@
 cabal-version:  3.0
 name:           if-instance
-version:        0.5.2.0
+version:        0.6.0.0
 synopsis:       Branch on whether a constraint is satisfied
 license:        BSD-3-Clause
 build-type:     Simple
@@ -86,6 +86,8 @@ library
   build-depends:
     ghc-tcplugin-api
       >= 0.13 && < 0.15,
+    transformers
+      >= 0.5.6.2 && < 0.7
 
   exposed-modules:
     Data.Constraint.If

src/IfSat/Plugin.hs

@@ -1,5 +1,6 @@
-{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE CPP #-}
+
+{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE PatternSynonyms #-}
@@ -10,10 +11,8 @@ module IfSat.Plugin
   where
 
 -- base
-import Control.Monad
-  ( filterM )
 import Data.Foldable
-  ( for_ )
+  ( traverse_ )
 import Data.Maybe
   ( catMaybes, mapMaybe )
 
@@ -35,12 +34,8 @@ import GHC.Tc.Types
   ( TcM )
 import GHC.Tc.Types.Constraint
   ( isEmptyWC, CtEvidence (..), ctEvEvId )
-import GHC.Tc.Utils.TcType
-  ( MetaDetails(..), metaTyVarRef
-  , tyCoVarsOfTypeList
-  )
 import GHC.Tc.Utils.TcMType
-  ( isUnfilledMetaTyVar, newTcEvBinds )
+  ( newTcEvBinds )
 
 -- ghc-tcplugin-api
 import GHC.TcPlugin.API
@@ -49,7 +44,9 @@ import GHC.TcPlugin.API.Internal
 
 -- if-instance
 import IfSat.Plugin.Compat
-  ( wrapTcS, getRestoreTcS )
+  ( wrapTcS, getRestoreTcS
+  , unfilledRefsOfType, unfillMutableRef
+  )
 
 --------------------------------------------------------------------------------
 -- Plugin definition.
@@ -151,9 +148,7 @@ solveWanted defs@( PluginDefs { orClass } ) givens wanted
 
       -- Keep track of the current solver state in order to backtrack
       -- in the event that our attempt at solving 'ct_l' fails.
-      ct_l_unfilled_metas <- wrapTcS
-                           $ filterM isUnfilledMetaTyVar
-                           $ tyCoVarsOfTypeList ct_l_ty
+      ct_l_unfilled <- wrapTcS $ unfilledRefsOfType ct_l_ty
       restoreTcS <- getRestoreTcS
 
       -- Try to solve 'ct_l', using both Givens and top-level instances.
@@ -183,11 +178,8 @@ solveWanted defs@( PluginDefs { orClass } ) givens wanted
           -- Reset the solver state to before we attempted to solve 'ct_l',
           -- and undo any type variable unifications that happened.
           restoreTcS
-          wrapTcS $ for_ ct_l_unfilled_metas \ meta ->
-            writeTcRef ( metaTyVarRef meta ) Flexi
-          ct_r_unfilled_metas <- wrapTcS
-                               $ filterM isUnfilledMetaTyVar
-                               $ tyCoVarsOfTypeList ct_r_ty
+          wrapTcS $ traverse_ unfillMutableRef ct_l_unfilled
+          ct_r_unfilled <- wrapTcS $ unfilledRefsOfType ct_r_ty
 
           -- Try to solve 'ct_r', using both Givens and top-level instances.
           residual_ct_r <- solveSimpleWanteds ( unitBag ct_r )
@@ -212,8 +204,7 @@ solveWanted defs@( PluginDefs { orClass } ) givens wanted
               -- Reset the solver state to before we attempted to solve 'ct_r',
               -- and undo any type variable unifications that happened.
               restoreTcS
-              wrapTcS $ for_ ct_r_unfilled_metas \ meta ->
-                writeTcRef ( metaTyVarRef meta ) Flexi
+              wrapTcS $ traverse_ unfillMutableRef ct_r_unfilled
 
               pure Nothing
       pure $ ( , wanted ) <$> mb_wanted_evTerm
@@ -356,8 +347,8 @@ usedGivenCoercions givens ev = mapMaybe dep_cv givens
 -- selector warnings.
 wantedEvDest :: HasDebugCallStack => CtEvidence -> TcEvDest
 wantedEvDest ( CtWanted { ctev_dest = dst } ) = dst
-wantedEvDest g@( CtGiven {} ) =
-  pprPanic "wantedEvDest called on CtGiven" (ppr g)
+wantedEvDest non_wtd =
+  pprPanic "wantedEvDest, but not a Wanted Ct" (ppr non_wtd)
 
 --------------------------------------------------------------------------------
 
@@ -383,9 +374,7 @@ isSatRewriter ( PluginDefs { isSatTyCon } ) givens [ct_ty] = do
 
     -- Keep track of the current solver state in order to undo any
     -- side-effects after calling 'solveSimpleWanteds' on 'ct'.
-    ct_unfilled_metas <- wrapTcS
-                       $ filterM isUnfilledMetaTyVar
-                       $ tyCoVarsOfTypeList ct_ty
+    ct_unfilled <- wrapTcS $ unfilledRefsOfType ct_ty
     restoreTcS <- getRestoreTcS
 
     -- Try to solve 'ct', using both Givens and top-level instances.
@@ -396,8 +385,7 @@ isSatRewriter ( PluginDefs { isSatTyCon } ) givens [ct_ty] = do
     -- Reset the solver state to before we attempted to solve 'ct',
     -- and undo any type variable unifications that happened.
     restoreTcS
-    wrapTcS $ for_ ct_unfilled_metas \ meta ->
-      writeTcRef ( metaTyVarRef meta ) Flexi
+    wrapTcS $ traverse_ unfillMutableRef ct_unfilled
 
     let
       is_sat :: Bool

src/IfSat/Plugin/Compat.hs

@@ -1,38 +1,61 @@
-{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE CPP #-}
 
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE TypeApplications #-}
+
 {-# OPTIONS_GHC -Wno-unused-top-binds #-}
 
 module IfSat.Plugin.Compat
-  ( wrapTcS, getRestoreTcS )
+  ( -- * Dealing with TcS state
+    wrapTcS, getRestoreTcS
+    -- * Dealing with mutable references
+  , UnfilledRef(..), unfilledRefsOfType, unfillMutableRef
+  )
   where
 
 -- base
+import Control.Monad
+  ( when )
 import Unsafe.Coerce
   ( unsafeCoerce )
 
 -- ghc
+import GHC.Core.Coercion
+  ( mkCoVarCo, mkHoleCo, coVarName, coHoleCoVar )
+import GHC.Core.Type
+  ( TyCoMapper(..), mapTyCo )
 #if MIN_VERSION_ghc(9,4,0)
 import GHC.Tc.Solver.InertSet
-  ( WorkList, InertSet )
+  ( WorkList )
 #endif
-import GHC.Tc.Solver.Monad
-  ( TcS
 #if MIN_VERSION_ghc(9,1,0)
-  , TcLevel, wrapTcS
+import GHC.Tc.Solver.Monad
+  ( wrapTcS )
 #endif
 #if !MIN_VERSION_ghc(9,4,0)
-  , WorkList, InertSet
+import GHC.Tc.Solver.Monad
+  ( WorkList )
 #endif
-  )
 import GHC.Tc.Types
   ( TcM, TcRef )
+
+import GHC.Tc.Utils.TcMType
+  ( isUnfilledMetaTyVar )
+import GHC.Tc.Utils.TcType
+  ( MetaDetails(..), metaTyVarRef )
 import GHC.Tc.Types.Evidence
   ( EvBindsVar(..) )
 
+-- transformers
+import Control.Monad.Trans.Class ( lift )
+import Control.Monad.Trans.Writer.CPS ( WriterT )
+import qualified Control.Monad.Trans.Writer.CPS as Writer
+
 -- ghc-tcplugin-api
 import GHC.TcPlugin.API
-  ( readTcRef, writeTcRef )
+import GHC.Types.Name.Env
+import GHC.Types.Var (tyVarName)
 
 --------------------------------------------------------------------------------
 
@@ -107,3 +130,61 @@ data ShimTcSEnv
 #endif
   , shim_tcs_worklist           :: TcRef WorkList
   }
+
+--------------------------------------------------------------------------------
+
+-- | A mutable reference that was originally unfilled
+data UnfilledRef
+  -- | A metavariable that was originally unfilled
+  = UnfilledMeta   !( TcRef MetaDetails )
+  -- | A coercion hole that was originally unfilled
+  | UnfilledCoHole !( TcRef ( Maybe Coercion ) )
+
+-- | Gather all the unfilled mutable references of a type: unfilled
+-- metavariables and unfilled coercion holes.
+unfilledRefsOfType :: TcType -> TcM [ UnfilledRef ]
+unfilledRefsOfType ty0
+  = fmap
+#if MIN_VERSION_ghc(9,3,0)
+      nonDetNameEnvElts
+#else
+      nameEnvElts
+#endif
+  $ Writer.execWriterT
+  $ go_ty ty0
+  where
+    (go_ty, _go_tys, _go_co, _go_cos) =
+      mapTyCo @( WriterT ( NameEnv UnfilledRef ) TcM ) $
+        TyCoMapper
+          { tcm_tyvar = \ _ tv -> do
+              unfilled_meta <- lift $ isUnfilledMetaTyVar tv
+              when unfilled_meta do
+                let nm = tyVarName tv
+                    ref = UnfilledMeta $ metaTyVarRef tv
+                Writer.tell $ unitNameEnv nm ref
+              return $ mkTyVarTy tv
+          , tcm_tycobinder =
+#if MIN_VERSION_ghc(9,7,0)
+              \ () tv _ftf k -> k () tv
+#else
+              \ () tv _af -> return ( (), tv )
+#endif
+          , tcm_tycon = return
+          , tcm_covar = \ _ cv -> return $ mkCoVarCo cv
+          , tcm_hole = \ _ hole@(CoercionHole { ch_ref = hole_ref }) -> do
+              hole_contents <- lift $ readTcRef hole_ref
+              case hole_contents of
+                Nothing -> do
+                  let nm = coVarName $ coHoleCoVar hole
+                      ref = UnfilledCoHole hole_ref
+                  Writer.tell $ unitNameEnv nm ref
+                Just {} ->
+                  return ()
+              return $ mkHoleCo hole
+          }
+
+-- | Restore a mutable reference to the unfilled state.
+unfillMutableRef :: UnfilledRef -> TcM ()
+unfillMutableRef = \case
+  UnfilledMeta   ref  -> writeTcRef ref  Flexi
+  UnfilledCoHole hole -> writeTcRef hole Nothing