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.6 && < 0.7
 
   exposed-modules:
     Data.Constraint.If

src/IfSat/Plugin.hs

@@ -1,19 +1,22 @@
 {-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeApplications #-}
 
 module IfSat.Plugin
   ( plugin )
   where
 
 -- base
 import Control.Monad
-  ( filterM )
+  ( when )
 import Data.Foldable
-  ( for_ )
+  ( traverse_ )
 import Data.Maybe
   ( catMaybes, mapMaybe )
 
@@ -35,18 +38,23 @@ import GHC.Tc.Types
   ( TcM )
 import GHC.Tc.Types.Constraint
   ( isEmptyWC, CtEvidence (..), ctEvEvId )
+import GHC.Tc.Utils.Monad
+  ( TcRef )
 import GHC.Tc.Utils.TcType
-  ( MetaDetails(..), metaTyVarRef
-  , tyCoVarsOfTypeList
-  )
+  ( MetaDetails (..), metaTyVarRef )
 import GHC.Tc.Utils.TcMType
-  ( isUnfilledMetaTyVar, newTcEvBinds )
+  ( newTcEvBinds, isUnfilledMetaTyVar )
 
 -- ghc-tcplugin-api
 import GHC.TcPlugin.API
 import GHC.TcPlugin.API.Internal
   ( unsafeLiftTcM )
 
+-- transformers
+import Control.Monad.Trans.Class ( lift )
+import Control.Monad.Trans.Writer.CPS ( WriterT )
+import qualified Control.Monad.Trans.Writer.CPS as Writer
+
 -- if-instance
 import IfSat.Plugin.Compat
   ( wrapTcS, getRestoreTcS )
@@ -151,9 +159,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 +189,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 +215,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
@@ -306,6 +308,45 @@ dispatchFalseEvTerm defs@( PluginDefs { orClass } ) givens ct_l_ty ct_r_ty ct_r_
         )
       ]
 
+-- | 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 = Writer.execWriterT . go_ty
+  where
+    (go_ty, _go_tys, _go_co, _go_cos) =
+      mapTyCo @( WriterT [UnfilledRef] TcM ) $
+        TyCoMapper
+          { tcm_tyvar = \ _ tv -> do
+              unfilled_meta <- lift $ isUnfilledMetaTyVar tv
+              when unfilled_meta $
+                Writer.tell $ [ UnfilledMeta $ metaTyVarRef tv ]
+              return $ mkTyVarTy tv
+          , tcm_tycobinder = \ _ tcv _ftf k -> k () tcv
+          , 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 ->
+                  Writer.tell $ [ UnfilledCoHole hole_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
+
 -- The type @IsSat ct ~ b@.
 sat_eqTy :: PluginDefs -> Type -> Bool -> Type
 sat_eqTy ( PluginDefs { isSatTyCon } ) ct_ty booly
@@ -383,9 +424,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 +435,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