database: improve BTreeCappedSet

Signed-off-by: Yuki Kishimoto <[email protected]>

CHANGELOG.md

@@ -33,6 +33,7 @@
 * pool: changes in `RelayPool::remove_relay` behavior ([Yuki Kishimoto])
 * pool: deprecate `RelayPool::remove_all_relays` ([Yuki Kishimoto])
 * sdk: deprecate `Client::get_events_of` and `Client::get_events_from` methods ([Yuki Kishimoto])
+* database: improve `BTreeCappedSet` ([Yuki Kishimoto])
 * lmdb: not save event deletion if contains not owned events ([Yuki Kishimoto])
 * lmdb: return iterator instead of vector in `Lmdb::single_filter_query` ([Yuki Kishimoto])
 * signer: bootstrap NIP-46 signer on demand ([Yuki Kishimoto])

crates/nostr-database/src/tree.rs

@@ -5,7 +5,8 @@
 #![allow(dead_code)]
 
 use std::borrow::Borrow;
-use std::collections::btree_set::Iter;
+use std::cmp::Ordering;
+use std::collections::btree_set::{IntoIter, Iter};
 use std::collections::BTreeSet;
 
 /// Represents the possible options for removing a value.
@@ -18,7 +19,7 @@ pub enum OverCapacityPolicy {
     Last,
 }
 
-#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
 pub enum Capacity {
     #[default]
     Unbounded,
@@ -28,6 +29,25 @@ pub enum Capacity {
     },
 }
 
+impl PartialOrd for Capacity {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Capacity {
+    fn cmp(&self, other: &Self) -> Ordering {
+        match (self, other) {
+            (Self::Unbounded, Self::Unbounded) => Ordering::Equal,
+            (Self::Unbounded, Self::Bounded { .. }) => Ordering::Greater,
+            (Self::Bounded { .. }, Self::Unbounded) => Ordering::Less,
+            (Self::Bounded { max: this_max, .. }, Self::Bounded { max: other_max, .. }) => {
+                this_max.cmp(other_max)
+            }
+        }
+    }
+}
+
 impl Capacity {
     #[inline]
     pub fn bounded(max: usize) -> Self {
@@ -117,6 +137,11 @@ where
         self.set.len()
     }
 
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.set.is_empty()
+    }
+
     #[inline]
     pub fn contains<Q>(&self, value: &Q) -> bool
     where
@@ -173,6 +198,24 @@ where
         }
     }
 
+    /// Extend with values
+    pub fn extend<I>(&mut self, values: I)
+    where
+        I: IntoIterator<Item = T>,
+    {
+        match self.capacity {
+            Capacity::Bounded { .. } => {
+                // TODO: find more efficient way
+                for value in values.into_iter() {
+                    self.insert(value);
+                }
+            }
+            Capacity::Unbounded => {
+                self.set.extend(values);
+            }
+        }
+    }
+
     #[inline]
     pub fn remove<Q>(&mut self, value: &Q) -> bool
     where
@@ -182,12 +225,48 @@ where
         self.set.remove(value)
     }
 
+    /// Get first value
+    #[inline]
+    pub fn first(&self) -> Option<&T>
+    where
+        T: Ord,
+    {
+        self.set.first()
+    }
+
+    /// Get last value
+    #[inline]
+    pub fn last(&self) -> Option<&T>
+    where
+        T: Ord,
+    {
+        self.set.last()
+    }
+
     #[inline]
     pub fn iter(&self) -> Iter<'_, T> {
         self.set.iter()
     }
 }
 
+impl<T> From<BTreeSet<T>> for BTreeCappedSet<T> {
+    fn from(set: BTreeSet<T>) -> Self {
+        Self {
+            set,
+            capacity: Capacity::Unbounded,
+        }
+    }
+}
+
+impl<T> IntoIterator for BTreeCappedSet<T> {
+    type Item = T;
+    type IntoIter = IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.set.into_iter()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -285,4 +364,11 @@ mod tests {
         assert_eq!(iter.next(), Some(&2));
         assert_eq!(iter.next(), Some(&3));
     }
+
+    #[test]
+    fn test_cmp_capacity() {
+        assert!(Capacity::Unbounded > Capacity::bounded(1000));
+        assert!(Capacity::bounded(1) < Capacity::bounded(1000));
+        assert_eq!(Capacity::Unbounded, Capacity::Unbounded);
+    }
 }