refactor

XiangpengHao · Nov 19, 2024 · 6339c15 · 6339c15
1 parent 52b8d6c
commit 6339c15
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 34 deletions.
diff --git a/src/node_ptr.rs b/src/node_ptr.rs
@@ -3,9 +3,41 @@ use std::ptr::NonNull;
 use crate::{
     base_node::{BaseNode, Node},
     node_256::Node256,
+    utils::LastLevelKey,
 };
 
-pub(crate) struct LastLevelProof {}
+mod private {
+    use super::{ChildIsPayload, LastLevelKey};
+
+    pub trait LastLevelProofInner {}
+
+    impl LastLevelProofInner for LastLevelKey<'_> {}
+
+    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 LastLevelProof for LastLevelKey<'_> {}
+
+pub(crate) struct ChildIsPayload<'a> {
+    _marker: std::marker::PhantomData<&'a ()>,
+}
+
+impl<'a> ChildIsPayload<'a> {
+    pub(crate) fn try_new<const K_LEN: usize>(level: usize) -> Option<Self> {
+        if level == (K_LEN - 1) {
+            Some(Self {
+                _marker: std::marker::PhantomData,
+            })
+        } else {
+            None
+        }
+    }
+}
 
 #[derive(Clone, Copy)]
 pub(crate) union NodePtr {
@@ -39,7 +71,7 @@ impl NodePtr {
         unsafe { self.payload }
     }
 
-    pub(crate) fn as_payload(&self, _proof: &LastLevelProof) -> usize {
+    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() }
     }

diff --git a/src/range_scan.rs b/src/range_scan.rs
@@ -1,6 +1,6 @@
 use crate::error::ArtError;
 use crate::node_256::Node256;
-use crate::node_ptr::LastLevelProof;
+use crate::utils::LastLevelKey;
 use crate::{base_node::BaseNode, lock::ReadGuard, node_ptr::NodePtr, utils::KeyTracker};
 use std::cmp;
 use std::ptr::NonNull;
@@ -41,9 +41,8 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
         self.start < self.end
     }
 
