Switched awaitMessage to continuation-passing style.

1 files changed, 4 insertions, 5 deletions

diff --git a/Mainline.hs b/Mainline.hs
index 291a196f..860372dc 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -419,11 +419,10 @@ showPacket f addr flow bs = L8.unpack $ L8.unlines es
 
 -- Add detailed printouts for every packet.
 addVerbosity tr =
-    tr { awaitMessage = do
-            m <- awaitMessage tr
+    tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
             forM_ m $ mapM_ $ \(msg,addr) -> do
                 hPutStrLn stderr (showPacket id addr " --> " msg)
-            return m
+            kont m
        , sendMessage = \addr msg -> do
             hPutStrLn stderr (showPacket id addr " <-- " msg)
             sendMessage tr addr msg
@@ -566,7 +565,7 @@ newClient addr = do
         -- recursive since 'updateRouting' does not invoke 'awaitMessage' which
         -- which was modified by 'onInbound'.  However, I'm going to avoid the
         -- mutual reference just to be safe.
-        outgoingClient = client { clientNet = net { awaitMessage = return Nothing } }
+        outgoingClient = client { clientNet = net { awaitMessage = ($ Nothing) } }
 
         dispatch = DispatchMethods
             { classifyInbound = classify -- :: x -> MessageClass err meth tid
@@ -587,7 +586,7 @@ newClient addr = do
         gen cnt = (TransactionId $ S.encode cnt, cnt+1)
 
         client = Client
-            { clientNet           = net
+            { clientNet           = addHandler (handleMessage client) net
             , clientDispatcher    = dispatch
             , clientErrorReporter = ignoreErrors -- printErrors stderr
             , clientPending       = map_var