From 49d1c9d1ba238647568fece048fa45573bb49a4c Mon Sep 17 00:00:00 2001
From: Nathan Sobo <nathan@zed.dev>
Date: Fri, 10 Dec 2021 23:33:15 -0700
Subject: [PATCH] Introduce sum_tree::TreeMap<K, V>

I think this will be useful to avoid cloning HashMaps in certain cases such as snapshots.
---
 crates/sum_tree/src/sum_tree.rs |   2 +
 crates/sum_tree/src/tree_map.rs | 109 ++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 crates/sum_tree/src/tree_map.rs

diff --git a/crates/sum_tree/src/sum_tree.rs b/crates/sum_tree/src/sum_tree.rs
index 63fb379d53..67c056d858 100644
--- a/crates/sum_tree/src/sum_tree.rs
+++ b/crates/sum_tree/src/sum_tree.rs
@@ -1,9 +1,11 @@
 mod cursor;
+mod tree_map;
 
 use arrayvec::ArrayVec;
 pub use cursor::{Cursor, FilterCursor, Iter};
 use std::marker::PhantomData;
 use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
+pub use tree_map::TreeMap;
 
 #[cfg(test)]
 const TREE_BASE: usize = 2;
diff --git a/crates/sum_tree/src/tree_map.rs b/crates/sum_tree/src/tree_map.rs
new file mode 100644
index 0000000000..d5ca7c1327
--- /dev/null
+++ b/crates/sum_tree/src/tree_map.rs
@@ -0,0 +1,109 @@
+use std::{cmp::Ordering, fmt::Debug};
+
+use crate::{Bias, Dimension, Item, KeyedItem, SeekTarget, SumTree, Summary};
+
+pub struct TreeMap<K, V>(SumTree<MapEntry<K, V>>)
+where
+    K: Clone + Debug + Default,
+    V: Clone + Debug + Default;
+
+#[derive(Clone)]
+pub struct MapEntry<K, V> {
+    key: K,
+    value: V,
+}
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
+pub struct MapKey<K>(K);
+
+#[derive(Clone, Debug, Default)]
+pub struct MapKeyRef<'a, K>(Option<&'a K>);
+
+impl<K: Clone + Debug + Default + Ord, V: Clone + Debug + Default> TreeMap<K, V> {
+    pub fn get<'a>(&self, key: &'a K) -> Option<&V> {
+        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
+        let key = MapKeyRef(Some(key));
+        cursor.seek(&key, Bias::Left, &());
+        if key.cmp(cursor.start(), &()) == Ordering::Equal {
+            Some(&cursor.item().unwrap().value)
+        } else {
+            None
+        }
+    }
+
+    pub fn insert(&mut self, key: K, value: V) {
+        self.0.insert_or_replace(MapEntry { key, value }, &());
+    }
+
+    pub fn remove<'a>(&mut self, key: &'a K) -> Option<V> {
+        let mut removed = None;
+        let mut cursor = self.0.cursor::<MapKeyRef<'_, K>>();
+        let key = MapKeyRef(Some(key));
+        let mut new_tree = cursor.slice(&key, Bias::Left, &());
+        if key.cmp(cursor.start(), &()) == Ordering::Equal {
+            removed = Some(cursor.item().unwrap().value.clone());
+            cursor.next(&());
+        }
+        new_tree.push_tree(cursor.suffix(&()), &());
+        drop(cursor);
+        self.0 = new_tree;
+        removed
+    }
+}
+
+impl<K, V> Item for MapEntry<K, V>
+where
+    K: Clone + Debug + Default + Clone,
+    V: Clone,
+{
+    type Summary = MapKey<K>;
+
+    fn summary(&self) -> Self::Summary {
+        todo!()
+    }
+}
+
+impl<K, V> KeyedItem for MapEntry<K, V>
+where
+    K: Clone + Debug + Default + Ord,
+    V: Clone,
+{
+    type Key = MapKey<K>;
+
+    fn key(&self) -> Self::Key {
+        MapKey(self.key.clone())
+    }
+}
+
+impl<K> Summary for MapKey<K>
+where
+    K: Clone + Debug + Default,
+{
+    type Context = ();
+
+    fn add_summary(&mut self, summary: &Self, cx: &()) {
+        *self = summary.clone()
+    }
+}
+
+impl<'a, K> Dimension<'a, MapKey<K>> for MapKeyRef<'a, K>
+where
+    K: Clone + Debug + Default + Ord,
+{
+    fn add_summary(&mut self, summary: &'a MapKey<K>, _: &()) {
+        self.0 = Some(&summary.0)
+    }
+}
+
+impl<'a, K> SeekTarget<'a, MapKey<K>, MapKeyRef<'a, K>> for MapKeyRef<'_, K>
+where
+    K: Clone + Debug + Default + Ord,
+{
+    fn cmp(&self, cursor_location: &MapKeyRef<K>, cx: &()) -> Ordering {
+        if let Some(key) = cursor_location.0 {
+            self.0.cmp(&cursor_location.0)
+        } else {
+            Ordering::Greater
+        }
+    }
+}