2 files changed, 16 insertions, 3 deletions

diff --git a/src/Data/MinMaxPSQ.hs b/src/Data/MinMaxPSQ.hs
index f41da4a4..f385f258 100644
--- a/src/Data/MinMaxPSQ.hs
+++ b/src/Data/MinMaxPSQ.hs
@@ -55,9 +55,12 @@ maxView :: (PSQKey k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k
 maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> p, MinMaxPSQ (PSQ.delete k nq) xq'))
                             $ PSQ.minView xq
 
+-- | Maintains a bounded 'MinMaxPSQ' by deleting the maximum element if the
+-- insertion would cause the queue to have too many elements.
 insertTake :: (PSQKey k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
 insertTake n k p q = take n $ insert k p q
 
+-- | Truncate the 'MinMaxPSQ' to the given number of lowest priority elements.
 take :: (PSQKey k, Ord p) => Int -> MinMaxPSQ k p -> MinMaxPSQ k p
 take !n !q | (size q <= n) = q
            | null q        = q
diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
index 91a1079b..c562b988 100644
--- a/src/Network/BitTorrent/DHT/Search.hs
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -33,8 +33,12 @@ data IterativeSearch nid addr ni r = IterativeSearch
     , searchSpace        :: KademliaSpace nid ni
     , searchNodeAddress  :: ni -> addr
     , searchQuery        :: ni -> IO ([ni], [r])
+      -- | The number of pending queries.  Incremented before any query is sent
+      -- and decremented when we get a reply.
     , searchPendingCount :: TVar Int
+      -- | Nodes scheduled to be queried.
     , searchQueued       :: TVar (MinMaxPSQ ni nid)
+      -- | The nearest K nodes that issued a reply.
     , searchInformant    :: TVar (MinMaxPSQ ni nid)
     , searchVisited      :: TVar (Set addr)
     , searchResults      :: TVar (Set r)
@@ -127,11 +131,17 @@ search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do
             found <- MM.minView <$> readTVar searchQueued
             case found of
                 Just (ni :-> d, q)
-                  | (MM.size informants < searchK) && (cnt > 0 || not (MM.null q))
+                  | -- If there's fewer than /k/ informants and there's any
+                    -- node we haven't yet got a response from.
+                    (MM.size informants < searchK) && (cnt > 0 || not (MM.null q))
+                    -- Or if the closest scheduled node is nearer than the
+                    -- nearest /k/ informants.
                     || (PSQ.prio (fromJust $ MM.findMax informants) > d)
-                  -> do writeTVar searchQueued q
+                  -> -- Then the search continues, send a query.
+                     do writeTVar searchQueued q
                         modifyTVar searchVisited $ Set.insert (searchNodeAddress ni)
                         modifyTVar searchPendingCount succ
                         return $ withAsync g (sendQuery s (ni :-> d)) (const again)
-                _ -> do check (cnt == 0)
+                _ -> -- Otherwise, we are finished.
+                     do check (cnt == 0)
                         return $ return ()