Remote GHCi: separate out message types

Summary:
From a suggestion by @goldfire: clean up the message types, so that
rather than one Message type with all the messages, we have a separate
THMessage type for messages sent back to GHC during TH execution.  At
the same time I also removed the QDone/QFailed/QException messages
into their own type, and made the result type of RunTH more accurate.

Test Plan: validate

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

Subscribers: thomie, goldfire

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

4 files changed, 129 insertions, 96 deletions

diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index bb9cfb32d0..69cacd58e1 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -926,6 +926,7 @@ finishTH = do
           liftIO $ withForeignRef fhv $ \rhv ->
             writeIServ i (putMessage (FinishTH rhv))
           () <- runRemoteTH i []
+          () <- readQResult i
           writeTcRef (tcg_th_remote_state tcg) Nothing
 
 runTHExp :: ForeignHValue -> TcM TH.Exp
@@ -959,22 +960,20 @@ 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 []
+        runRemoteTH i []
+        bs <- readQResult i
         return $! runGet get (LB.fromStrict bs)
 
--- | communicate with a remotely-running TH computation until it
--- finishes and returns a result.
+
+-- | communicate with a remotely-running TH computation until it finishes
 runRemoteTH
-  :: Binary a
-  => IServ
+  :: IServ
   -> [Messages]   --  saved from nested calls to qRecover
-  -> TcM a
+  -> TcM ()
 runRemoteTH iserv recovers = do
-  Msg msg <- liftIO $ readIServ iserv getMessage
+  THMsg msg <- liftIO $ readIServ iserv getTHMessage
   case msg of
-    QDone -> liftIO $ readIServ iserv get
-    QException str -> liftIO $ throwIO (ErrorCall str)
-    QFail str -> fail str
+    RunTHDone -> return ()
     StartRecover -> do -- Note [TH recover with -fexternal-interpreter]
       v <- getErrsVar
       msgs <- readTcRef v
@@ -994,6 +993,15 @@ runRemoteTH iserv recovers = do
       liftIO $ writeIServ iserv (put r)
       runRemoteTH iserv recovers
 
