1mod cursor;
2
3use arrayvec::ArrayVec;
4pub use cursor::{Cursor, FilterCursor, Iter};
5use std::marker::PhantomData;
6use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
7
8#[cfg(test)]
9const TREE_BASE: usize = 2;
10#[cfg(not(test))]
11const TREE_BASE: usize = 6;
12
13pub trait Item: Clone {
14 type Summary: Summary;
15
16 fn summary(&self) -> Self::Summary;
17}
18
19pub trait KeyedItem: Item {
20 type Key: for<'a> Dimension<'a, Self::Summary> + Ord;
21
22 fn key(&self) -> Self::Key;
23}
24
25pub trait Summary: Default + Clone + fmt::Debug {
26 type Context;
27
28 fn add_summary(&mut self, summary: &Self, cx: &Self::Context);
29}
30
31pub trait Dimension<'a, S: Summary>: Clone + fmt::Debug + Default {
32 fn add_summary(&mut self, _summary: &'a S, _: &S::Context);
33
34 fn from_summary(summary: &'a S, cx: &S::Context) -> Self {
35 let mut dimension = Self::default();
36 dimension.add_summary(summary, cx);
37 dimension
38 }
39}
40
41impl<'a, T: Summary> Dimension<'a, T> for T {
42 fn add_summary(&mut self, summary: &'a T, cx: &T::Context) {
43 Summary::add_summary(self, summary, cx);
44 }
45}
46
47pub trait SeekTarget<'a, S: Summary, D: Dimension<'a, S>>: fmt::Debug {
48 fn cmp(&self, cursor_location: &D, cx: &S::Context) -> Ordering;
49}
50
51impl<'a, S: Summary, D: Dimension<'a, S> + Ord> SeekTarget<'a, S, D> for D {
52 fn cmp(&self, cursor_location: &Self, _: &S::Context) -> Ordering {
53 Ord::cmp(self, cursor_location)
54 }
55}
56
57impl<'a, T: Summary> Dimension<'a, T> for () {
58 fn add_summary(&mut self, _: &'a T, _: &T::Context) {}
59}
60
61impl<'a, T: Summary, D1: Dimension<'a, T>, D2: Dimension<'a, T>> Dimension<'a, T> for (D1, D2) {
62 fn add_summary(&mut self, summary: &'a T, cx: &T::Context) {
63 self.0.add_summary(summary, cx);
64 self.1.add_summary(summary, cx);
65 }
66}
67
68impl<'a, S: Summary, D1: SeekTarget<'a, S, D1> + Dimension<'a, S>, D2: Dimension<'a, S>>
69 SeekTarget<'a, S, (D1, D2)> for D1
70{
71 fn cmp(&self, cursor_location: &(D1, D2), cx: &S::Context) -> Ordering {
72 self.cmp(&cursor_location.0, cx)
73 }
74}
75
76struct End<D>(PhantomData<D>);
77
78impl<D> End<D> {
79 fn new() -> Self {
80 Self(PhantomData)
81 }
82}
83
84impl<'a, S: Summary, D: Dimension<'a, S>> SeekTarget<'a, S, D> for End<D> {
85 fn cmp(&self, _: &D, _: &S::Context) -> Ordering {
86 Ordering::Greater
87 }
88}
89
90impl<D> fmt::Debug for End<D> {
91 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92 f.debug_tuple("End").finish()
93 }
94}
95
96#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash)]
97pub enum Bias {
98 Left,
99 Right,
100}
101
102impl PartialOrd for Bias {
103 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
104 Some(self.cmp(other))
105 }
106}
107
108impl Ord for Bias {
109 fn cmp(&self, other: &Self) -> Ordering {
110 match (self, other) {
111 (Self::Left, Self::Left) => Ordering::Equal,
112 (Self::Left, Self::Right) => Ordering::Less,
113 (Self::Right, Self::Right) => Ordering::Equal,
114 (Self::Right, Self::Left) => Ordering::Greater,
115 }
116 }
117}
118
119#[derive(Debug, Clone)]
120pub struct SumTree<T: Item>(Arc<Node<T>>);
121
122impl<T: Item> SumTree<T> {
123 pub fn new() -> Self {
124 SumTree(Arc::new(Node::Leaf {
125 summary: T::Summary::default(),
126 items: ArrayVec::new(),
127 item_summaries: ArrayVec::new(),
128 }))
129 }
130
131 pub fn from_item(item: T, cx: &<T::Summary as Summary>::Context) -> Self {
132 let mut tree = Self::new();
133 tree.push(item, cx);
134 tree
135 }
136
137 pub fn from_iter<I: IntoIterator<Item = T>>(
138 iter: I,
139 cx: &<T::Summary as Summary>::Context,
140 ) -> Self {
141 let mut tree = Self::new();
142 tree.extend(iter, cx);
143 tree
144 }
145
146 #[allow(unused)]
147 pub fn items(&self, cx: &<T::Summary as Summary>::Context) -> Vec<T> {
148 let mut items = Vec::new();
149 let mut cursor = self.cursor::<()>();
150 cursor.next(cx);
151 while let Some(item) = cursor.item() {
152 items.push(item.clone());
153 cursor.next(cx);
154 }
155 items
156 }
157
158 pub fn iter(&self) -> Iter<T> {
159 Iter::new(self)
160 }
161
162 pub fn cursor<'a, S>(&'a self) -> Cursor<T, S>
163 where
164 S: Dimension<'a, T::Summary>,
165 {
166 Cursor::new(self)
167 }
168
169 pub fn filter<'a, F, U>(
170 &'a self,
171 filter_node: F,
172 cx: &<T::Summary as Summary>::Context,
173 ) -> FilterCursor<F, T, U>
174 where
175 F: FnMut(&T::Summary) -> bool,
176 U: Dimension<'a, T::Summary>,
177 {
178 FilterCursor::new(self, filter_node, cx)
179 }
180
181 #[allow(dead_code)]
182 pub fn first(&self) -> Option<&T> {
183 self.leftmost_leaf().0.items().first()
184 }
185
186 pub fn last(&self) -> Option<&T> {
187 self.rightmost_leaf().0.items().last()
188 }
189
190 pub fn update_last(&mut self, f: impl FnOnce(&mut T), cx: &<T::Summary as Summary>::Context) {
191 self.update_last_recursive(f, cx);
192 }
193
194 fn update_last_recursive(
195 &mut self,
196 f: impl FnOnce(&mut T),
197 cx: &<T::Summary as Summary>::Context,
198 ) -> Option<T::Summary> {
199 match Arc::make_mut(&mut self.0) {
200 Node::Internal {
201 summary,
202 child_summaries,
203 child_trees,
204 ..
205 } => {
206 let last_summary = child_summaries.last_mut().unwrap();
207 let last_child = child_trees.last_mut().unwrap();
208 *last_summary = last_child.update_last_recursive(f, cx).unwrap();
209 *summary = sum(child_summaries.iter(), cx);
210 Some(summary.clone())
211 }
212 Node::Leaf {
213 summary,
214 items,
215 item_summaries,
216 } => {
217 if let Some((item, item_summary)) = items.last_mut().zip(item_summaries.last_mut())
218 {
219 (f)(item);
220 *item_summary = item.summary();
221 *summary = sum(item_summaries.iter(), cx);
222 Some(summary.clone())
223 } else {
224 None
225 }
226 }
227 }
228 }
229
230 pub fn extent<'a, D: Dimension<'a, T::Summary>>(
231 &'a self,
232 cx: &<T::Summary as Summary>::Context,
233 ) -> D {
234 let mut extent = D::default();
235 match self.0.as_ref() {
236 Node::Internal { summary, .. } | Node::Leaf { summary, .. } => {
237 extent.add_summary(summary, cx);
238 }
239 }
240 extent
241 }
242
243 pub fn summary(&self) -> T::Summary {
244 match self.0.as_ref() {
245 Node::Internal { summary, .. } => summary.clone(),
246 Node::Leaf { summary, .. } => summary.clone(),
247 }
248 }
249
250 pub fn is_empty(&self) -> bool {
251 match self.0.as_ref() {
252 Node::Internal { .. } => false,
253 Node::Leaf { items, .. } => items.is_empty(),
254 }
255 }
256
257 pub fn extend<I>(&mut self, iter: I, cx: &<T::Summary as Summary>::Context)
258 where
259 I: IntoIterator<Item = T>,
260 {
261 let mut leaf: Option<Node<T>> = None;
262
263 for item in iter {
264 if leaf.is_some() && leaf.as_ref().unwrap().items().len() == 2 * TREE_BASE {
265 self.push_tree(SumTree(Arc::new(leaf.take().unwrap())), cx);
266 }
267
268 if leaf.is_none() {
269 leaf = Some(Node::Leaf::<T> {
270 summary: T::Summary::default(),
271 items: ArrayVec::new(),
272 item_summaries: ArrayVec::new(),
273 });
274 }
275
276 if let Some(Node::Leaf {
277 summary,
278 items,
279 item_summaries,
280 }) = leaf.as_mut()
281 {
282 let item_summary = item.summary();
283 <T::Summary as Summary>::add_summary(summary, &item_summary, cx);
284 items.push(item);
285 item_summaries.push(item_summary);
286 } else {
287 unreachable!()
288 }
289 }
290
291 if leaf.is_some() {
292 self.push_tree(SumTree(Arc::new(leaf.take().unwrap())), cx);
293 }
294 }
295
296 pub fn push(&mut self, item: T, cx: &<T::Summary as Summary>::Context) {
297 let summary = item.summary();
298 self.push_tree(
299 SumTree(Arc::new(Node::Leaf {
300 summary: summary.clone(),
301 items: ArrayVec::from_iter(Some(item)),
302 item_summaries: ArrayVec::from_iter(Some(summary)),
303 })),
304 cx,
305 );
306 }
307
308 pub fn push_tree(&mut self, other: Self, cx: &<T::Summary as Summary>::Context) {
309 if !other.0.is_leaf() || other.0.items().len() > 0 {
310 if self.0.height() < other.0.height() {
311 for tree in other.0.child_trees() {
312 self.push_tree(tree.clone(), cx);
313 }
314 } else if let Some(split_tree) = self.push_tree_recursive(other, cx) {
315 *self = Self::from_child_trees(self.clone(), split_tree, cx);
316 }
317 }
318 }
319
320 fn push_tree_recursive(
321 &mut self,
322 other: SumTree<T>,
323 cx: &<T::Summary as Summary>::Context,
324 ) -> Option<SumTree<T>> {
325 match Arc::make_mut(&mut self.0) {
326 Node::Internal {
327 height,
328 summary,
329 child_summaries,
330 child_trees,
331 ..
332 } => {
333 let other_node = other.0.clone();
334 <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
335
336 let height_delta = *height - other_node.height();
337 let mut summaries_to_append = ArrayVec::<T::Summary, { 2 * TREE_BASE }>::new();
338 let mut trees_to_append = ArrayVec::<SumTree<T>, { 2 * TREE_BASE }>::new();
339 if height_delta == 0 {
340 summaries_to_append.extend(other_node.child_summaries().iter().cloned());
341 trees_to_append.extend(other_node.child_trees().iter().cloned());
342 } else if height_delta == 1 && !other_node.is_underflowing() {
343 summaries_to_append.push(other_node.summary().clone());
344 trees_to_append.push(other)
345 } else {
346 let tree_to_append = child_trees
347 .last_mut()
348 .unwrap()
349 .push_tree_recursive(other, cx);
350 *child_summaries.last_mut().unwrap() =
351 child_trees.last().unwrap().0.summary().clone();
352
353 if let Some(split_tree) = tree_to_append {
354 summaries_to_append.push(split_tree.0.summary().clone());
355 trees_to_append.push(split_tree);
356 }
357 }
358
359 let child_count = child_trees.len() + trees_to_append.len();
360 if child_count > 2 * TREE_BASE {
361 let left_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
362 let right_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
363 let left_trees;
364 let right_trees;
365
366 let midpoint = (child_count + child_count % 2) / 2;
367 {
368 let mut all_summaries = child_summaries
369 .iter()
370 .chain(summaries_to_append.iter())
371 .cloned();
372 left_summaries = all_summaries.by_ref().take(midpoint).collect();
373 right_summaries = all_summaries.collect();
374 let mut all_trees =
375 child_trees.iter().chain(trees_to_append.iter()).cloned();
376 left_trees = all_trees.by_ref().take(midpoint).collect();
377 right_trees = all_trees.collect();
378 }
379 *summary = sum(left_summaries.iter(), cx);
380 *child_summaries = left_summaries;
381 *child_trees = left_trees;
382
383 Some(SumTree(Arc::new(Node::Internal {
384 height: *height,
385 summary: sum(right_summaries.iter(), cx),
386 child_summaries: right_summaries,
387 child_trees: right_trees,
388 })))
389 } else {
390 child_summaries.extend(summaries_to_append);
391 child_trees.extend(trees_to_append);
392 None
393 }
394 }
395 Node::Leaf {
396 summary,
397 items,
398 item_summaries,
399 } => {
400 let other_node = other.0;
401
402 let child_count = items.len() + other_node.items().len();
403 if child_count > 2 * TREE_BASE {
404 let left_items;
405 let right_items;
406 let left_summaries;
407 let right_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>;
408
409 let midpoint = (child_count + child_count % 2) / 2;
410 {
411 let mut all_items = items.iter().chain(other_node.items().iter()).cloned();
412 left_items = all_items.by_ref().take(midpoint).collect();
413 right_items = all_items.collect();
414
415 let mut all_summaries = item_summaries
416 .iter()
417 .chain(other_node.child_summaries())
418 .cloned();
419 left_summaries = all_summaries.by_ref().take(midpoint).collect();
420 right_summaries = all_summaries.collect();
421 }
422 *items = left_items;
423 *item_summaries = left_summaries;
424 *summary = sum(item_summaries.iter(), cx);
425 Some(SumTree(Arc::new(Node::Leaf {
426 items: right_items,
427 summary: sum(right_summaries.iter(), cx),
428 item_summaries: right_summaries,
429 })))
430 } else {
431 <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
432 items.extend(other_node.items().iter().cloned());
433 item_summaries.extend(other_node.child_summaries().iter().cloned());
434 None
435 }
436 }
437 }
438 }
439
440 fn from_child_trees(
441 left: SumTree<T>,
442 right: SumTree<T>,
443 cx: &<T::Summary as Summary>::Context,
444 ) -> Self {
445 let height = left.0.height() + 1;
446 let mut child_summaries = ArrayVec::new();
447 child_summaries.push(left.0.summary().clone());
448 child_summaries.push(right.0.summary().clone());
449 let mut child_trees = ArrayVec::new();
450 child_trees.push(left);
451 child_trees.push(right);
452 SumTree(Arc::new(Node::Internal {
453 height,
454 summary: sum(child_summaries.iter(), cx),
455 child_summaries,
456 child_trees,
457 }))
458 }
459
460 fn leftmost_leaf(&self) -> &Self {
461 match *self.0 {
462 Node::Leaf { .. } => self,
463 Node::Internal {
464 ref child_trees, ..
465 } => child_trees.first().unwrap().leftmost_leaf(),
466 }
467 }
468
469 fn rightmost_leaf(&self) -> &Self {
470 match *self.0 {
471 Node::Leaf { .. } => self,
472 Node::Internal {
473 ref child_trees, ..
474 } => child_trees.last().unwrap().rightmost_leaf(),
475 }
476 }
477}
478
479impl<T: KeyedItem> SumTree<T> {
480 pub fn insert_or_replace(&mut self, item: T, cx: &<T::Summary as Summary>::Context) -> bool {
481 let mut replaced = false;
482 *self = {
483 let mut cursor = self.cursor::<T::Key>();
484 let mut new_tree = cursor.slice(&item.key(), Bias::Left, cx);
485 if cursor
486 .item()
487 .map_or(false, |cursor_item| cursor_item.key() == item.key())
488 {
489 cursor.next(cx);
490 replaced = true;
491 }
492 new_tree.push(item, cx);
493 new_tree.push_tree(cursor.suffix(cx), cx);
494 new_tree
495 };
496 replaced
497 }
498
499 pub fn edit(
500 &mut self,
501 mut edits: Vec<Edit<T>>,
502 cx: &<T::Summary as Summary>::Context,
503 ) -> Vec<T> {
504 if edits.is_empty() {
505 return Vec::new();
506 }
507
508 let mut removed = Vec::new();
509 edits.sort_unstable_by_key(|item| item.key());
510
511 *self = {
512 let mut cursor = self.cursor::<T::Key>();
513 let mut new_tree = SumTree::new();
514 let mut buffered_items = Vec::new();
515
516 cursor.seek(&T::Key::default(), Bias::Left, cx);
517 for edit in edits {
518 let new_key = edit.key();
519 let mut old_item = cursor.item();
520
521 if old_item
522 .as_ref()
523 .map_or(false, |old_item| old_item.key() < new_key)
524 {
525 new_tree.extend(buffered_items.drain(..), cx);
526 let slice = cursor.slice(&new_key, Bias::Left, cx);
527 new_tree.push_tree(slice, cx);
528 old_item = cursor.item();
529 }
530
531 if let Some(old_item) = old_item {
532 if old_item.key() == new_key {
533 removed.push(old_item.clone());
534 cursor.next(cx);
535 }
536 }
537
538 match edit {
539 Edit::Insert(item) => {
540 buffered_items.push(item);
541 }
542 Edit::Remove(_) => {}
543 }
544 }
545
546 new_tree.extend(buffered_items, cx);
547 new_tree.push_tree(cursor.suffix(cx), cx);
548 new_tree
549 };
550
551 removed
552 }
553
554 pub fn get(&self, key: &T::Key, cx: &<T::Summary as Summary>::Context) -> Option<&T> {
555 let mut cursor = self.cursor::<T::Key>();
556 if cursor.seek(key, Bias::Left, cx) {
557 cursor.item()
558 } else {
559 None
560 }
561 }
562}
563
564impl<T: Item> Default for SumTree<T> {
565 fn default() -> Self {
566 Self::new()
567 }
568}
569
570#[derive(Clone, Debug)]
571pub enum Node<T: Item> {
572 Internal {
573 height: u8,
574 summary: T::Summary,
575 child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
576 child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }>,
577 },
578 Leaf {
579 summary: T::Summary,
580 items: ArrayVec<T, { 2 * TREE_BASE }>,
581 item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
582 },
583}
584
585impl<T: Item> Node<T> {
586 fn is_leaf(&self) -> bool {
587 match self {
588 Node::Leaf { .. } => true,
589 _ => false,
590 }
591 }
592
593 fn height(&self) -> u8 {
594 match self {
595 Node::Internal { height, .. } => *height,
596 Node::Leaf { .. } => 0,
597 }
598 }
599
600 fn summary(&self) -> &T::Summary {
601 match self {
602 Node::Internal { summary, .. } => summary,
603 Node::Leaf { summary, .. } => summary,
604 }
605 }
606
607 fn child_summaries(&self) -> &[T::Summary] {
608 match self {
609 Node::Internal {
610 child_summaries, ..
611 } => child_summaries.as_slice(),
612 Node::Leaf { item_summaries, .. } => item_summaries.as_slice(),
613 }
614 }
615
616 fn child_trees(&self) -> &ArrayVec<SumTree<T>, { 2 * TREE_BASE }> {
617 match self {
618 Node::Internal { child_trees, .. } => child_trees,
619 Node::Leaf { .. } => panic!("Leaf nodes have no child trees"),
620 }
621 }
622
623 fn items(&self) -> &ArrayVec<T, { 2 * TREE_BASE }> {
624 match self {
625 Node::Leaf { items, .. } => items,
626 Node::Internal { .. } => panic!("Internal nodes have no items"),
627 }
628 }
629
630 fn is_underflowing(&self) -> bool {
631 match self {
632 Node::Internal { child_trees, .. } => child_trees.len() < TREE_BASE,
633 Node::Leaf { items, .. } => items.len() < TREE_BASE,
634 }
635 }
636}
637
638#[derive(Debug)]
639pub enum Edit<T: KeyedItem> {
640 Insert(T),
641 Remove(T::Key),
642}
643
644impl<T: KeyedItem> Edit<T> {
645 fn key(&self) -> T::Key {
646 match self {
647 Edit::Insert(item) => item.key(),
648 Edit::Remove(key) => key.clone(),
649 }
650 }
651}
652
653fn sum<'a, T, I>(iter: I, cx: &T::Context) -> T
654where
655 T: 'a + Summary,
656 I: Iterator<Item = &'a T>,
657{
658 let mut sum = T::default();
659 for value in iter {
660 sum.add_summary(value, cx);
661 }
662 sum
663}
664
665#[cfg(test)]
666mod tests {
667 use super::*;
668 use rand::{distributions, prelude::*};
669 use std::cmp;
670
671 #[test]
672 fn test_extend_and_push_tree() {
673 let mut tree1 = SumTree::new();
674 tree1.extend(0..20, &());
675
676 let mut tree2 = SumTree::new();
677 tree2.extend(50..100, &());
678
679 tree1.push_tree(tree2, &());
680 assert_eq!(
681 tree1.items(&()),
682 (0..20).chain(50..100).collect::<Vec<u8>>()
683 );
684 }
685
686 #[test]
687 fn test_random() {
688 let mut starting_seed = 0;
689 if let Ok(value) = std::env::var("SEED") {
690 starting_seed = value.parse().expect("invalid SEED variable");
691 }
692 let mut num_iterations = 100;
693 if let Ok(value) = std::env::var("ITERATIONS") {
694 num_iterations = value.parse().expect("invalid ITERATIONS variable");
695 }
696
697 for seed in starting_seed..(starting_seed + num_iterations) {
698 let mut rng = StdRng::seed_from_u64(seed);
699
700 let rng = &mut rng;
701 let mut tree = SumTree::<u8>::new();
702 let count = rng.gen_range(0..10);
703 tree.extend(rng.sample_iter(distributions::Standard).take(count), &());
704
705 for _ in 0..5 {
706 let splice_end = rng.gen_range(0..tree.extent::<Count>(&()).0 + 1);
707 let splice_start = rng.gen_range(0..splice_end + 1);
708 let count = rng.gen_range(0..3);
709 let tree_end = tree.extent::<Count>(&());
710 let new_items = rng
711 .sample_iter(distributions::Standard)
712 .take(count)
713 .collect::<Vec<u8>>();
714
715 let mut reference_items = tree.items(&());
716 reference_items.splice(splice_start..splice_end, new_items.clone());
717
718 tree = {
719 let mut cursor = tree.cursor::<Count>();
720 let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right, &());
721 new_tree.extend(new_items, &());
722 cursor.seek(&Count(splice_end), Bias::Right, &());
723 new_tree.push_tree(cursor.slice(&tree_end, Bias::Right, &()), &());
724 new_tree
725 };
726
727 assert_eq!(tree.items(&()), reference_items);
728 assert_eq!(
729 tree.iter().collect::<Vec<_>>(),
730 tree.cursor::<()>().collect::<Vec<_>>()
731 );
732
733 let mut filter_cursor =
734 tree.filter::<_, Count>(|summary| summary.contains_even, &());
735 let mut reference_filter = tree
736 .items(&())
737 .into_iter()
738 .enumerate()
739 .filter(|(_, item)| (item & 1) == 0);
740 while let Some(actual_item) = filter_cursor.item() {
741 let (reference_index, reference_item) = reference_filter.next().unwrap();
742 assert_eq!(actual_item, &reference_item);
743 assert_eq!(filter_cursor.start().0, reference_index);
744 filter_cursor.next(&());
745 }
746 assert!(reference_filter.next().is_none());
747
748 let mut pos = rng.gen_range(0..tree.extent::<Count>(&()).0 + 1);
749 let mut before_start = false;
750 let mut cursor = tree.cursor::<Count>();
751 cursor.seek(&Count(pos), Bias::Right, &());
752
753 for i in 0..10 {
754 assert_eq!(cursor.start().0, pos);
755
756 if pos > 0 {
757 assert_eq!(cursor.prev_item().unwrap(), &reference_items[pos - 1]);
758 } else {
759 assert_eq!(cursor.prev_item(), None);
760 }
761
762 if pos < reference_items.len() && !before_start {
763 assert_eq!(cursor.item().unwrap(), &reference_items[pos]);
764 } else {
765 assert_eq!(cursor.item(), None);
766 }
767
768 if i < 5 {
769 cursor.next(&());
770 if pos < reference_items.len() {
771 pos += 1;
772 before_start = false;
773 }
774 } else {
775 cursor.prev(&());
776 if pos == 0 {
777 before_start = true;
778 }
779 pos = pos.saturating_sub(1);
780 }
781 }
782 }
783
784 for _ in 0..10 {
785 let end = rng.gen_range(0..tree.extent::<Count>(&()).0 + 1);
786 let start = rng.gen_range(0..end + 1);
787 let start_bias = if rng.gen() { Bias::Left } else { Bias::Right };
788 let end_bias = if rng.gen() { Bias::Left } else { Bias::Right };
789
790 let mut cursor = tree.cursor::<Count>();
791 cursor.seek(&Count(start), start_bias, &());
792 let slice = cursor.slice(&Count(end), end_bias, &());
793
794 cursor.seek(&Count(start), start_bias, &());
795 let summary = cursor.summary::<_, Sum>(&Count(end), end_bias, &());
796
797 assert_eq!(summary.0, slice.summary().sum);
798 }
799 }
800 }
801
802 #[test]
803 fn test_cursor() {
804 // Empty tree
805 let tree = SumTree::<u8>::new();
806 let mut cursor = tree.cursor::<IntegersSummary>();
807 assert_eq!(
808 cursor.slice(&Count(0), Bias::Right, &()).items(&()),
809 Vec::<u8>::new()
810 );
811 assert_eq!(cursor.item(), None);
812 assert_eq!(cursor.prev_item(), None);
813 assert_eq!(cursor.start().sum, 0);
814
815 // Single-element tree
816 let mut tree = SumTree::<u8>::new();
817 tree.extend(vec![1], &());
818 let mut cursor = tree.cursor::<IntegersSummary>();
819 assert_eq!(
820 cursor.slice(&Count(0), Bias::Right, &()).items(&()),
821 Vec::<u8>::new()
822 );
823 assert_eq!(cursor.item(), Some(&1));
824 assert_eq!(cursor.prev_item(), None);
825 assert_eq!(cursor.start().sum, 0);
826
827 cursor.next(&());
828 assert_eq!(cursor.item(), None);
829 assert_eq!(cursor.prev_item(), Some(&1));
830 assert_eq!(cursor.start().sum, 1);
831
832 cursor.prev(&());
833 assert_eq!(cursor.item(), Some(&1));
834 assert_eq!(cursor.prev_item(), None);
835 assert_eq!(cursor.start().sum, 0);
836
837 let mut cursor = tree.cursor::<IntegersSummary>();
838 assert_eq!(cursor.slice(&Count(1), Bias::Right, &()).items(&()), [1]);
839 assert_eq!(cursor.item(), None);
840 assert_eq!(cursor.prev_item(), Some(&1));
841 assert_eq!(cursor.start().sum, 1);
842
843 cursor.seek(&Count(0), Bias::Right, &());
844 assert_eq!(
845 cursor
846 .slice(&tree.extent::<Count>(&()), Bias::Right, &())
847 .items(&()),
848 [1]
849 );
850 assert_eq!(cursor.item(), None);
851 assert_eq!(cursor.prev_item(), Some(&1));
852 assert_eq!(cursor.start().sum, 1);
853
854 // Multiple-element tree
855 let mut tree = SumTree::new();
856 tree.extend(vec![1, 2, 3, 4, 5, 6], &());
857 let mut cursor = tree.cursor::<IntegersSummary>();
858
859 assert_eq!(cursor.slice(&Count(2), Bias::Right, &()).items(&()), [1, 2]);
860 assert_eq!(cursor.item(), Some(&3));
861 assert_eq!(cursor.prev_item(), Some(&2));
862 assert_eq!(cursor.start().sum, 3);
863
864 cursor.next(&());
865 assert_eq!(cursor.item(), Some(&4));
866 assert_eq!(cursor.prev_item(), Some(&3));
867 assert_eq!(cursor.start().sum, 6);
868
869 cursor.next(&());
870 assert_eq!(cursor.item(), Some(&5));
871 assert_eq!(cursor.prev_item(), Some(&4));
872 assert_eq!(cursor.start().sum, 10);
873
874 cursor.next(&());
875 assert_eq!(cursor.item(), Some(&6));
876 assert_eq!(cursor.prev_item(), Some(&5));
877 assert_eq!(cursor.start().sum, 15);
878
879 cursor.next(&());
880 cursor.next(&());
881 assert_eq!(cursor.item(), None);
882 assert_eq!(cursor.prev_item(), Some(&6));
883 assert_eq!(cursor.start().sum, 21);
884
885 cursor.prev(&());
886 assert_eq!(cursor.item(), Some(&6));
887 assert_eq!(cursor.prev_item(), Some(&5));
888 assert_eq!(cursor.start().sum, 15);
889
890 cursor.prev(&());
891 assert_eq!(cursor.item(), Some(&5));
892 assert_eq!(cursor.prev_item(), Some(&4));
893 assert_eq!(cursor.start().sum, 10);
894
895 cursor.prev(&());
896 assert_eq!(cursor.item(), Some(&4));
897 assert_eq!(cursor.prev_item(), Some(&3));
898 assert_eq!(cursor.start().sum, 6);
899
900 cursor.prev(&());
901 assert_eq!(cursor.item(), Some(&3));
902 assert_eq!(cursor.prev_item(), Some(&2));
903 assert_eq!(cursor.start().sum, 3);
904
905 cursor.prev(&());
906 assert_eq!(cursor.item(), Some(&2));
907 assert_eq!(cursor.prev_item(), Some(&1));
908 assert_eq!(cursor.start().sum, 1);
909
910 cursor.prev(&());
911 assert_eq!(cursor.item(), Some(&1));
912 assert_eq!(cursor.prev_item(), None);
913 assert_eq!(cursor.start().sum, 0);
914
915 cursor.prev(&());
916 assert_eq!(cursor.item(), None);
917 assert_eq!(cursor.prev_item(), None);
918 assert_eq!(cursor.start().sum, 0);
919
920 cursor.next(&());
921 assert_eq!(cursor.item(), Some(&1));
922 assert_eq!(cursor.prev_item(), None);
923 assert_eq!(cursor.start().sum, 0);
924
925 let mut cursor = tree.cursor::<IntegersSummary>();
926 assert_eq!(
927 cursor
928 .slice(&tree.extent::<Count>(&()), Bias::Right, &())
929 .items(&()),
930 tree.items(&())
931 );
932 assert_eq!(cursor.item(), None);
933 assert_eq!(cursor.prev_item(), Some(&6));
934 assert_eq!(cursor.start().sum, 21);
935
936 cursor.seek(&Count(3), Bias::Right, &());
937 assert_eq!(
938 cursor
939 .slice(&tree.extent::<Count>(&()), Bias::Right, &())
940 .items(&()),
941 [4, 5, 6]
942 );
943 assert_eq!(cursor.item(), None);
944 assert_eq!(cursor.prev_item(), Some(&6));
945 assert_eq!(cursor.start().sum, 21);
946
947 // Seeking can bias left or right
948 cursor.seek(&Count(1), Bias::Left, &());
949 assert_eq!(cursor.item(), Some(&1));
950 cursor.seek(&Count(1), Bias::Right, &());
951 assert_eq!(cursor.item(), Some(&2));
952
953 // Slicing without resetting starts from where the cursor is parked at.
954 cursor.seek(&Count(1), Bias::Right, &());
955 assert_eq!(
956 cursor.slice(&Count(3), Bias::Right, &()).items(&()),
957 vec![2, 3]
958 );
959 assert_eq!(
960 cursor.slice(&Count(6), Bias::Left, &()).items(&()),
961 vec![4, 5]
962 );
963 assert_eq!(
964 cursor.slice(&Count(6), Bias::Right, &()).items(&()),
965 vec![6]
966 );
967 }
968
969 #[test]
970 fn test_edit() {
971 let mut tree = SumTree::<u8>::new();
972
973 let removed = tree.edit(vec![Edit::Insert(1), Edit::Insert(2), Edit::Insert(0)], &());
974 assert_eq!(tree.items(&()), vec![0, 1, 2]);
975 assert_eq!(removed, Vec::<u8>::new());
976 assert_eq!(tree.get(&0, &()), Some(&0));
977 assert_eq!(tree.get(&1, &()), Some(&1));
978 assert_eq!(tree.get(&2, &()), Some(&2));
979 assert_eq!(tree.get(&4, &()), None);
980
981 let removed = tree.edit(vec![Edit::Insert(2), Edit::Insert(4), Edit::Remove(0)], &());
982 assert_eq!(tree.items(&()), vec![1, 2, 4]);
983 assert_eq!(removed, vec![0, 2]);
984 assert_eq!(tree.get(&0, &()), None);
985 assert_eq!(tree.get(&1, &()), Some(&1));
986 assert_eq!(tree.get(&2, &()), Some(&2));
987 assert_eq!(tree.get(&4, &()), Some(&4));
988 }
989
990 #[derive(Clone, Default, Debug)]
991 pub struct IntegersSummary {
992 count: usize,
993 sum: usize,
994 contains_even: bool,
995 max: u8,
996 }
997
998 #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
999 struct Count(usize);
1000
1001 #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1002 struct Sum(usize);
1003
1004 impl Item for u8 {
1005 type Summary = IntegersSummary;
1006
1007 fn summary(&self) -> Self::Summary {
1008 IntegersSummary {
1009 count: 1,
1010 sum: *self as usize,
1011 contains_even: (*self & 1) == 0,
1012 max: *self,
1013 }
1014 }
1015 }
1016
1017 impl KeyedItem for u8 {
1018 type Key = u8;
1019
1020 fn key(&self) -> Self::Key {
1021 *self
1022 }
1023 }
1024
1025 impl Summary for IntegersSummary {
1026 type Context = ();
1027
1028 fn add_summary(&mut self, other: &Self, _: &()) {
1029 self.count += other.count;
1030 self.sum += other.sum;
1031 self.contains_even |= other.contains_even;
1032 self.max = cmp::max(self.max, other.max);
1033 }
1034 }
1035
1036 impl<'a> Dimension<'a, IntegersSummary> for u8 {
1037 fn add_summary(&mut self, summary: &IntegersSummary, _: &()) {
1038 *self = summary.max;
1039 }
1040 }
1041
1042 impl<'a> Dimension<'a, IntegersSummary> for Count {
1043 fn add_summary(&mut self, summary: &IntegersSummary, _: &()) {
1044 self.0 += summary.count;
1045 }
1046 }
1047
1048 impl<'a> SeekTarget<'a, IntegersSummary, IntegersSummary> for Count {
1049 fn cmp(&self, cursor_location: &IntegersSummary, _: &()) -> Ordering {
1050 self.0.cmp(&cursor_location.count)
1051 }
1052 }
1053
1054 impl<'a> Dimension<'a, IntegersSummary> for Sum {
1055 fn add_summary(&mut self, summary: &IntegersSummary, _: &()) {
1056 self.0 += summary.sum;
1057 }
1058 }
1059}