TcPlugin: access to irreducible givens + fix passed ev_binds_var

6 files changed, 116 insertions, 0 deletions

diff --git a/testsuite/tests/tcplugins/Common.hs b/testsuite/tests/tcplugins/Common.hs
index f2f425381d..d5eb1767d3 100644
--- a/testsuite/tests/tcplugins/Common.hs
+++ b/testsuite/tests/tcplugins/Common.hs
@@ -67,6 +67,7 @@ data PluginDefs =
     , zero     :: !TyCon
     , succ     :: !TyCon
     , add      :: !TyCon
+    , ctIdFam  :: !TyCon
     }
 
 definitionsModule :: TcPluginM Module
@@ -87,6 +88,7 @@ lookupDefs = do
   ( promoteDataCon    -> zero ) <- tcLookupDataCon =<< lookupOrig defs ( mkDataOcc "Zero"     )
   ( promoteDataCon    -> succ ) <- tcLookupDataCon =<< lookupOrig defs ( mkDataOcc "Succ"     )
   add                           <- tcLookupTyCon   =<< lookupOrig defs ( mkTcOcc   "Add"      )
+  ctIdFam                       <- tcLookupTyCon   =<< lookupOrig defs ( mkTcOcc   "CtId"     )
   pure ( PluginDefs { .. } )
 
 mkPlugin :: ( [String] -> PluginDefs -> EvBindsVar -> [Ct] -> [Ct] -> TcPluginM TcPluginSolveResult )
