1 files changed, 28 insertions, 14 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 7f0e1c94..e2bc5aa9 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -638,33 +638,33 @@ deriving instance ( Show ip
 --  [ /tbl'/ ]         The updated routing 'Table'.
 --
 updateForInbound :: Ord ni =>
-    (nid -> Word -> Bool)
-    -> (ni -> nid)
+    KademliaSpace nid ni
     -> Timestamp -> ni -> Table ni nid -> (Bool, [ni], Table ni nid)
-updateForInbound testIdBit nodeId tm ni tbl =
+updateForInbound space tm ni tbl =
     maybe (False, [],tbl) (\(ps,tbl') -> (True, ps, tbl'))
-                   $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
-                                            (nodeId ni)
-                                            (updateBucketForInbound tm ni)
-                                            tbl
+                   $ modifyBucket (kademliaTestBit space)
+                                   (\ni -> kademliaTestBit space $ kademliaLocation space ni)
+                                   (kademliaLocation space ni)
+                                   (updateBucketForInbound tm ni)
+                                   tbl
 
 -- | Update the routing table with the results of a ping.
 --
 -- Each (a,(tm,b)) in the returned list indicates that the node /a/ was deleted from the
 -- routing table and the  node /b/, with timestamp /tm/, has taken its place.
 updateForPingResult :: Ord ni =>
-    (nid -> Word -> Bool)
-    -> (ni -> nid)
+    KademliaSpace nid ni
     -> ni           -- ^ The pinged node.
     -> Bool         -- ^ True if we got a reply, False if it timed out.
     -> Table ni nid -- ^ The routing table.
     -> ( [(ni,(Timestamp, ni))], Table ni nid )
-updateForPingResult testIdBit nodeId ni got_reply tbl =
+updateForPingResult space ni got_reply tbl =
     fromMaybe ([],tbl)
-                   $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
-                                            (nodeId ni)
-                                            (updateBucketForPingResult ni got_reply)
-                                            tbl
+                   $ modifyBucket (kademliaTestBit space)
+                                  (\ni -> kademliaTestBit space $ kademliaLocation space ni)
+                                  (kademliaLocation space ni)
+                                  (updateBucketForPingResult ni got_reply)
+                                  tbl
 
 
 {-----------------------------------------------------------------------
@@ -685,3 +685,17 @@ toBucketList (One   b t) = b : toBucketList t
 toList :: Table ni nid -> [[TableEntry ni]]
 toList = L.map (L.map tableEntry . PSQ.toList . bktNodes) . toBucketList
 
+data KademliaSpace nid ni = KademliaSpace
+    { -- | Given a node record (probably including IP address), yields a
+      -- kademlia xor-metric location.
+      kademliaLocation :: ni -> nid
+      -- | Used when comparing locations.  This is similar to
+      -- 'Data.Bits.testBit' except that the ordering of bits is reversed, so
+      -- that 0 is the most significant bit.
+    , kademliaTestBit :: nid -> Word -> Bool
+    }
+
+contramapKS f ks = ks
+    { kademliaLocation = kademliaLocation ks . f
+    }
+