2 files changed, 28 insertions, 33 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 293d58ab..73fe358a 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -14,6 +14,7 @@
 --   <http://www.bittorrent.org/beps/bep_0005.html#routing-table>
 --
 {-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ViewPatterns    #-}
 {-# LANGUAGE TypeOperators   #-}
 {-# LANGUAGE DeriveGeneric   #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
@@ -39,8 +40,8 @@ module Network.BitTorrent.DHT.Routing
          -- * Lookup
        , K
        , defaultK
-       , Network.BitTorrent.DHT.Routing.kclosest
-       , Network.BitTorrent.DHT.Routing.kclosestHash
+       , TableKey (..)
+       , kclosest
 
          -- * Construction
        , Network.BitTorrent.DHT.Routing.nullTable
@@ -383,22 +384,24 @@ type K = Int
 defaultK :: K
 defaultK = 8
 
--- | Get a list of /K/ closest nodes using XOR metric. Used in
--- 'find_node' queries.
-kclosest :: Eq ip => K -> NodeId -> Table ip -> [NodeInfo ip]
-kclosest k nid = L.take k . rank nid
-               . L.map key . PSQ.toList . fromMaybe PSQ.empty
-               . lookupBucket nid
-
-coerceId :: (Serialize a, Serialize b) => a -> b
-coerceId = either (error msg) id . S.decode . S.encode
-  where
-    msg = "coerceId: impossible"
+class TableKey k where
+  toNodeId :: k -> NodeId
+
+instance TableKey NodeId where
+  toNodeId = id
 
--- | Get a list of /K/ nodes with node id closest to the specific
--- infohash. Used in 'get_peers' queries.
-kclosestHash :: Eq a => K -> InfoHash -> Table a -> [NodeInfo a]
-kclosestHash k nid t = kclosest k (coerceId nid) t
+instance TableKey InfoHash where
+  toNodeId = either (error msg) id . S.decode . S.encode
+    where -- TODO unsafe coerse?
+      msg = "tableKey: impossible"
+
+-- | 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 k (toNodeId -> nid)
+  = L.take k . rank nid
+  . L.map PSQ.key . PSQ.toList . fromMaybe PSQ.empty
+  . lookupBucket nid
 
 {-----------------------------------------------------------------------
 --  Routing
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 149b7dd2..5ceca4a3 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -26,7 +26,6 @@ module Network.BitTorrent.DHT.Session
        , getTable
        , getNodeId
        , getClosest
-       , getClosestHash
        , insertNode
 
          -- * Peer storage
@@ -327,21 +326,10 @@ getNodeId = asks thisNodeId
 --
 --   This operation used for 'find_nodes' query.
 --
-getClosest :: Eq ip => NodeId -> DHT ip [NodeInfo ip]
-getClosest nid = do
+getClosest :: Eq ip => TableKey k => k -> DHT ip [NodeInfo ip]
+getClosest node = do
   k <- asks (optK . options)
-  kclosest k nid <$> getTable
-
--- | Find a set of closest nodes from routing table of this node. (in
--- no particular order)
---
---   This operation used as failback in 'get_peers' query, see
---   'getPeerList'.
---
-getClosestHash :: Eq ip => InfoHash -> DHT ip [NodeInfo ip]
-getClosestHash ih = do
-  k <- asks (optK . options)
-  kclosestHash k ih <$> getTable
+  kclosest k node <$> getTable
 
 -- | This operation do not block but acquire exclusive access to
 --   routing table.
@@ -378,11 +366,15 @@ lookupPeers ih = do
 
 type PeerList ip = Either [NodeInfo ip] [PeerAddr ip]
 
+-- |
+--
+--   This operation used 'getClosest' as failback.
+--
 getPeerList :: Eq ip => InfoHash -> DHT ip (PeerList ip)
 getPeerList ih = do
   ps <- lookupPeers ih
   if L.null ps
-    then Left <$> getClosestHash ih
+    then Left <$> getClosest ih
     else return (Right ps)
 
 {-----------------------------------------------------------------------