refactor

XiangpengHao · Nov 20, 2024 · 72c5fc2 · 72c5fc2
1 parent 3d62451
commit 72c5fc2
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 162 deletions.
diff --git a/src/node_ptr.rs b/src/node_ptr.rs
@@ -3,26 +3,8 @@ use std::ptr::NonNull;
 use crate::{
     base_node::{BaseNode, Node},
     node_256::Node256,
-    utils::LastLevelKey,
 };
 
-mod private {
-    use super::{ChildIsPayload, LastLevelKey};
-
-    pub trait LastLevelProofInner {}
-
-    impl<const K_LEN: usize> LastLevelProofInner for LastLevelKey<'_, K_LEN> {}
-
-    impl LastLevelProofInner for ChildIsPayload<'_> {}
-}
-
-/// A proof that the NodePtr is the payload rather than a sub node
-pub(crate) trait LastLevelProof: private::LastLevelProofInner {}
-
-impl LastLevelProof for ChildIsPayload<'_> {}
-
-impl<const K_LEN: usize> LastLevelProof for LastLevelKey<'_, K_LEN> {}
-
 pub(crate) struct ChildIsPayload<'a> {
     _marker: std::marker::PhantomData<&'a ()>,
 }
@@ -84,11 +66,6 @@ impl NodePtr {
         unsafe { self.payload }
     }
 
