1 files changed, 22 insertions, 14 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index cf4a4de3..6cf7f122 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -83,8 +83,10 @@ import Text.PrettyPrint.HughesPJClass (pPrint,Pretty)
 import qualified Data.ByteString as BS
 import Data.Bits
 
+import Network.KRPC.Message (KMessageOf)
 import Data.Torrent
 import Network.BitTorrent.Address
+import Network.DHT.Mainline
 
 {-----------------------------------------------------------------------
 --  Routing monad
@@ -180,7 +182,7 @@ runRouting ping_node find_nodes timestamper = go
 -- other words: new nodes are used only when older nodes disappear.
 
 -- | Timestamp - last time this node is pinged.
-type NodeEntry ip = Binding (NodeInfo ip) Timestamp
+type NodeEntry ip = Binding (NodeInfo KMessageOf ip ()) Timestamp
 
 -- TODO instance Pretty where
 
@@ -211,7 +213,7 @@ fromQ embed project QueueMethods{..} =
                  }
 -}
 
-seqQ :: QueueMethods Identity (NodeInfo ip) (Seq.Seq (NodeInfo ip))
+seqQ :: QueueMethods Identity (NodeInfo KMessageOf ip ()) (Seq.Seq (NodeInfo KMessageOf ip ()))
 seqQ = QueueMethods
     { pushBack = \e fifo -> pure (fifo Seq.|> e)
     , popFront = \fifo -> case Seq.viewl fifo of
@@ -220,9 +222,9 @@ seqQ = QueueMethods
     , emptyQueue = pure Seq.empty
     }
 
-type BucketQueue ip = Seq.Seq (NodeInfo ip)
+type BucketQueue ip = Seq.Seq (NodeInfo KMessageOf ip ())
 
-bucketQ :: QueueMethods Identity (NodeInfo ip) (BucketQueue ip)
+bucketQ :: QueueMethods Identity (NodeInfo KMessageOf ip ()) (BucketQueue ip)
 bucketQ = seqQ
 
 -- | Bucket is also limited in its length — thus it's called k-bucket.
@@ -232,7 +234,7 @@ bucketQ = seqQ
 --   very unlikely that all nodes in bucket fail within an hour of
 --   each other.
 --
-data Bucket ip = Bucket { bktNodes :: !(PSQ (NodeInfo ip) Timestamp)
+data Bucket ip = Bucket { bktNodes :: !(PSQ (NodeInfo KMessageOf ip ()) Timestamp)
                         , bktQ     :: !(BucketQueue ip)
                         } deriving (Show,Generic)
 
@@ -303,7 +305,7 @@ insertBucket curTime (PingResult bad_node got_response) bucket
                  pure $ PSQ.insert info curTime nodes'
         | otherwise        = id
 
-updateStamps :: Eq ip => Timestamp -> [NodeInfo ip] -> PSQ (NodeInfo ip) Timestamp -> PSQ (NodeInfo ip) Timestamp
+updateStamps :: Eq ip => Timestamp -> [NodeInfo KMessageOf ip ()] -> PSQ (NodeInfo KMessageOf ip ()) Timestamp -> PSQ (NodeInfo KMessageOf ip ()) Timestamp
 updateStamps curTime stales nodes = foldl' (\q n -> PSQ.insert n curTime q) nodes stales
 
 
@@ -330,6 +332,12 @@ split i b = (Bucket ns qs, Bucket ms rs)
   where
     (ns,ms) = (PSQ.fromList *** PSQ.fromList) . partition (spanBit . key) . PSQ.toList $ bktNodes b
     (qs,rs) = runIdentity $ partitionQ bucketQ spanBit $ bktQ b
+    {-
+    spanBit :: forall (dht :: * -> *) addr u.
+               FiniteBits (Network.RPC.NodeId dht) =>
+               NodeInfo dht addr u -> Bool
+    -}
+    spanBit :: NodeInfo KMessageOf addr () -> Bool
     spanBit entry = testIdBit (nodeId entry) i
 
 {-----------------------------------------------------------------------
@@ -458,7 +466,7 @@ compatibleNodeId tbl = genBucketSample prefix br
  where
     br = bucketRange (L.length (shape tbl) - 1) True
     bs = BS.pack $ take nodeIdSize $ tablePrefix tbl ++ repeat 0
-    prefix = asNodeId bs
+    prefix = either error id $ S.decode bs
 
 tablePrefix :: Table ip -> [Word8]
 tablePrefix = map (packByte . take 8 . (++repeat False))
@@ -503,7 +511,7 @@ instance TableKey InfoHash where
 
 -- | Get a list of /K/ closest nodes using XOR metric. Used in
 -- 'find_node' and 'get_peers' queries.
-kclosest :: Eq ip => TableKey a => K -> a -> Table ip -> [NodeInfo ip]
+kclosest :: Eq ip => TableKey a => K -> a -> Table ip -> [NodeInfo KMessageOf ip ()]
 kclosest k (toNodeId -> nid) tbl = take k $ rank nodeId nid (L.concat bucket)
                                           ++ rank nodeId nid (L.concat everyone)
  where
@@ -547,19 +555,19 @@ modifyBucket nodeId f = go (0 :: BitIx)
                            <|>  go i (splitTip nid n i bucket)
 
 -- | Triggering event for atomic table update
-data Event ip = TryInsert { foreignNode :: NodeInfo ip }
-              | PingResult { foreignNode :: NodeInfo ip
+data Event ip = TryInsert { foreignNode :: NodeInfo KMessageOf ip () }
+              | PingResult { foreignNode :: NodeInfo KMessageOf ip ()
                            , ponged :: Bool
                            }
-    deriving (Eq,Ord,Show)
+    deriving (Eq,Show) -- Ord
 
 eventId :: Event ip -> NodeId
 eventId (TryInsert NodeInfo{..})    = nodeId
 eventId (PingResult NodeInfo{..} _) = nodeId
 
 -- | Actions requested by atomic table update
-data CheckPing ip = CheckPing [NodeInfo ip]
-    deriving (Eq,Ord,Show)
+data CheckPing ip = CheckPing [NodeInfo KMessageOf ip ()]
+    deriving (Eq,Show) -- Ord
 
 
 -- | Atomic 'Table' update
@@ -571,7 +579,7 @@ insert tm event tbl = pure $ fromMaybe ([],tbl) $ modifyBucket (eventId event) (
 --  Conversion
 -----------------------------------------------------------------------}
 
-type TableEntry ip = (NodeInfo ip, Timestamp)
+type TableEntry ip = (NodeInfo KMessageOf ip (), Timestamp)
 
 tableEntry :: NodeEntry ip -> TableEntry ip
 tableEntry (a :-> b) = (a, b)