From a2958250c5dfd34d061f5ac8929c90299d5f346d Mon Sep 17 00:00:00 2001
From: Kerollmops <renault.cle@gmail.com>
Date: Sun, 13 May 2018 19:04:39 +0200
Subject: [PATCH] feat: Implement a custom capped `BTreeMap`

---
 src/capped_btree_map.rs | 93 +++++++++++++++++++++++++++++++++++++++++
 src/lib.rs              |  3 +-
 2 files changed, 95 insertions(+), 1 deletion(-)
 create mode 100644 src/capped_btree_map.rs

diff --git a/src/capped_btree_map.rs b/src/capped_btree_map.rs
new file mode 100644
index 000000000..d2ec99d44
--- /dev/null
+++ b/src/capped_btree_map.rs
@@ -0,0 +1,93 @@
+use std::collections::BTreeMap;
+
+#[derive(Debug)]
+pub struct CappedBTreeMap<K, V> {
+    inner: BTreeMap<K, V>,
+    capacity: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum Insertion<V> {
+    OldValue(V),
+    Evicted(V), // FIXME give (key *and* value)
+    Nothing,
+}
+
+impl<K: Ord, V> CappedBTreeMap<K, V> {
+    pub fn new(capacity: usize) -> Self {
+        assert!(capacity != 0, "It is invalid to set the capacity to zero.");
+        Self {
+            inner: BTreeMap::new(),
+            capacity: capacity,
+        }
+    }
+
+    pub fn clear(&mut self) {
+        self.inner.clear()
+    }
+
+    /// This method insert the `key`, `value` pair in the tree *but* will
+    /// remove the _smallest_ one, if the capacity is already reached,
+    /// before insertion.
+    ///
+    /// The _smallest_ `value` is not removed if the `key` inserted is already
+    /// present in the tree, in this case, the old replaced `value` is returned.
+    ///
+    /// ```
+    /// # extern crate raptor;
+    /// use raptor::{CappedBTreeMap, Insertion};
+    ///
+    /// let mut tree = CappedBTreeMap::new(3);
+    ///
+    /// let res = tree.insert(1, "a");
+    /// assert_eq!(res, Insertion::Nothing);
+    ///
+    /// tree.insert(2, "b");
+    /// tree.insert(3, "c");
+    ///
+    /// assert_eq!(tree.insert(4, "d"), Insertion::Evicted("c"));
+    ///
+    /// assert_eq!(tree.insert(1, "d"), Insertion::OldValue("a"));
+    /// ```
+    ///
+    pub fn insert(&mut self, key: K, value: V) -> Insertion<V>
+    where K: Clone,
+    {
+        if self.len() == self.capacity {
+            if self.inner.contains_key(&key) {
+                let value = self.inner.insert(key, value).unwrap();
+                Insertion::OldValue(value)
+            }
+            else {
+                let evicted_value = {
+                    // it is not possible to panic because we have reached
+                    // the capacity and the capacity cannot be set to zero.
+
+                    // FIXME remove this clone, find a way to remove
+                    //       the smallest key/value avoid borrowing problems
+                    let key = self.inner.keys().next_back().unwrap().clone();
+                    self.inner.remove(&key).unwrap()
+                };
+
+                self.inner.insert(key, value);
+                Insertion::Evicted(evicted_value)
+            }
+        }
+        else {
+            self.inner.insert(key, value);
+            Insertion::Nothing
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    pub fn capacity(&self) -> usize {
+        self.capacity
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index dde6a0ab1..3f436a907 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,6 +5,7 @@ extern crate levenshtein_automata;
 extern crate serde;
 
 pub mod map;
+pub mod capped_btree_map;
 mod levenshtein;
 
 pub use self::map::{Map, MapBuilder, Values};
@@ -12,7 +13,7 @@ pub use self::map::{
     OpBuilder, IndexedValues,
     OpWithStateBuilder, IndexedValuesWithState,
 };
-
+pub use self::capped_btree_map::{CappedBTreeMap, Insertion};
 pub use self::levenshtein::LevBuilder;
 
 #[derive(Debug, Serialize, Deserialize)]