From dc4d1ea990843538ecdaa9ad33e70ebffe92498c Mon Sep 17 00:00:00 2001
From: Fedor Sakharov <fedor.sakharov@gmail.com>
Date: Mon, 8 Jul 2019 11:57:49 +0300
Subject: [PATCH] Adds a store_mut method to kademlia. (#1192)

* Adds a retain method to kademlia.

* Add a store_mut getter to Kademlia

* Removes a blank line

* Changes store_mut comment appropriately

* Fixes build

* Return a type, not a trait
---
 protocols/kad/src/behaviour.rs | 5 +++++
 protocols/kad/src/record.rs    | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index 501b8845..15f71982 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -276,6 +276,11 @@ where
         }
     }
 
+    /// Get a mutable reference to internal record store.
+    pub fn store_mut(&mut self) -> &mut TStore {
+        &mut self.records
+    }
+
     /// Bootstraps the local node to join the DHT.
     ///
     /// Bootstrapping is a multi-step operation that starts with a lookup of the local node's
diff --git a/protocols/kad/src/record.rs b/protocols/kad/src/record.rs
index 611acaa2..4296cf61 100644
--- a/protocols/kad/src/record.rs
+++ b/protocols/kad/src/record.rs
@@ -78,6 +78,15 @@ impl Default for MemoryRecordStorage {
     }
 }
 
+impl MemoryRecordStorage {
+    /// Retain the elements by a predicate.
+    pub fn retain<F>(&mut self, f: F)
+        where F: FnMut(&Multihash, &mut Record) -> bool
+    {
+        self.records.retain(f);
+    }
+}
+
 impl RecordStore for MemoryRecordStorage {
     fn get(&self, k: &Multihash) -> Option<Cow<Record>> {
         match self.records.get(k) {