TcPlugins: `newWanted` uses the provided `CtLoc`

The `GHC.Tc.Plugin.newWanted` function takes a `CtLoc` as an argument,
but it used to discard the location information, keeping only
the `CtOrigin`. It would then retrieve the source location from the
`TcM` environment using `getCtLocM`.

This patch changes this so that `GHC.Tc.Plugin.newWanted` passes on
the full `CtLoc`. This means that authors of type-checking plugins
no longer need to manually set the `CtLoc` environment in the `TcM`
monad if they want to create a new Wanted constraint with the given
`CtLoc` (in particular, for setting the `SrcSpan` of an emitted
constraint). This makes the `newWanted` function consistent with
`newGiven`, which always used the full `CtLoc` instead of using
the environment.

Fixes #20895

4 files changed, 122 insertions, 0 deletions

diff --git a/testsuite/tests/tcplugins/EmitWantedPlugin.hs b/testsuite/tests/tcplugins/EmitWantedPlugin.hs
new file mode 100644
index 0000000000..d5175dc13e
--- /dev/null
+++ b/testsuite/tests/tcplugins/EmitWantedPlugin.hs
@@ -0,0 +1,77 @@
+{-# LANGUAGE RecordWildCards #-}
+
+module EmitWantedPlugin where
+
+-- base
+import Data.Maybe
+  ( catMaybes )
+
+-- ghc
+import GHC.Builtin.Types
+  ( unitTy )
+import GHC.Core
+  ( Expr(Type) )
+import GHC.Core.Class
+  ( Class(..) )
+import GHC.Core.Coercion
+  ( mkSymCo, mkSubCo, mkPrimEqPred )
+import GHC.Core.DataCon
+  ( classDataCon )
+import GHC.Core.Make
+  ( mkCoreConApps, unitExpr )
+import GHC.Core.Type
+  ( eqType )
+import GHC.Core.Utils
+  ( mkCast )
+import GHC.Plugins
+  ( Plugin )
+import GHC.Tc.Plugin
+  ( TcPluginM, newWanted )
+import GHC.Tc.Types
+  ( TcPluginSolveResult(..) )
+import GHC.Tc.Types.Constraint
+  ( Ct(..), ctLoc, ctEvCoercion, mkNonCanonical )
+import GHC.Tc.Types.Evidence
+  ( EvBindsVar, EvTerm(EvExpr) )
+
+-- common
+import Common
+  ( PluginDefs(..), mkPlugin, don'tRewrite )
+
+--------------------------------------------------------------------------------
+
+-- This plugin emits a new Wanted constraint @ty ~# ()@ whenever it encounters
+-- a Wanted constraint of the form @MyClass ty@, and uses the coercion hole
+-- from the @ty ~# ()@ constraint to solve the @MyClass ty@ constraint.
+--
+-- This is used to check that unsolved Wanted constraints are reported
+-- with the correct source location information.
+
+plugin :: Plugin
+plugin = mkPlugin solver don'tRewrite
+
+-- Find Wanteds of the form @MyClass ty@ for some type @ty@,
+-- emits a new Wanted equality @ty ~ ()@, and solves the
+-- @MyClass ty@ constraint using it.
+solver :: [String]
+       -> PluginDefs -> EvBindsVar -> [Ct] -> [Ct] -> [Ct]
+       -> TcPluginM TcPluginSolveResult
+solver args defs _ev _gs _ds ws = do
+  (solved, new) <- unzip . catMaybes <$> traverse ( solveCt defs ) ws
+  pure $ TcPluginOk solved new
+
+solveCt :: PluginDefs -> Ct -> TcPluginM ( Maybe ( (EvTerm, Ct), Ct ) )
+solveCt ( PluginDefs {..} ) ct@( CDictCan { cc_class, cc_tyargs } )
+  | className cc_class == className myClass
+  , [tyArg] <- cc_tyargs
+  = do
+      new_wanted_ctev <- newWanted (ctLoc ct) (mkPrimEqPred tyArg unitTy)
+      let
+        -- co :: tyArg ~# ()
+        co = ctEvCoercion new_wanted_ctev
+        new_wanted_ct = mkNonCanonical new_wanted_ctev
+        ev_term = EvExpr $
+          mkCoreConApps ( classDataCon myClass )
+            [ Type tyArg, mkCast unitExpr (mkSubCo $ mkSymCo co) ]
+      pure $ Just ( ( ev_term, ct ), new_wanted_ct )
+solveCt _ ct = pure Nothing
diff --git a/testsuite/tests/tcplugins/TcPlugin_EmitWanted.hs b/testsuite/tests/tcplugins/TcPlugin_EmitWanted.hs
new file mode 100644
index 0000000000..1b9af19902
--- /dev/null
+++ b/testsuite/tests/tcplugins/TcPlugin_EmitWanted.hs
@@ -0,0 +1,22 @@
+{-# OPTIONS_GHC -dcore-lint #-}
+{-# OPTIONS_GHC -fplugin EmitWantedPlugin #-}
+
+module TcPlugin_EmitWanted where
+
+import Definitions
+  ( MyClass(methC) )
+
+foo :: MyClass a => a
+foo = methC
+
+bar :: Bool
+bar = foo
+  -- We need to solve [W] MyClass a.
+  -- The plugin emits [W] a ~# (), and solves [W] MyClass a
+  -- using the coercion hole.
+  -- We then report an error for the unsolved a ~# () constraint,
+  -- where we get to see whether the source location of the newly
+  -- emitted Wanted constraint is as expected.
+  --
+  -- The crucial thing is that the error message should have
+  -- the correct SrcSpan, in this case line 13 column 7.
diff --git a/testsuite/tests/tcplugins/TcPlugin_EmitWanted.stderr b/testsuite/tests/tcplugins/TcPlugin_EmitWanted.stderr
new file mode 100644
index 0000000000..d85d8bf4ed
--- /dev/null
+++ b/testsuite/tests/tcplugins/TcPlugin_EmitWanted.stderr
@@ -0,0 +1,9 @@
+[1 of 4] Compiling Common           ( Common.hs, Common.o )
+[2 of 4] Compiling Definitions      ( Definitions.hs, Definitions.o )
+[3 of 4] Compiling EmitWantedPlugin ( EmitWantedPlugin.hs, EmitWantedPlugin.o )
+[4 of 4] Compiling TcPlugin_EmitWanted ( TcPlugin_EmitWanted.hs, TcPlugin_EmitWanted.o )
+
+TcPlugin_EmitWanted.hs:13:7: error:
+    • Couldn't match type ‘Bool’ with ‘()’ arising from a use of ‘foo’
+    • In the expression: foo
+      In an equation for ‘bar’: bar = foo
diff --git a/testsuite/tests/tcplugins/all.T b/testsuite/tests/tcplugins/all.T
index 216d46de57..52264e83db 100644
--- a/testsuite/tests/tcplugins/all.T
+++ b/testsuite/tests/tcplugins/all.T
@@ -70,3 +70,17 @@ test('TcPlugin_RewritePerf'
     , [ 'TcPlugin_RewritePerf.hs'
       , '-dynamic -package ghc' if have_dynamic() else '-package ghc' ]
     )
+
+# See EmitWantedPlugin.hs for a description of this plugin.
+test('TcPlugin_EmitWanted'
+    , [ extra_files(
+        [ 'Definitions.hs'
+        , 'Common.hs'
+        , 'EmitWantedPlugin.hs'
+        , 'TcPlugin_EmitWanted.hs'
+        ])
+      ]
+    , multimod_compile_fail
+    , [ 'TcPlugin_EmitWanted.hs'
+      , '-dynamic -package ghc' if have_dynamic() else '-package ghc ' ]
+    )