From 33d7ae4995561fd11c5b1014a6c7cfc0a5763dac Mon Sep 17 00:00:00 2001
From: Xiangpeng Hao <haoxiangpeng123@gmail.com>
Date: Wed, 20 Nov 2024 18:20:40 -0600
Subject: [PATCH] refactor

---
 src/base_node.rs  |  7 -------
 src/node_ptr.rs   |  8 --------
 src/range_scan.rs | 24 +++++++++++++++---------
 src/utils.rs      | 24 +++++++++++++-----------
 4 files changed, 28 insertions(+), 35 deletions(-)

diff --git a/src/base_node.rs b/src/base_node.rs
index 58da4b4..6880e43 100644
--- a/src/base_node.rs
+++ b/src/base_node.rs
@@ -247,13 +247,6 @@ impl BaseNode {
         Self::read_lock_inner(node)
     }
 
-    pub(crate) fn read_lock_deprecated<'a, const MAX_LEVEL: usize>(
-        node: NodePtr,
-        current_level: usize,
-    ) -> Result<ReadGuard<'a>, ArtError> {
-        Self::read_lock_inner(node.as_ptr_safe::<MAX_LEVEL>(current_level))
-    }
-
     pub(crate) fn read_lock_root<'a>(node: NonNull<Node256>) -> Result<ReadGuard<'a>, ArtError> {
         Self::read_lock_inner(unsafe {
             std::mem::transmute::<NonNull<Node256>, NonNull<BaseNode>>(node)
diff --git a/src/node_ptr.rs b/src/node_ptr.rs
index 4ca9c93..d81ca4f 100644
--- a/src/node_ptr.rs
+++ b/src/node_ptr.rs
@@ -67,14 +67,6 @@ impl NodePtr {
         unsafe { self.payload }
     }
 
-    pub(crate) fn as_ptr_safe<const MAX_LEVEL: usize>(
-        &self,
-        current_level: usize,
-    ) -> NonNull<BaseNode> {
-        debug_assert!(current_level < MAX_LEVEL);
-        unsafe { self.sub_node }
-    }
-
     pub(crate) fn downcast<const K_LEN: usize>(&self, current_level: usize) -> PtrType {
         if current_level == (K_LEN - 1) {
             PtrType::Payload(unsafe { self.as_payload_unchecked() })
diff --git a/src/range_scan.rs b/src/range_scan.rs
index c0301a7..f199ae2 100644
--- a/src/range_scan.rs
+++ b/src/range_scan.rs
@@ -93,7 +93,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
                                     if k == start_level {
                                         self.find_start(sub_node, &node, key_tracker.clone())?;
                                     } else if k > start_level && k < end_level {
-                                        let cur_key = KeyTracker::append_prefix(n, &key_tracker);
+                                        let cur_key = KeyTracker::append_prefix(n, &key_tracker)?;
                                         self.copy_node_recursive(n, &cur_key)?;
                                     } else if k == end_level {
                                         self.find_end(sub_node, &node, key_tracker.clone())?;
@@ -120,11 +120,17 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
                         }
                         key_tracker.push(start_level);
 
-                        let next_node =
-                            BaseNode::read_lock_deprecated::<K_LEN>(next_node_tmp, level)?;
-                        parent_node = Some(node);
-                        node = next_node;
-                        continue;
+                        match next_node_tmp.downcast_key_tracker::<K_LEN>(&key_tracker) {
+                            PtrType::Payload(_payload) => {
+                                unreachable!()
+                            }
+                            PtrType::SubNode(sub_node) => {
+                                let next_node = BaseNode::read_lock(sub_node)?;
+                                parent_node = Some(node);
+                                node = next_node;
+                                continue;
+                            }
+                        }
                     }
                     return Ok(self.result_found);
                 }
@@ -176,7 +182,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
                             if k == end_level {
                                 self.find_end(sub_node, &node, key_tracker.clone())?;
                             } else if k < end_level {
-                                let cur_key = KeyTracker::append_prefix(n, &key_tracker);
+                                let cur_key = KeyTracker::append_prefix(n, &key_tracker)?;
                                 self.copy_node_recursive(n, &cur_key)?;
                             }
                         }
@@ -234,7 +240,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
                             if k == start_level {
                                 self.find_start(sub_node, &node, key_tracker.clone())?;
                             } else if k > start_level {
-                                let cur_key = KeyTracker::append_prefix(n, &key_tracker);
+                                let cur_key = KeyTracker::append_prefix(n, &key_tracker)?;
                                 self.copy_node_recursive(n, &cur_key)?;
                             }
                         }
@@ -284,7 +290,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
 
                     key_tracker.push(k);
 
-                    let cur_key = KeyTracker::append_prefix(c, &key_tracker);
+                    let cur_key = KeyTracker::append_prefix(c, &key_tracker)?;
                     self.copy_node_recursive(c, &cur_key)?;
 
                     if self.to_continue {
diff --git a/src/utils.rs b/src/utils.rs
index a68a86e..96cdef4 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,5 +1,6 @@
 use crate::base_node::BaseNode;
-use crate::node_ptr::NodePtr;
+use crate::error::ArtError;
+use crate::node_ptr::{NodePtr, PtrType};
 use core::cell::Cell;
 use core::fmt;
 
@@ -131,17 +132,18 @@ impl<const K_LEN: usize> KeyTracker<K_LEN> {
     pub(crate) fn append_prefix(
         node: NodePtr,
         key_tracker: &KeyTracker<K_LEN>,
-    ) -> KeyTracker<K_LEN> {
-        let mut cur_key = key_tracker.clone();
-        if key_tracker.len() == K_LEN {
-            cur_key
-        } else {
-            let node_ref = BaseNode::read_lock_deprecated::<K_LEN>(node, 0).unwrap();
-            let n_prefix = node_ref.as_ref().prefix().iter().skip(key_tracker.len());
-            for i in n_prefix {
-                cur_key.push(*i);
+    ) -> Result<KeyTracker<K_LEN>, ArtError> {
+        match node.downcast_key_tracker::<K_LEN>(key_tracker) {
+            PtrType::Payload(_payload) => Ok(key_tracker.clone()),
+            PtrType::SubNode(sub_node) => {
+                let node_ref = BaseNode::read_lock(sub_node)?;
+                let n_prefix = node_ref.as_ref().prefix().iter().skip(key_tracker.len());
+                let mut cur_key = key_tracker.clone();
+                for i in n_prefix {
+                    cur_key.push(*i);
+                }
+                Ok(cur_key)
             }
-            cur_key
         }
     }