diff --git a/testsuite/tests/tcplugins/CtIdPlugin.hs b/testsuite/tests/tcplugins/CtIdPlugin.hs
new file mode 100644
index 0000000000..2511c902d3
--- /dev/null
+++ b/testsuite/tests/tcplugins/CtIdPlugin.hs
@@ -0,0 +1,80 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE ViewPatterns #-}
+
+module CtIdPlugin where
+
+-- base
+import Data.Maybe
+import Data.Traversable
+
+-- ghc
+import GHC.Core.Class
+import GHC.Core.Coercion
+import GHC.Core.DataCon
+import GHC.Core.Make
+import GHC.Core.Predicate
+import GHC.Core.TyCo.Rep
+import GHC.Plugins
+import GHC.Tc.Plugin
+import GHC.Tc.Types
+import GHC.Tc.Types.Constraint
+import GHC.Tc.Types.Evidence
+
+-- common
+import Common
+
+--------------------------------------------------------------------------------
+
+-- This plugin simplifies Given and Wanted 'CtId' constraints.
+-- To do this, we just look through the Givens and Wanteds,
+-- find any irreducible constraint whose TyCon matches that of 'CtId',
+-- in which case we substitute it for its argument:
+-- We create a new Given or Wanted and remove the old one using a cast.
+
+plugin :: Plugin
+plugin = mkPlugin solver don'tRewrite
+
+-- Solve "CtId".
+solver :: [String]
+       -> PluginDefs -> EvBindsVar -> [Ct] -> [Ct]
+       -> TcPluginM TcPluginSolveResult
+solver _args defs ev givens wanteds = do
+  let pluginCo = mkUnivCo (PluginProv "CtIdPlugin") Representational
+  let substEvidence ct ct' =
+        evCast (ctEvExpr $ ctEvidence ct') $ pluginCo (ctPred ct') (ctPred ct)
+
+  if null wanteds
+    then do
+      newGivenPredTypes <- traverse (solveCt defs) givens
+      newGivens <- for (zip newGivenPredTypes givens) \case
+        (Nothing, _) -> return Nothing
+        (Just pred, ct) ->
+          let EvExpr expr =
+                evCast (ctEvExpr $ ctEvidence ct) $ pluginCo (ctPred ct) pred
+          in Just . mkNonCanonical <$> newGiven ev (ctLoc ct) pred expr
+      let removedGivens =
+            [ (substEvidence ct ct', ct)
+            | (Just ct', ct) <- zip newGivens givens
+            ]
+      pure $ TcPluginOk removedGivens (catMaybes newGivens)
+    else do
+      newWantedPredTypes <- traverse (solveCt defs) wanteds
+      newWanteds <- for (zip newWantedPredTypes wanteds) \case
+        (Nothing, _) -> return Nothing
+        (Just pred, ct) -> do
+          evidence <- newWanted (ctLoc ct) pred
+          return $ Just (mkNonCanonical evidence)
+      let removedWanteds =
+            [ (substEvidence ct ct', ct)
+            | (Just ct', ct) <- zip newWanteds wanteds
+            ]
+      pure $ TcPluginOk removedWanteds (catMaybes newWanteds)
+
+solveCt :: PluginDefs -> Ct -> TcPluginM (Maybe PredType)
+solveCt (PluginDefs {..}) ct@(classifyPredType . ctPred -> IrredPred pred)
+  | Just (tyCon, [arg]) <- splitTyConApp_maybe pred
+  , tyCon == ctIdFam
+  = pure $ Just arg
+solveCt _ ct = pure Nothing
diff --git a/testsuite/tests/tcplugins/Definitions.hs b/testsuite/tests/tcplugins/Definitions.hs
index 5a84967c07..70d04b0296 100644
--- a/testsuite/tests/tcplugins/Definitions.hs
+++ b/testsuite/tests/tcplugins/Definitions.hs
@@ -28,6 +28,9 @@ class MyClass a where
 type MyTyFam :: Type -> Type -> Type
 type family MyTyFam a b where
 
+type CtId :: Constraint -> Constraint
+type family CtId a where
+
 data Nat = Zero | Succ Nat
 
 type Add :: Nat -> Nat -> Nat
diff --git a/testsuite/tests/tcplugins/TcPlugin_CtId.hs b/testsuite/tests/tcplugins/TcPlugin_CtId.hs
new file mode 100644
index 0000000000..14698a9f3f
--- /dev/null
+++ b/testsuite/tests/tcplugins/TcPlugin_CtId.hs
@@ -0,0 +1,13 @@
+{-# OPTIONS_GHC -dcore-lint #-}
+{-# OPTIONS_GHC -fplugin CtIdPlugin #-}
+
+module TcPlugin_CtId where
+
+import Definitions
+  ( CtId )
+
+foo :: CtId (Num a) => a
+foo = 5
+
+bar :: Int
+bar = foo
diff --git a/testsuite/tests/tcplugins/TcPlugin_CtId.stderr b/testsuite/tests/tcplugins/TcPlugin_CtId.stderr
new file mode 100644
index 0000000000..f7e7913f9f
--- /dev/null
+++ b/testsuite/tests/tcplugins/TcPlugin_CtId.stderr
@@ -0,0 +1,4 @@
+[1 of 4] Compiling Common           ( Common.hs, Common.o )
+[2 of 4] Compiling CtIdPlugin       ( CtIdPlugin.hs, CtIdPlugin.o )
+[3 of 4] Compiling Definitions      ( Definitions.hs, Definitions.o )
+[4 of 4] Compiling TcPlugin_CtId    ( TcPlugin_CtId.hs, TcPlugin_CtId.o )
diff --git a/testsuite/tests/tcplugins/all.T b/testsuite/tests/tcplugins/all.T
index 52264e83db..c371deaaa8 100644
--- a/testsuite/tests/tcplugins/all.T
+++ b/testsuite/tests/tcplugins/all.T
@@ -84,3 +84,17 @@ test('TcPlugin_EmitWanted'
     , [ 'TcPlugin_EmitWanted.hs'
       , '-dynamic -package ghc' if have_dynamic() else '-package ghc ' ]
     )
+
+# See TcPlugin_CtId.hs for a description of this plugin.
+test('TcPlugin_CtId'
+    , [ extra_files(
+        [ 'Definitions.hs'
+        , 'Common.hs'
+        , 'CtIdPlugin.hs'
+        , 'TcPlugin_CtId.hs'
+        ])
+      ]
+    , multimod_compile
+    , [ 'TcPlugin_CtId.hs'
+      , '-dynamic -package ghc' if have_dynamic() else '-package ghc' ]
+    )