-    pub(crate) fn as_payload<P: LastLevelProof>(&self, _proof: &P) -> usize {
-        // Safety: We have a proof that the node is at the last level
-        unsafe { self.as_payload_unchecked() }
-    }
-
     pub(crate) fn as_ptr_safe<const MAX_LEVEL: usize>(
         &self,
         current_level: usize,

diff --git a/src/range_scan.rs b/src/range_scan.rs
@@ -1,5 +1,6 @@
 use crate::error::ArtError;
 use crate::node_256::Node256;
+use crate::node_ptr::PtrType;
 use crate::utils::LastLevelKey;
 use crate::{base_node::BaseNode, lock::ReadGuard, node_ptr::NodePtr, utils::KeyTracker};
 use std::cmp;
@@ -239,37 +240,41 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
         node: NodePtr,
         key_tracker: &KeyTracker<K_LEN>,
     ) -> Result<(), ArtError> {
-        if let Some(last_level_key) = key_tracker.as_last_level() {
-            if self.key_in_range(&last_level_key) {
-                if self.result_found == self.result.len() {
-                    self.to_continue = node.as_payload(&last_level_key);
-                    return Ok(());
-                }
-                self.result[self.result_found] =
-                    (key_tracker.get_key(), node.as_payload(&last_level_key));
-                self.result_found += 1;
-            };
-        } else {
-            let node = BaseNode::read_lock_deprecated::<K_LEN>(node, key_tracker.len())?;
-            let mut key_tracker = key_tracker.clone();
+        match node.downcast::<K_LEN>(key_tracker.len() - 1) {
+            PtrType::Payload(payload) => {
+                let last_level_key = unsafe { key_tracker.as_last_level_unchecked() };
+                if self.key_in_range(&last_level_key) {
+                    if self.result_found == self.result.len() {
+                        self.to_continue = payload;
+                        return Ok(());
+                    }
+                    self.result[self.result_found] = (key_tracker.get_key(), payload);
+                    self.result_found += 1;
+                };
+            }
+            PtrType::SubNode(sub_node) => {
+                let node = BaseNode::read_lock(sub_node)?;
+                let mut key_tracker = key_tracker.clone();
 
-            let children = node.as_ref().get_children(0, 255);
+                let children = node.as_ref().get_children(0, 255);
 
-            for (k, c) in children {
-                node.check_version()?;
+                for (k, c) in children {
+                    node.check_version()?;
 
-                key_tracker.push(k);
+                    key_tracker.push(k);
 
-                let cur_key = KeyTracker::append_prefix(c, &key_tracker);
-                self.copy_node(c, &cur_key)?;
+                    let cur_key = KeyTracker::append_prefix(c, &key_tracker);
+                    self.copy_node(c, &cur_key)?;
 
-                if self.to_continue != 0 {
-                    break;
-                }
+                    if self.to_continue != 0 {
+                        break;
+                    }
 
-                key_tracker.pop();
+                    key_tracker.pop();
+                }
             }
-        }
+        };
+
         Ok(())
     }
 

diff --git a/src/utils.rs b/src/utils.rs
@@ -118,14 +118,11 @@ impl<const K_LEN: usize> KeyTracker<K_LEN> {
         v
     }
 
-    pub(crate) fn as_last_level(&self) -> Option<LastLevelKey<K_LEN>> {
-        if self.len == K_LEN {
-            Some(LastLevelKey { key: self })
-        } else {
-            None
-        }
+    pub(crate) unsafe fn as_last_level_unchecked(&self) -> LastLevelKey<K_LEN> {
+        LastLevelKey { key: self }
     }
 
+
     pub(crate) fn get_key<const N: usize>(&self) -> [u8; N] {
         assert!(self.len == N);
         self.data[..N].try_into().unwrap()

diff --git a/tests/art_usize.rs b/tests/art_usize.rs
@@ -184,112 +184,3 @@ fn fuzz_0() {
 
     test_runner(&ops);
 }
-
-#[cfg(feature = "db_extension")]
-#[test]
-fn compute_if_present() {
-    let tree = Congee::default();
-    let guard = tree.pin();
-    _ = tree.insert(1, 42, &guard);
-    let (old_v, new_v) = tree
-        .compute_if_present(
-            &1,
-            |v| {
-                assert_eq!(v, 42);
-                Some(v + 1)
-            },
-            &guard,
-        )
-        .unwrap();
-    assert_eq!(old_v, 42);
-    assert_eq!(new_v, Some(43));
-
-    let mut tmp_v = 0;
-    let (old_v, new_v) = tree
-        .compute_if_present(
-            &1,
-            |v| {
-                assert_eq!(v, 43);
-                tmp_v = v;
-                Some(v + 1)
-            },
-            &guard,
-        )
-        .unwrap();
-    assert_eq!(old_v, 43);
-    assert_eq!(new_v, Some(44));
-    assert_eq!(tmp_v, 43);
-}
-
-#[cfg(feature = "db_extension")]
-#[test]
-fn random_value() {
-    let tree = Congee::default();
-    let guard = tree.pin();
-    _ = tree.insert(1, 42, &guard);
-    let mut rng = rand::thread_rng();
-    let (key, old_v, new_v) = tree
-        .compute_on_random(
-            &mut rng,
-            |k, v| {
-                assert_eq!(k, 1);
-                assert_eq!(v, 42);
-                v + 1
-            },
-            &guard,
-        )
-        .unwrap();
-    assert_eq!(key, 1);
-    assert_eq!(old_v, 42);
-    assert_eq!(new_v, 43);
-
-    let mut tmp_v = 0;
-    let (key, old_v, new_v) = tree
-        .compute_on_random(
-            &mut rng,
-            |k, v| {
-                assert_eq!(k, 1);
-                assert_eq!(v, 43);
-                tmp_v = v;
-                v + 1
-            },
-            &guard,
-        )
-        .unwrap();
-    assert_eq!(key, 1);
-    assert_eq!(old_v, 43);
-    assert_eq!(new_v, 44);
-    assert_eq!(tmp_v, 43);
-}
-
-#[cfg(feature = "db_extension")]
-#[test]
-fn compare_exchange() {
-    let tree = Congee::default();
-    let guard = tree.pin();
-    _ = tree.insert(1, 42, &guard);
-
-    let v = tree.compare_exchange(&1, &42, Some(43), &guard).unwrap();
-    assert_eq!(v, Some(43));
-
-    if let Err(v) = tree.compare_exchange(&1, &42, Some(45), &guard) {
-        assert_eq!(v.unwrap(), 43);
-    } else {
-        panic!("should have failed");
-    }
-
-    match tree.compare_exchange(&1, &43, Some(43), &guard) {
-        Ok(v) => assert_eq!(v, Some(43)),
-        Err(_v) => panic!("should have succeeded because the old value matches."),
-    }
-    match tree.compare_exchange(&1, &42, Some(43), &guard) {
-        Ok(_) => panic!("should failed because the old value is wrong"),
-        Err(v) => assert_eq!(v.unwrap(), 43),
-    }
-
-    if let Err(v) = tree.compare_exchange(&0, &43, Some(42), &guard) {
-        assert!(v.is_none());
-    } else {
-        panic!("should have failed");
-    }
-}