DHT rewrite: kclosest

1 files changed, 36 insertions, 10 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 34d8385f..bfaaf022 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -62,6 +62,7 @@ module Network.DHT.Routing
 
          -- * Routing
        , Timestamp
+       , getTimestamp
        ) -} where
 
 import Control.Applicative as A
@@ -83,6 +84,7 @@ import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.HughesPJClass (pPrint,Pretty)
 import qualified Data.ByteString as BS
 import Data.Bits
+import Data.Ord
 
 import Network.Address
 
@@ -103,6 +105,11 @@ import Network.Address
 --
 type Timestamp = POSIXTime
 
+getTimestamp :: IO Timestamp
+getTimestamp = do
+  utcTime <- getCurrentTime
+  return $ utcTimeToPOSIXSeconds utcTime
+
 
 
 {-----------------------------------------------------------------------
@@ -435,8 +442,11 @@ size = L.sum . shape
 depth :: Table -> BucketCount
 depth = L.length . shape
 
-lookupBucket :: ( FiniteBits (NodeId)
-                ) => NodeId -> Table -> [Bucket]
+#endif
+
+lookupBucket :: ( FiniteBits nid
+                , Ord nid
+                ) => nid -> Table ni nid -> [Bucket ni]
 lookupBucket nid = go 0 []
   where
     go i bs (Zero table bucket)
@@ -447,7 +457,6 @@ lookupBucket nid = go 0 []
       |     otherwise        = bucket : toBucketList table ++ bs
     go _ bs (Tip _ _ bucket) = bucket : bs
 
-#endif
 
 compatibleNodeId :: forall ni nid.
                     ( Serialize nid, FiniteBits nid) =>
@@ -498,14 +507,31 @@ class TableKey dht k where
 instance TableKey dht (NodeId) where
   toNodeId = id
 
+#endif
+
+-- | In Kademlia, the distance metric is XOR and the result is
+-- interpreted as an unsigned integer.
+newtype NodeDistance nodeid = NodeDistance nodeid
+  deriving (Eq, Ord)
+
+-- | distance(A,B) = |A xor B| Smaller values are closer.
+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))
+
+
 -- | Get a list of /K/ closest nodes using XOR metric. Used in
 -- 'find_node' and 'get_peers' queries.
-kclosest :: ( Eq ip
-            , Ord (NodeId)
-            , FiniteBits (NodeId)
-            ) => TableKey dht a => K -> a -> Table -> [NodeInfo]
-kclosest k (toNodeId -> nid) tbl = take k $ rank nodeId nid (L.concat bucket)
-                                          ++ rank nodeId nid (L.concat everyone)
+kclosest :: ( FiniteBits nid
+            , Ord nid
+            ) => (ni -> nid) -> Int -> nid -> Table ni nid -> [ni]
+kclosest nodeId k nid tbl = take k $ rank nodeId nid (L.concat bucket)
+                              ++ rank nodeId nid (L.concat everyone)
  where
     (bucket,everyone) =
         L.splitAt 1
@@ -513,7 +539,7 @@ kclosest k (toNodeId -> nid) tbl = take k $ rank nodeId nid (L.concat bucket)
         . lookupBucket nid
         $ tbl
 
-#endif
+
 
 {-----------------------------------------------------------------------
 --  Routing