Introduce a new `TreeSet` struct in `sum_tree`

Antonio Scandurra created 3 years ago

This is just a special case of `TreeMap<K, V>` where `V = ()`.

Change summary

crates/sum_tree/src/sum_tree.rs |  2 +-
crates/sum_tree/src/tree_map.rs | 31 +++++++++++++++++++++++++++++++
2 files changed, 32 insertions(+), 1 deletion(-)

Detailed changes

crates/sum_tree/src/sum_tree.rs

@@ -5,7 +5,7 @@ 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;
+pub use tree_map::{TreeMap, TreeSet};
 
 #[cfg(test)]
 const TREE_BASE: usize = 2;

crates/sum_tree/src/tree_map.rs

@@ -20,6 +20,11 @@ pub struct MapKey<K>(K);
 #[derive(Clone, Debug, Default)]
 pub struct MapKeyRef<'a, K>(Option<&'a K>);
 
+#[derive(Clone)]
+pub struct TreeSet<K>(TreeMap<K, ()>)
+where
+    K: Clone + Debug + Default + Ord;
+
 impl<K: Clone + Debug + Default + Ord, V: Clone + Debug> TreeMap<K, V> {
     pub fn from_ordered_entries(entries: impl IntoIterator<Item = (K, V)>) -> Self {
         let tree = SumTree::from_iter(
@@ -136,6 +141,32 @@ where
     }
 }
 
+impl<K> Default for TreeSet<K>
+where
+    K: Clone + Debug + Default + Ord,
+{
+    fn default() -> Self {
+        Self(Default::default())
+    }
+}
+
+impl<K> TreeSet<K>
+where
+    K: Clone + Debug + Default + Ord,
+{
+    pub fn insert(&mut self, key: K) {
+        self.0.insert(key, ());
+    }
+
+    pub fn contains(&self, key: &K) -> bool {
+        self.0.get(key).is_some()
+    }
+
+    pub fn iter<'a>(&'a self) -> impl 'a + Iterator<Item = &'a K> {
+        self.0.iter().map(|(k, _)| k)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;