From bfd30033467ae95957618acd94594f7d2a649a50 Mon Sep 17 00:00:00 2001 From: sharnoff Date: Mon, 19 Sep 2022 20:53:41 -0700 Subject: [PATCH] Drain pt 1: `tree/iter` refactor --- src/tree/iter.rs | 192 ++++++++++++++++++++++------------------------- 1 file changed, 90 insertions(+), 102 deletions(-) diff --git a/src/tree/iter.rs b/src/tree/iter.rs index af28bb6..15900f0 100644 --- a/src/tree/iter.rs +++ b/src/tree/iter.rs @@ -4,7 +4,7 @@ use crate::param::RleTreeConfig; use crate::public_traits::{Index, Slice}; use crate::range::{EndBound, RangeBounds, StartBound}; -use crate::Cursor; +use crate::{Cursor, PathComponent}; use std::any::TypeId; use std::fmt::Debug; use std::panic::{RefUnwindSafe, UnwindSafe}; @@ -33,7 +33,8 @@ where P: RleTreeConfig, { start: I, - end: IncludedOrExcludedBound, + end: I, + end_bound_is_exclusive: bool, state: Option>, } @@ -177,12 +178,6 @@ impl<'t, I, S, P: RleTreeConfig, const M: usize> Debug for IterStack<'t } } -#[derive(Debug, Copy, Clone)] -enum IncludedOrExcludedBound { - Included(I), - Excluded(I), -} - #[track_caller] fn panic_internal_error_or_bad_index() -> ! { if crate::public_traits::perfect_index_impls().contains(&TypeId::of::()) { @@ -192,6 +187,67 @@ fn panic_internal_error_or_bad_index() -> ! { } } +/// Custom function for iteration that's like [`search_step`], but goes to the key or child to the +/// target's immediate left if `excluded` is true +fn bounded_search_step<'t, I, S, P, const M: usize>( + node: NodeHandle, I, S, P, M>, + hint: Option, + target: I, + excluded: bool, +) -> ChildOrKey<(u8, I), (u8, I)> +where + I: Index, + P: RleTreeConfig, +{ + let result = if target == node.leaf().subtree_size() { + debug_assert!(excluded); + ChildOrKey::Key((node.leaf().len(), target)) + } else { + search_step(node, hint, target) + }; + + if !excluded { + return result; + } + + // Because we have an excluded bound, check for key/child positions exactly matching `target`, + // and move to the previous key or child + match result { + ChildOrKey::Child((c_idx, c_pos)) if c_pos == target => { + if c_idx == 0 { + panic_internal_error_or_bad_index::(); + } + + let k_idx = c_idx - 1; + // SAFETY: child indexes returned from `search_step` are guaranteed to be valid + // children, so the key to their left (at `c_idx - 1`) is a valid key + let k_pos = unsafe { node.leaf().key_pos(k_idx) }; + ChildOrKey::Key((k_idx, k_pos)) + } + ChildOrKey::Key((k_idx, k_pos)) if k_pos == target => match node.typed_ref() { + Type::Leaf(n) => { + if k_idx == 0 { + panic_internal_error_or_bad_index::(); + } + + // SAFETY: key indexes returned from `search_step` are guaranteed to be in bounds, + // and our special case is also non-zero; we checked above that `k_idx != 0`, so + // subtracting one won't overlflow: `k_idx - 1` is still in bounds. + let lhs_k_pos = unsafe { n.leaf().key_pos(k_idx - 1) }; + ChildOrKey::Key((k_idx - 1, lhs_k_pos)) + } + Type::Internal(n) => { + let c_idx = k_idx; + // SAFETY: key indexes returned from `search_step` are guaranteed to be in bounds, + // so a child at the same index (i.e. to its left) will also be valid. + let c_pos = k_pos.sub_right(unsafe { n.child_size(c_idx) }); + ChildOrKey::Child((c_idx, c_pos)) + } + }, + r => r, + } +} + impl<'t, C, I, S, P, const M: usize> Iter<'t, C, I, S, P, M> where C: Cursor, @@ -220,13 +276,15 @@ where StartBound::Unbounded => I::ZERO, }; - let (end, strict_end) = match end_bound { - EndBound::Included(i) => (IncludedOrExcludedBound::Included(i), EndBound::Included(i)), - EndBound::Excluded(i) => (IncludedOrExcludedBound::Excluded(i), EndBound::Excluded(i)), - EndBound::Unbounded => ( - IncludedOrExcludedBound::Excluded(tree_size), - EndBound::Excluded(tree_size), - ), + let (end, end_bound_is_exclusive) = match end_bound { + EndBound::Included(i) => (i, false), + EndBound::Excluded(i) => (i, true), + EndBound::Unbounded => (tree_size, true), + }; + + let strict_end = match end_bound_is_exclusive { + true => EndBound::Excluded(end), + false => EndBound::Included(end), }; if range.starts_after_end() { @@ -241,18 +299,14 @@ where // Special case for empty iterators, so we don't attempt to traverse the tree when we won't // end up with anything. - let bounded_end_bound = match end { - IncludedOrExcludedBound::Included(i) => EndBound::Included(i), - IncludedOrExcludedBound::Excluded(i) => EndBound::Excluded(i), - }; - - if (StartBound::Included(start), bounded_end_bound).is_empty_naive() { + if (StartBound::Included(start), strict_end).is_empty_naive() { root = None; } Iter { start, end, + end_bound_is_exclusive, state: root.map(|(root, store)| IterState { store, root, @@ -266,80 +320,18 @@ where fn make_stack( root: NodeHandle, I, S, P, M>, cursor: Option, - target: IncludedOrExcludedBound, + mut target: I, + excluded: bool, ) -> IterStack<'t, I, S, P, M> { let mut stack = Vec::new(); let mut cursor = cursor.map(|c| c.into_path()); let mut head_node = root; let mut head_pos = I::ZERO; - let (mut target, excluded) = match target { - IncludedOrExcludedBound::Included(i) => (i, false), - IncludedOrExcludedBound::Excluded(i) => (i, true), - }; // Search downward with the cursor loop { let hint = cursor.as_mut().and_then(|c| c.next()); - - let result = if target != head_node.leaf().subtree_size() { - match search_step(head_node, hint, target) { - ChildOrKey::Key((k_idx, k_pos)) if excluded && k_pos == target => { - match head_node.into_typed() { - Type::Leaf(n) => { - if k_idx == 0 { - panic_internal_error_or_bad_index::(); - } - - // SAFETY: key indexes returned from `search_step` are guaranteed - // to be in bounds; we checked above that `k_idx != 0`, so - // subtracting one won't overlflow: `k_idx - 1` is still in bounds. - let lhs_k_pos = unsafe { n.leaf().key_pos(k_idx - 1) }; - ChildOrKey::Key((k_idx - 1, lhs_k_pos)) - } - Type::Internal(n) => { - let c_idx = k_idx; - // SAFETY: key indexes returned from `search_step` are guaranteed - // to be in bounds, so a child at the same index (i.e. to its left) - // will also be valid. - let c_pos = k_pos.sub_right(unsafe { n.child_size(c_idx) }); - ChildOrKey::Child((c_idx, c_pos)) - } - } - } - ChildOrKey::Child((c_idx, c_pos)) if excluded && c_pos == target => { - if c_idx == 0 { - panic_internal_error_or_bad_index::(); - } - - let k_idx = c_idx - 1; - // SAFETY: child indexes returned from `search_step` are guaranteed to be - // valid children, so the key to their left (at `c_idx - 1`) is a valid key - let k_pos = unsafe { head_node.leaf().key_pos(k_idx) }; - ChildOrKey::Key((k_idx, k_pos)) - } - res => res, - } - } else { - let len = head_node.leaf().len(); - - // Use the key or child immediately to the left. - match head_node.into_typed() { - Type::Leaf(_) => { - let i = len - 1; - // SAFETY: `len` is guaranteed to be `>= 1`, so `len - 1` can't overflow, - // and key indexes must be `< len`. - let p = unsafe { head_node.leaf().key_pos(i) }; - ChildOrKey::Key((i, p)) - } - Type::Internal(node) => { - // The child immediately left of the key has the same index - let c_idx = len; - // SAFETY: the index `len` is always a valid child index - let c_pos = unsafe { node.child_pos(c_idx) }; - ChildOrKey::Child((c_idx, c_pos)) - } - } - }; + let result = bounded_search_step(head_node, hint, target, excluded); match result { ChildOrKey::Key((k_idx, k_pos)) => { @@ -586,16 +578,11 @@ where // Update `fwd` to the next entry: let next_start = fwd.end_pos; // ... but first, check that it'll be within bounds: - match self.end { - IncludedOrExcludedBound::Included(i) if i < next_start => { - self.state = None; - return None; - } - IncludedOrExcludedBound::Excluded(i) if i <= next_start => { - self.state = None; - return None; - } - _ => (), + if (self.end_bound_is_exclusive && self.end <= next_start) + || (!self.end_bound_is_exclusive && self.end < next_start) + { + self.state = None; + return None; } fwd.fwd_step(); @@ -603,11 +590,7 @@ where (fwd, next_start) } None => { - let s = Self::make_stack( - state.root, - state.cursor.take(), - IncludedOrExcludedBound::Included(self.start), - ); + let s = Self::make_stack(state.root, state.cursor.take(), self.start, false); // We can't use `self.start` for the start position because it might be in the // middle of a slice @@ -656,7 +639,12 @@ where bkwd } None => { - let s = Self::make_stack(state.root, state.cursor.take(), self.end); + let s = Self::make_stack( + state.root, + state.cursor.take(), + self.end, + self.end_bound_is_exclusive, + ); state.bkwd.insert(s) }