use std::cmp::{min, Ordering};
use std::fmt::Debug;
use std::mem;
use std::ops::RangeBounds;
use super::normalize_range;
pub trait SpanWeight<S>: PartialEq + Default + Debug + Clone {
fn alloc(_span: &S) -> Self;
fn and(self, b: &Self) -> Self;
fn union(&mut self, other: &Self);
fn maybe_sub(&mut self, other: &Self) -> bool;
}
impl<T> SpanWeight<T> for () {
fn alloc(_t: &T) -> Self {}
fn and(self, _b: &Self) -> Self {}
fn union(&mut self, _other: &Self) {}
fn maybe_sub(&mut self, _other: &Self) -> bool {
true
}
}
pub(crate) const B: usize = 16;
#[derive(Clone, Debug)]
pub struct SpanTree<T, W = ()>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
{
root_node: Option<TreeNode<T, W>>,
}
#[derive(Clone, Debug, PartialEq)]
struct TreeNode<T: Clone + Debug + Default, W: SpanWeight<T>> {
elements: Vec<T>,
children: Vec<TreeNode<T, W>>,
length: usize,
weight: W,
}
#[derive(Clone, Debug, PartialEq)]
pub struct SubCursor<'a, T: Clone + Debug + Default, W: SpanWeight<T>> {
pub index: usize,
pub weight: W,
pub element: &'a T,
}
impl<T: Clone + Debug + Default, W: SpanWeight<T> + Copy> Copy for SubCursor<'_, T, W> {}
impl<'a, T: Clone + Debug + Default, W: SpanWeight<T>> SubCursor<'a, T, W> {
fn new(index: usize, weight: W, element: &'a T) -> Self {
Self {
index,
weight,
element,
}
}
}
impl<T: Clone + Debug + Default, W: SpanWeight<T>> SpanTree<T, W> {
pub fn new2(element: T) -> Self {
let mut t = Self::default();
t.push(element);
t
}
pub fn load<I: IntoIterator<Item = T>>(elements: I) -> Self {
let mut t = Self::default();
for e in elements.into_iter() {
t.push(e);
}
if t.is_empty() {
t.push(T::default());
}
t
}
pub(crate) fn resume(&self, state: SpanTreeIterState<W>) -> SpanTreeIter<'_, T, W> {
SpanTreeIter {
inner: Some(self),
state,
}
}
pub fn weight(&self) -> Option<&W> {
self.root_node.as_ref().map(|n| n.weight())
}
pub fn last_weight(&self) -> W {
self.root_node
.as_ref()
.map(|n| n.last_weight())
.unwrap_or_default()
}
pub fn len(&self) -> usize {
self.root_node.as_ref().map_or(0, |n| n.len())
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn to_vec(&self) -> Vec<T> {
self.iter().cloned().collect()
}
pub fn iter(&self) -> SpanTreeIter<'_, T, W> {
SpanTreeIter {
inner: Some(self),
state: SpanTreeIterState {
index: 0,
weight: Default::default(),
},
}
}
pub fn splice<R, I>(&mut self, range: R, values: I)
where
R: RangeBounds<usize>,
I: IntoIterator<Item = T>,
{
let mut values = values.into_iter();
let (start, end) = normalize_range(range);
let end = min(end, self.len());
let mut index = start;
let mut to_delete = end - start;
while to_delete > 0 {
if let Some(val) = values.next() {
self.replace(index, val);
index += 1;
} else {
self.remove(index);
}
to_delete -= 1;
}
for val in values {
self.insert(index, val);
index += 1;
}
}
pub fn insert(&mut self, index: usize, element: T) {
let old_len = self.len();
#[cfg(feature = "slow_path_assertions")]
let old_weight = self.weight().cloned().unwrap_or_default();
let weight = W::alloc(&element);
#[cfg(feature = "slow_path_assertions")]
let weight_chk = weight.clone();
if let Some(root) = self.root_node.as_mut() {
#[cfg(feature = "slow_path_assertions")]
{
debug_assert_eq!(root.check_weight(), *root.weight());
root.check();
}
if root.is_full() {
let original_len = root.len();
let new_root = TreeNode::new();
let old_root = mem::replace(root, new_root);
assert_eq!(root.length, 0);
assert_eq!(root.weight, W::default());
root.length += old_root.len();
root.weight.union(old_root.weight());
root.children.push(old_root);
root.split_child(0);
assert_eq!(original_len, root.len());
let first_child_len = root.children[0].len();
let (child, insertion_index) = if first_child_len < index {
(&mut root.children[1], index - (first_child_len + 1))
} else {
(&mut root.children[0], index)
};
root.length += 1;
root.weight.union(&weight);
child.insert_into_non_full_node(insertion_index, element)
} else {
root.insert_into_non_full_node(index, element)
}
} else {
self.root_node = Some(TreeNode {
elements: vec![element],
children: Vec::new(),
length: 1,
weight,
})
}
assert_eq!(self.len(), old_len + 1, "{:#?}", self);
#[cfg(feature = "slow_path_assertions")]
{
let w = self.weight().cloned().unwrap_or_default();
assert_eq!(w, weight_chk.and(&old_weight));
assert_eq!(Some(w), self.check_weight());
}
}
pub fn push(&mut self, element: T) {
let l = self.len();
self.insert(l, element)
}
pub fn get(&self, index: usize) -> Option<&T> {
self.root_node.as_ref().and_then(|n| n.get(index))
}
pub fn get_cursor(&self, index: usize) -> Option<SubCursor<'_, T, W>> {
let acc = Default::default();
self.root_node
.as_ref()
.and_then(|n| n.get_cursor(0, index, acc))
}
pub fn last(&self) -> Option<&T> {
self.root_node.as_ref().and_then(|n| n.last())
}
pub fn iter_where<F>(&self, f: F) -> SpanTreeFnIter<'_, T, W, F>
where
F: Fn(&W, &W) -> bool,
{
let mut iter = SpanTreeFnIter::new(f);
if let Some(node) = self.root_node.as_ref() {
iter.stack.push(NodeWalker::new(node, 0, W::default()))
}
iter
}
pub fn get_where<F>(&self, f: F) -> Option<SubCursor<'_, T, W>>
where
F: Fn(&W, &W) -> bool,
{
self.iter_where(f).next()
}
pub fn get_last_cursor(&self) -> SubCursor<'_, T, W> {
assert!(!self.is_empty());
let element = self.last().unwrap();
let weight = self.last_weight();
let index = self.len() - 1;
SubCursor {
index,
weight,
element,
}
}
pub fn get_where_or_last<F>(&self, f: F) -> SubCursor<'_, T, W>
where
F: Fn(&W, &W) -> bool,
{
self.iter_where(f)
.next()
.unwrap_or_else(|| self.get_last_cursor())
}
pub fn remove(&mut self, index: usize) -> T {
if let Some(root) = self.root_node.as_mut() {
#[cfg(feature = "slow_path_assertions")]
let len = root.check();
#[cfg(feature = "slow_path_assertions")]
debug_assert_eq!(&root.check_weight(), root.weight());
let old = root.remove(index);
if root.elements.is_empty() {
if root.is_leaf() {
self.root_node = None;
} else {
self.root_node = Some(root.children.remove(0));
}
}
#[cfg(feature = "slow_path_assertions")]
{
debug_assert_eq!(len, self.root_node.as_ref().map_or(0, |r| r.check()) + 1);
debug_assert_eq!(self.check_weight().as_ref(), self.weight());
}
old
} else {
panic!("remove from empty tree")
}
}
pub fn replace(&mut self, index: usize, element: T) -> T {
if let Some(root) = self.root_node.as_mut() {
#[cfg(feature = "slow_path_assertions")]
let len = root.check();
let old = root.replace(index, element);
#[cfg(feature = "slow_path_assertions")]
{
debug_assert_eq!(len, self.root_node.as_ref().map_or(0, |r| r.check()));
debug_assert_eq!(self.check_weight().as_ref(), self.weight());
}
old
} else {
panic!("remove from empty tree")
}
}
#[cfg(feature = "slow_path_assertions")]
fn check_weight(&self) -> Option<W> {
self.root_node.as_ref().map(|root| root.check_weight())
}
}
impl<T: Clone + Debug + Default, W: SpanWeight<T>> TreeNode<T, W> {
fn new() -> Self {
Self {
elements: Vec::new(),
children: Vec::new(),
length: 0,
weight: Default::default(),
}
}
fn weight(&self) -> &W {
&self.weight
}
fn last_weight(&self) -> W {
let mut weight = Default::default();
if self.is_leaf() {
let mut iter = self.elements.iter().peekable();
while let Some(e) = iter.next() {
if iter.peek().is_none() {
return weight;
}
weight.union(&W::alloc(e));
}
} else {
for e in self.elements.iter() {
weight.union(&W::alloc(e));
}
let mut iter = self.children.iter().peekable();
while let Some(c) = iter.next() {
if iter.peek().is_none() {
return weight.and(&c.last_weight());
}
weight.union(c.weight());
}
}
panic!()
}
fn len(&self) -> usize {
self.length
}
fn recompute_weight(&mut self) {
let acc = W::default();
let acc = self.elements.iter().fold(acc, |a, e| a.and(&W::alloc(e)));
let acc = self.children.iter().fold(acc, |a, c| a.and(c.weight()));
self.weight = acc;
debug_assert_eq!(&self.check_weight(), self.weight());
}
fn recompute_len(&mut self) {
self.length = self.elements.len() + self.children.iter().map(|c| c.len()).sum::<usize>();
self.recompute_weight()
}
fn is_leaf(&self) -> bool {
self.children.is_empty()
}
fn is_full(&self) -> bool {
self.elements.len() >= 2 * B - 1
}
fn find_child_index(&self, index: usize) -> (usize, usize) {
let mut cumulative_len = 0;
for (child_index, child) in self.children.iter().enumerate() {
if cumulative_len + child.len() >= index {
return (child_index, index - cumulative_len);
} else {
cumulative_len += child.len() + 1;
}
}
panic!("index not found in node")
}
fn insert_into_non_full_node(&mut self, index: usize, element: T) {
assert!(!self.is_full());
if self.is_leaf() {
self.length += 1;
self.weight.union(&W::alloc(&element));
self.elements.insert(index, element);
} else {
let (child_index, sub_index) = self.find_child_index(index);
let child = &mut self.children[child_index];
let element_weight = W::alloc(&element);
if child.is_full() {
self.split_child(child_index);
let (child_index, sub_index) = self.find_child_index(index);
let child = &mut self.children[child_index];
child.insert_into_non_full_node(sub_index, element);
} else {
child.insert_into_non_full_node(sub_index, element);
}
self.length += 1;
self.weight.union(&element_weight);
}
debug_assert_eq!(&self.check_weight(), self.weight());
}
fn split_child(&mut self, full_child_index: usize) {
debug_assert_eq!(&self.check_weight(), self.weight());
let original_len_self = self.len();
let original_weight_self = self.weight().clone();
let mut successor_sibling = TreeNode::new();
let full_child = &mut self.children[full_child_index];
let original_len = full_child.len();
let original_weight = full_child.weight().clone();
assert!(full_child.is_full());
successor_sibling.elements = full_child.elements.split_off(B);
if !full_child.is_leaf() {
successor_sibling.children = full_child.children.split_off(B);
}
let middle = full_child.elements.pop().unwrap();
full_child.recompute_len();
successor_sibling.recompute_len();
let z_len = successor_sibling.len();
let z_weight = successor_sibling.weight().clone();
let full_child_len = full_child.len();
let full_child_weight = full_child.weight().clone();
self.children
.insert(full_child_index + 1, successor_sibling);
let middle_weight = W::alloc(&middle);
self.elements.insert(full_child_index, middle);
assert_eq!(full_child_len + z_len + 1, original_len, "{:#?}", self);
debug_assert_eq!(
full_child_weight.and(&z_weight).and(&middle_weight),
original_weight,
"{:#?}",
self
);
assert_eq!(original_len_self, self.len());
assert_eq!(&original_weight_self, self.weight());
debug_assert_eq!(&self.check_weight(), self.weight());
}
fn remove_from_leaf(&mut self, index: usize) -> T {
let s = self.elements.remove(index);
self.length -= 1;
self.subtract_weight(&W::alloc(&s));
s
}
fn remove_element_from_non_leaf(&mut self, index: usize, element_index: usize) -> T {
let result = if self.children[element_index].elements.len() >= B {
let total_index = self.cumulative_index(element_index);
let predecessor = self.children[element_index].remove(index - 1 - total_index);
mem::replace(&mut self.elements[element_index], predecessor)
} else if self.children[element_index + 1].elements.len() >= B {
let total_index = self.cumulative_index(element_index + 1);
let successor = self.children[element_index + 1].remove(index + 1 - total_index);
mem::replace(&mut self.elements[element_index], successor)
} else {
let middle_element = self.elements.remove(element_index);
let successor_child = self.children.remove(element_index + 1);
self.children[element_index].merge(middle_element, successor_child);
let total_index = self.cumulative_index(element_index);
self.children[element_index].remove(index - total_index)
};
self.length -= 1;
self.subtract_weight(&W::alloc(&result));
result
}
fn cumulative_index(&self, child_index: usize) -> usize {
self.children[0..child_index]
.iter()
.map(|c| c.len() + 1)
.sum()
}
fn remove_from_internal_child(&mut self, index: usize, mut child_index: usize) -> T {
if self.children[child_index].elements.len() < B
&& if child_index > 0 {
self.children[child_index - 1].elements.len() < B
} else {
true
}
&& if child_index + 1 < self.children.len() {
self.children[child_index + 1].elements.len() < B
} else {
true
}
{
if child_index > 0 {
let middle = self.elements.remove(child_index - 1);
let successor = self.children.remove(child_index);
child_index -= 1;
self.children[child_index].merge(middle, successor);
} else {
let middle = self.elements.remove(child_index);
let successor = self.children.remove(child_index + 1);
self.children[child_index].merge(middle, successor);
}
} else if self.children[child_index].elements.len() < B {
if child_index > 0
&& self
.children
.get(child_index - 1)
.is_some_and(|c| c.elements.len() >= B)
{
let last_element = self.children[child_index - 1].elements.pop().unwrap();
assert!(!self.children[child_index - 1].elements.is_empty());
self.children[child_index - 1].length -= 1;
self.children[child_index - 1].subtract_weight(&W::alloc(&last_element));
let parent_element =
mem::replace(&mut self.elements[child_index - 1], last_element);
self.children[child_index].length += 1;
self.children[child_index]
.weight
.union(&W::alloc(&parent_element));
self.children[child_index]
.elements
.insert(0, parent_element);
if let Some(last_child) = self.children[child_index - 1].children.pop() {
self.children[child_index - 1].length -= last_child.len();
self.children[child_index - 1].subtract_weight(last_child.weight());
self.children[child_index].length += last_child.len();
self.children[child_index].weight.union(last_child.weight());
self.children[child_index].children.insert(0, last_child);
}
} else if self
.children
.get(child_index + 1)
.is_some_and(|c| c.elements.len() >= B)
{
let first_element = self.children[child_index + 1].elements.remove(0);
self.children[child_index + 1].length -= 1;
self.children[child_index + 1].subtract_weight(&W::alloc(&first_element));
assert!(!self.children[child_index + 1].elements.is_empty());
let parent_element = mem::replace(&mut self.elements[child_index], first_element);
self.children[child_index].length += 1;
self.children[child_index]
.weight
.union(&W::alloc(&parent_element));
self.children[child_index].elements.push(parent_element);
if !self.children[child_index + 1].is_leaf() {
let first_child = self.children[child_index + 1].children.remove(0);
self.children[child_index + 1].length -= first_child.len();
self.children[child_index + 1].subtract_weight(first_child.weight());
self.children[child_index].length += first_child.len();
self.children[child_index]
.weight
.union(first_child.weight());
self.children[child_index].children.push(first_child);
}
}
}
let total_index = self.cumulative_index(child_index);
let v = self.children[child_index].remove(index - total_index);
self.length -= 1;
self.subtract_weight(&W::alloc(&v));
v
}
fn subtract_weight(&mut self, weight: &W) {
if !self.weight.maybe_sub(weight) {
self.recompute_weight();
}
}
fn check_weight(&self) -> W {
let m = W::default();
let m = self.elements.iter().fold(m, |acc, e| acc.and(&W::alloc(e)));
let m = self.children.iter().fold(m, |acc, c| acc.and(c.weight()));
m
}
fn check(&self) -> usize {
let l = self.elements.len() + self.children.iter().map(|c| c.check()).sum::<usize>();
assert_eq!(self.len(), l, "{:#?}", self);
l
}
fn remove(&mut self, index: usize) -> T {
let original_len = self.len();
if self.is_leaf() {
let v = self.remove_from_leaf(index);
assert_eq!(original_len, self.len() + 1);
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight());
v
} else {
let mut total_index = 0;
for (child_index, child) in self.children.iter().enumerate() {
match (total_index + child.len()).cmp(&index) {
Ordering::Less => {
total_index += child.len() + 1;
continue;
}
Ordering::Equal => {
let v = self.remove_element_from_non_leaf(
index,
min(child_index, self.elements.len() - 1),
);
assert_eq!(original_len, self.len() + 1);
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight());
return v;
}
Ordering::Greater => {
let v = self.remove_from_internal_child(index, child_index);
assert_eq!(original_len, self.len() + 1);
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight());
return v;
}
}
}
panic!(
"index not found to remove {} {} {} {}",
index,
total_index,
self.len(),
self.check()
);
}
}
fn replace(&mut self, index: usize, element: T) -> T {
let original_len = self.len();
debug_assert_eq!(&self.check_weight(), self.weight()); if self.is_leaf() {
self.weight.union(&W::alloc(&element));
let v = mem::replace(&mut self.elements[index], element);
self.subtract_weight(&W::alloc(&v));
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight()); v
} else {
let mut total_index = 0;
for (child_index, child) in self.children.iter().enumerate() {
match (total_index + child.len()).cmp(&index) {
Ordering::Less => {
total_index += child.len() + 1;
continue;
}
Ordering::Equal => {
let child_index = min(child_index, self.elements.len() - 1);
self.weight.union(&W::alloc(&element));
let v = mem::replace(&mut self.elements[child_index], element);
self.subtract_weight(&W::alloc(&v));
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight());
return v;
}
Ordering::Greater => {
self.weight.union(&W::alloc(&element));
let v = self.children[child_index].replace(index - total_index, element);
self.subtract_weight(&W::alloc(&v));
assert_eq!(original_len, self.len());
debug_assert_eq!(self.check(), self.len());
debug_assert_eq!(&self.check_weight(), self.weight());
return v;
}
}
}
panic!(
"index not found to remove {} {} {} {}",
index,
total_index,
self.len(),
self.check()
);
}
}
fn merge(&mut self, middle: T, successor_sibling: TreeNode<T, W>) {
self.length += successor_sibling.length + 1;
self.weight.union(&successor_sibling.weight);
self.weight.union(&W::alloc(&middle));
self.elements.push(middle);
self.elements.extend(successor_sibling.elements);
self.children.extend(successor_sibling.children);
debug_assert_eq!(&self.check_weight(), self.weight());
assert!(self.is_full());
}
fn last(&self) -> Option<&T> {
if self.is_leaf() {
self.elements.last()
} else {
self.children.last().and_then(|c| c.last())
}
}
fn get(&self, index: usize) -> Option<&T> {
if self.is_leaf() {
return self.elements.get(index);
} else {
let mut cumulative_len = 0;
for (child_index, child) in self.children.iter().enumerate() {
match (cumulative_len + child.len()).cmp(&index) {
Ordering::Less => {
cumulative_len += child.len() + 1;
}
Ordering::Equal => return self.elements.get(child_index),
Ordering::Greater => {
return child.get(index - cumulative_len);
}
}
}
}
None
}
fn get_cursor(
&self,
mut current: usize,
index: usize,
mut acc: W,
) -> Option<SubCursor<'_, T, W>> {
if self.is_leaf() {
let iter = self.elements.iter();
for e in iter {
if current == index {
return Some(SubCursor::new(index, acc, e));
}
current += 1;
acc.union(&W::alloc(e));
}
} else {
for i in 0..self.children.len() {
if let Some(child) = self.children.get(i) {
if current + child.len() > index {
return child.get_cursor(current, index, acc);
}
current += child.len();
acc.union(child.weight());
}
if let Some(e) = self.elements.get(i) {
if current == index {
return Some(SubCursor::new(index, acc, e));
}
current += 1;
acc.union(&W::alloc(e));
}
}
}
None
}
}
impl<T: Clone + Debug + Default, W: SpanWeight<T>> Default for SpanTree<T, W> {
fn default() -> SpanTree<T, W> {
Self { root_node: None }
}
}
impl<'a, T: Clone + Debug + Default, W: SpanWeight<T>> IntoIterator for &'a SpanTree<T, W> {
type Item = &'a T;
type IntoIter = SpanTreeIter<'a, T, W>;
fn into_iter(self) -> Self::IntoIter {
SpanTreeIter {
inner: Some(self),
state: SpanTreeIterState {
index: 0,
weight: Default::default(),
},
}
}
}
#[derive(Debug, Clone, Default)]
pub struct SpanTreeIter<'a, T: Clone + Debug + Default, W: SpanWeight<T>> {
inner: Option<&'a SpanTree<T, W>>,
state: SpanTreeIterState<W>,
}
#[derive(Debug, Clone, Default)]
pub(crate) struct SpanTreeIterState<W> {
index: usize,
weight: W,
}
impl<'a, T: Clone + Debug + Default, W: SpanWeight<T>> SpanTreeIter<'a, T, W> {
pub(crate) fn new(
tree: &'a SpanTree<T, W>,
cursor: SubCursor<'_, T, W>,
) -> SpanTreeIter<'a, T, W> {
Self {
inner: Some(tree),
state: SpanTreeIterState {
index: cursor.index + 1,
weight: cursor.weight.and(&W::alloc(cursor.element)),
},
}
}
pub(crate) fn suspend(&self) -> SpanTreeIterState<W> {
self.state.clone()
}
pub fn weight(&self) -> &W {
&self.state.weight
}
pub fn total_weight(&self) -> Option<&W> {
self.inner.as_ref().and_then(|t| t.weight())
}
pub fn index(&self) -> usize {
self.state.index
}
pub(crate) fn span_tree(&self) -> Option<&'a SpanTree<T, W>> {
self.inner
}
pub(crate) fn peek(&self) -> Option<&'a T> {
self.inner?.get(self.state.index + 1)
}
pub(crate) fn current(&self) -> Option<&'a T> {
if self.state.index > 0 {
Some(self.inner?.get(self.state.index - 1)?)
} else {
None
}
}
}
impl<T: Clone + Debug + Default, W: SpanWeight<T> + Copy> Copy for SpanTreeIter<'_, T, W> {}
impl<W: Copy> Copy for SpanTreeIterState<W> {}
impl<'a, T: Clone + Debug + Default, W: SpanWeight<T>> Iterator for SpanTreeIter<'a, T, W> {
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
let element = self.inner?.get(self.state.index)?;
self.state.index += 1;
self.state.weight.union(&W::alloc(element)); Some(element)
}
}
#[cfg(test)]
mod legacy_tests {
use proptest::prelude::*;
use super::*;
#[test]
fn push_back() {
let mut t = SpanTree::<u32>::default();
t.push(1);
t.push(2);
t.push(3);
t.push(4);
t.push(5);
t.push(6);
t.push(8);
t.push(100);
}
#[test]
fn insert() {
let mut t = SpanTree::<u32>::default();
t.insert(0, 1);
t.insert(1, 1);
t.insert(0, 1);
t.insert(0, 1);
t.insert(0, 1);
t.insert(3, 1);
t.insert(4, 1);
}
#[test]
fn insert_book() {
let mut t = SpanTree::<()>::default();
for i in 0..100 {
t.insert(i % 2, ());
}
}
#[test]
fn insert_book_vec() {
let mut t = SpanTree::<()>::default();
let mut v = Vec::new();
for i in 0..100 {
t.insert(i % 3, ());
v.insert(i % 3, ());
assert_eq!(v, t.iter().copied().collect::<Vec<_>>())
}
}
fn arb_indices() -> impl Strategy<Value = Vec<usize>> {
proptest::collection::vec(any::<usize>(), 0..1000).prop_map(|v| {
let mut len = 0;
v.into_iter()
.map(|i| {
len += 1;
i % len
})
.collect::<Vec<_>>()
})
}
proptest! {
#[test]
fn proptest_insert(indices in arb_indices()) {
let mut t = SpanTree::<usize,usize>::default();
let mut v = Vec::new();
for i in indices{
if i <= v.len() {
t.insert(i % 3, i);
v.insert(i % 3, i);
} else {
return Err(proptest::test_runner::TestCaseError::reject("index out of bounds"))
}
assert_eq!(v, t.iter().copied().collect::<Vec<_>>())
}
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(20))]
#[test]
fn proptest_remove(inserts in arb_indices(), removes in arb_indices()) {
let mut t = SpanTree::<usize,usize>::default();
let mut v = Vec::new();
for i in inserts {
if i <= v.len() {
t.insert(i , i);
v.insert(i , i);
} else {
return Err(proptest::test_runner::TestCaseError::reject("index out of bounds"))
}
assert_eq!(v, t.iter().copied().collect::<Vec<_>>())
}
for i in removes {
if i < v.len() {
let tr = t.remove(i);
let vr = v.remove(i);
assert_eq!(tr, vr);
} else {
return Err(proptest::test_runner::TestCaseError::reject("index out of bounds"))
}
assert_eq!(v, t.iter().copied().collect::<Vec<_>>())
}
}
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::{ColumnData, UIntCursor};
use std::ops::Range;
#[derive(Default, Clone, Debug, PartialEq)]
struct TestWidth {
index: usize,
weight: usize,
}
impl SpanWeight<TestWidth> for usize {
fn alloc(t: &TestWidth) -> usize {
t.weight
}
fn and(self, b: &usize) -> usize {
self + *b
}
fn union(&mut self, other: &usize) {
*self += *other
}
fn maybe_sub(&mut self, other: &Self) -> bool {
*self -= *other;
true
}
}
impl SpanWeight<usize> for usize {
fn alloc(t: &usize) -> usize {
*t
}
fn and(self, b: &usize) -> usize {
self + *b
}
fn union(&mut self, other: &usize) {
*self += *other
}
fn maybe_sub(&mut self, other: &usize) -> bool {
*self -= *other;
true
}
}
#[test]
fn test_basic_get_where() {
let mut tree = SpanTree::<usize, usize>::default();
tree.push(10);
tree.push(20);
tree.push(30);
tree.push(40);
tree.push(50);
assert_eq!(*tree.weight().unwrap(), 150);
assert_eq!(tree.get(0), Some(&10));
assert_eq!(tree.get(1), Some(&20));
assert_eq!(tree.get(2), Some(&30));
assert_eq!(tree.get(3), Some(&40));
assert_eq!(tree.get(4), Some(&50));
assert_eq!(
tree.get_where(|acc, next| 0 < *acc + *next),
Some(SubCursor::new(0, 0, &10))
);
assert_eq!(
tree.get_where(|acc, next| 5 < *acc + *next),
Some(SubCursor::new(0, 0, &10))
);
assert_eq!(
tree.get_where(|acc, next| 6 < *acc + *next),
Some(SubCursor::new(0, 0, &10))
);
assert_eq!(
tree.get_where(|acc, next| 10 < *acc + *next),
Some(SubCursor::new(1, 10, &20))
);
assert_eq!(
tree.get_where(|acc, next| 15 < *acc + *next),
Some(SubCursor::new(1, 10, &20))
);
assert_eq!(
tree.get_where(|acc, next| 16 < *acc + *next),
Some(SubCursor::new(1, 10, &20))
);
assert_eq!(
tree.get_where(|acc, next| 29 < *acc + *next),
Some(SubCursor::new(1, 10, &20))
);
assert_eq!(
tree.get_where(|acc, next| 30 < *acc + *next),
Some(SubCursor::new(2, 30, &30))
);
assert_eq!(
tree.get_where(|acc, next| 40 < *acc + *next),
Some(SubCursor::new(2, 30, &30))
);
assert_eq!(
tree.get_where(|acc, next| 50 < *acc + *next),
Some(SubCursor::new(2, 30, &30))
);
assert_eq!(
tree.get_where(|acc, next| 60 < *acc + *next),
Some(SubCursor::new(3, 60, &40))
);
}
#[test]
fn really_big_fn_iter_test() {
let col: ColumnData<UIntCursor> = (10..2_000_010).collect();
let found = col.find_by_value(2_000_009).collect::<Vec<_>>();
assert_eq!(found, vec![1999999]);
}
#[test]
fn test_advanced_get_where() {
let mut tree = SpanTree::<TestWidth, usize>::default();
const MAX: usize = 1000;
for index in 0..MAX {
tree.push(TestWidth { index, weight: 10 });
}
assert_eq!(*tree.weight().unwrap(), MAX * 10);
for i in 0..MAX * 10 {
let index = i / 10;
let weight = 10;
assert_eq!(
tree.get_where(|acc, next| i < *acc + *next),
Some(SubCursor::new(
index,
index * weight,
&TestWidth { index, weight }
))
);
}
tree.replace(
10,
TestWidth {
index: 10,
weight: 100,
},
);
assert_eq!(*tree.weight().unwrap(), MAX * 10 + 90);
assert_eq!(
tree.get_where(|acc, next| 201 < *acc + *next),
Some(SubCursor::new(
11,
200,
&TestWidth {
index: 11,
weight: 10
}
))
);
tree.replace(
20,
TestWidth {
index: 10,
weight: 1000,
},
);
assert_eq!(*tree.weight().unwrap(), MAX * 10 + 90 + 990);
let result = tree.iter().map(usize::alloc).sum::<usize>();
assert_eq!(tree.weight(), Some(&result));
}
fn test_splice(
tree: &mut SpanTree<usize, usize>,
vec: &mut Vec<usize>,
range: Range<usize>,
values: Vec<usize>,
) {
vec.splice(range.clone(), values.clone());
tree.splice(range, values);
assert_eq!(&tree.to_vec(), vec);
}
#[test]
fn test_tree_splice() {
let mut tree = SpanTree::<usize, usize>::default();
let mut base = vec![];
test_splice(&mut tree, &mut base, 0..0, vec![1, 2, 3]);
test_splice(&mut tree, &mut base, 1..2, vec![7]);
test_splice(&mut tree, &mut base, 1..2, vec![10, 11, 12, 13, 14]);
test_splice(&mut tree, &mut base, 2..5, vec![50, 60, 70]);
}
}
#[derive(Clone, Debug)]
struct NodeWalker<'a, T, W>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
{
children: std::slice::Iter<'a, TreeNode<T, W>>,
elements: std::slice::Iter<'a, T>,
index: usize,
acc: W,
}
impl<'a, T, W> NodeWalker<'a, T, W>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
{
fn new(node: &'a TreeNode<T, W>, index: usize, acc: W) -> Self {
NodeWalker {
children: node.children.iter(),
elements: node.elements.iter(),
index,
acc,
}
}
fn next_walker<F>(&mut self, f: &F) -> Option<NodeWalker<'a, T, W>>
where
F: Fn(&W, &W) -> bool,
{
let child = self.children.next()?;
let next = child.weight();
let index = self.index;
let acc = self.acc.clone();
self.index += child.len();
self.acc.union(next);
if f(&acc, next) {
Some(NodeWalker::new(child, index, acc))
} else {
None
}
}
fn next_cursor<F>(&mut self, f: &F) -> Option<Option<SubCursor<'a, T, W>>>
where
F: Fn(&W, &W) -> bool,
{
let e = self.elements.next()?;
let next = W::alloc(e);
let index = self.index;
let acc = self.acc.clone();
self.index += 1;
self.acc.union(&next);
if f(&acc, &next) {
Some(Some(SubCursor::new(index, acc, e)))
} else {
Some(None)
}
}
}
#[derive(Clone, Debug)]
pub struct SpanTreeFnIter<'a, T, W, F>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
F: Fn(&W, &W) -> bool,
{
stack: Vec<NodeWalker<'a, T, W>>,
func: F,
}
impl<'a, T, W, F> SpanTreeFnIter<'a, T, W, F>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
F: Fn(&W, &W) -> bool,
{
fn new(func: F) -> Self {
Self {
stack: vec![],
func,
}
}
fn pop_cursor(&mut self) -> Option<SubCursor<'a, T, W>> {
if let Some(result) = self.stack.last_mut()?.next_cursor(&self.func) {
result
} else {
self.stack.pop();
self.pop_cursor()
}
}
}
impl<'a, T, W, F> Iterator for SpanTreeFnIter<'a, T, W, F>
where
T: Clone + Debug + Default,
W: SpanWeight<T>,
F: Fn(&W, &W) -> bool,
{
type Item = SubCursor<'a, T, W>;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(walker) = self.stack.last_mut()?.next_walker(&self.func) {
self.stack.push(walker);
} else if let Some(result) = self.pop_cursor() {
return Some(result);
}
}
}
}