partitionTransport utility.

1 files changed, 30 insertions, 0 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 11fe7c32..1346174f 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -69,6 +69,36 @@ layerTransport parse encode tr =
                 sendMessage tr addr msg
        }
 
+-- | Paritions a 'Transport' into two higher-level transports.  Note: An 'MVar'
+-- is used to share the same underlying socket, so be sure to fork a thread for
+-- both returned 'Transport's to avoid hanging.
+partitionTransport :: ((b,a) -> Either (x,xaddr) (y,yaddr))
+                      -> ((x,xaddr) -> (b,a))
+                      -> ((y,yaddr) -> (b,a))
+                      -> Transport err a b
+                      -> IO (Transport err xaddr x, Transport err yaddr y)
+partitionTransport parse encodex encodey tr = do
+    mvar <- newEmptyMVar
+    let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
+                    awaitMessage tr $ \m -> case m of
+                        Just (Right msg) -> either (kont . Just . Right)
+                                                   (\y -> putMVar mvar y >> again)
+                                                $ parse msg
+                        Just (Left e)    -> kont $ Just (Left e)
+                        Nothing          -> kont Nothing
+                 , sendMessage = \addr' msg' -> do
+                    let (msg,addr) = encodex (msg',addr')
+                    sendMessage tr addr msg
+                 }
+        ytr = Transport
+                { awaitMessage = \kont -> takeMVar mvar >>= kont . Just . Right
+                , sendMessage = \addr' msg' -> do
+                    let (msg,addr) = encodey (msg',addr')
+                    sendMessage tr addr msg
+                , closeTransport = return ()
+                }
+    return (xtr, ytr)
+
 addHandler :: (addr -> x -> IO (Maybe (x -> x))) -> Transport err addr x -> Transport err addr x
 addHandler f tr = tr
     { awaitMessage = \kont -> fix $ \eat -> awaitMessage tr $ \m -> do