Add constraint creation functions to TcPluginM API

Summary:
This extends the TcPluginM API with functions to create new constraints,
as described here:
https://ghc.haskell.org/trac/ghc/wiki/Plugins/TypeChecker#Post-7.10changestoTcPluginMAPI

Test Plan: validate and hope

Reviewers: austin, yav, christiaanb

Reviewed By: christiaanb

Subscribers: bgamari, thomie

Differential Revision: https://phabricator.haskell.org/D909

4 files changed, 81 insertions, 15 deletions

diff --git a/compiler/typecheck/TcPluginM.hs b/compiler/typecheck/TcPluginM.hs
index 5acf1b898c..ecf8ed9e45 100644
--- a/compiler/typecheck/TcPluginM.hs
+++ b/compiler/typecheck/TcPluginM.hs
@@ -31,12 +31,24 @@ module TcPluginM (
         matchFam,
 
         -- * Type variables
+        newUnique,
         newFlexiTyVar,
         isTouchableTcPluginM,
 
         -- * Zonking
         zonkTcType,
-        zonkCt
+        zonkCt,
+
+        -- * Creating constraints
+        newWanted,
+        newDerived,
+        newGiven,
+
+        -- * Manipulating evidence bindings
+        newEvVar,
+        setEvBind,
+        getEvBindsTcPluginM,
+        getEvBindsTcPluginM_maybe
 #endif
     ) where
 
@@ -51,11 +63,14 @@ import qualified IfaceEnv
 import qualified Finder
 
 import FamInstEnv ( FamInstEnv )
-import TcRnMonad  ( TcGblEnv, TcLclEnv, Ct, TcPluginM
-                  , unsafeTcPluginTcM, liftIO, traceTc )
+import TcRnMonad  ( TcGblEnv, TcLclEnv, Ct, CtLoc, TcPluginM
+                  , unsafeTcPluginTcM, getEvBindsTcPluginM_maybe
+                  , liftIO, traceTc )
 import TcMType    ( TcTyVar, TcType )
 import TcEnv      ( TcTyThing )
-import TcEvidence ( TcCoercion )
+import TcEvidence ( TcCoercion, EvTerm, EvBind, EvBindsVar, mkGivenEvBind )
+import TcRnTypes  ( CtEvidence(..) )
+import Var        ( EvVar )
 
 import Module
 import Name
@@ -68,6 +83,8 @@ import Type
 import Id
 import InstEnv
 import FastString
+import Maybes
+import Unique
 
 
 -- | Perform some IO, typically to interact with an external tool.
@@ -123,6 +140,9 @@ matchFam :: TyCon -> [Type] -> TcPluginM (Maybe (TcCoercion, TcType))
 matchFam tycon args = unsafeTcPluginTcM $ TcSMonad.matchFamTcM tycon args
 
 
+newUnique :: TcPluginM Unique
+newUnique = unsafeTcPluginTcM TcRnMonad.newUnique
+
 newFlexiTyVar :: Kind -> TcPluginM TcTyVar
 newFlexiTyVar = unsafeTcPluginTcM . TcMType.newFlexiTyVar
 
@@ -135,4 +155,43 @@ zonkTcType = unsafeTcPluginTcM . TcMType.zonkTcType
 
 zonkCt :: Ct -> TcPluginM Ct
 zonkCt = unsafeTcPluginTcM . TcMType.zonkCt
+
+
+-- | Create a new wanted constraint.
+newWanted  :: CtLoc -> PredType -> TcPluginM CtEvidence
+newWanted loc pty = do
+    new_ev <- newEvVar pty
+    return CtWanted { ctev_pred = pty, ctev_evar = new_ev, ctev_loc = loc }
+
+-- | Create a new derived constraint.
+newDerived :: CtLoc -> PredType -> TcPluginM CtEvidence
+newDerived loc pty = return CtDerived { ctev_pred = pty, ctev_loc = loc }
+
+-- | Create a new given constraint, with the supplied evidence.  This
+-- must not be invoked from 'tcPluginInit' or 'tcPluginStop', or it
+-- will panic.
+newGiven :: CtLoc -> PredType -> EvTerm -> TcPluginM CtEvidence
+newGiven loc pty evtm = do
+   new_ev <- newEvVar pty
+   setEvBind $ mkGivenEvBind new_ev evtm
+   return CtGiven { ctev_pred = pty, ctev_evar = new_ev, ctev_loc = loc }
+
+-- | Create a fresh evidence variable.
+newEvVar :: PredType -> TcPluginM EvVar
+newEvVar = unsafeTcPluginTcM . TcMType.newEvVar
+
+-- | Bind an evidence variable.  This must not be invoked from
+-- 'tcPluginInit' or 'tcPluginStop', or it will panic.
+setEvBind :: EvBind -> TcPluginM ()
+setEvBind ev_bind = do
+    tc_evbinds <- getEvBindsTcPluginM
+    unsafeTcPluginTcM $ TcMType.addTcEvBind tc_evbinds ev_bind
+
+-- | Access the 'EvBindsVar' carried by the 'TcPluginM' during
+-- constraint solving.  This must not be invoked from 'tcPluginInit'
+-- or 'tcPluginStop', or it will panic.
+getEvBindsTcPluginM :: TcPluginM EvBindsVar
+getEvBindsTcPluginM = fmap (expectJust oops) getEvBindsTcPluginM_maybe
+  where
+    oops = "plugin attempted to read EvBindsVar outside the constraint solver"
 #endif
