1 files changed, 37 insertions, 8 deletions

diff --git a/src/Network/Kademlia/Search.hs b/src/Network/Kademlia/Search.hs
index 770d2f13..5f024cd0 100644
--- a/src/Network/Kademlia/Search.hs
+++ b/src/Network/Kademlia/Search.hs
@@ -25,7 +25,7 @@ import System.IO.Error
 import qualified Data.MinMaxPSQ   as MM
          ;import Data.MinMaxPSQ   (MinMaxPSQ, MinMaxPSQ')
 import qualified Data.Wrapper.PSQ as PSQ
-         ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ, PSQKey)
+         ;import Data.Wrapper.PSQ (pattern (:->), Binding, pattern Binding, Binding', PSQ, PSQKey)
 import Network.Address hiding (NodeId)
 import Network.Kademlia.Routing as R
 #ifdef THREAD_DEBUG
@@ -94,6 +94,29 @@ newSearch (Search space nAddr qry) target ns = do
     return -- (Search space nAddr qry) , r , target
            ( SearchState c q i v )
 
+-- | Discard a value from a key-priority-value tuple. This is useful for
+-- swaping items from a "MinMaxPSQ'" to a "MinMaxPSQ".
+stripValue :: Binding' k p v -> Binding k p
+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)
+       -> Search nid addr1 tok1 ni r1
+       -> nid
+       -> SearchState nid addr tok ni r
+       -> STM ()
+reset bkts 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
+    priorInformants <- map stripValue . MM.toList <$> readTVar (searchInformant  st)
+    writeTVar (searchQueued       st) $ MM.fromList $ priorInformants ++ bktNodes
+    writeTVar (searchInformant    st) MM.empty
+    writeTVar (searchVisited      st) Set.empty
+    writeTVar (searchPendingCount st) 0
+
 searchAlpha :: Int
 searchAlpha = 8
 
@@ -108,7 +131,7 @@ sendQuery :: forall addr nid tok ni r.
             ) =>
             Search nid addr tok ni r
             -> nid
-            -> (r -> STM Bool)
+            -> (r -> STM Bool) -- ^ return False to terminate the search.
             -> SearchState nid addr tok ni r
             -> Binding ni nid
             -> IO ()
@@ -124,16 +147,18 @@ sendQuery Search{..} searchTarget searchResult sch@SearchState{..} (ni :-> d) =
     go (ns,rs,tok) = do
         vs <- readTVar searchVisited
         -- We only queue a node if it is not yet visited
-        let insertFoundNode :: ni
+        let insertFoundNode :: Int
+                           -> ni
                            -> MinMaxPSQ ni nid
                            -> MinMaxPSQ ni nid
-            insertFoundNode n q
+            insertFoundNode k n q
              | searchNodeAddress n `Set.member` vs
                          = q
-             | otherwise = MM.insertTake searchK n ( kademliaXor searchSpace searchTarget
-                                                     $ kademliaLocation searchSpace n )
+             | otherwise = MM.insertTake k n ( kademliaXor searchSpace searchTarget
+                                               $ kademliaLocation searchSpace n )
                                                    q
-        modifyTVar searchQueued    $ \q -> foldr insertFoundNode q ns
+        qsize <- MM.size <$> readTVar searchQueued
+        modifyTVar searchQueued    $ \q -> foldr (insertFoundNode qsize) q ns
         modifyTVar searchInformant $ MM.insertTake' searchK ni tok d
         flip fix rs $ \loop -> \case
             r:rs' -> do
@@ -175,7 +200,11 @@ search sch buckets target result = do
     return st
 
 searchLoop :: ( Ord addr, Ord nid, Ord ni, Show nid, Hashable nid, Hashable ni )
-           => Search nid addr tok ni r -> nid -> (r -> STM Bool) -> SearchState nid addr tok ni r -> IO ()
+           => Search nid addr tok ni r      -- ^ Query and distance methods.
+           -> nid                           -- ^ The target we are searching for.
+           -> (r -> STM Bool)               -- ^ Invoked on each result.  Return False to quit searching.
+           -> SearchState nid addr tok ni r -- ^ Search-related state.
+           -> IO ()
 searchLoop sch@Search{..} target result s@SearchState{..} = do
     myThreadId >>= flip labelThread ("search."++show target)
     withTaskGroup ("search.g."++show target) searchAlpha $ \g -> fix $ \again -> do