1 files changed, 10 insertions, 2 deletions

diff --git a/src/Network/DHT/Routing.hs b/src/Network/DHT/Routing.hs
index 58e0cbd3..0eb1078b 100644
--- a/src/Network/DHT/Routing.hs
+++ b/src/Network/DHT/Routing.hs
@@ -23,6 +23,7 @@
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE StandaloneDeriving, FlexibleContexts, MultiParamTypeClasses, FlexibleInstances #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 module Network.DHT.Routing
@@ -304,7 +305,8 @@ updateBucketForInbound curTime info bucket
     -- map_q f = bucket { bktQ = runIdentity \$ f (bktQ bucket) }
 
 updateBucketForPingResult bad_node got_response bucket
-    = pure ( replacements
+    = pure ( map (,Nothing) forgotten
+             ++ map (second Just) replacements
            , Bucket (foldr replace
                            (bktNodes bucket)
                            replacements)
@@ -312,6 +314,7 @@ updateBucketForPingResult bad_node got_response bucket
            )
  where
     (top, popped) = runIdentity $ popFront bucketQ (bktQ bucket)
+
     replacements | got_response     = [] -- Timestamp was already updated by TryInsert.
                  | Just info <- top = do
                                          -- Insert only if there's a removal.
@@ -319,6 +322,9 @@ updateBucketForPingResult bad_node got_response bucket
                                          return (bad_node, info)
                  | otherwise        = []
 
+    forgotten | got_response = maybeToList $ fmap snd top
+
+
     replace (bad_node, (tm, info)) =
                  PSQ.insert (coerce info) tm
                  . PSQ.delete (coerce bad_node)
@@ -447,6 +453,8 @@ instance Pretty (BucketList ni) where
       ss = shape t
 
 -- | Empty table with specified /spine/ node id.
+--
+-- XXX: The comparison function argument is awkward here.
 nullTable :: (ni -> ni -> Ordering) -> ni -> Int -> BucketList ni
 nullTable cmp ni n =
     reify (Compare cmp)
@@ -704,7 +712,7 @@ updateForPingResult ::
     -> ni            -- ^ The pinged node.
     -> Bool          -- ^ True if we got a reply, False if it timed out.
     -> BucketList ni -- ^ The routing table.
-    -> ( [(ni,(Timestamp, ni))], BucketList ni )
+    -> ( [(ni,Maybe (Timestamp, ni))], BucketList ni )
 updateForPingResult space ni got_reply tbl =
     fromMaybe ([],tbl)
                    $ modifyBucket space