1 files changed, 18 insertions, 11 deletions

diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 79cc9489..854f26c7 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 module Network.BitTorrent.DHT.Search where
 
 import Control.Concurrent
@@ -21,20 +22,23 @@ import qualified Data.MinMaxPSQ   as MM
          ;import Data.MinMaxPSQ   (MinMaxPSQ)
 import qualified Data.Wrapper.PSQ as PSQ
          ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ)
-import Network.BitTorrent.Address
+import Network.BitTorrent.Address hiding (NodeId)
+import Network.RPC
+import Network.KRPC.Message (KMessageOf)
+import Network.DHT.Mainline ()
 
 data IterativeSearch ip r = IterativeSearch
-    { searchTarget       :: NodeId
-    , searchQuery        :: NodeInfo ip -> IO ([NodeInfo ip], [r])
+    { searchTarget       :: NodeId KMessageOf
+    , searchQuery        :: NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [r])
     , searchPendingCount :: TVar Int
-    , searchQueued       :: TVar (MinMaxPSQ (NodeInfo ip) NodeDistance)
-    , searchInformant    :: TVar (MinMaxPSQ (NodeInfo ip) NodeDistance)
+    , searchQueued       :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf)))
+    , searchInformant    :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf)))
     , searchVisited      :: TVar (Set (NodeAddr ip))
     , searchResults      :: TVar (Set r)
     }
 
-newSearch :: Eq ip => (NodeInfo ip -> IO ([NodeInfo ip], [r]))
-                      -> NodeId -> [NodeInfo ip] -> IO (IterativeSearch ip r)
+newSearch :: Eq ip => (NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [r]))
+                      -> NodeId KMessageOf -> [NodeInfo KMessageOf ip ()] -> IO (IterativeSearch ip r)
 newSearch qry target ns = atomically $ do
     c <- newTVar 0
     q <- newTVar $ MM.fromList $ map (\n -> n :-> distance target (nodeId n)) ns
@@ -49,9 +53,9 @@ searchAlpha = 3
 searchK :: Int
 searchK = 8
 
-sendQuery :: (Ord a, Ord t) =>
-            IterativeSearch t a
-            -> Binding (NodeInfo t) NodeDistance
+sendQuery :: forall a ip. (Ord a, Ord ip) =>
+            IterativeSearch ip a
+            -> Binding (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf))
             -> IO ()
 sendQuery IterativeSearch{..} (ni :-> d) = do
     (ns,rs) <- handle (\(SomeException e) -> return ([],[]))
@@ -60,7 +64,10 @@ sendQuery IterativeSearch{..} (ni :-> d) = do
         modifyTVar searchPendingCount pred
         vs <- readTVar searchVisited
         -- We only queue a node if it is not yet visited
-        let insertFoundNode n q
+        let insertFoundNode :: NodeInfo KMessageOf ip u
+                           -> MinMaxPSQ (NodeInfo KMessageOf ip u) (NodeDistance (NodeId KMessageOf))
+                           -> MinMaxPSQ (NodeInfo KMessageOf ip u) (NodeDistance (NodeId KMessageOf))
+            insertFoundNode n q
              | nodeAddr n `Set.member` vs = q
              | otherwise                  = MM.insertTake searchK n (distance searchTarget $ nodeId n) q
         modifyTVar searchQueued    $ \q -> foldr insertFoundNode q ns