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

2 files changed, 31 insertions, 22 deletions

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