1 files changed, 29 insertions, 17 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 396c4b1d..46ebe472 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -536,19 +536,31 @@ depth = L.length . shape
 #endif
 
 lookupBucket :: forall ni nid x.
-                ( FiniteBits nid
-                , Ord nid
-                ) => (ni -> nid) -> nid -> (forall s. Reifies s (Compare ni) => [Bucket s ni] -> x) -> BucketList ni -> x
-lookupBucket nodeId nid kont (BucketList self _ bkts) = kont $ go 0 [] bkts
+                ( -- FiniteBits nid
+                 Ord nid
+                ) => KademliaSpace nid ni -> nid -> (forall s. Reifies s (Compare ni) => [Bucket s ni] -> x) -> BucketList ni -> x
+lookupBucket space nid kont (BucketList self _ bkts) = kont $ go 0 [] bkts
   where
-    d = nid `xor` nodeId self
+    d = kademliaXor space nid (kademliaLocation space self)
 
     go :: Word -> [Bucket s ni] -> [Bucket s ni] -> [Bucket s ni]
     go i bs (bucket : buckets)
-      |  testIdBit d i = go (succ i) (bucket:bs) buckets
-      |     otherwise  = bucket : buckets ++ bs
+      |  kademliaTestBit space d i = bucket : buckets ++ bs
+      |     otherwise              = go (succ i) (bucket:bs) buckets
     go _ bs []         = bs
 
+bucketNumber :: forall ni nid.
+                KademliaSpace nid ni -> nid -> BucketList ni -> Int
+bucketNumber space nid (BucketList self _ bkts) = fromIntegral $ go 0 bkts
+  where
+    d = kademliaXor space nid (kademliaLocation space self)
+
+    go :: Word -> [Bucket s ni] -> Word
+    go i (bucket : buckets)
+      |  kademliaTestBit space d i = i
+      |     otherwise              = go (succ i) buckets
+    go i []                        = i
+
 
 compatibleNodeId :: forall ni nid.
                     ( Serialize nid, FiniteBits nid) =>
@@ -614,23 +626,23 @@ distance :: Bits nid => nid -> nid -> NodeDistance nid
 distance a b = NodeDistance $ xor a b
 
 -- | Order by closeness: nearest nodes first.
-rank :: ( FiniteBits nid
-        , Ord nid
-        ) => (x -> nid) -> nid -> [x] -> [x]
-rank f nid = L.sortBy (comparing (distance nid . f))
+rank :: ( Ord nid
+        ) => KademliaSpace nid ni -> nid -> [ni] -> [ni]
+rank space nid = L.sortBy (comparing (kademliaXor space nid . kademliaLocation space))
 
 
 -- | Get a list of /K/ closest nodes using XOR metric. Used in
 -- 'find_node' and 'get_peers' queries.
-kclosest :: ( FiniteBits nid
-            , Ord nid
-            ) => (ni -> nid) -> Int -> nid -> BucketList ni -> [ni]
-kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
-                              ++ rank nodeId nid (L.concat everyone)
+kclosest :: ( -- FiniteBits nid
+             Ord nid
+            ) =>
+            KademliaSpace nid ni -> Int -> nid -> BucketList ni -> [ni]
+kclosest space k nid tbl = take k $ rank space nid (L.concat bucket)
+                              ++ rank space nid (L.concat everyone)
  where
     (bucket,everyone) =
         L.splitAt 1
-        . lookupBucket nodeId nid (L.map (coerce . L.map PSQ.key . PSQ.toList . bktNodes))
+        . lookupBucket space nid (L.map (coerce . L.map PSQ.key . PSQ.toList . bktNodes))
         $ tbl