use crate::bit_vec::fast_rs_vec::SelectIntoIter;
use crate::trees::mmt::MinMaxTree;
use crate::trees::{IsAncestor, LevelTree, SubtreeSize, Tree};
use crate::{BitVec, RsVec};
use std::cmp::{max, min};
use std::iter::FusedIterator;
const DEFAULT_BLOCK_SIZE: usize = 512;
const OPEN_PAREN: u64 = 1;
const CLOSE_PAREN: u64 = 0;
mod builder;
pub use builder::BpBuilder;
#[cfg(feature = "bp_u16_lookup")]
mod lookup;
#[cfg(feature = "bp_u16_lookup")]
use lookup::{process_block_bwd, process_block_fwd, LOOKUP_BLOCK_SIZE};
#[cfg(not(feature = "bp_u16_lookup"))]
mod lookup_query;
#[cfg(not(feature = "bp_u16_lookup"))]
use lookup_query::{process_block_bwd, process_block_fwd, LOOKUP_BLOCK_SIZE};
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "mem_dbg", derive(mem_dbg::MemSize, mem_dbg::MemDbg))]
pub struct BpTree<const BLOCK_SIZE: usize = DEFAULT_BLOCK_SIZE> {
vec: RsVec,
min_max_tree: MinMaxTree,
}
impl<const BLOCK_SIZE: usize> BpTree<BLOCK_SIZE> {
#[must_use]
pub fn from_bit_vector(bv: BitVec) -> Self {
let min_max_tree = MinMaxTree::excess_tree(&bv, BLOCK_SIZE);
let vec = bv.into();
Self { vec, min_max_tree }
}
pub fn fwd_search(&self, index: usize, mut relative_excess: i64) -> Option<usize> {
if index >= (self.vec.len() - 1) {
return None;
}
let block_index = (index + 1) / BLOCK_SIZE;
self.fwd_search_block(index, block_index, &mut relative_excess)
.map_or_else(
|()| {
let block = self.min_max_tree.fwd_search(block_index, relative_excess);
block.and_then(|(block, mut relative_excess)| {
self.fwd_search_block(block * BLOCK_SIZE - 1, block, &mut relative_excess)
.ok()
})
},
Some,
)
}
#[inline(always)]
fn fwd_search_block(
&self,
start_index: usize,
block_index: usize,
relative_excess: &mut i64,
) -> Result<usize, ()> {
let block_boundary = min((block_index + 1) * BLOCK_SIZE, self.vec.len());
let lookup_boundary = min(
(start_index + 1).div_ceil(LOOKUP_BLOCK_SIZE as usize) * LOOKUP_BLOCK_SIZE as usize,
block_boundary,
);
for i in start_index + 1..lookup_boundary {
let bit = self.vec.get_unchecked(i);
*relative_excess -= if bit == 1 { 1 } else { -1 };
if *relative_excess == 0 {
return Ok(i);
}
}
let upper_lookup_boundary = max(
lookup_boundary,
(block_boundary / LOOKUP_BLOCK_SIZE as usize) * LOOKUP_BLOCK_SIZE as usize,
);
for i in (lookup_boundary..upper_lookup_boundary).step_by(LOOKUP_BLOCK_SIZE as usize) {
if let Ok(idx) = process_block_fwd(
self.vec
.get_bits_unchecked(i, LOOKUP_BLOCK_SIZE as usize)
.try_into()
.unwrap(),
relative_excess,
) {
return Ok(i + idx as usize);
}
}
for i in upper_lookup_boundary..block_boundary {
let bit = self.vec.get_unchecked(i);
*relative_excess -= if bit == 1 { 1 } else { -1 };
if *relative_excess == 0 {
return Ok(i);
}
}
Err(())
}
pub fn bwd_search(&self, index: usize, mut relative_excess: i64) -> Option<usize> {
if index >= self.vec.len() {
return None;
}
if index == 0 {
return None;
}
let block_index = (index - 1) / BLOCK_SIZE;
self.bwd_search_block(index, block_index, &mut relative_excess)
.map_or_else(
|()| {
let block = self.min_max_tree.bwd_search(block_index, relative_excess);
block.and_then(|(block, mut relative_excess)| {
self.bwd_search_block((block + 1) * BLOCK_SIZE, block, &mut relative_excess)
.ok()
})
},
Some,
)
}
#[inline(always)]
fn bwd_search_block(
&self,
start_index: usize,
block_index: usize,
relative_excess: &mut i64,
) -> Result<usize, ()> {
let block_boundary = min(block_index * BLOCK_SIZE, self.vec.len());
let lookup_boundary = max(
((start_index - 1) / LOOKUP_BLOCK_SIZE as usize) * LOOKUP_BLOCK_SIZE as usize,
block_boundary,
);
for i in (lookup_boundary..start_index).rev() {
let bit = self.vec.get_unchecked(i);
*relative_excess += if bit == 1 { 1 } else { -1 };
if *relative_excess == 0 {
return Ok(i);
}
}
for i in (block_boundary..lookup_boundary)
.step_by(LOOKUP_BLOCK_SIZE as usize)
.rev()
{
if let Ok(idx) = process_block_bwd(
self.vec
.get_bits_unchecked(i, LOOKUP_BLOCK_SIZE as usize)
.try_into()
.unwrap(),
relative_excess,
) {
return Ok(i + idx as usize);
}
}
Err(())
}
#[must_use]
pub fn close(&self, index: usize) -> Option<usize> {
if index >= self.vec.len() {
return None;
}
self.fwd_search(index, -1)
}
#[must_use]
pub fn open(&self, index: usize) -> Option<usize> {
if index >= self.vec.len() {
return None;
}
self.bwd_search(index, -1)
}
#[must_use]
pub fn enclose(&self, index: usize) -> Option<usize> {
if index >= self.vec.len() {
return None;
}
self.bwd_search(
index,
if self.vec.get_unchecked(index) == 1 {
-1
} else {
-2
},
)
}
#[must_use]
pub fn excess(&self, index: usize) -> i64 {
debug_assert!(index < self.vec.len(), "Index out of bounds");
self.vec.rank1(index + 1) as i64 - self.vec.rank0(index + 1) as i64
}
pub fn iter(
&self,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
self.dfs_iter()
}
pub fn dfs_iter(
&self,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
self.vec.iter1()
}
pub fn dfs_post_iter(
&self,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
self.vec.iter0().map(|n| self.open(n).unwrap())
}
pub fn subtree_iter(
&self,
node: <BpTree<BLOCK_SIZE> as Tree>::NodeHandle,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
let index = self.vec.rank1(node);
let close = self.close(node).unwrap_or(node);
let subtree_size = self.vec.rank1(close) - index;
self.vec.iter1().skip(index).take(subtree_size)
}
pub fn subtree_post_iter(
&self,
node: <BpTree<BLOCK_SIZE> as Tree>::NodeHandle,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
let index = self.vec.rank0(node);
let close = self.close(node).unwrap_or(node);
let subtree_size = self.vec.rank0(close) + 1 - index;
self.vec
.iter0()
.skip(index)
.take(subtree_size)
.map(|n| self.open(n).unwrap())
}
pub fn children(
&self,
node: <BpTree<BLOCK_SIZE> as Tree>::NodeHandle,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
ChildrenIter::<BLOCK_SIZE, true>::new(self, node)
}
pub fn rev_children(
&self,
node: <BpTree<BLOCK_SIZE> as Tree>::NodeHandle,
) -> impl Iterator<Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle> + use<'_, BLOCK_SIZE> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
ChildrenIter::<BLOCK_SIZE, false>::new(self, node)
}
#[must_use]
pub fn into_parentheses_vec(self) -> RsVec {
self.vec
}
#[must_use]
pub fn heap_size(&self) -> usize {
self.vec.heap_size() + self.min_max_tree.heap_size()
}
}
impl<const BLOCK_SIZE: usize> Tree for BpTree<BLOCK_SIZE> {
type NodeHandle = usize;
fn root(&self) -> Option<Self::NodeHandle> {
if self.vec.is_empty() {
None
} else {
Some(0)
}
}
fn parent(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.enclose(node)
}
fn first_child(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
if let Some(bit) = self.vec.get(node + 1) {
if bit == OPEN_PAREN {
return Some(node + 1);
}
}
None
}
fn next_sibling(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.close(node).and_then(|i| {
self.vec
.get(i + 1)
.and_then(|bit| if bit == OPEN_PAREN { Some(i + 1) } else { None })
})
}
fn previous_sibling(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
if node == 0 {
None
} else {
self.vec.get(node - 1).and_then(|bit| {
if bit == CLOSE_PAREN {
self.open(node - 1)
} else {
None
}
})
}
}
fn last_child(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.vec.get(node + 1).and_then(|bit| {
if bit == OPEN_PAREN {
if let Some(i) = self.close(node) {
self.open(i - 1)
} else {
None
}
} else {
None
}
})
}
fn node_index(&self, node: Self::NodeHandle) -> usize {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.vec.rank1(node)
}
fn node_handle(&self, index: usize) -> Self::NodeHandle {
self.vec.select1(index)
}
fn is_leaf(&self, node: Self::NodeHandle) -> bool {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.vec.get(node + 1) == Some(CLOSE_PAREN)
}
fn depth(&self, node: Self::NodeHandle) -> u64 {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
let excess: u64 = self.excess(node).try_into().unwrap_or(0);
excess.saturating_sub(1)
}
fn size(&self) -> usize {
self.vec.rank1(self.vec.len())
}
fn is_empty(&self) -> bool {
self.vec.is_empty()
}
}
impl<const BLOCK_SIZE: usize> IsAncestor for BpTree<BLOCK_SIZE> {
fn is_ancestor(
&self,
ancestor: Self::NodeHandle,
descendant: Self::NodeHandle,
) -> Option<bool> {
debug_assert!(
self.vec.get(ancestor) == Some(OPEN_PAREN),
"Node handle is invalid"
);
debug_assert!(
self.vec.get(descendant) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.close(ancestor)
.map(|closing| ancestor <= descendant && descendant < closing)
}
}
impl<const BLOCK_SIZE: usize> LevelTree for BpTree<BLOCK_SIZE> {
fn level_ancestor(&self, node: Self::NodeHandle, level: u64) -> Option<Self::NodeHandle> {
if level == 0 {
return Some(node);
}
#[allow(clippy::cast_possible_wrap)]
self.bwd_search(node, -(level as i64))
}
fn level_next(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
self.fwd_search(self.close(node)?, 1)
}
fn level_prev(&self, node: Self::NodeHandle) -> Option<Self::NodeHandle> {
self.open(self.bwd_search(node, 1)?)
}
fn level_leftmost(&self, level: u64) -> Option<Self::NodeHandle> {
if level == 0 {
return Some(0);
}
#[allow(clippy::cast_possible_wrap)]
self.fwd_search(0, level as i64)
}
fn level_rightmost(&self, level: u64) -> Option<Self::NodeHandle> {
#[allow(clippy::cast_possible_wrap)]
self.open(self.bwd_search(self.size() * 2 - 1, level as i64)?)
}
}
impl<const BLOCK_SIZE: usize> SubtreeSize for BpTree<BLOCK_SIZE> {
fn subtree_size(&self, node: Self::NodeHandle) -> Option<usize> {
debug_assert!(
self.vec.get(node) == Some(OPEN_PAREN),
"Node handle is invalid"
);
self.close(node)
.map(|c| self.vec.rank1(c) - self.vec.rank1(node))
}
}
impl<const BLOCK_SIZE: usize> IntoIterator for BpTree<BLOCK_SIZE> {
type Item = <BpTree<BLOCK_SIZE> as Tree>::NodeHandle;
type IntoIter = SelectIntoIter<false>;
fn into_iter(self) -> Self::IntoIter {
self.vec.into_iter1()
}
}
impl<const BLOCK_SIZE: usize> From<BitVec> for BpTree<BLOCK_SIZE> {
fn from(bv: BitVec) -> Self {
Self::from_bit_vector(bv)
}
}
impl<const BLOCK_SIZE: usize> From<BpTree<BLOCK_SIZE>> for BitVec {
fn from(value: BpTree<BLOCK_SIZE>) -> Self {
value.into_parentheses_vec().into_bit_vec()
}
}
impl<const BLOCK_SIZE: usize> From<BpTree<BLOCK_SIZE>> for RsVec {
fn from(value: BpTree<BLOCK_SIZE>) -> Self {
value.into_parentheses_vec()
}
}
struct ChildrenIter<'a, const BLOCK_SIZE: usize, const FORWARD: bool> {
tree: &'a BpTree<BLOCK_SIZE>,
current_sibling: Option<usize>,
}
impl<'a, const BLOCK_SIZE: usize, const FORWARD: bool> ChildrenIter<'a, BLOCK_SIZE, FORWARD> {
fn new(tree: &'a BpTree<BLOCK_SIZE>, node: usize) -> Self {
Self {
tree,
current_sibling: if FORWARD {
tree.first_child(node)
} else {
tree.last_child(node)
},
}
}
}
impl<const BLOCK_SIZE: usize, const FORWARD: bool> Iterator
for ChildrenIter<'_, BLOCK_SIZE, FORWARD>
{
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
let current = self.current_sibling?;
let next = if FORWARD {
self.tree.next_sibling(current)
} else {
self.tree.previous_sibling(current)
};
self.current_sibling = next;
Some(current)
}
}
impl<const BLOCK_SIZE: usize, const FORWARD: bool> FusedIterator
for ChildrenIter<'_, BLOCK_SIZE, FORWARD>
{
}
#[cfg(test)]
mod tests;