From df33556693a78cc0977de80c12eddefc037ad964 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Mon, 18 Jul 2022 13:40:30 +0200
Subject: [PATCH] Introduce a new `TreeSet` struct in `sum_tree`

This is just a special case of `TreeMap<K, V>` where `V = ()`.
---
 crates/sum_tree/src/sum_tree.rs |  2 +-
 crates/sum_tree/src/tree_map.rs | 31 +++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/crates/sum_tree/src/sum_tree.rs b/crates/sum_tree/src/sum_tree.rs
index 193786112b46f5dec84c0343ca8e3c9d5c24926a..efd5c43480174b777674eeed331013466841fa8a 100644
--- a/crates/sum_tree/src/sum_tree.rs
+++ b/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;
diff --git a/crates/sum_tree/src/tree_map.rs b/crates/sum_tree/src/tree_map.rs
index 80143aad6971d9cde4164236e77d142060a24b14..5218d2b4dbad5e4c3b59573c4db9e3ca3aab2f8a 100644
--- a/crates/sum_tree/src/tree_map.rs
+++ b/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::*;