chore: add feature flag for btree version

surrealdb · Dec 30, 2024 · 0952aae · 0952aae
1 parent 435bb15
commit 0952aae
Show file tree

Hide file tree

Showing 3 changed files with 121 additions and 14 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -24,6 +24,10 @@ rand = "0.8.5"
 criterion = "0.5.1"
 divan = "0.1.14"
 
+[features]
+default = ["vec_store"]
+vec_store = []
+
 [[bench]]
 name = "vart_bench"
 path = "benches/vart_bench.rs"

diff --git a/src/node.rs b/src/node.rs
@@ -1,7 +1,14 @@
 use std::slice::from_ref;
 use std::sync::Arc;
 
-use crate::{art::QueryType, version::BTree, KeyTrait};
+use crate::{
+    art::QueryType,
+    version::{VecStore, VersionStore},
+    KeyTrait,
+};
+
+#[cfg(not(feature = "vec_store"))]
+use crate::version::BTreeStore;
 
 /*
     Immutable nodes
@@ -22,7 +29,10 @@ pub(crate) trait NodeTrait<N> {
 pub(crate) struct TwigNode<K: KeyTrait, V: Clone> {
     pub(crate) prefix: K,
     pub(crate) key: K,
-    pub(crate) values: BTree<V>,
+    #[cfg(feature = "vec_store")]
+    pub(crate) values: VecStore<V>,
+    #[cfg(not(feature = "vec_store"))]
+    pub(crate) values: BTreeStore<V>,
     pub(crate) version: u64, // Version for the twig node
 }
 
@@ -52,28 +62,20 @@ impl<K: KeyTrait, V: Clone> TwigNode<K, V> {
         TwigNode {
             prefix,
             key,
-            values: BTree::new(),
+            values: VersionStore::new(),
             version: 0,
         }
     }
 
     pub(crate) fn version(&self) -> u64 {
-        self.values
-            .iter()
-            .map(|value| value.version)
-            .max()
-            .unwrap_or(self.version)
+        self.values.get_max_version().unwrap_or(self.version)
     }
 
     pub(crate) fn insert(&self, value: V, version: u64, ts: u64) -> TwigNode<K, V> {
         let mut new_values = self.values.clone();
         new_values.insert(value, version, ts);
 
-        let new_version = new_values
-            .iter()
-            .map(|value| value.version)
-            .max()
-            .unwrap_or(self.version);
+        let new_version = new_values.get_max_version().unwrap_or(self.version);
 
         TwigNode {
             prefix: self.prefix.clone(),

diff --git a/src/version.rs b/src/version.rs
@@ -15,6 +15,7 @@ pub(crate) struct BNode<V: Clone> {
     is_leaf: bool,
 }
 
+#[allow(unused)]
 impl<V: Clone> BNode<V> {
     fn new(is_leaf: bool) -> Self {
         BNode {
@@ -34,6 +35,7 @@ impl<V: Clone> BNode<V> {
     }
 }
 
+#[allow(unused)]
 impl<V: Clone> BTree<V> {
     pub fn new() -> Self {
         BTree { root: None }
@@ -220,7 +222,7 @@ impl<'a, V: Clone> Iterator for BNodeIterator<'a, V> {
     }
 }
 
-impl<'a, V: Clone> DoubleEndedIterator for BNodeIterator<'a, V> {
+impl<V: Clone> DoubleEndedIterator for BNodeIterator<'_, V> {
     fn next_back(&mut self) -> Option<Self::Item> {
         while let Some((node, idx)) = self.reverse_stack.last_mut() {
             if *idx > 0 {
@@ -241,6 +243,105 @@ impl<'a, V: Clone> DoubleEndedIterator for BNodeIterator<'a, V> {
     }
 }
 
+pub(crate) trait VersionStore<V: Clone> {
+    fn new() -> Self;
+    fn insert(&mut self, value: V, version: u64, ts: u64);
+    fn clear(&mut self);
+    fn iter(&self) -> Box<dyn DoubleEndedIterator<Item = &Arc<LeafValue<V>>> + '_>;
+    fn get_max_version(&self) -> Option<u64>;
+}
+
+// Implement for Vec
+#[cfg(feature = "vec_store")]
+#[derive(Clone)]
+pub(crate) struct VecStore<V: Clone> {
+    values: Vec<Arc<LeafValue<V>>>,
+}
+
+#[cfg(feature = "vec_store")]
+impl<V: Clone> VersionStore<V> for VecStore<V> {
+    fn new() -> Self {
+        Self { values: Vec::new() }
+    }
+
+    fn insert(&mut self, value: V, version: u64, ts: u64) {
+        let new_leaf_value = LeafValue::new(value, version, ts);
+        match self
+            .values
+            .binary_search_by(|v| v.version.cmp(&new_leaf_value.version))
+        {
+            Ok(index) => {
+                if self.values[index].ts == ts {
+                    self.values[index] = Arc::new(new_leaf_value);
+                } else {
+                    let mut insert_position = index;
+                    if self.values[index].ts < ts {
+                        insert_position += self.values[index..]
+                            .iter()
+                            .take_while(|v| v.ts <= ts)
+                            .count();
+                    } else {
+                        insert_position -= self.values[..index]
+                            .iter()
+                            .rev()
+                            .take_while(|v| v.ts >= ts)
+                            .count();
+                    }
+                    self.values
+                        .insert(insert_position, Arc::new(new_leaf_value));
+                }
+            }
+            Err(index) => {
+                self.values.insert(index, Arc::new(new_leaf_value));
+            }
+        }
+    }
+
+    fn clear(&mut self) {
+        self.values.clear();
+    }
+
+    fn iter(&self) -> Box<dyn DoubleEndedIterator<Item = &Arc<LeafValue<V>>> + '_> {
+        Box::new(self.values.iter())
+    }
+
+    fn get_max_version(&self) -> Option<u64> {
+        self.values.iter().map(|value| value.version).max()
+    }
+}
+
+// Implement for BTree
+#[cfg(not(feature = "vec_store"))]
+#[derive(Clone)]
+pub(crate) struct BTreeStore<V: Clone> {
+    values: BTree<V>,
+}
+
+#[cfg(not(feature = "vec_store"))]
+impl<V: Clone> VersionStore<V> for BTreeStore<V> {
+    fn new() -> Self {
+        Self {
+            values: BTree::new(),
+        }
+    }
+
+    fn insert(&mut self, value: V, version: u64, ts: u64) {
+        self.values.insert(value, version, ts);
+    }
+
+    fn clear(&mut self) {
+        self.values.clear();
+    }
+
+    fn iter(&self) -> Box<dyn DoubleEndedIterator<Item = &Arc<LeafValue<V>>> + '_> {
+        Box::new(self.values.iter())
+    }
+
+    fn get_max_version(&self) -> Option<u64> {
+        self.values.iter().map(|value| value.version).max()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;