1 files changed, 48 insertions, 42 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 631761ac..2521f7b9 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -227,9 +227,7 @@ delta = 15 * 60
 -- | Should maintain a set of stable long running nodes.
 --
 -- Note: pings are triggerd only when a bucket is full.
-insertBucket :: (Alternative f, Ord ni) => -- (Eq ip, Alternative f, Ord (NodeId))  =>
-    Timestamp -> Event ni -> Bucket ni -> f ([CheckPing ni], Bucket ni)
-insertBucket curTime (TryInsert info) bucket
+updateBucketForInbound curTime info bucket
   -- Just update timestamp if a node is already in bucket.
   --
   -- Note PingResult events should only occur for nodes we requested a ping for,
@@ -247,7 +245,7 @@ insertBucket curTime (TryInsert info) bucket
   -- then one is replaced by the new node in the next insertBucket
   -- iteration.
   | not (L.null stales)
-    = pure ( [CheckPing stales]
+    = pure ( stales
            , bucket { -- Update timestamps so that we don't redundantly ping.
                       bktNodes = updateStamps curTime stales $ bktNodes bucket
                       -- Update queue with the pending NodeInfo in case of ping fail.
@@ -272,15 +270,24 @@ insertBucket curTime (TryInsert info) bucket
     map_ns f = bucket { bktNodes = f (bktNodes bucket) }
     -- map_q f = bucket { bktQ = runIdentity \$ f (bktQ bucket) }
 
-insertBucket curTime (PingResult bad_node got_response) bucket
-    = pure ([], Bucket (upd $ bktNodes bucket) popped)
+updateBucketForPingResult curTime bad_node got_response bucket
+    = pure ( replacements
+           , Bucket (foldr replace
+                           (bktNodes bucket)
+                           replacements)
+                    popped
+           )
  where
     (top, popped) = runIdentity $ popFront bucketQ (bktQ bucket)
-    upd | got_response     = id -- Timestamp was already updated by TryInsert.
-        | Just info <- top = \nodes ->
-            fromMaybe nodes $ do
-                 _ <- PSQ.lookup bad_node nodes -- Insert only if there's a removal.
-                 let nodes' = PSQ.delete bad_node nodes
+    replacements | got_response     = [] -- Timestamp was already updated by TryInsert.
+                 | Just info <- top = do
+                                         -- Insert only if there's a removal.
+                                         _ <- maybeToList $ PSQ.lookup bad_node (bktNodes bucket)
+                                         return (bad_node, info)
+                 | otherwise        = []
+
+    replace (bad_node, info) =
+                 PSQ.insert info curTime
                  -- XXX: curTime is the time somebody pinged out, not the time
                  -- of which the popped node was originally queued.  That's not
                  -- quite right.
@@ -290,8 +297,7 @@ insertBucket curTime (PingResult bad_node got_response) bucket
                  --
                  -- This means that initial time-outs for replacement nodes are up to
                  -- twice as long as for the originals.
-                 pure $ PSQ.insert info curTime nodes'
-        | otherwise        = id
+                 . PSQ.delete bad_node
 
 
 updateStamps :: Ord ni => Timestamp -> [ni] -> PSQ ni Timestamp -> PSQ ni Timestamp
@@ -627,48 +633,48 @@ deriving instance ( Show ip
 
 #endif
 
--- | Atomic 'Table' update
+
+-- | Call on every inbound packet (including requested ping results).
+-- Returns a triple (was_inserted, to_ping, tbl') where
 --
--- Deprecated, use these two functions instead:
+--  [ /was_inserted/ ] True if the node was added to the routing table.
 --
---  [ 'updateForInbound' ]
+--  [ /to_ping/ ]      A list of nodes to ping and then run 'updateForPingResult'.
+--                     This will be empty if /was_inserted/, but a non-inserted node
+--                     may be added to a replacement queue and will be inserted if
+--                     one of the items in this list time out.
 --
---  [ 'updateForPingResult' ]
+--  [ /tbl'/ ]         The updated routing 'Table'.
 --
-insert :: -- ( Eq ip , Applicative m , Ord (NodeId) , FiniteBits (NodeId)) =>
-    (Applicative m, Ord ni) =>
-    (nid -> Word -> Bool)
-    -> (ni -> nid)
-    -> Timestamp -> Event ni -> Table ni nid -> m ([CheckPing ni], Table ni nid)
-insert testIdBit nodeId tm event tbl
-    = pure $ fromMaybe ([],tbl)
-           $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
-                                    (eventId nodeId event)
-                                    (insertBucket tm event)
-                                    tbl
-
-
--- | Call on every inbound packet (including requested ping results).  Returns
--- a list of nodes to ping and an updated routing table.  The caller should
--- also invoke 'updateForPingResult' on each node in the to-ping list to update
--- the routing table with that information as it becomes available.
 updateForInbound :: Ord ni =>
     (nid -> Word -> Bool)
     -> (ni -> nid)
-    -> Timestamp -> ni -> Table ni nid -> ([ni], Table ni nid)
+    -> Timestamp -> ni -> Table ni nid -> (Bool, [ni], Table ni nid)
 updateForInbound testIdBit nodeId tm ni tbl =
-    let Identity (ps, tbl') = Network.DHT.Routing.insert testIdBit nodeId tm (TryInsert ni) tbl
-        decon (CheckPing ns) = ns
-    in (concatMap decon ps, tbl')
+    maybe (False, [],tbl) (\(ps,tbl') -> (True, ps, tbl'))
+                   $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
+                                            (nodeId ni)
+                                            (updateBucketForInbound tm ni)
+                                            tbl
 
+-- | Update the routing table with the results of a ping.
+--
+-- Each (a,b) in the returned list indicates that the node /a/ was deleted from the
+-- routing table and the  node /b/ has taken its place.
 updateForPingResult :: Ord ni =>
     (nid -> Word -> Bool)
     -> (ni -> nid)
-    -> Timestamp -> ni -> Bool -> Table ni nid -> Table ni nid
+    -> Timestamp
+    -> ni           -- ^ The pinged node.
+    -> Bool         -- ^ True if we got a reply, False if it timed out.
+    -> Table ni nid -- ^ The routing table.
+    -> ( [(ni,ni)], Table ni nid )
 updateForPingResult testIdBit nodeId tm ni got_reply tbl =
-    -- The _ should always be empty in this case.
-    let Identity (_, tbl') = Network.DHT.Routing.insert testIdBit nodeId tm (PingResult ni got_reply) tbl
-    in tbl'
+    fromMaybe ([],tbl)
+                   $ modifyBucket testIdBit (\ni -> testIdBit $ nodeId ni)
+                                            (nodeId ni)
+                                            (updateBucketForPingResult tm ni got_reply)
+                                            tbl
 
 
 {-----------------------------------------------------------------------