-    fn key_in_range(&self, key: &KeyTracker, last_level_proof: &LastLevelProof) -> bool {
-        debug_assert_eq!(key.len(), 8);
-        let cur_key = key.to_usize_key(last_level_proof);
+    fn key_in_range(&self, key: &LastLevelKey) -> bool {
+        let cur_key = key.to_usize_key();
 
         let start_key = unsafe { *(self.start.as_ptr() as *const usize) }.swap_bytes();
         let end_key = unsafe { *(self.end.as_ptr() as *const usize) }.swap_bytes();
@@ -243,13 +242,14 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
     }
 
     fn copy_node(&mut self, node: NodePtr, key_tracker: &KeyTracker) -> Result<(), ArtError> {
-        if let Some(proof) = key_tracker.is_last_level::<K_LEN>() {
-            if self.key_in_range(key_tracker, &proof) {
+        if let Some(last_level_key) = key_tracker.as_last_level::<K_LEN>() {
+            if self.key_in_range(&last_level_key) {
                 if self.result_found == self.result.len() {
-                    self.to_continue = node.as_payload(&proof);
+                    self.to_continue = node.as_payload(&last_level_key);
                     return Ok(());
                 }
-                self.result[self.result_found] = (key_tracker.get_key(), node.as_payload(&proof));
+                self.result[self.result_found] =
+                    (key_tracker.get_key(), node.as_payload(&last_level_key));
                 self.result_found += 1;
             };
         } else {

diff --git a/src/tree.rs b/src/tree.rs
@@ -8,7 +8,7 @@ use crate::{
     lock::ReadGuard,
     node_256::Node256,
     node_4::Node4,
-    node_ptr::{LastLevelProof, NodePtr},
+    node_ptr::{ChildIsPayload, NodePtr},
     range_scan::RangeScan,
     utils::Backoff,
     Allocator, DefaultAllocator,
@@ -103,12 +103,8 @@ impl<const K_LEN: usize, A: Allocator + Clone + Send> RawCongee<K_LEN, A> {
         }
     }
 
-    fn is_last_level(current_level: usize) -> Option<LastLevelProof> {
-        if current_level == (K_LEN - 1) {
-            Some(LastLevelProof {})
-        } else {
-            None
-        }
+    fn is_last_level<'a>(current_level: usize) -> Option<ChildIsPayload<'a>> {
+        ChildIsPayload::try_new::<K_LEN>(current_level)
     }
 
     #[inline]
@@ -522,25 +518,28 @@ impl<const K_LEN: usize, A: Allocator + Clone + Send> RawCongee<K_LEN, A> {
 
             key_tracker.push(k);
 
-            if let Some(proof) = key_tracker.is_last_level::<K_LEN>() {
+            if let Some(last_level_key) = key_tracker.as_last_level::<K_LEN>() {
                 let new_v = f(
-                    key_tracker.to_usize_key(&proof),
-                    child_node.as_payload(&proof),
+                    last_level_key.to_usize_key(),
+                    child_node.as_payload(&last_level_key),
                 );
-                if new_v == child_node.as_payload(&proof) {
+                if new_v == child_node.as_payload(&last_level_key) {
                     // Don't acquire the lock if the value is not changed
-                    return Ok(Some((key_tracker.to_usize_key(&proof), new_v, new_v)));
+                    return Ok(Some((last_level_key.to_usize_key(), new_v, new_v)));
                 }
 
                 let mut write_n = node.upgrade().map_err(|(_n, v)| v)?;
 
                 let old_v = write_n.as_mut().change(k, NodePtr::from_payload(new_v));
 
-                debug_assert_eq!(old_v.as_payload(&proof), child_node.as_payload(&proof));
+                debug_assert_eq!(
+                    old_v.as_payload(&last_level_key),
+                    child_node.as_payload(&last_level_key)
+                );
 
                 return Ok(Some((
-                    key_tracker.to_usize_key(&proof),
-                    child_node.as_payload(&proof),
+                    last_level_key.to_usize_key(),
+                    child_node.as_payload(&last_level_key),
                     new_v,
                 )));
             }

diff --git a/src/utils.rs b/src/utils.rs
@@ -1,5 +1,5 @@
 use crate::base_node::{BaseNode, MAX_KEY_LEN};
-use crate::node_ptr::{LastLevelProof, NodePtr};
+use crate::node_ptr::NodePtr;
 use core::cell::Cell;
 use core::fmt;
 
@@ -77,6 +77,17 @@ impl Default for Backoff {
     }
 }
 
+pub(crate) struct LastLevelKey<'a> {
+    key: &'a KeyTracker,
+}
+
+impl<'a> LastLevelKey<'a> {
+    pub(crate) fn to_usize_key(&self) -> usize {
+        let val = unsafe { *((&self.key.data) as *const [u8; 8] as *const usize) };
+        val.swap_bytes()
+    }
+}
+
 #[derive(Default, Clone)]
 pub(crate) struct KeyTracker {
     len: usize,
@@ -101,20 +112,14 @@ impl KeyTracker {
         v
     }
 
-    pub(crate) fn is_last_level<const K_LEN: usize>(&self) -> Option<LastLevelProof> {
+    pub(crate) fn as_last_level<const K_LEN: usize>(&self) -> Option<LastLevelKey> {
         if self.len == K_LEN {
-            Some(LastLevelProof {})
+            Some(LastLevelKey { key: self })
         } else {
             None
         }
     }
 
-    #[inline]
-    pub(crate) fn to_usize_key(&self, _last_level_proof: &LastLevelProof) -> usize {
-        let val = unsafe { *((&self.data) as *const [u8; 8] as *const usize) };
-        val.swap_bytes()
-    }
-
     pub(crate) fn get_key<const N: usize>(&self) -> [u8; N] {
         assert!(self.len == N);
         self.data[..N].try_into().unwrap()