Support for qRecover in TH with -fexternal-interpreter

Summary: This completes the support for TH with -fexternal-interpreter.

Test Plan: validate

Reviewers: bgamari, ezyang, austin, niteria, goldfire, erikd

Subscribers: thomie

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

GHC Trac Issues: #11100

5 files changed, 92 insertions, 33 deletions

diff --git a/compiler/main/ErrUtils.hs b/compiler/main/ErrUtils.hs
index 0677240522..11b30fd13c 100644
--- a/compiler/main/ErrUtils.hs
+++ b/compiler/main/ErrUtils.hs
@@ -14,6 +14,7 @@ module ErrUtils (
         -- * Messages
         MsgDoc, ErrMsg, ErrDoc, errDoc, WarnMsg,
         Messages, ErrorMessages, WarningMessages,
+        unionMessages,
         errMsgSpan, errMsgContext,
         errorsFound, isEmptyMessages,
 
@@ -48,7 +49,7 @@ module ErrUtils (
 
 #include "HsVersions.h"
 
-import Bag              ( Bag, bagToList, isEmptyBag, emptyBag )
+import Bag
 import Exception
 import Outputable
 import Panic
@@ -100,6 +101,10 @@ type Messages        = (WarningMessages, ErrorMessages)
 type WarningMessages = Bag WarnMsg
 type ErrorMessages   = Bag ErrMsg
 
+unionMessages :: Messages -> Messages -> Messages
+unionMessages (warns1, errs1) (warns2, errs2) =
+  (warns1 `unionBags` warns2, errs1 `unionBags` errs2)
+
 data ErrMsg = ErrMsg {
         errMsgSpan        :: SrcSpan,
         errMsgContext     :: PrintUnqualified,
diff --git a/compiler/typecheck/TcRnMonad.hs b/compiler/typecheck/TcRnMonad.hs
index 7ce60bc852..f55f5dd548 100644
--- a/compiler/typecheck/TcRnMonad.hs
+++ b/compiler/typecheck/TcRnMonad.hs
@@ -722,18 +722,17 @@ warnIf True  msg = addWarn msg
 warnIf False _   = return ()
 
 addMessages :: Messages -> TcRn ()
-addMessages (m_warns, m_errs)
+addMessages msgs1
   = do { errs_var <- getErrsVar ;
-         (warns, errs) <- readTcRef errs_var ;
-         writeTcRef errs_var (warns `unionBags` m_warns,
-                               errs  `unionBags` m_errs) }
+         msgs0 <- readTcRef errs_var ;
+         writeTcRef errs_var (unionMessages msgs0 msgs1) }
 
 discardWarnings :: TcRn a -> TcRn a
 -- Ignore warnings inside the thing inside;
 -- used to ignore-unused-variable warnings inside derived code
 discardWarnings thing_inside
   = do  { errs_var <- getErrsVar
-        ; (old_warns, _) <- readTcRef errs_var ;
+        ; (old_warns, _) <- readTcRef errs_var
 
         ; result <- thing_inside
 
diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index cdb47901c0..d24de8bae0 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -915,7 +915,7 @@ finishTH = do
         Just fhv -> do
           liftIO $ withForeignRef fhv $ \rhv ->
             writeIServ i (putMessage (FinishTH rhv))
-          () <- runRemoteTH i
+          () <- runRemoteTH i []
           writeTcRef (tcg_th_remote_state tcg) Nothing
 
 runTHExp :: ForeignHValue -> TcM TH.Exp
@@ -949,22 +949,68 @@ runTH ty fhv = do
           withForeignRef rstate $ \state_hv ->
           withForeignRef fhv $ \q_hv ->
             writeIServ i (putMessage (RunTH state_hv q_hv ty (Just loc)))
-        bs <- runRemoteTH i
+        bs <- runRemoteTH i []
         return $! runGet get (LB.fromStrict bs)
 
 -- | communicate with a remotely-running TH computation until it
 -- finishes and returns a result.
-runRemoteTH :: Binary a => IServ -> TcM a
-runRemoteTH iserv = do
+runRemoteTH
+  :: Binary a
+  => IServ
+  -> [Messages]   --  saved from nested calls to qRecover
+  -> TcM a
+runRemoteTH iserv recovers = do
   Msg msg <- liftIO $ readIServ iserv getMessage
   case msg of
     QDone -> liftIO $ readIServ iserv get
     QException str -> liftIO $ throwIO (ErrorCall str)
     QFail str -> fail str
+    StartRecover -> do -- Note [TH recover with -fexternal-interpreter]
+      v <- getErrsVar
+      msgs <- readTcRef v
+      writeTcRef v emptyMessages
+      runRemoteTH iserv (msgs : recovers)
+    EndRecover caught_error -> do
+      v <- getErrsVar
+      let (prev_msgs, rest) = case recovers of
+             [] -> panic "EndRecover"
+             a : b -> (a,b)
+      if caught_error
+        then writeTcRef v prev_msgs
+        else updTcRef v (unionMessages prev_msgs)
+      runRemoteTH iserv rest
     _other -> do
       r <- handleTHMessage msg
       liftIO $ writeIServ iserv (put r)
-      runRemoteTH iserv
+      runRemoteTH iserv recovers
+
+{- Note [TH recover with -fexternal-interpreter]
+
+Recover is slightly tricky to implement.
+
+The meaning of "recover a b" is
+ - Do a
+   - If it finished successfully, then keep the messages it generated
+   - If it failed, discard any messages it generated, and do b
+
+The messages are managed by GHC in the TcM monad, whereas the
+exception-handling is done in the ghc-iserv process, so we have to
+coordinate between the two.
+
+On the server:
+  - emit a StartRecover message
+  - run "a" inside a catch
+    - if it finishes, emit EndRecover False
+    - if it fails, emit EndRecover True, then run "b"
+
+Back in GHC, when we receive:
+
+  StartRecover
+    save the current messages and start with an empty set.
+  EndRecover caught_error
+    Restore the previous messages,
+    and merge in the new messages if caught_error is false.
+-}
 
 getTHState :: IServ -> TcM (ForeignRef (IORef QState))
 getTHState i = do
diff --git a/libraries/ghci/GHCi/Message.hs b/libraries/ghci/GHCi/Message.hs
index 59d6483089..4bc2d25b66 100644
--- a/libraries/ghci/GHCi/Message.hs
+++ b/libraries/ghci/GHCi/Message.hs
@@ -191,6 +191,9 @@ data Message a where
   IsExtEnabled :: Extension -> Message (THResult Bool)
   ExtsEnabled :: Message (THResult [Extension])
 
+  StartRecover :: Message ()
+  EndRecover :: Bool -> Message ()
+
   -- Template Haskell return values
 
   -- | RunTH finished successfully; return value follows
@@ -347,8 +350,10 @@ getMessage = do
       43 -> Msg <$> AddTopDecls <$> get
       44 -> Msg <$> (IsExtEnabled <$> get)
       45 -> Msg <$> return ExtsEnabled
-      46 -> Msg <$> return QDone
-      47 -> Msg <$> QException <$> get
+      46 -> Msg <$> return StartRecover
+      47 -> Msg <$> EndRecover <$> get
+      48 -> Msg <$> return QDone
+      49 -> Msg <$> QException <$> get
       _  -> Msg <$> QFail <$> get
 
 putMessage :: Message a -> Put
@@ -399,9 +404,11 @@ putMessage m = case m of
   AddTopDecls a               -> putWord8 43 >> put a
   IsExtEnabled a              -> putWord8 44 >> put a
   ExtsEnabled                 -> putWord8 45
-  QDone                       -> putWord8 46
-  QException a                -> putWord8 47 >> put a
-  QFail a                     -> putWord8 48 >> put a
+  StartRecover                -> putWord8 46
+  EndRecover a                -> putWord8 47 >> put a
+  QDone                       -> putWord8 48
+  QException a                -> putWord8 49 >> put a
+  QFail a                     -> putWord8 50 >> put a
 
 -- -----------------------------------------------------------------------------
 -- Reading/writing messages
diff --git a/libraries/ghci/GHCi/TH.hs b/libraries/ghci/GHCi/TH.hs
index 799bd6261b..2c7a50172c 100644
--- a/libraries/ghci/GHCi/TH.hs
+++ b/libraries/ghci/GHCi/TH.hs
@@ -81,27 +81,28 @@ ghcCmd m = GHCiQ $ \s -> do
 instance TH.Quasi GHCiQ where
   qNewName str = ghcCmd (NewName str)
   qReport isError msg = ghcCmd (Report isError msg)
-  qRecover = undefined
-{-
-  qRecover (GHCiQ h) (GHCiQ a) = GHCiQ $ \s -> do
-    let r :: Bool -> IO ()
-        r b = do EndRecover' <- sendRequest (EndRecover b)
-                 return ()
-    StartRecover' <- sendRequest StartRecover
-    (a s >>= \s' -> r False >> return s') `E.catch`
-      \(GHCiQException s' _ _) -> r True >> h s
--}
+
+  -- See Note [TH recover with -fexternal-interpreter] in TcSplice
+  qRecover (GHCiQ h) (GHCiQ a) = GHCiQ $ \s -> (do
+    remoteCall (qsPipe s) StartRecover
+    (r, s') <- a s
+    remoteCall (qsPipe s) (EndRecover False)
+    return (r,s'))
+      `catch`
+       \GHCiQException{} -> remoteCall (qsPipe s) (EndRecover True) >> h s
   qLookupName isType occ = ghcCmd (LookupName isType occ)
   qReify name = ghcCmd (Reify name)
   qReifyFixity name = ghcCmd (ReifyFixity name)
   qReifyInstances name tys = ghcCmd (ReifyInstances name tys)
   qReifyRoles name = ghcCmd (ReifyRoles name)
 
-  -- To reify annotations, we send GHC the AnnLookup and also the TypeRep of the
-  -- thing we're looking for, to avoid needing to serialize irrelevant annotations.
+  -- To reify annotations, we send GHC the AnnLookup and also the
+  -- TypeRep of the thing we're looking for, to avoid needing to
+  -- serialize irrelevant annotations.
   qReifyAnnotations :: forall a . Data a => TH.AnnLookup -> GHCiQ [a]
   qReifyAnnotations lookup =
-    map (deserializeWithData . B.unpack) <$> ghcCmd (ReifyAnnotations lookup typerep)
+    map (deserializeWithData . B.unpack) <$>
+      ghcCmd (ReifyAnnotations lookup typerep)
     where typerep = typeOf (undefined :: a)
 
   qReifyModule m = ghcCmd (ReifyModule m)
@@ -149,11 +150,12 @@ runTH pipe rstate rhv ty mb_loc = do
     THAnnWrapper -> do
       hv <- unsafeCoerce <$> localRef rhv
       case hv :: AnnotationWrapper of
-        AnnotationWrapper thing ->
-          return $! LB.toStrict (runPut (put (toSerialized serializeWithData thing)))
+        AnnotationWrapper thing -> return $!
+          LB.toStrict (runPut (put (toSerialized serializeWithData thing)))
 
-runTHQ :: Binary a => Pipe -> RemoteRef (IORef QState) -> Maybe TH.Loc -> TH.Q a
-       -> IO ByteString
+runTHQ
+  :: Binary a => Pipe -> RemoteRef (IORef QState) -> Maybe TH.Loc -> TH.Q a
+  -> IO ByteString
 runTHQ pipe@Pipe{..} rstate mb_loc ghciq = do
   qstateref <- localRef rstate
   qstate <- readIORef qstateref