+-- | Read a value of type QResult from the iserv
+readQResult :: Binary a => IServ -> TcM a
+readQResult i = do
+  qr <- liftIO $ readIServ i get
+  case qr of
+    QDone a -> return a
+    QException str -> liftIO $ throwIO (ErrorCall str)
+    QFail str -> fail str
+
 {- Note [TH recover with -fexternal-interpreter]
 
 Recover is slightly tricky to implement.
@@ -1041,7 +1049,7 @@ wrapTHResult tcm = do
     Left e -> return (THException (show e))
     Right a -> return (THComplete a)
 
-handleTHMessage :: Message a -> TcM a
+handleTHMessage :: THMessage a -> TcM a
 handleTHMessage msg = case msg of
   NewName a -> wrapTHResult $ TH.qNewName a
   Report b str -> wrapTHResult $ TH.qReport b str
diff --git a/iserv/src/Main.hs b/iserv/src/Main.hs
index 46ae82b464..2e4555b017 100644
--- a/iserv/src/Main.hs
+++ b/iserv/src/Main.hs
@@ -58,21 +58,17 @@ serv verbose pipe@Pipe{..} restore = loop
   wrapRunTH :: forall a. (Binary a, Show a) => IO a -> IO ()
   wrapRunTH io = do
     r <- try io
+    writePipe pipe (putTHMessage RunTHDone)
     case r of
       Left e
         | Just (GHCiQException _ err) <- fromException e  -> do
-           when verbose $ putStrLn "iserv: QFail"
-           writePipe pipe (putMessage (QFail err))
-           loop
+           reply (QFail err :: QResult a)
         | otherwise -> do
-           when verbose $ putStrLn "iserv: QException"
            str <- showException e
-           writePipe pipe (putMessage (QException str))
-           loop
+           reply (QException str :: QResult a)
       Right a -> do
         when verbose $ putStrLn "iserv: QDone"
-        writePipe pipe (putMessage QDone)
-        reply a
+        reply (QDone a)
 
   -- carefully when showing an exception, there might be other exceptions
   -- lurking inside it.  If so, we return the inner exception instead.
diff --git a/libraries/ghci/GHCi/Message.hs b/libraries/ghci/GHCi/Message.hs
index b8f9fccd64..b46030f7ea 100644
--- a/libraries/ghci/GHCi/Message.hs
+++ b/libraries/ghci/GHCi/Message.hs
@@ -4,13 +4,15 @@
 
 module GHCi.Message
   ( Message(..), Msg(..)
+  , THMessage(..), THMsg(..)
+  , QResult(..)
   , EvalStatus_(..), EvalStatus, EvalResult(..), EvalOpts(..), EvalExpr(..)
   , SerializableException(..)
   , THResult(..), THResultType(..)
   , ResumeContext(..)
   , QState(..)
-  , getMessage, putMessage
-  , Pipe(..), remoteCall, readPipe, writePipe
+  , getMessage, putMessage, getTHMessage, putTHMessage
+  , Pipe(..), remoteCall, remoteTHCall, readPipe, writePipe
   ) where
 
 import GHCi.RemoteTypes
@@ -162,7 +164,7 @@ data Message a where
   StartTH :: Message (RemoteRef (IORef QState))
 
   -- | Run TH module finalizers, and free the HValueRef
-  FinishTH :: RemoteRef (IORef QState) -> Message ()
+  FinishTH :: RemoteRef (IORef QState) -> Message (QResult ())
 
   -- | Evaluate a TH computation.
   --
@@ -176,39 +178,99 @@ data Message a where
    -> HValueRef {- e.g. TH.Q TH.Exp -}
    -> THResultType
    -> Maybe TH.Loc
-   -> Message ByteString {- e.g. TH.Exp -}
-
-  -- Template Haskell Quasi monad operations
-  NewName :: String -> Message (THResult TH.Name)
-  Report :: Bool -> String -> Message (THResult ())
-  LookupName :: Bool -> String -> Message (THResult (Maybe TH.Name))
-  Reify :: TH.Name -> Message (THResult TH.Info)
-  ReifyFixity :: TH.Name -> Message (THResult (Maybe TH.Fixity))
-  ReifyInstances :: TH.Name -> [TH.Type] -> Message (THResult [TH.Dec])
-  ReifyRoles :: TH.Name -> Message (THResult [TH.Role])
-  ReifyAnnotations :: TH.AnnLookup -> TypeRep -> Message (THResult [ByteString])
-  ReifyModule :: TH.Module -> Message (THResult TH.ModuleInfo)
-  ReifyConStrictness :: TH.Name -> Message (THResult [TH.DecidedStrictness])
-
-  AddDependentFile :: FilePath -> Message (THResult ())
-  AddTopDecls :: [TH.Dec] -> Message (THResult ())
-  IsExtEnabled :: Extension -> Message (THResult Bool)
-  ExtsEnabled :: Message (THResult [Extension])
-
-  StartRecover :: Message ()
-  EndRecover :: Bool -> Message ()
-
-  -- Template Haskell return values
-
-  -- | RunTH finished successfully; return value follows
-  QDone :: Message ()
-  -- | RunTH threw an exception
-  QException :: String -> Message ()
-  -- | RunTH called 'fail'
-  QFail :: String -> Message ()
+   -> Message (QResult ByteString)
+
 
 deriving instance Show (Message a)
 
+
+-- | Template Haskell return values
+data QResult a
+  = QDone a
+    -- ^ RunTH finished successfully; return value follows
+  | QException String
+    -- ^ RunTH threw an exception
+  | QFail String
+    -- ^ RunTH called 'fail'
+  deriving (Generic, Show)
+
+instance Binary a => Binary (QResult a)
+
+
+-- | Messages sent back to GHC from GHCi.TH, to implement the methods
+-- of 'Quasi'.
+data THMessage a where
+  NewName :: String -> THMessage (THResult TH.Name)
+  Report :: Bool -> String -> THMessage (THResult ())
+  LookupName :: Bool -> String -> THMessage (THResult (Maybe TH.Name))
+  Reify :: TH.Name -> THMessage (THResult TH.Info)
+  ReifyFixity :: TH.Name -> THMessage (THResult (Maybe TH.Fixity))
+  ReifyInstances :: TH.Name -> [TH.Type] -> THMessage (THResult [TH.Dec])
+  ReifyRoles :: TH.Name -> THMessage (THResult [TH.Role])
+  ReifyAnnotations :: TH.AnnLookup -> TypeRep
+    -> THMessage (THResult [ByteString])
+  ReifyModule :: TH.Module -> THMessage (THResult TH.ModuleInfo)
+  ReifyConStrictness :: TH.Name -> THMessage (THResult [TH.DecidedStrictness])
+
+  AddDependentFile :: FilePath -> THMessage (THResult ())
+  AddTopDecls :: [TH.Dec] -> THMessage (THResult ())
+  IsExtEnabled :: Extension -> THMessage (THResult Bool)
+  ExtsEnabled :: THMessage (THResult [Extension])
+
+  StartRecover :: THMessage ()
+  EndRecover :: Bool -> THMessage ()
+
+  -- | Indicates that this RunTH is finished, and the next message
+  -- will be the result of RunTH (a QResult).
+  RunTHDone :: THMessage ()
+
+deriving instance Show (THMessage a)
+
+data THMsg = forall a . (Binary a, Show a) => THMsg (THMessage a)
+
+getTHMessage :: Get THMsg
+getTHMessage = do
+  b <- getWord8
+  case b of
+    0  -> THMsg <$> NewName <$> get
+    1  -> THMsg <$> (Report <$> get <*> get)
+    2  -> THMsg <$> (LookupName <$> get <*> get)
+    3  -> THMsg <$> Reify <$> get
+    4  -> THMsg <$> ReifyFixity <$> get
+    5  -> THMsg <$> (ReifyInstances <$> get <*> get)
+    6  -> THMsg <$> ReifyRoles <$> get
+    7  -> THMsg <$> (ReifyAnnotations <$> get <*> get)
+    8  -> THMsg <$> ReifyModule <$> get
+    9  -> THMsg <$> ReifyConStrictness <$> get
+    10 -> THMsg <$> AddDependentFile <$> get
+    11 -> THMsg <$> AddTopDecls <$> get
+    12 -> THMsg <$> (IsExtEnabled <$> get)
+    13 -> THMsg <$> return ExtsEnabled
+    14 -> THMsg <$> return StartRecover
+    15 -> THMsg <$> EndRecover <$> get
+    _  -> return (THMsg RunTHDone)
+
+putTHMessage :: THMessage a -> Put
+putTHMessage m = case m of
+  NewName a                   -> putWord8 0  >> put a
+  Report a b                  -> putWord8 1  >> put a >> put b
+  LookupName a b              -> putWord8 2  >> put a >> put b
+  Reify a                     -> putWord8 3  >> put a
+  ReifyFixity a               -> putWord8 4  >> put a
+  ReifyInstances a b          -> putWord8 5  >> put a >> put b
+  ReifyRoles a                -> putWord8 6  >> put a
+  ReifyAnnotations a b        -> putWord8 7  >> put a >> put b
+  ReifyModule a               -> putWord8 8  >> put a
+  ReifyConStrictness a        -> putWord8 9  >> put a
+  AddDependentFile a          -> putWord8 10 >> put a
+  AddTopDecls a               -> putWord8 11 >> put a
+  IsExtEnabled a              -> putWord8 12 >> put a
+  ExtsEnabled                 -> putWord8 13
+  StartRecover                -> putWord8 14
+  EndRecover a                -> putWord8 15 >> put a
+  RunTHDone                   -> putWord8 16
+
+
 data EvalOpts = EvalOpts
   { useSandboxThread :: Bool
   , singleStep :: Bool
@@ -341,26 +403,7 @@ getMessage = do
       30 -> Msg <$> (GetBreakpointVar <$> get <*> get)
       31 -> Msg <$> return StartTH
       32 -> Msg <$> FinishTH <$> get
-      33 -> Msg <$> (RunTH <$> get <*> get <*> get <*> get)
-      34 -> Msg <$> NewName <$> get
-      35 -> Msg <$> (Report <$> get <*> get)
-      36 -> Msg <$> (LookupName <$> get <*> get)
-      37 -> Msg <$> Reify <$> get
-      38 -> Msg <$> ReifyFixity <$> get
-      39 -> Msg <$> (ReifyInstances <$> get <*> get)
-      40 -> Msg <$> ReifyRoles <$> get
-      41 -> Msg <$> (ReifyAnnotations <$> get <*> get)
-      42 -> Msg <$> ReifyModule <$> get
-      43 -> Msg <$> ReifyConStrictness <$> get
-      44 -> Msg <$> AddDependentFile <$> get
-      45 -> Msg <$> AddTopDecls <$> get
-      46 -> Msg <$> (IsExtEnabled <$> get)
-      47 -> Msg <$> return ExtsEnabled
-      48 -> Msg <$> return StartRecover
-      49 -> Msg <$> EndRecover <$> get
-      50 -> Msg <$> return QDone
-      51 -> Msg <$> QException <$> get
-      _  -> Msg <$> QFail <$> get
+      _  -> Msg <$> (RunTH <$> get <*> get <*> get <*> get)
 
 putMessage :: Message a -> Put
 putMessage m = case m of
@@ -398,25 +441,6 @@ putMessage m = case m of
   StartTH                     -> putWord8 31
   FinishTH val                -> putWord8 32 >> put val
   RunTH st q loc ty           -> putWord8 33 >> put st >> put q >> put loc >> put ty
-  NewName a                   -> putWord8 34 >> put a
-  Report a b                  -> putWord8 35 >> put a >> put b
-  LookupName a b              -> putWord8 36 >> put a >> put b
-  Reify a                     -> putWord8 37 >> put a
-  ReifyFixity a               -> putWord8 38 >> put a
-  ReifyInstances a b          -> putWord8 39 >> put a >> put b
-  ReifyRoles a                -> putWord8 40 >> put a
-  ReifyAnnotations a b        -> putWord8 41 >> put a >> put b
-  ReifyModule a               -> putWord8 42 >> put a
-  ReifyConStrictness a        -> putWord8 43 >> put a
-  AddDependentFile a          -> putWord8 44 >> put a
-  AddTopDecls a               -> putWord8 45 >> put a
-  IsExtEnabled a              -> putWord8 46 >> put a
-  ExtsEnabled                 -> putWord8 47
-  StartRecover                -> putWord8 48
-  EndRecover a                -> putWord8 49 >> put a
-  QDone                       -> putWord8 50
-  QException a                -> putWord8 51 >> put a
-  QFail a                     -> putWord8 52  >> put a
 
 -- -----------------------------------------------------------------------------
 -- Reading/writing messages
@@ -432,6 +456,11 @@ remoteCall pipe msg = do
   writePipe pipe (putMessage msg)
   readPipe pipe get
 
+remoteTHCall :: Binary a => Pipe -> THMessage a -> IO a
+remoteTHCall pipe msg = do
+  writePipe pipe (putTHMessage msg)
+  readPipe pipe get
+
 writePipe :: Pipe -> Put -> IO ()
 writePipe Pipe{..} put
   | LB.null bs = return ()
diff --git a/libraries/ghci/GHCi/TH.hs b/libraries/ghci/GHCi/TH.hs
index 69f114cfd2..6d6158ffdb 100644
--- a/libraries/ghci/GHCi/TH.hs
+++ b/libraries/ghci/GHCi/TH.hs
@@ -75,9 +75,9 @@ putState s = GHCiQ $ \_ -> return ((),s)
 noLoc :: TH.Loc
 noLoc = TH.Loc "<no file>" "<no package>" "<no module>" (0,0) (0,0)
 
-ghcCmd :: Binary a => Message (THResult a) -> GHCiQ a
+ghcCmd :: Binary a => THMessage (THResult a) -> GHCiQ a
 ghcCmd m = GHCiQ $ \s -> do
-  r <- remoteCall (qsPipe s) m
+  r <- remoteTHCall (qsPipe s) m
   case r of
     THException str -> throwIO (GHCiQException s str)
     THComplete res -> return (res, s)
@@ -88,12 +88,12 @@ instance TH.Quasi GHCiQ where
 
   -- See Note [TH recover with -fexternal-interpreter] in TcSplice
   qRecover (GHCiQ h) (GHCiQ a) = GHCiQ $ \s -> (do
-    remoteCall (qsPipe s) StartRecover
+    remoteTHCall (qsPipe s) StartRecover
     (r, s') <- a s
-    remoteCall (qsPipe s) (EndRecover False)
+    remoteTHCall (qsPipe s) (EndRecover False)
     return (r,s'))
       `catch`
-       \GHCiQException{} -> remoteCall (qsPipe s) (EndRecover True) >> h s
+       \GHCiQException{} -> remoteTHCall (qsPipe s) (EndRecover True) >> h s
   qLookupName isType occ = ghcCmd (LookupName isType occ)
   qReify name = ghcCmd (Reify name)
   qReifyFixity name = ghcCmd (ReifyFixity name)