1 files changed, 15 insertions, 7 deletions

diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 854f26c7..844b4575 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -24,21 +25,28 @@ import qualified Data.Wrapper.PSQ as PSQ
          ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ)
 import Network.BitTorrent.Address hiding (NodeId)
 import Network.RPC
-import Network.KRPC.Message (KMessageOf)
+#ifdef VERSION_bencoding
 import Network.DHT.Mainline ()
+import Network.KRPC.Message (KMessageOf)
+type Ann = ()
+#else
+import Data.Tox as Tox
+type KMessageOf = Tox.Message
+type Ann = Bool
+#endif
 
 data IterativeSearch ip r = IterativeSearch
     { searchTarget       :: NodeId KMessageOf
-    , searchQuery        :: NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [r])
+    , searchQuery        :: NodeInfo KMessageOf ip Ann -> IO ([NodeInfo KMessageOf ip Ann], [r])
     , searchPendingCount :: TVar Int
-    , searchQueued       :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf)))
-    , searchInformant    :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf)))
+    , searchQueued       :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf)))
+    , searchInformant    :: TVar (MinMaxPSQ (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf)))
     , searchVisited      :: TVar (Set (NodeAddr ip))
     , searchResults      :: TVar (Set r)
     }
 
-newSearch :: Eq ip => (NodeInfo KMessageOf ip () -> IO ([NodeInfo KMessageOf ip ()], [r]))
-                      -> NodeId KMessageOf -> [NodeInfo KMessageOf ip ()] -> IO (IterativeSearch ip r)
+newSearch :: Eq ip => (NodeInfo KMessageOf ip Ann -> IO ([NodeInfo KMessageOf ip Ann], [r]))
+                      -> NodeId KMessageOf -> [NodeInfo KMessageOf ip Ann] -> 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
@@ -55,7 +63,7 @@ searchK = 8
 
 sendQuery :: forall a ip. (Ord a, Ord ip) =>
             IterativeSearch ip a
-            -> Binding (NodeInfo KMessageOf ip ()) (NodeDistance (NodeId KMessageOf))
+            -> Binding (NodeInfo KMessageOf ip Ann) (NodeDistance (NodeId KMessageOf))
             -> IO ()
 sendQuery IterativeSearch{..} (ni :-> d) = do
     (ns,rs) <- handle (\(SomeException e) -> return ([],[]))