Implement FoldMap's folds using a SumTree

This required passing a `Context` object to `Summary` and introducing a
new `SeekDimension` trait that allows comparing two dimensions and pass
an additional context object.
This commit is contained in:
Antonio Scandurra
2021-05-06 16:20:23 +02:00
parent 7fb7a6adfc
commit 901525bf77
9 changed files with 354 additions and 151 deletions
+88 -33
View File
@@ -382,22 +382,48 @@ where
impl<'a, T, S, U> Cursor<'a, T, S, U>
where
T: Item,
S: Dimension<'a, T::Summary> + Ord,
S: SeekDimension<'a, T::Summary>,
U: Dimension<'a, T::Summary>,
{
pub fn seek(&mut self, pos: &S, bias: SeekBias) -> bool {
self.seek_with_ctx(pos, bias, None)
}
pub fn seek_with_ctx(
&mut self,
pos: &S,
bias: SeekBias,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> bool {
self.reset();
self.seek_internal::<()>(pos, bias, &mut SeekAggregate::None)
self.seek_internal::<()>(pos, bias, &mut SeekAggregate::None, ctx)
}
#[allow(unused)]
pub fn seek_forward(&mut self, pos: &S, bias: SeekBias) -> bool {
self.seek_internal::<()>(pos, bias, &mut SeekAggregate::None)
self.seek_forward_with_ctx(pos, bias, None)
}
pub fn seek_forward_with_ctx(
&mut self,
pos: &S,
bias: SeekBias,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> bool {
self.seek_internal::<()>(pos, bias, &mut SeekAggregate::None, ctx)
}
pub fn slice(&mut self, end: &S, bias: SeekBias) -> SumTree<T> {
self.slice_with_ctx(end, bias, None)
}
pub fn slice_with_ctx(
&mut self,
end: &S,
bias: SeekBias,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> SumTree<T> {
let mut slice = SeekAggregate::Slice(SumTree::new());
self.seek_internal::<()>(end, bias, &mut slice);
self.seek_internal::<()>(end, bias, &mut slice, ctx);
if let SeekAggregate::Slice(slice) = slice {
slice
} else {
@@ -406,9 +432,16 @@ where
}
pub fn suffix(&mut self) -> SumTree<T> {
self.suffix_with_ctx(None)
}
pub fn suffix_with_ctx(
&mut self,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> SumTree<T> {
let extent = self.tree.extent::<S>();
let mut slice = SeekAggregate::Slice(SumTree::new());
self.seek_internal::<()>(&extent, SeekBias::Right, &mut slice);
self.seek_internal::<()>(&extent, SeekBias::Right, &mut slice, ctx);
if let SeekAggregate::Slice(slice) = slice {
slice
} else {
@@ -417,11 +450,23 @@ where
}
pub fn summary<D>(&mut self, end: &S, bias: SeekBias) -> D
where
D: Dimension<'a, T::Summary>,
{
self.summary_with_ctx(end, bias, None)
}
pub fn summary_with_ctx<D>(
&mut self,
end: &S,
bias: SeekBias,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> D
where
D: Dimension<'a, T::Summary>,
{
let mut summary = SeekAggregate::Summary(D::default());
self.seek_internal(end, bias, &mut summary);
self.seek_internal(end, bias, &mut summary, ctx);
if let SeekAggregate::Summary(summary) = summary {
summary
} else {
@@ -434,11 +479,12 @@ where
target: &S,
bias: SeekBias,
aggregate: &mut SeekAggregate<T, D>,
ctx: Option<&'a <T::Summary as Summary>::Context>,
) -> bool
where
D: Dimension<'a, T::Summary>,
{
debug_assert!(target >= &self.seek_dimension);
debug_assert!(target.cmp(&self.seek_dimension, ctx) >= Ordering::Equal);
let mut containing_subtree = None;
if self.did_seek {
@@ -458,7 +504,7 @@ where
let mut child_end = self.seek_dimension.clone();
child_end.add_summary(&child_summary);
let comparison = target.cmp(&child_end);
let comparison = target.cmp(&child_end, ctx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == SeekBias::Right)
{
@@ -467,7 +513,7 @@ where
match aggregate {
SeekAggregate::None => {}
SeekAggregate::Slice(slice) => {
slice.push_tree(child_tree.clone());
slice.push_tree_with_ctx(child_tree.clone(), ctx);
}
SeekAggregate::Summary(summary) => {
summary.add_summary(child_summary);
@@ -500,7 +546,7 @@ where
let mut item_end = self.seek_dimension.clone();
item_end.add_summary(item_summary);
let comparison = target.cmp(&item_end);
let comparison = target.cmp(&item_end, ctx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == SeekBias::Right)
{
@@ -514,7 +560,7 @@ where
slice_items_summary
.as_mut()
.unwrap()
.add_summary(item_summary);
.add_summary_with_ctx(item_summary, ctx);
}
SeekAggregate::Summary(summary) => {
summary.add_summary(item_summary);
@@ -523,11 +569,14 @@ where
entry.index += 1;
} else {
if let SeekAggregate::Slice(slice) = aggregate {
slice.push_tree(SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})));
slice.push_tree_with_ctx(
SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})),
ctx,
);
}
break 'outer;
}
@@ -535,11 +584,14 @@ where
if let SeekAggregate::Slice(slice) = aggregate {
if !slice_items.is_empty() {
slice.push_tree(SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})));
slice.push_tree_with_ctx(
SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})),
ctx,
);
}
}
}
@@ -568,7 +620,7 @@ where
let mut child_end = self.seek_dimension.clone();
child_end.add_summary(child_summary);
let comparison = target.cmp(&child_end);
let comparison = target.cmp(&child_end, ctx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == SeekBias::Right)
{
@@ -577,7 +629,7 @@ where
match aggregate {
SeekAggregate::None => {}
SeekAggregate::Slice(slice) => {
slice.push_tree(child_trees[index].clone());
slice.push_tree_with_ctx(child_trees[index].clone(), ctx);
}
SeekAggregate::Summary(summary) => {
summary.add_summary(child_summary);
@@ -614,7 +666,7 @@ where
let mut child_end = self.seek_dimension.clone();
child_end.add_summary(item_summary);
let comparison = target.cmp(&child_end);
let comparison = target.cmp(&child_end, ctx);
if comparison == Ordering::Greater
|| (comparison == Ordering::Equal && bias == SeekBias::Right)
{
@@ -627,7 +679,7 @@ where
slice_items_summary
.as_mut()
.unwrap()
.add_summary(item_summary);
.add_summary_with_ctx(item_summary, ctx);
slice_item_summaries.push(item_summary.clone());
}
SeekAggregate::Summary(summary) => {
@@ -647,11 +699,14 @@ where
if let SeekAggregate::Slice(slice) = aggregate {
if !slice_items.is_empty() {
slice.push_tree(SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})));
slice.push_tree_with_ctx(
SumTree(Arc::new(Node::Leaf {
summary: slice_items_summary.unwrap(),
items: slice_items,
item_summaries: slice_item_summaries,
})),
ctx,
);
}
}
}
@@ -672,9 +727,9 @@ where
if let Some(summary) = self.item_summary() {
end.add_summary(summary);
}
*target == end
target.cmp(&end, ctx) == Ordering::Equal
} else {
*target == self.seek_dimension
target.cmp(&self.seek_dimension, ctx) == Ordering::Equal
}
}
}
+76 -28
View File
@@ -3,7 +3,7 @@ mod cursor;
use arrayvec::ArrayVec;
pub use cursor::Cursor;
pub use cursor::FilterCursor;
use std::{fmt, iter::FromIterator, sync::Arc};
use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
#[cfg(test)]
const TREE_BASE: usize = 2;
@@ -23,17 +23,36 @@ pub trait KeyedItem: Item {
}
pub trait Summary: Default + Clone + fmt::Debug {
fn add_summary(&mut self, summary: &Self);
type Context;
fn add_summary(&mut self, _summary: &Self) {
unimplemented!();
}
fn add_summary_with_ctx(&mut self, summary: &Self, ctx: Option<&Self::Context>) {
assert!(ctx.is_none());
self.add_summary(summary);
}
}
pub trait Dimension<'a, Summary: Default>: Clone + fmt::Debug + Default {
fn add_summary(&mut self, summary: &'a Summary);
pub trait Dimension<'a, S: Summary>: Clone + fmt::Debug + Default {
fn add_summary(&mut self, _summary: &'a S);
}
impl<'a, T: Default> Dimension<'a, T> for () {
impl<'a, T: Summary> Dimension<'a, T> for () {
fn add_summary(&mut self, _: &'a T) {}
}
pub trait SeekDimension<'a, T: Summary>: Dimension<'a, T> {
fn cmp(&self, other: &Self, ctx: Option<&T::Context>) -> Ordering;
}
impl<'a, S: Summary, T: Dimension<'a, S> + Ord> SeekDimension<'a, S> for T {
fn cmp(&self, other: &Self, _ctx: Option<&S::Context>) -> Ordering {
Ord::cmp(self, other)
}
}
#[derive(Copy, Clone, Eq, PartialEq)]
pub enum SeekBias {
Left,
@@ -94,7 +113,7 @@ impl<T: Item> SumTree<T> {
let mut extent = D::default();
match self.0.as_ref() {
Node::Internal { summary, .. } | Node::Leaf { summary, .. } => {
extent.add_summary(summary)
extent.add_summary(summary);
}
}
extent
@@ -154,30 +173,50 @@ impl<T: Item> SumTree<T> {
}
pub fn push(&mut self, item: T) {
self.push_with_ctx(item, None);
}
pub fn push_with_ctx(&mut self, item: T, ctx: Option<&<T::Summary as Summary>::Context>) {
let summary = item.summary();
self.push_tree(SumTree::from_child_trees(vec![SumTree(Arc::new(
Node::Leaf {
summary: summary.clone(),
items: ArrayVec::from_iter(Some(item)),
item_summaries: ArrayVec::from_iter(Some(summary)),
},
))]))
self.push_tree_with_ctx(
SumTree::from_child_trees(
vec![SumTree(Arc::new(Node::Leaf {
summary: summary.clone(),
items: ArrayVec::from_iter(Some(item)),
item_summaries: ArrayVec::from_iter(Some(summary)),
}))],
ctx,
),
ctx,
)
}
pub fn push_tree(&mut self, other: Self) {
self.push_tree_with_ctx(other, None);
}
pub fn push_tree_with_ctx(
&mut self,
other: Self,
ctx: Option<&<T::Summary as Summary>::Context>,
) {
let other_node = other.0.clone();
if !other_node.is_leaf() || other_node.items().len() > 0 {
if self.0.height() < other_node.height() {
for tree in other_node.child_trees() {
self.push_tree(tree.clone());
self.push_tree_with_ctx(tree.clone(), ctx);
}
} else if let Some(split_tree) = self.push_tree_recursive(other) {
*self = Self::from_child_trees(vec![self.clone(), split_tree]);
} else if let Some(split_tree) = self.push_tree_recursive(other, ctx) {
*self = Self::from_child_trees(vec![self.clone(), split_tree], ctx);
}
}
}
fn push_tree_recursive(&mut self, other: SumTree<T>) -> Option<SumTree<T>> {
fn push_tree_recursive(
&mut self,
other: SumTree<T>,
ctx: Option<&<T::Summary as Summary>::Context>,
) -> Option<SumTree<T>> {
match Arc::make_mut(&mut self.0) {
Node::Internal {
height,
@@ -187,7 +226,7 @@ impl<T: Item> SumTree<T> {
..
} => {
let other_node = other.0.clone();
summary.add_summary(other_node.summary());
summary.add_summary_with_ctx(other_node.summary(), ctx);
let height_delta = *height - other_node.height();
let mut summaries_to_append = ArrayVec::<[T::Summary; 2 * TREE_BASE]>::new();
@@ -199,7 +238,10 @@ impl<T: Item> SumTree<T> {
summaries_to_append.push(other_node.summary().clone());
trees_to_append.push(other)
} else {
let tree_to_append = child_trees.last_mut().unwrap().push_tree_recursive(other);
let tree_to_append = child_trees
.last_mut()
.unwrap()
.push_tree_recursive(other, ctx);
*child_summaries.last_mut().unwrap() =
child_trees.last().unwrap().0.summary().clone();
@@ -229,13 +271,13 @@ impl<T: Item> SumTree<T> {
left_trees = all_trees.by_ref().take(midpoint).collect();
right_trees = all_trees.collect();
}
*summary = sum(left_summaries.iter());
*summary = sum(left_summaries.iter(), ctx);
*child_summaries = left_summaries;
*child_trees = left_trees;
Some(SumTree(Arc::new(Node::Internal {
height: *height,
summary: sum(right_summaries.iter()),
summary: sum(right_summaries.iter(), ctx),
child_summaries: right_summaries,
child_trees: right_trees,
})))
@@ -274,14 +316,14 @@ impl<T: Item> SumTree<T> {
}
*items = left_items;
*item_summaries = left_summaries;
*summary = sum(item_summaries.iter());
*summary = sum(item_summaries.iter(), ctx);
Some(SumTree(Arc::new(Node::Leaf {
items: right_items,
summary: sum(right_summaries.iter()),
summary: sum(right_summaries.iter(), ctx),
item_summaries: right_summaries,
})))
} else {
summary.add_summary(other_node.summary());
summary.add_summary_with_ctx(other_node.summary(), ctx);
items.extend(other_node.items().iter().cloned());
item_summaries.extend(other_node.child_summaries().iter().cloned());
None
@@ -290,13 +332,16 @@ impl<T: Item> SumTree<T> {
}
}
fn from_child_trees(child_trees: Vec<SumTree<T>>) -> Self {
fn from_child_trees(
child_trees: Vec<SumTree<T>>,
ctx: Option<&<T::Summary as Summary>::Context>,
) -> Self {
let height = child_trees[0].0.height() + 1;
let mut child_summaries = ArrayVec::new();
for child in &child_trees {
child_summaries.push(child.0.summary().clone());
}
let summary = sum(child_summaries.iter());
let summary = sum(child_summaries.iter(), ctx);
SumTree(Arc::new(Node::Internal {
height,
summary,
@@ -486,14 +531,14 @@ impl<T: KeyedItem> Edit<T> {
}
}
fn sum<'a, T, I>(iter: I) -> T
fn sum<'a, T, I>(iter: I, ctx: Option<&T::Context>) -> T
where
T: 'a + Summary,
I: Iterator<Item = &'a T>,
{
let mut sum = T::default();
for value in iter {
sum.add_summary(value);
sum.add_summary_with_ctx(value, ctx);
}
sum
}
@@ -840,7 +885,10 @@ mod tests {
*self
}
}
impl Summary for IntegersSummary {
type Context = ();
fn add_summary(&mut self, other: &Self) {
self.count.0 += &other.count.0;
self.sum.0 += &other.sum.0;