From 3d52a6b02155b4bff6e05a3434e52649036eb91f Mon Sep 17 00:00:00 2001
From: folex <0xdxdy@gmail.com>
Date: Thu, 26 Mar 2020 19:06:44 +0300
Subject: [PATCH] Sort weighted bucket iter() by weight

---
 protocols/kad/src/behaviour/test.rs   |  2 +-
 protocols/kad/src/kbucket/bucket.rs   |  6 ------
 protocols/kad/src/kbucket/weighted.rs | 16 ++++++++++++----
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
index 59fa41c4..fcaa6875 100644
--- a/protocols/kad/src/behaviour/test.rs
+++ b/protocols/kad/src/behaviour/test.rs
@@ -74,7 +74,7 @@ fn build_nodes_with_config(num: usize, cfg: KademliaConfig) -> (u64, Vec<(ed2551
 
         let local_id = local_public_key.clone().into_peer_id();
         let store = MemoryStore::new(local_id.clone());
-        let trust = TrustGraph::new(HashMap::new());
+        let trust = TrustGraph::new(Vec::new());
         let behaviour = Kademlia::with_config(ed25519_key.clone(), local_id.clone(), store, cfg.clone(), trust);
         result.push((ed25519_key, Swarm::new(transport, behaviour, local_id)));
     }
diff --git a/protocols/kad/src/kbucket/bucket.rs b/protocols/kad/src/kbucket/bucket.rs
index ccba5ef4..4b385fab 100644
--- a/protocols/kad/src/kbucket/bucket.rs
+++ b/protocols/kad/src/kbucket/bucket.rs
@@ -178,12 +178,6 @@ where
     /// the node that was replaced. `None` indicates that the nodes in the
     /// bucket remained unchanged.
     pub fn apply_pending(&mut self) -> Vec<AppliedPending<TKey, TVal>> {
-        // self.weighted
-        //     .apply_pending()
-        //     .iter()
-        //     .chain(self.swamp.apply_pending().iter())
-        //     .collect()
-
         Iterator::chain(
             self.weighted.apply_pending().into_iter(),
             self.swamp.apply_pending().into_iter(),
diff --git a/protocols/kad/src/kbucket/weighted.rs b/protocols/kad/src/kbucket/weighted.rs
index 680888cc..7c6dc749 100644
--- a/protocols/kad/src/kbucket/weighted.rs
+++ b/protocols/kad/src/kbucket/weighted.rs
@@ -311,10 +311,18 @@ where
         }
     }
 
-    pub fn iter(&self) -> impl Iterator<Item = (&Node<TKey, TVal>, NodeStatus)> {
-        self.map
-            .values()
-            .map(|bucket| bucket.iter().map(|(n, s)| (&n.inner, s)))
+    pub fn iter(&self) -> impl Iterator<Item = (&Node<TKey, TVal>, NodeStatus)> + '_ {
+        let mut keys = self.map.keys().collect::<Vec<_>>();
+        keys.sort();
+        let map: &HashMap<u32, _> = &self.map;
+
+        keys.into_iter()
+            .map(move |w| {
+                map.get(&w)
+                    .into_iter()
+                    .map(|bucket| bucket.iter().map(|(n, s)| (&n.inner, s)))
+                    .flatten()
+            })
             .flatten()
     }