example: Add SearchArgument for EvenCount as an example on how to enh… (

#29)

* example: Add SearchArgument for EvenCount as an example on how to enhaunce query

* rename example name

* refine

* more comment

examples/augment_even_count.rs

@@ -0,0 +1,95 @@
+use sweep_bptree::argument::SearchArgument;
+use sweep_bptree::{argument::Argument, BPlusTreeMap};
+
+/// An argumentation that counts even numbers
+#[derive(Default, Clone, Debug)]
+struct EvenCount(usize);
+
+fn value_is_even(v: i64) -> bool {
+    v % 2 == 0
+}
+
+impl Argument<i64> for EvenCount {
+    fn from_leaf(keys: &[i64]) -> Self {
+        // For leafs, we count all keys that are even
+        Self(keys.iter().filter(|i| value_is_even(**i)).count())
+    }
+
+    fn from_inner(_keys: &[i64], arguments: &[Self]) -> Self {
+        // For inner nodes, we aggregate all the EvenCount
+        Self(arguments.iter().map(|a| a.0).sum::<usize>())
+    }
+}
+
+/// This implementation enables get key by 'nth' even number. This effectively makes
+/// `EvenCount` a secondary index
+impl SearchArgument<i64> for EvenCount {
+    /// offset of even number
+    type Query = usize;
+
+    fn locate_in_leaf(mut offset: usize, keys: &[i64]) -> Option<usize> {
+        for (idx, key) in keys.iter().enumerate() {
+            if value_is_even(*key) {
+                if offset == 0 {
+                    return Some(idx);
+                }
+
+                offset -= 1;
+            }
+        }
+
+        None
+    }
+
+    fn locate_in_inner(offset: usize, _keys: &[i64], arguments: &[Self]) -> Option<(usize, usize)> {
+        let mut relative_offset = offset;
+        for (child_idx, argument) in arguments.iter().enumerate() {
+            if argument.0 > relative_offset {
+                return Some((child_idx, relative_offset));
+            } else {
+                relative_offset -= argument.0
+            }
+        }
+
+        // not found
+        None
+    }
+}
+
+fn main() {
+    // create a tree with the argument
+    let mut tree = BPlusTreeMap::<i64, i64, EvenCount>::new();
+
+    // insert 100000 numbers
+    for i in 0..100000 {
+        tree.insert(i, i);
+    }
+
+    // check we get the correct count
+    assert_eq!(dbg!(tree.root_argument()).0, 50000);
+
+    // then remove some keys
+    for i in 0..100 {
+        tree.remove(&(i * 2));
+    }
+    assert_eq!(dbg!(tree.root_argument()).0, 49900);
+
+    // remove odd numbers should not affect the even count
+    for i in 0..100 {
+        tree.remove(&(i * 2 + 1));
+    }
+    assert_eq!(dbg!(tree.root_argument()).0, 49900);
+
+    // able to get nth even value easily
+    for i in 0..tree.root_argument().0 {
+        let Some((k, _)) = tree.get_by_argument::<usize>(i) else {
+            panic!("should got a value");
+        };
+        assert_eq!(k % 2, 0);
+    }
+
+    // offset = length - 1, get(length) should be None
+    assert!(tree
+        .get_by_argument::<usize>(tree.root_argument().0)
+        .is_none());
+}

src/argument/group.rs

@@ -6,7 +6,7 @@ use super::{Argument, RankArgument, SearchArgument};
 
 /// Argument to count the number of groups in a set of keys
 /// Note, the group must be ordered
-/// This Argument basicly provides two capabilities:
+/// This Argument basically provides two capabilities:
 /// 1. Get the group count
 /// 2. Query inside group by offset
 #[derive(Clone, Debug)]

src/argument/mod.rs

@@ -3,7 +3,9 @@ use crate::Key;
 pub mod count;
 pub mod group;
 
-/// Augument trait, it is used to store augumentation, like 'size'
+/// Argument trait, it is used to store argumentation, like 'size'
+/// NOTE: Since the lib has no control on how value changes, so argument only calculated from keys
+/// e.g: Map<i64, Arc<Mutex<i64>>
 pub trait Argument<K: Key>: Clone + Default {
     fn is_zst() -> bool {
         false
@@ -23,6 +25,8 @@ pub trait Argument<K: Key>: Clone + Default {
 }
 
 /// Whether the argumentation able to locate element
+/// `SearchArgument` acts like a secondary index, it is able to locate
+/// the record.
 pub trait SearchArgument<K: Key>: Argument<K> {
     type Query;
 
@@ -68,12 +72,8 @@ impl<K: Key> Argument<K> for () {
     }
 
     #[inline(always)]
-    fn from_leaf(_: &[K]) -> Self {
-
-    }
+    fn from_leaf(_: &[K]) -> Self {}
 
     #[inline(always)]
-    fn from_inner(_: &[K], _: &[Self]) -> Self {
-
-    }
+    fn from_inner(_: &[K], _: &[Self]) -> Self {}
 }

src/tree/cursor.rs

@@ -4,6 +4,8 @@ use crate::Key;
 use crate::NodeStore;
 
 /// `Cursor` points to a key value pair in the tree. Not like Iterator, it can move to next or prev.
+/// One key feature of this Cursor is, it didn't borrow the tree, so multiple Cursor can be created
+/// and the underlying tree can be updated without invalidate existing Cursor.
 #[derive(Debug, Clone, Copy)]
 pub struct Cursor<K: Key> {
     /// The key this cursor points to. It is possible the `k` doesn't exist in the tree.

src/tree/leaf_node.rs

@@ -325,7 +325,7 @@ impl<K: Key, V> LeafNode<K, V> {
         match self.locate_slot(k) {
             Ok(idx) => {
                 // exact match, go to right child.
-                // if the child split, then the new key should inserted idx + 1
+                // if the child split, then the new key should insert idx + 1
                 (idx, {
                     let v = unsafe { self.value_area(idx).assume_init_ref() };
                     Some(v)

src/tree/mod.rs

@@ -612,12 +612,13 @@ where
         iterator::Iter::new(self)
     }
 
-    /// Create an cursor from first elem
+    /// Create a `Cursor` from first elem if exists
+    ///
     pub fn cursor_first(&self) -> Option<Cursor<S::K>> {
         Cursor::first(self).map(|c| c.0)
     }
 
-    /// Create an cursor for k
+    /// Create a cursor and value for k if k exists.
     pub fn get_cursor(&self, k: &S::K) -> Option<(Cursor<S::K>, Option<&S::V>)> {
         let node_id = self.root;
         let leaf_id = match node_id {
@@ -737,7 +738,7 @@ where
         Self::remove_by_ref(entry_ref.into_detached().into_ref(self))
     }
 
-    /// Get the (&K, &V) pair for referece
+    /// Get the (&K, &V) pair for `EntryRef`
     fn get_by_ref(entry_ref: EntryRef<&Self>) -> Option<(&S::K, &S::V)> {
         let leaf = entry_ref.tree.node_store.get_leaf(entry_ref.leaf_id);
         let slot = entry_ref.offset;