diff --git a/compiler/typecheck/TcRnDriver.hs b/compiler/typecheck/TcRnDriver.hs
index ec22699ed1..1df1ca3b8a 100644
--- a/compiler/typecheck/TcRnDriver.hs
+++ b/compiler/typecheck/TcRnDriver.hs
@@ -2157,13 +2157,13 @@ withTcPlugins hsc_env m =
                 -- error occurs during compilation (Fix of #10078)
                 eitherRes <- tryM $ do
                   updGblEnv (\e -> e { tcg_tc_plugins = solvers }) m
-                mapM_ runTcPluginM stops
+                mapM_ (flip runTcPluginM Nothing) stops
                 case eitherRes of
                   Left _ -> failM
                   Right res -> return res
   where
   startPlugin (TcPlugin start solve stop) =
-    do s <- runTcPluginM start
+    do s <- runTcPluginM start Nothing
        return (solve s, stop s)
 
 loadTcPlugins :: HscEnv -> IO [TcPlugin]
diff --git a/compiler/typecheck/TcRnTypes.hs b/compiler/typecheck/TcRnTypes.hs
index 3014755858..5262e186f0 100644
--- a/compiler/typecheck/TcRnTypes.hs
+++ b/compiler/typecheck/TcRnTypes.hs
@@ -84,6 +84,7 @@ module TcRnTypes(
         -- Constraint solver plugins
         TcPlugin(..), TcPluginResult(..), TcPluginSolver,
         TcPluginM, runTcPluginM, unsafeTcPluginTcM,
+        getEvBindsTcPluginM_maybe,
 
         CtFlavour(..), ctEvFlavour,
 
@@ -2209,7 +2210,7 @@ type TcPluginSolver = [Ct]    -- given
                    -> [Ct]    -- wanted
                    -> TcPluginM TcPluginResult
 
-newtype TcPluginM a = TcPluginM (TcM a)
+newtype TcPluginM a = TcPluginM (Maybe EvBindsVar -> TcM a)
 
 instance Functor     TcPluginM where
   fmap = liftM
@@ -2219,21 +2220,27 @@ instance Applicative TcPluginM where
   (<*>) = ap
 
 instance Monad TcPluginM where
-  return x = TcPluginM (return x)
-  fail x   = TcPluginM (fail x)
+  return x = TcPluginM (const $ return x)
+  fail x   = TcPluginM (const $ fail x)
   TcPluginM m >>= k =
-    TcPluginM (do a <- m
-                  let TcPluginM m1 = k a
-                  m1)
+    TcPluginM (\ ev -> do a <- m ev
+                          runTcPluginM (k a) ev)
 
-runTcPluginM :: TcPluginM a -> TcM a
+runTcPluginM :: TcPluginM a -> Maybe EvBindsVar -> TcM a
 runTcPluginM (TcPluginM m) = m
 
 -- | This function provides an escape for direct access to
 -- the 'TcM` monad.  It should not be used lightly, and
 -- the provided 'TcPluginM' API should be favoured instead.
 unsafeTcPluginTcM :: TcM a -> TcPluginM a
-unsafeTcPluginTcM = TcPluginM
+unsafeTcPluginTcM = TcPluginM . const
+
+-- | Access the 'EvBindsVar' carried by the 'TcPluginM' during
+-- constraint solving.  Returns 'Nothing' if invoked during
+-- 'tcPluginInit' or 'tcPluginStop'.
+getEvBindsTcPluginM_maybe :: TcPluginM (Maybe EvBindsVar)
+getEvBindsTcPluginM_maybe = TcPluginM return
+
 
 data TcPlugin = forall s. TcPlugin
   { tcPluginInit  :: TcPluginM s
diff --git a/compiler/typecheck/TcSMonad.hs b/compiler/typecheck/TcSMonad.hs
index 39b01e7d69..3a3f91246f 100644
--- a/compiler/typecheck/TcSMonad.hs
+++ b/compiler/typecheck/TcSMonad.hs
@@ -1241,7 +1241,7 @@ traceTcS :: String -> SDoc -> TcS ()
 traceTcS herald doc = wrapTcS (TcM.traceTc herald doc)
 
 runTcPluginTcS :: TcPluginM a -> TcS a
-runTcPluginTcS = wrapTcS . runTcPluginM
+runTcPluginTcS m = wrapTcS . runTcPluginM m . Just =<< getTcEvBinds
 
 instance HasDynFlags TcS where
     getDynFlags = wrapTcS getDynFlags