use super::cid::Cid;
use super::error::Error;
use super::node::Node;
use super::store::Store;
use super::tree::Tree;
use super::Prolly;
#[cfg(feature = "async-store")]
use super::{store::AsyncStore, AsyncProlly};
#[cfg(feature = "async-store")]
use futures_util::stream::{self, Stream};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
type RangeItem = Result<(Vec<u8>, Vec<u8>), Error>;
type LeafEntry = (Vec<u8>, Vec<u8>);
type OptionalLeafEntry = Result<Option<LeafEntry>, Error>;
pub(crate) const RANGE_CHILD_PREFETCH_PARALLELISM: usize = 16;
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RangeCursor {
after_key: Option<Vec<u8>>,
}
impl RangeCursor {
pub fn start() -> Self {
Self { after_key: None }
}
pub fn after_key(key: impl Into<Vec<u8>>) -> Self {
Self {
after_key: Some(key.into()),
}
}
pub fn after(&self) -> Option<&[u8]> {
self.after_key.as_deref()
}
pub fn is_start(&self) -> bool {
self.after_key.is_none()
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct RangePage {
pub entries: Vec<(Vec<u8>, Vec<u8>)>,
pub next_cursor: Option<RangeCursor>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReverseCursor {
before_key: Option<Vec<u8>>,
}
impl ReverseCursor {
pub fn end() -> Self {
Self { before_key: None }
}
pub fn before_key(key: impl Into<Vec<u8>>) -> Self {
Self {
before_key: Some(key.into()),
}
}
pub fn before(&self) -> Option<&[u8]> {
self.before_key.as_deref()
}
pub fn is_end(&self) -> bool {
self.before_key.is_none()
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ReversePage {
pub entries: Vec<(Vec<u8>, Vec<u8>)>,
pub next_cursor: Option<ReverseCursor>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct CursorWindow {
pub position_key: Option<Vec<u8>>,
pub position_value: Option<Vec<u8>>,
pub found: bool,
pub entries: Vec<(Vec<u8>, Vec<u8>)>,
pub next_cursor: Option<RangeCursor>,
}
#[cfg(feature = "async-store")]
pub type AsyncRangePage = RangePage;
#[cfg(feature = "async-store")]
pub type AsyncReversePage = ReversePage;
pub fn create_range_iter<'a, S: Store>(
prolly: &'a Prolly<S>,
tree: &Tree,
start: &[u8],
end: Option<&[u8]>,
) -> Result<RangeIter<'a, S>, Error> {
if end.is_some_and(|end| end <= start) {
return Ok(RangeIter::new(prolly, Vec::new(), start, end));
}
let path = prolly.find_path_arcs(tree, start)?;
Ok(RangeIter::new(prolly, path, start, end))
}
pub fn create_range_after_iter<'a, S: Store>(
prolly: &'a Prolly<S>,
tree: &Tree,
after_key: &[u8],
end: Option<&[u8]>,
) -> Result<RangeIter<'a, S>, Error> {
if end.is_some_and(|end| end <= after_key) {
return Ok(RangeIter::new_after(prolly, Vec::new(), after_key, end));
}
let path = prolly.find_path_arcs(tree, after_key)?;
Ok(RangeIter::new_after(prolly, path, after_key, end))
}
#[cfg(feature = "async-store")]
pub async fn create_async_range_iter<'a, S>(
prolly: &'a AsyncProlly<S>,
tree: &Tree,
start: &[u8],
end: Option<&[u8]>,
) -> Result<AsyncRangeIter<'a, S>, Error>
where
S: AsyncStore,
S::Error: Send + Sync,
{
if end.is_some_and(|end| end <= start) {
return Ok(AsyncRangeIter::new(prolly, Vec::new(), start, end));
}
let path = prolly.find_path_arcs(tree, start).await?;
Ok(AsyncRangeIter::new(prolly, path, start, end))
}
#[cfg(feature = "async-store")]
pub async fn create_async_range_after_iter<'a, S>(
prolly: &'a AsyncProlly<S>,
tree: &Tree,
after_key: &[u8],
end: Option<&[u8]>,
) -> Result<AsyncRangeIter<'a, S>, Error>
where
S: AsyncStore,
S::Error: Send + Sync,
{
if end.is_some_and(|end| end <= after_key) {
return Ok(AsyncRangeIter::new_after(
prolly,
Vec::new(),
after_key,
end,
));
}
let path = prolly.find_path_arcs(tree, after_key).await?;
Ok(AsyncRangeIter::new_after(prolly, path, after_key, end))
}
pub struct RangeIter<'a, S: Store> {
prolly: &'a Prolly<S>,
stack: Vec<(Arc<Node>, usize)>,
end: Option<Vec<u8>>,
started: bool,
start_key: Vec<u8>,
skip_start_key: bool,
last_key: Option<Vec<u8>>,
}
impl<'a, S: Store> RangeIter<'a, S> {
pub(crate) fn new(
prolly: &'a Prolly<S>,
stack: Vec<(Arc<Node>, usize)>,
start: &[u8],
end: Option<&[u8]>,
) -> Self {
Self {
prolly,
stack,
end: end.map(|e| e.to_vec()),
started: false,
start_key: start.to_vec(),
skip_start_key: false,
last_key: None,
}
}
pub(crate) fn new_after(
prolly: &'a Prolly<S>,
stack: Vec<(Arc<Node>, usize)>,
after_key: &[u8],
end: Option<&[u8]>,
) -> Self {
Self {
prolly,
stack,
end: end.map(|e| e.to_vec()),
started: false,
start_key: after_key.to_vec(),
skip_start_key: true,
last_key: None,
}
}
pub fn resume_cursor(&self) -> RangeCursor {
self.last_key
.clone()
.map(RangeCursor::after_key)
.unwrap_or_else(RangeCursor::start)
}
fn position_at_start(&mut self) -> Option<RangeItem> {
self.started = true;
if self.stack.is_empty() {
return None;
}
let (node, idx) = self.stack.last_mut()?;
if node.leaf {
let start_idx = match node.search(&self.start_key) {
Ok(i) if self.skip_start_key => i.saturating_add(1),
Ok(i) => i, Err(i) => i, };
*idx = start_idx;
if *idx >= node.len() {
return self.advance_to_next_leaf();
}
match leaf_entry_before_end(node, *idx, self.end.as_deref()) {
Ok(Some(entry)) => {
*idx += 1;
self.last_key = Some(entry.0.clone());
return Some(Ok(entry));
}
Ok(None) => return None,
Err(e) => return Some(Err(e)),
}
}
self.descend_to_leaf()
}
fn descend_to_leaf(&mut self) -> Option<RangeItem> {
loop {
let (node, idx) = self.stack.last()?;
if node.leaf {
if *idx >= node.len() {
return self.advance_to_next_leaf();
}
match leaf_entry_before_end(node, *idx, self.end.as_deref()) {
Ok(Some(entry)) => {
if let Some((_, idx)) = self.stack.last_mut() {
*idx += 1;
}
self.last_key = Some(entry.0.clone());
return Some(Ok(entry));
}
Ok(None) => return None,
Err(e) => return Some(Err(e)),
}
}
if *idx >= node.len() {
return self.advance_to_next_leaf();
}
match child_starts_at_or_after_end(self.end.as_deref(), node, *idx) {
Ok(true) => return None,
Ok(false) => {}
Err(e) => return Some(Err(e)),
}
match self.load_child_for_descent(node, *idx) {
Ok(child) => {
self.stack.push((child, 0));
}
Err(e) => return Some(Err(e)),
}
}
}
fn advance_to_next_leaf(&mut self) -> Option<RangeItem> {
loop {
self.stack.pop();
if self.stack.is_empty() {
return None;
}
if let Some((parent, parent_idx)) = self.stack.last_mut() {
*parent_idx += 1;
if *parent_idx < parent.len() {
match child_starts_at_or_after_end(self.end.as_deref(), parent, *parent_idx) {
Ok(true) => return None,
Ok(false) => {}
Err(e) => return Some(Err(e)),
}
return self.descend_to_leaf();
}
if parent.keys.len() != parent.vals.len() {
return Some(Err(Error::InvalidNode));
}
}
}
}
fn load_child_for_descent(&self, node: &Node, idx: usize) -> Result<Arc<Node>, Error> {
let child_cid = child_cid_at(node, idx)?;
if !self.prolly.store().prefers_batch_reads() {
return self.prolly.load_arc(&child_cid);
}
if let Some(child) = self.prolly.cached_node_arc(&child_cid) {
return Ok(child);
}
let max_child_idx = node
.len()
.min(idx.saturating_add(RANGE_CHILD_PREFETCH_PARALLELISM));
let mut child_cids = Vec::with_capacity(max_child_idx.saturating_sub(idx));
child_cids.push(child_cid);
for child_idx in idx + 1..max_child_idx {
match child_starts_at_or_after_end(self.end.as_deref(), node, child_idx) {
Ok(true) => break,
Ok(false) => {}
Err(_) => break,
}
match child_cid_at(node, child_idx) {
Ok(cid) => child_cids.push(cid),
Err(_) => break,
}
}
if child_cids.len() == 1 {
return self.prolly.load_arc(&child_cids[0]);
}
let children = self
.prolly
.load_many_ordered_with_parallelism(&child_cids, RANGE_CHILD_PREFETCH_PARALLELISM)?;
children.into_iter().next().ok_or(Error::InvalidNode)
}
}
impl<'a, S: Store> Iterator for RangeIter<'a, S> {
type Item = Result<(Vec<u8>, Vec<u8>), Error>;
fn next(&mut self) -> Option<Self::Item> {
if !self.started {
return self.position_at_start();
}
loop {
let (node, idx) = self.stack.last_mut()?;
if !node.leaf && node.keys.len() != node.vals.len() {
return Some(Err(Error::InvalidNode));
}
if *idx >= node.len() {
return self.advance_to_next_leaf();
}
if node.leaf {
match leaf_entry_before_end(node, *idx, self.end.as_deref()) {
Ok(Some(entry)) => {
*idx += 1;
self.last_key = Some(entry.0.clone());
return Some(Ok(entry));
}
Ok(None) => return None,
Err(e) => return Some(Err(e)),
}
}
match child_starts_at_or_after_end(self.end.as_deref(), node, *idx) {
Ok(true) => return None,
Ok(false) => {}
Err(e) => return Some(Err(e)),
}
let child = {
let (node, idx) = self.stack.last()?;
self.load_child_for_descent(node, *idx)
};
match child {
Ok(child) => {
self.stack.push((child, 0));
}
Err(e) => return Some(Err(e)),
}
}
}
}
#[cfg(feature = "async-store")]
pub struct AsyncRangeIter<'a, S: AsyncStore> {
prolly: &'a AsyncProlly<S>,
stack: Vec<(Arc<Node>, usize)>,
end: Option<Vec<u8>>,
started: bool,
start_key: Vec<u8>,
skip_start_key: bool,
last_key: Option<Vec<u8>>,
}
#[cfg(feature = "async-store")]
impl<'a, S> AsyncRangeIter<'a, S>
where
S: AsyncStore,
S::Error: Send + Sync,
{
pub(crate) fn new(
prolly: &'a AsyncProlly<S>,
stack: Vec<(Arc<Node>, usize)>,
start: &[u8],
end: Option<&[u8]>,
) -> Self {
Self {
prolly,
stack,
end: end.map(|e| e.to_vec()),
started: false,
start_key: start.to_vec(),
skip_start_key: false,
last_key: None,
}
}
pub(crate) fn new_after(
prolly: &'a AsyncProlly<S>,
stack: Vec<(Arc<Node>, usize)>,
after_key: &[u8],
end: Option<&[u8]>,
) -> Self {
Self {
prolly,
stack,
end: end.map(|e| e.to_vec()),
started: false,
start_key: after_key.to_vec(),
skip_start_key: true,
last_key: None,
}
}
pub async fn next(&mut self) -> Option<RangeItem> {
self.position_at_start();
loop {
let (node, idx) = self.stack.last_mut()?;
if !node.leaf && node.keys.len() != node.vals.len() {
return Some(Err(Error::InvalidNode));
}
if *idx >= node.len() {
match self.advance_to_next_sibling() {
Ok(true) => continue,
Ok(false) => return None,
Err(e) => return Some(Err(e)),
}
}
if node.leaf {
match leaf_entry_before_end(node, *idx, self.end.as_deref()) {
Ok(Some(entry)) => {
*idx += 1;
self.last_key = Some(entry.0.clone());
return Some(Ok(entry));
}
Ok(None) => return None,
Err(e) => return Some(Err(e)),
}
}
match child_starts_at_or_after_end(self.end.as_deref(), node, *idx) {
Ok(true) => return None,
Ok(false) => {}
Err(e) => return Some(Err(e)),
}
let child = {
let (node, idx) = self.stack.last()?;
self.load_child_for_descent(node, *idx).await
};
match child {
Ok(child) => self.stack.push((child, 0)),
Err(e) => return Some(Err(e)),
}
}
}
pub async fn collect(mut self) -> Result<Vec<LeafEntry>, Error> {
let mut entries = Vec::new();
while let Some(item) = self.next().await {
entries.push(item?);
}
Ok(entries)
}
pub fn resume_cursor(&self) -> RangeCursor {
self.last_key
.clone()
.map(RangeCursor::after_key)
.unwrap_or_else(RangeCursor::start)
}
pub fn into_stream(self) -> impl Stream<Item = RangeItem> + 'a {
stream::unfold(self, |mut iter| async move {
iter.next().await.map(|item| (item, iter))
})
}
fn position_at_start(&mut self) {
if self.started {
return;
}
self.started = true;
let Some((node, idx)) = self.stack.last_mut() else {
return;
};
if node.leaf {
*idx = match node.search(&self.start_key) {
Ok(i) if self.skip_start_key => i.saturating_add(1),
Ok(i) | Err(i) => i,
};
}
}
fn advance_to_next_sibling(&mut self) -> Result<bool, Error> {
loop {
self.stack.pop();
let Some((parent, parent_idx)) = self.stack.last_mut() else {
return Ok(false);
};
*parent_idx += 1;
if *parent_idx < parent.len() {
if child_starts_at_or_after_end(self.end.as_deref(), parent, *parent_idx)? {
return Ok(false);
}
return Ok(true);
}
if parent.keys.len() != parent.vals.len() {
return Err(Error::InvalidNode);
}
}
}
async fn load_child_for_descent(&self, node: &Node, idx: usize) -> Result<Arc<Node>, Error> {
let child_cid = child_cid_at(node, idx)?;
if !self.prolly.store().prefers_batch_reads() {
return self.prolly.load_arc(&child_cid).await;
}
if let Some(child) = self.prolly.cached_node_arc(&child_cid) {
return Ok(child);
}
let max_child_idx = node
.len()
.min(idx.saturating_add(RANGE_CHILD_PREFETCH_PARALLELISM));
let mut child_cids = Vec::with_capacity(max_child_idx.saturating_sub(idx));
child_cids.push(child_cid);
for child_idx in idx + 1..max_child_idx {
if child_starts_at_or_after_end(self.end.as_deref(), node, child_idx).unwrap_or(true) {
break;
}
match child_cid_at(node, child_idx) {
Ok(cid) => child_cids.push(cid),
Err(_) => break,
}
}
if child_cids.len() == 1 {
return self.prolly.load_arc(&child_cids[0]).await;
}
let children = self.prolly.load_child_frontier_ordered(&child_cids).await?;
children.into_iter().next().ok_or(Error::InvalidNode)
}
}
fn leaf_entry_before_end(node: &Node, idx: usize, end: Option<&[u8]>) -> OptionalLeafEntry {
let key = node.keys.get(idx).ok_or(Error::InvalidNode)?;
if let Some(end) = end {
if key.as_slice() >= end {
return Ok(None);
}
}
let val = node.vals.get(idx).ok_or(Error::InvalidNode)?;
Ok(Some((key.clone(), val.clone())))
}
fn child_starts_at_or_after_end(
end: Option<&[u8]>,
node: &Node,
child_index: usize,
) -> Result<bool, Error> {
let Some(end) = end else {
return Ok(false);
};
let first_key = node.keys.get(child_index).ok_or(Error::InvalidNode)?;
Ok(first_key.as_slice() >= end)
}
fn child_cid_at(node: &Node, idx: usize) -> Result<Cid, Error> {
let child = node.vals.get(idx).ok_or(Error::InvalidNode)?;
Ok(Cid(child
.as_slice()
.try_into()
.map_err(|_| Error::InvalidNode)?))
}