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

7 files changed, 149 insertions, 10 deletions

diff --git a/compiler/GHC/Tc/Plugin.hs b/compiler/GHC/Tc/Plugin.hs
index 2edee72207..8984bffce9 100644
--- a/compiler/GHC/Tc/Plugin.hs
+++ b/compiler/GHC/Tc/Plugin.hs
@@ -66,7 +66,7 @@ import GHC.Core.FamInstEnv     ( FamInstEnv )
 import GHC.Tc.Utils.Monad      ( TcGblEnv, TcLclEnv, TcPluginM
                                , unsafeTcPluginTcM
                                , liftIO, traceTc )
-import GHC.Tc.Types.Constraint ( Ct, CtLoc, CtEvidence(..), ctLocOrigin )
+import GHC.Tc.Types.Constraint ( Ct, CtLoc, CtEvidence(..) )
 import GHC.Tc.Utils.TcMType    ( TcTyVar, TcType )
 import GHC.Tc.Utils.Env        ( TcTyThing )
 import GHC.Tc.Types.Evidence   ( CoercionHole, EvTerm(..)
@@ -162,10 +162,10 @@ zonkTcType = unsafeTcPluginTcM . TcM.zonkTcType
 zonkCt :: Ct -> TcPluginM Ct
 zonkCt = unsafeTcPluginTcM . TcM.zonkCt
 
--- | Create a new wanted constraint.
+-- | Create a new Wanted constraint with the given 'CtLoc'.
 newWanted :: CtLoc -> PredType -> TcPluginM CtEvidence
 newWanted loc pty
-  = unsafeTcPluginTcM (TcM.newWanted (ctLocOrigin loc) Nothing pty)
+  = unsafeTcPluginTcM (TcM.newWantedWithLoc loc pty)
 
 -- | Create a new derived constraint.
 newDerived :: CtLoc -> PredType -> TcPluginM CtEvidence
diff --git a/compiler/GHC/Tc/Utils/TcMType.hs b/compiler/GHC/Tc/Utils/TcMType.hs
index 699bf5f69a..4565dabab9 100644
--- a/compiler/GHC/Tc/Utils/TcMType.hs
+++ b/compiler/GHC/Tc/Utils/TcMType.hs
@@ -36,7 +36,7 @@ module GHC.Tc.Utils.TcMType (
   --------------------------------
   -- Creating new evidence variables
   newEvVar, newEvVars, newDict,
-  newWanted, newWanteds, cloneWanted, cloneWC,
+  newWantedWithLoc, newWanted, newWanteds, cloneWanted, cloneWC,
   emitWanted, emitWantedEq, emitWantedEvVar, emitWantedEvVars,
   emitDerivedEqs,
   newTcEvBinds, newNoTcEvBinds, addTcEvBind,
@@ -186,17 +186,26 @@ newEvVar :: TcPredType -> TcRnIf gbl lcl EvVar
 newEvVar ty = do { name <- newSysName (predTypeOccName ty)
                  ; return (mkLocalIdOrCoVar name Many ty) }
 
-newWanted :: CtOrigin -> Maybe TypeOrKind -> PredType -> TcM CtEvidence
--- Deals with both equality and non-equality predicates
-newWanted orig t_or_k pty
-  = do loc <- getCtLocM orig t_or_k
-       d <- if isEqPrimPred pty then HoleDest  <$> newCoercionHole pty
+-- | Create a new Wanted constraint with the given 'CtLoc'.
+newWantedWithLoc :: CtLoc -> PredType -> TcM CtEvidence
+newWantedWithLoc loc pty
+  = do d <- if isEqPrimPred pty then HoleDest  <$> newCoercionHole pty
                                 else EvVarDest <$> newEvVar pty
        return $ CtWanted { ctev_dest = d
                          , ctev_pred = pty
                          , ctev_nosh = WDeriv
-                         , ctev_loc = loc }
+                         , ctev_loc  = loc }
+
+-- | Create a new Wanted constraint with the given 'CtOrigin', and
+-- location information taken from the 'TcM' environment.
+newWanted :: CtOrigin -> Maybe TypeOrKind -> PredType -> TcM CtEvidence
+-- Deals with both equality and non-equality predicates
+newWanted orig t_or_k pty
+  = do loc <- getCtLocM orig t_or_k
+       newWantedWithLoc loc pty
 
+-- | Create new Wanted constraints with the given 'CtOrigin',
+-- and location information taken from the 'TcM' environment.
 newWanteds :: CtOrigin -> ThetaType -> TcM [CtEvidence]
 newWanteds orig = mapM (newWanted orig Nothing)
 
diff --git a/docs/users_guide/9.4.1-notes.rst b/docs/users_guide/9.4.1-notes.rst
index bb70761f1d..31f09fa6bc 100644
--- a/docs/users_guide/9.4.1-notes.rst
+++ b/docs/users_guide/9.4.1-notes.rst
@@ -109,8 +109,16 @@ Compiler
   for computing the ``Type`` of an ``HsExpr GhcTc`` in a pure fashion.
   The ``hsLitType`` and ``hsPatType`` functions that previously lived in
   ``GHC.Tc.Utils.Zonk`` have been moved to this module.
+
 - A ``Typeable`` constraint has been added to ``fromStaticPtr`` in the
   class ``GHC.StaticPtr.IsStatic``. GHC automatically wraps each use of
   the ``static`` keyword with ``fromStaticPtr``. Because ``static`` requires
   its argument to be an instance of ``Typeable``, ``fromStaticPtr`` can
   safely carry this constraint as well.
+
+- The ``newWanted`` function exported by ``GHC.Tc.Plugin`` now passes on
+  the full ``CtLoc`` instead of reconstituting it from the type-checking
+  environment. This makes ``newWanted`` consistent with ``newGiven``.
+  For authors of type-checking plugins, this means you don't need to wrap
+  a call to ``newWanted`` in ``setCtLocM`` to create a new Wanted constraint
+  with the provided ``CtLoc``.
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 ' ]
+    )