Skip to content

Commit

Permalink
refactor scan
Browse files Browse the repository at this point in the history
  • Loading branch information
XiangpengHao committed Nov 20, 2024
1 parent d7a2a74 commit 3d62451
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 40 deletions.
8 changes: 4 additions & 4 deletions src/node_ptr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ mod private {

pub trait LastLevelProofInner {}

impl LastLevelProofInner for LastLevelKey<'_> {}
impl<const K_LEN: usize> LastLevelProofInner for LastLevelKey<'_, K_LEN> {}

impl LastLevelProofInner for ChildIsPayload<'_> {}
}
Expand All @@ -21,13 +21,13 @@ pub(crate) trait LastLevelProof: private::LastLevelProofInner {}

impl LastLevelProof for ChildIsPayload<'_> {}

impl LastLevelProof for LastLevelKey<'_> {}
impl<const K_LEN: usize> LastLevelProof for LastLevelKey<'_, K_LEN> {}

pub(crate) struct ChildIsPayload<'a> {
_marker: std::marker::PhantomData<&'a ()>,
}

impl<'a> ChildIsPayload<'a> {
impl ChildIsPayload<'_> {
pub(crate) fn new() -> Self {
Self {
_marker: std::marker::PhantomData,
Expand All @@ -39,7 +39,7 @@ pub(crate) struct ChildIsSubNode<'a> {
_marker: std::marker::PhantomData<&'a ()>,
}

impl<'a> ChildIsSubNode<'a> {
impl ChildIsSubNode<'_> {
pub(crate) fn new() -> Self {
Self {
_marker: std::marker::PhantomData,
Expand Down
37 changes: 17 additions & 20 deletions src/range_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,16 +41,8 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
self.start < self.end
}

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();

if start_key <= cur_key && cur_key < end_key {
return true;
}
false
fn key_in_range(&self, key: &LastLevelKey<K_LEN>) -> bool {
self.start <= key.key() && key.key() < self.end
}

pub(crate) fn scan(&mut self) -> Result<usize, ArtError> {
Expand All @@ -59,7 +51,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
self.to_continue = 0;
self.result_found = 0;

let mut key_tracker = KeyTracker::default();
let mut key_tracker = KeyTracker::empty();

loop {
let prefix_check_result = self.check_prefix_equals(node.as_ref(), &mut key_tracker);
Expand Down Expand Up @@ -122,7 +114,8 @@ 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)?;
let next_node =
BaseNode::read_lock_deprecated::<K_LEN>(next_node_tmp, level)?;
parent_node = Some(node);
node = next_node;
continue;
Expand All @@ -144,9 +137,9 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
&mut self,
node: NodePtr,
parent_node: &ReadGuard,
mut key_tracker: KeyTracker,
mut key_tracker: KeyTracker<K_LEN>,
) -> Result<(), ArtError> {
debug_assert!(key_tracker.len() != 8);
debug_assert!(key_tracker.len() != K_LEN);

let node = BaseNode::read_lock_deprecated::<K_LEN>(node, key_tracker.len())?;
let prefix_result =
Expand Down Expand Up @@ -194,9 +187,9 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
&mut self,
node: NodePtr,
parent_node: &ReadGuard,
mut key_tracker: KeyTracker,
mut key_tracker: KeyTracker<K_LEN>,
) -> Result<(), ArtError> {
debug_assert!(key_tracker.len() != 8);
debug_assert!(key_tracker.len() != K_LEN);

let node = BaseNode::read_lock_deprecated::<K_LEN>(node, key_tracker.len())?;
let prefix_result =
Expand Down Expand Up @@ -241,8 +234,12 @@ 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(last_level_key) = key_tracker.as_last_level::<K_LEN>() {
fn copy_node(
&mut self,
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);
Expand Down Expand Up @@ -281,7 +278,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
n: &BaseNode,
k: &[u8; K_LEN],
fill_key: u8,
key_tracker: &mut KeyTracker,
key_tracker: &mut KeyTracker<K_LEN>,
) -> cmp::Ordering {
let n_prefix = n.prefix();
if !n_prefix.is_empty() {
Expand Down Expand Up @@ -324,7 +321,7 @@ impl<'a, const K_LEN: usize> RangeScan<'a, K_LEN> {
fn check_prefix_equals(
&self,
n: &BaseNode,
key_tracker: &mut KeyTracker,
key_tracker: &mut KeyTracker<K_LEN>,
) -> PrefixCheckEqualsResult {
let n_prefix = n.prefix();

Expand Down
41 changes: 25 additions & 16 deletions src/utils.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use crate::base_node::{BaseNode, MAX_KEY_LEN};
use crate::base_node::BaseNode;
use crate::node_ptr::NodePtr;
use core::cell::Cell;
use core::fmt;
Expand Down Expand Up @@ -77,27 +77,33 @@ impl Default for Backoff {
}
}

pub(crate) struct LastLevelKey<'a> {
key: &'a KeyTracker,
pub(crate) struct LastLevelKey<'a, const K_LEN: usize> {
key: &'a KeyTracker<K_LEN>,
}

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()
impl<const K_LEN: usize> LastLevelKey<'_, K_LEN> {
pub(crate) fn key(&self) -> &[u8; K_LEN] {
&self.key.data
}
}

#[derive(Default, Clone)]
pub(crate) struct KeyTracker {
#[derive(Clone)]
pub(crate) struct KeyTracker<const K_LEN: usize> {
len: usize,
data: [u8; 8],
data: [u8; K_LEN],
}

impl KeyTracker {
impl<const K_LEN: usize> KeyTracker<K_LEN> {
pub(crate) fn empty() -> Self {
Self {
len: 0,
data: [0; K_LEN],
}
}

#[inline]
pub(crate) fn push(&mut self, key: u8) {
debug_assert!(self.len <= 8);
debug_assert!(self.len <= K_LEN);

self.data[self.len] = key;
self.len += 1;
Expand All @@ -112,7 +118,7 @@ impl KeyTracker {
v
}

pub(crate) fn as_last_level<const K_LEN: usize>(&self) -> Option<LastLevelKey> {
pub(crate) fn as_last_level(&self) -> Option<LastLevelKey<K_LEN>> {
if self.len == K_LEN {
Some(LastLevelKey { key: self })
} else {
Expand All @@ -126,12 +132,15 @@ impl KeyTracker {
}

#[inline]
pub(crate) fn append_prefix(node: NodePtr, key_tracker: &KeyTracker) -> KeyTracker {
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() == MAX_KEY_LEN {
if key_tracker.len() == K_LEN {
cur_key
} else {
let node_ref = BaseNode::read_lock_deprecated::<MAX_KEY_LEN>(node, 0).unwrap();
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);
Expand Down

0 comments on commit 3d62451

Please sign in to comment.