use crate::bplus_tree::{BPlusTree, Node};
use crate::codecs::{IdCodec, KeyCodec};
use crate::errors::Error;
use crate::pager::Pager;
use crate::{
types::{CursorStepRes, EntryRefRes, FramePred, NodeWithNextRes},
views::key_view::KeyRef,
views::node_view::{NodeTag, NodeView},
views::value_view::ValueRef,
};
use core::cmp::Ordering;
use core::marker::PhantomData;
use core::ops::Bound;
use core::ops::Bound::{Excluded, Included, Unbounded};
use std::sync::Arc;
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Direction {
Forward,
Reverse,
}
#[derive(Clone)]
pub struct ScanOpts<'a, KC>
where
KC: KeyCodec,
{
pub dir: Direction,
pub lower: Bound<&'a KC::Key>,
pub upper: Bound<&'a KC::Key>,
pub prefix: Option<&'a [u8]>,
pub frame_predicate: Option<FramePred>,
}
impl<'a, KC> Default for ScanOpts<'a, KC>
where
KC: KeyCodec,
{
fn default() -> Self {
Self {
dir: Direction::Forward,
lower: Unbounded,
upper: Unbounded,
prefix: None,
frame_predicate: None,
}
}
}
impl<'a, KC> ScanOpts<'a, KC>
where
KC: KeyCodec,
{
#[inline]
pub fn forward() -> Self {
Self {
dir: Direction::Forward,
..Default::default()
}
}
#[inline]
pub fn reverse() -> Self {
Self {
dir: Direction::Reverse,
..Default::default()
}
}
#[inline]
pub fn start_at(mut self, k: &'a KC::Key) -> Self {
self.lower = Included(k);
self
}
#[inline]
pub fn start_after(mut self, k: &'a KC::Key) -> Self {
self.lower = Excluded(k);
self
}
#[inline]
pub fn end_at(mut self, k: &'a KC::Key) -> Self {
self.upper = Included(k);
self
}
#[inline]
pub fn end_before(mut self, k: &'a KC::Key) -> Self {
self.upper = Excluded(k);
self
}
#[inline]
pub fn with_prefix(mut self, p: &'a [u8]) -> Self {
self.prefix = Some(p);
self
}
#[inline]
pub fn with_bounds(mut self, lo: Bound<&'a KC::Key>, up: Bound<&'a KC::Key>) -> Self {
self.lower = lo;
self.upper = up;
self
}
#[inline]
pub fn with_frame_predicate<F>(mut self, pred: F) -> Self
where
F: Fn(&[u8], &[u8]) -> bool + Send + Sync + 'static,
{
self.frame_predicate = Some(Arc::new(pred));
self
}
}
pub trait ValueResolver<'a, P, KC, IC>
where
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
fn read_node(&self, id: &P::Id) -> Result<Node<KC::Key, P::Id>, Error>;
fn read_one(&self, id: &P::Id) -> Result<P::Page, Error>;
fn root_id(&self) -> P::Id;
fn entry_at(&self, page: P::Page, view: &NodeView<P>, i: usize) -> EntryRefRes<P>;
}
pub struct PlainResolver<'a, P, KC, IC>
where
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
tree: &'a BPlusTree<P, KC, IC>,
}
impl<'a, P, KC, IC> PlainResolver<'a, P, KC, IC>
where
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
#[inline]
pub fn new(tree: &'a BPlusTree<P, KC, IC>) -> Self {
Self { tree }
}
}
impl<'a, P, KC, IC> ValueResolver<'a, P, KC, IC> for PlainResolver<'a, P, KC, IC>
where
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
#[inline]
fn read_node(&self, id: &P::Id) -> Result<Node<KC::Key, P::Id>, Error> {
self.tree.read_node(id)
}
#[inline]
fn read_one(&self, id: &P::Id) -> Result<P::Page, Error> {
self.tree.read_one(id)
}
#[inline]
fn root_id(&self) -> P::Id {
self.tree.root_id()
}
#[inline]
fn entry_at(
&self,
page: P::Page,
view: &NodeView<P>,
i: usize,
) -> Result<(KeyRef<P::Page>, ValueRef<P::Page>), Error> {
let (k_enc, _v_enc) = view.leaf_entry_slices(i);
let kref = KeyRef::from_subslice(page.clone(), k_enc);
let vr = view.leaf_entry_value_range(i);
let vref = ValueRef::new(page, vr);
Ok((kref, vref))
}
}
pub struct ScanIter<'a, R, P, KC, IC>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
resolver: R,
cur_page: Option<P::Page>,
cur_view: Option<NodeView<P>>,
cur_count: usize,
pos: usize,
pos_after: usize,
leaf_id: Option<P::Id>,
lower: Bound<&'a KC::Key>,
upper: Bound<&'a KC::Key>,
prefix: Option<&'a [u8]>,
frame_predicate: Option<FramePred>,
first_key_enc: Option<Vec<u8>>,
dir: Direction,
_p: PhantomData<(KC, IC)>,
}
impl<'a, R, P, KC, IC> ScanIter<'a, R, P, KC, IC>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
pub fn from_parts(resolver: R, opts: ScanOpts<'a, KC>) -> Result<Self, Error> {
{
let root = resolver.root_id();
if resolver.read_node(&root)?.entry_count() == 0 {
return Ok(Self {
resolver,
cur_page: None,
cur_view: None,
cur_count: 0,
pos: 0,
pos_after: 0,
leaf_id: None,
lower: opts.lower,
upper: opts.upper,
prefix: opts.prefix,
frame_predicate: opts.frame_predicate,
first_key_enc: None,
dir: opts.dir,
_p: PhantomData,
});
}
}
match opts.dir {
Direction::Forward => {
let (page, view, pos, leaf_id) =
descend_lower_pos::<R, P, KC, IC>(&resolver, opts.lower)?;
let count = view.count();
Ok(Self {
resolver,
cur_page: Some(page),
cur_view: Some(view),
cur_count: count,
pos,
pos_after: 0,
leaf_id,
lower: opts.lower,
upper: opts.upper,
prefix: opts.prefix,
frame_predicate: opts.frame_predicate,
first_key_enc: None,
dir: opts.dir,
_p: PhantomData,
})
}
Direction::Reverse => {
let (page, view, pos_after, leaf_id) =
descend_from_upper::<R, P, KC, IC>(&resolver, opts.upper)?;
let count = view.count();
Ok(Self {
resolver,
cur_page: Some(page),
cur_view: Some(view),
cur_count: count,
pos: 0,
pos_after,
leaf_id,
lower: opts.lower,
upper: opts.upper,
prefix: opts.prefix,
frame_predicate: opts.frame_predicate,
first_key_enc: None,
dir: opts.dir,
_p: PhantomData,
})
}
}
}
#[inline]
pub fn with_resolver(resolver: R, opts: ScanOpts<'a, KC>) -> Result<Self, Error> {
Self::from_parts(resolver, opts)
}
#[inline]
fn within_upper(&self, view: &NodeView<P>, i: usize) -> bool {
match self.upper {
Unbounded => true,
Included(up) => {
let (k_enc, _) = view.leaf_entry_slices(i);
KC::compare_encoded(k_enc, up) != Ordering::Greater
}
Excluded(up) => {
let (k_enc, _) = view.leaf_entry_slices(i);
KC::compare_encoded(k_enc, up) == Ordering::Less
}
}
}
#[inline]
fn within_lower(&self, view: &NodeView<P>, i: usize) -> bool {
match self.lower {
Unbounded => true,
Included(lo) => {
let (k_enc, _) = view.leaf_entry_slices(i);
KC::compare_encoded(k_enc, lo) != Ordering::Less
}
Excluded(lo) => {
let (k_enc, _) = view.leaf_entry_slices(i);
KC::compare_encoded(k_enc, lo) == Ordering::Greater
}
}
}
#[inline]
fn within_prefix(&self, view: &NodeView<P>, i: usize) -> bool {
if let Some(pref) = self.prefix {
let (k_enc, _) = view.leaf_entry_slices(i);
k_enc.starts_with(pref)
} else {
true
}
}
fn advance_leaf_forward(&mut self) -> Result<bool, Error> {
let mut nid = match self.leaf_id.take() {
Some(id) => id,
None => {
self.cur_page = None;
self.cur_view = None;
self.cur_count = 0;
self.pos = 0;
return Ok(false);
}
};
loop {
let page = self.resolver.read_one(&nid)?;
let view = NodeView::<P>::new(page.clone())?;
let count = view.count();
let next_id = decode_next_id::<P, IC>(&view);
self.cur_page = Some(page);
self.cur_view = Some(view);
self.cur_count = count;
self.pos = 0;
self.leaf_id = next_id;
if self.cur_count > 0 {
return Ok(true);
}
nid = match self.leaf_id.take() {
Some(id) => id,
None => {
self.cur_page = None;
self.cur_view = None;
self.cur_count = 0;
self.pos = 0;
return Ok(false);
}
};
}
}
fn step_reverse(&mut self) -> Result<bool, Error> {
if self.pos_after > 0 {
self.pos_after -= 1;
return Ok(true);
}
let view = match self.cur_view.as_ref() {
Some(v) => v,
None => return Ok(false),
};
if view.count() == 0 {
return Ok(false);
}
let (k_enc, _) = view.leaf_entry_slices(0);
let (page, v, pos_after, next_id) =
descend_lt_encoded::<R, P, KC, IC>(&self.resolver, k_enc)?;
self.cur_page = Some(page);
self.cur_view = Some(v);
self.cur_count = self.cur_view.as_ref().unwrap().count();
self.pos_after = pos_after;
self.leaf_id = next_id;
Ok(self.cur_count > 0 && self.pos_after > 0)
}
}
impl<'a, R, P, KC, IC> Iterator for ScanIter<'a, R, P, KC, IC>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
type Item = (KeyRef<P::Page>, ValueRef<P::Page>);
fn next(&mut self) -> Option<Self::Item> {
self.cur_page.as_ref()?;
match self.dir {
Direction::Forward => {
loop {
if self.pos >= self.cur_count && !(self.advance_leaf_forward().ok()?) {
return None;
}
let page = self.cur_page.as_ref()?.clone();
let view = self.cur_view.as_ref()?;
if !self.within_lower(view, self.pos) {
return None;
}
if !self.within_upper(view, self.pos) {
return None;
}
if !self.within_prefix(view, self.pos) {
self.pos += 1;
continue; }
if let Some(pred) = self.frame_predicate.as_ref() {
let (k_enc, _) = view.leaf_entry_slices(self.pos);
if let Some(head) = self.first_key_enc.as_ref() {
if !pred(head, k_enc) {
return None;
}
} else {
self.first_key_enc = Some(k_enc.to_vec());
}
}
let out = self.resolver.entry_at(page, view, self.pos).ok()?;
self.pos += 1;
return Some(out);
}
}
Direction::Reverse => {
loop {
if self.pos_after == 0 && !(self.step_reverse().ok()?) {
return None;
}
let page = self.cur_page.as_ref()?.clone();
let view = self.cur_view.as_ref()?;
let idx = self.pos_after - 1;
if !self.within_lower(view, idx) {
return None;
}
if !self.within_prefix(view, idx) {
if self.pos_after > 0 {
self.pos_after -= 1;
}
continue; }
if let Some(pred) = self.frame_predicate.as_ref() {
let (k_enc, _) = view.leaf_entry_slices(idx);
if let Some(head) = self.first_key_enc.as_ref() {
if !pred(head, k_enc) {
return None;
}
} else {
self.first_key_enc = Some(k_enc.to_vec());
}
}
let out = self.resolver.entry_at(page, view, idx).ok()?;
if self.pos_after > 0 {
self.pos_after -= 1;
}
return Some(out);
}
}
}
}
}
pub type BPlusTreeIter<'a, P, KC, IC> = ScanIter<'a, PlainResolver<'a, P, KC, IC>, P, KC, IC>;
impl<'a, P, KC, IC> BPlusTreeIter<'a, P, KC, IC>
where
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
#[inline]
pub fn new(tree: &'a BPlusTree<P, KC, IC>) -> Result<Self, Error> {
Self::with_opts(tree, ScanOpts::forward())
}
#[inline]
pub fn with_opts(
tree: &'a BPlusTree<P, KC, IC>,
opts: ScanOpts<'a, KC>,
) -> Result<Self, Error> {
let r = PlainResolver::new(tree);
ScanIter::with_resolver(r, opts)
}
}
#[inline]
fn decode_next_id<P, IC>(view: &NodeView<P>) -> Option<P::Id>
where
P: Pager,
IC: IdCodec<Id = P::Id>,
{
let aux = view.leaf_next_aux();
if aux.is_empty() {
None
} else {
Some(IC::decode_from(aux).ok()?.0)
}
}
fn descend_leftmost<'a, R, P, KC, IC>(r: &R) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
let mut id = r.root_id();
loop {
match r.read_node(&id)? {
Node::Internal { entries } => {
let (_, child) = entries.first().expect("internal without child").clone();
id = child;
}
Node::Leaf { .. } => {
let page = r.read_one(&id)?;
let view = NodeView::<P>::new(page.clone())?;
let next_id = decode_next_id::<P, IC>(&view);
return Ok((page, view, 0, next_id));
}
}
}
}
fn descend_rightmost<'a, R, P, KC, IC>(r: &R) -> NodeWithNextRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
let mut id = r.root_id();
loop {
match r.read_node(&id)? {
Node::Internal { entries } => {
let (_, child) = entries.last().expect("internal without child").clone();
id = child;
}
Node::Leaf { .. } => {
let page = r.read_one(&id)?;
let view = NodeView::<P>::new(page.clone())?;
let next_id = decode_next_id::<P, IC>(&view);
return Ok((page, view, next_id));
}
}
}
}
fn leaf_lower_bound<P, KC>(view: &NodeView<P>, key: &KC::Key) -> usize
where
P: Pager,
KC: KeyCodec,
{
let mut lo = 0isize;
let mut hi = view.count() as isize;
while lo < hi {
let mid = ((lo + hi) >> 1) as usize;
let (k_enc, _) = view.leaf_entry_slices(mid);
match KC::compare_encoded(k_enc, key) {
Ordering::Less => lo = mid as isize + 1,
_ => hi = mid as isize,
}
}
lo as usize
}
fn leaf_upper_pos<P, KC>(view: &NodeView<P>, key: &KC::Key, inclusive: bool) -> usize
where
P: Pager,
KC: KeyCodec,
{
let mut lo = 0isize;
let mut hi = view.count() as isize;
while lo < hi {
let mid = ((lo + hi) >> 1) as usize;
let (k_enc, _) = view.leaf_entry_slices(mid);
let ord = KC::compare_encoded(k_enc, key);
let is_le = if inclusive {
ord != Ordering::Greater
} else {
ord == Ordering::Less
};
if is_le {
lo = mid as isize + 1;
} else {
hi = mid as isize;
}
}
lo as usize
}
fn descend_ge<'a, R, P, KC, IC>(r: &R, key: &KC::Key) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
let mut id = r.root_id();
loop {
let node_page = r.read_one(&id)?;
let view = NodeView::<P>::new(node_page)?;
if let Ok(NodeTag::Leaf) = view.tag() {
let page = r.read_one(&id)?;
let view = NodeView::<P>::new(page.clone())?;
let pos = leaf_lower_bound::<P, KC>(&view, key);
let next_id = decode_next_id::<P, IC>(&view);
return Ok((page, view, pos, next_id));
}
let n = view.count();
let mut lo = 0usize;
let mut hi = n;
while lo < hi {
let mid = (lo + hi) / 2;
let (k_enc, _) = view.internal_entry_slices(mid);
if KC::compare_encoded(k_enc, key) == Ordering::Less {
lo = mid + 1;
} else {
hi = mid;
}
}
let idx = lo.min(n - 1);
let (_, child_raw) = view.internal_entry_slices(idx);
let (child_id, _) = IC::decode_from(child_raw)?;
id = child_id;
}
}
fn descend_upper_pos<'a, R, P, KC, IC>(r: &R, key: &KC::Key, inclusive: bool) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
let mut id = r.root_id();
loop {
let node_page = r.read_one(&id)?;
let view = NodeView::<P>::new(node_page)?;
if let Ok(NodeTag::Leaf) = view.tag() {
let page = r.read_one(&id)?;
let view = NodeView::<P>::new(page.clone())?;
let pos_after = leaf_upper_pos::<P, KC>(&view, key, inclusive);
let next_id = decode_next_id::<P, IC>(&view);
return Ok((page, view, pos_after, next_id));
}
let n = view.count();
let mut lo = 0usize;
let mut hi = n;
while lo < hi {
let mid = (lo + hi) / 2;
let (k_enc, _) = view.internal_entry_slices(mid);
let ord = KC::compare_encoded(k_enc, key);
let ok = if inclusive {
ord != Ordering::Greater
} else {
ord == Ordering::Less
};
if ok {
lo = mid + 1;
} else {
hi = mid;
}
}
let chosen_idx = lo.min(n - 1);
let (_, child_raw) = view.internal_entry_slices(chosen_idx);
let (child_id, _) = IC::decode_from(child_raw)?;
id = child_id;
}
}
fn descend_lt_encoded<'a, R, P, KC, IC>(r: &R, enc: &[u8]) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
let _hex8 = |b: &[u8]| -> String {
let mut s = String::new();
for (i, x) in b.iter().take(8).enumerate() {
if i > 0 {
s.push(' ');
}
use core::fmt::Write as _;
let _ = write!(&mut s, "{:02x}", x);
}
s
};
let key = KC::decode_from(enc)?;
let mut path: Vec<(P::Id, usize)> = Vec::new();
let mut id = r.root_id();
loop {
let node_page = r.read_one(&id)?;
let view = NodeView::<P>::new(node_page)?;
if let Ok(NodeTag::Leaf) = view.tag() {
let page = r.read_one(&id)?;
let view = NodeView::<P>::new(page.clone())?;
let pos_after = leaf_upper_pos::<P, KC>(&view, &key, false);
let next_id = decode_next_id::<P, IC>(&view);
if pos_after > 0 {
return Ok((page, view, pos_after, next_id));
}
while let Some((anc_id, idx)) = path.pop() {
if idx == 0 {
continue;
}
let anc_page = r.read_one(&anc_id)?;
let anc_view = NodeView::<P>::new(anc_page)?;
let (_, child_raw) = anc_view.internal_entry_slices(idx - 1);
let (mut child_id, _) = IC::decode_from(child_raw)?;
loop {
let c_page = r.read_one(&child_id)?;
let c_view = NodeView::<P>::new(c_page.clone())?;
if let Ok(NodeTag::Leaf) = c_view.tag() {
let next = decode_next_id::<P, IC>(&c_view);
let cnt = c_view.count();
return Ok((c_page, c_view, cnt, next));
}
let last = c_view.count().saturating_sub(1);
let (_, cr) = c_view.internal_entry_slices(last);
let (next_child, _) = IC::decode_from(cr)?;
child_id = next_child;
}
}
return Ok((page, view, 0, next_id));
}
let n = view.count();
let mut lo = 0usize;
let mut hi = n;
while lo < hi {
let mid = (lo + hi) / 2;
let (k_enc, _) = view.internal_entry_slices(mid);
let ord = KC::compare_encoded(k_enc, &key);
if ord == core::cmp::Ordering::Less {
lo = mid + 1;
} else {
hi = mid;
}
}
let chosen_idx = lo.min(n.saturating_sub(1));
let (_, child_raw) = view.internal_entry_slices(chosen_idx);
let (child_id, _) = IC::decode_from(child_raw)?;
path.push((id, chosen_idx));
id = child_id;
}
}
fn descend_lower_pos<'a, R, P, KC, IC>(r: &R, bound: Bound<&KC::Key>) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
match bound {
Unbounded => descend_leftmost::<R, P, KC, IC>(r),
Included(k) => descend_ge::<R, P, KC, IC>(r, k),
Excluded(k) => {
let (page, view, mut pos, next) =
descend_upper_pos::<R, P, KC, IC>(r, k, false)?;
if pos < view.count()
&& KC::compare_encoded(view.leaf_entry_slices(pos).0, k) == Ordering::Equal
{
pos += 1;
}
Ok((page, view, pos, next))
}
}
}
fn descend_from_upper<'a, R, P, KC, IC>(r: &R, bound: Bound<&KC::Key>) -> CursorStepRes<P>
where
R: ValueResolver<'a, P, KC, IC>,
P: Pager,
KC: KeyCodec,
IC: IdCodec<Id = P::Id>,
{
match bound {
Unbounded => {
let (p, v, n) = descend_rightmost::<R, P, KC, IC>(r)?;
let count = v.count();
Ok((p, v, count, n))
}
Included(k) => descend_upper_pos::<R, P, KC, IC>(r, k, true),
Excluded(k) => descend_upper_pos::<R, P, KC, IC>(r, k, false),
}
}