1 files changed, 3 insertions, 4 deletions

diff --git a/src/Network/Kademlia/Search.hs b/src/Network/Kademlia/Search.hs
index b7d07f7c..29a8fc61 100644
--- a/src/Network/Kademlia/Search.hs
+++ b/src/Network/Kademlia/Search.hs
@@ -101,16 +101,15 @@ stripValue (Binding ni _ nid) = (ni :-> nid)
 
 -- | Reset a 'SearchState' object to ready it for a repeated search.
 reset :: (Ord ni, Ord nid, Hashable ni, Hashable nid) =>
-       TVar (BucketList ni)
+       (nid -> STM [ni])
        -> Search nid addr1 tok1 ni r1
        -> nid
        -> SearchState nid addr tok ni r
        -> STM (SearchState nid addr tok ni r)
-reset bkts qsearch target st = do
+reset nearestNodes qsearch target st = do
     searchIsFinished st >>= check -- Wait for a search to finish before resetting.
     bktNodes <- map (\ni -> ni :-> kademliaLocation (searchSpace qsearch) ni)
-                . R.kclosest (searchSpace qsearch) searchK target
-                <$> readTVar bkts
+                <$> nearestNodes target
     priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant  st)
     writeTVar (searchQueued       st) $ MM.fromList $ priorInformants ++ bktNodes
     writeTVar (searchInformant    st) MM.empty