1 files changed, 38 insertions, 9 deletions

diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index ba686929..d8dc8bfa 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -21,8 +21,12 @@ import Control.Exception
 import Control.Monad
 import qualified Data.ByteString    as B
          ;import Data.ByteString    (ByteString)
+import Data.Dependent.Map as DMap
+import Data.Dependent.Sum
 import Data.Function
 import Data.Functor.Contravariant
+import Data.Functor.Identity
+import Data.GADT.Show
 import qualified Data.IntMap.Strict as IntMap
          ;import Data.IntMap.Strict (IntMap)
 import qualified Data.Map.Strict    as Map
@@ -40,6 +44,7 @@ import System.Endian
 import System.IO
 import System.IO.Error
 import System.Timeout
+
 import DPut
 import DebugTag
 import Data.TableMethods
@@ -122,22 +127,22 @@ partitionTransportM :: ((b,a) -> IO (Either (x,xaddr) (b,a)))
                       -> IO (Transport err xaddr x, Transport err a b)
 partitionTransportM parse encodex tr = do
     tchan <- atomically newTChan
-    let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
+    let ytr = tr { awaitMessage = \kont -> fix $ \again -> do
                     awaitMessage tr $ \m -> case m of
                         Arrival adr msg  -> parse (msg,adr) >>= \case
-                            Left (x,xaddr) -> kont $ Arrival xaddr x
-                            Right y        -> atomically (writeTChan tchan (Just y)) >> join (atomically again)
+                            Left x          -> atomically (writeTChan tchan (Just x)) >> join (atomically again)
+                            Right (y,yaddr) -> kont $ Arrival yaddr y
                         ParseError e     -> kont $ ParseError e
                         Terminated       -> atomically (writeTChan tchan Nothing) >> kont Terminated
-                 , sendMessage = \addr' msg' -> do
-                    msg_addr <- encodex (msg',addr')
-                    mapM_ (uncurry . flip $ sendMessage tr) msg_addr
+                 , sendMessage = sendMessage tr
                  }
-        ytr = Transport
+        xtr = Transport
                 { awaitMessage = \kont -> readTChan tchan >>= pure . kont . \case
                     Nothing        -> Terminated
-                    Just (y,yaddr) -> Arrival yaddr y
-                , sendMessage = sendMessage tr
+                    Just (x,xaddr) -> Arrival xaddr x
+                , sendMessage = \addr' msg' -> do
+                    msg_addr <- encodex (msg',addr')
+                    mapM_ (uncurry . flip $ sendMessage tr) msg_addr
                 , setActive = const $ return ()
                 }
     return (xtr, ytr)
@@ -687,3 +692,27 @@ testPairTransport = do
         b = SockAddrInet 2 2
     return ( chanTransport (const bchan) a achan aclosed
            , chanTransport (const achan) b bchan bclosed )
+
+newtype ByAddress err x addr = ByAddress (Transport err addr x)
+
+newtype Tagged x addr = Tagged x
+
+decorateAddr :: tag addr -> Arrival e addr x -> Arrival e (DSum tag Identity) x
+decorateAddr tag Terminated       = Terminated
+decorateAddr tag (ParseError e)   = ParseError e
+decorateAddr tag (Arrival addr x) = Arrival (tag ==> addr) x
+
+mergeTransports :: GCompare tag => DMap tag (ByAddress err x) -> IO (Transport err (DSum tag Identity) x)
+mergeTransports tmap = do
+    -- vmap <- traverseWithKey (\k v -> Tagged <$> newEmptyMVar) tmap
+    -- foldrWithKey (\k v n -> forkMergeBranch k v >> n) (return ()) vmap
+    return Transport
+        { awaitMessage = \kont ->
+            foldrWithKey (\k (ByAddress tr) n -> awaitMessage tr (kont . decorateAddr k) `orElse` n)
+                         retry
+                         tmap
+        , sendMessage = \(tag :=> Identity addr) x -> case DMap.lookup tag tmap of
+            Just (ByAddress tr) -> sendMessage tr addr x
+            Nothing             -> return ()
+        , setActive = \toggle -> foldrWithKey (\_ (ByAddress tr) next -> setActive tr toggle >> next) (return ()) tmap
+        }