1mod cursor;
2mod tree_map;
3
4use arrayvec::ArrayVec;
5pub use cursor::{Cursor, FilterCursor, Iter};
6use futures::{StreamExt, stream};
7use itertools::Itertools as _;
8use std::marker::PhantomData;
9use std::mem;
10use std::{cmp::Ordering, fmt, iter::FromIterator, sync::Arc};
11pub use tree_map::{MapSeekTarget, TreeMap, TreeSet};
12
13#[cfg(test)]
14pub const TREE_BASE: usize = 2;
15#[cfg(not(test))]
16pub const TREE_BASE: usize = 6;
17
18pub trait BackgroundSpawn {
19 type Task<R>: Future<Output = R> + Send + Sync
20 where
21 R: Send + Sync;
22 fn background_spawn<R>(
23 &self,
24 future: impl Future<Output = R> + Send + Sync + 'static,
25 ) -> Self::Task<R>
26 where
27 R: Send + Sync + 'static;
28}
29
30/// An item that can be stored in a [`SumTree`]
31///
32/// Must be summarized by a type that implements [`Summary`]
33pub trait Item: Clone {
34 type Summary: Summary;
35
36 fn summary(&self, cx: <Self::Summary as Summary>::Context<'_>) -> Self::Summary;
37}
38
39/// An [`Item`] whose summary has a specific key that can be used to identify it
40pub trait KeyedItem: Item {
41 type Key: for<'a> Dimension<'a, Self::Summary> + Ord;
42
43 fn key(&self) -> Self::Key;
44}
45
46/// A type that describes the Sum of all [`Item`]s in a subtree of the [`SumTree`]
47///
48/// Each Summary type can have multiple [`Dimension`]s that it measures,
49/// which can be used to navigate the tree
50pub trait Summary: Clone {
51 type Context<'a>: Copy;
52 fn zero<'a>(cx: Self::Context<'a>) -> Self;
53 fn add_summary<'a>(&mut self, summary: &Self, cx: Self::Context<'a>);
54}
55
56pub trait ContextLessSummary: Clone {
57 fn zero() -> Self;
58 fn add_summary(&mut self, summary: &Self);
59}
60
61impl<T: ContextLessSummary> Summary for T {
62 type Context<'a> = ();
63
64 fn zero<'a>((): ()) -> Self {
65 T::zero()
66 }
67
68 fn add_summary<'a>(&mut self, summary: &Self, (): ()) {
69 T::add_summary(self, summary)
70 }
71}
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub struct NoSummary;
75
76/// Catch-all implementation for when you need something that implements [`Summary`] without a specific type.
77/// We implement it on a `NoSummary` instead of re-using `()`, as that avoids blanket impl collisions with `impl<T: Summary> Dimension for T`
78/// (as we also need unit type to be a fill-in dimension)
79impl ContextLessSummary for NoSummary {
80 fn zero() -> Self {
81 NoSummary
82 }
83
84 fn add_summary(&mut self, _: &Self) {}
85}
86
87/// Each [`Summary`] type can have more than one [`Dimension`] type that it measures.
88///
89/// You can use dimensions to seek to a specific location in the [`SumTree`]
90///
91/// # Example:
92/// Zed's rope has a `TextSummary` type that summarizes lines, characters, and bytes.
93/// Each of these are different dimensions we may want to seek to
94pub trait Dimension<'a, S: Summary>: Clone {
95 fn zero(cx: S::Context<'_>) -> Self;
96
97 fn add_summary(&mut self, summary: &'a S, cx: S::Context<'_>);
98 #[must_use]
99 fn with_added_summary(mut self, summary: &'a S, cx: S::Context<'_>) -> Self {
100 self.add_summary(summary, cx);
101 self
102 }
103
104 fn from_summary(summary: &'a S, cx: S::Context<'_>) -> Self {
105 let mut dimension = Self::zero(cx);
106 dimension.add_summary(summary, cx);
107 dimension
108 }
109}
110
111impl<'a, T: Summary> Dimension<'a, T> for T {
112 fn zero(cx: T::Context<'_>) -> Self {
113 Summary::zero(cx)
114 }
115
116 fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
117 Summary::add_summary(self, summary, cx);
118 }
119}
120
121pub trait SeekTarget<'a, S: Summary, D: Dimension<'a, S>> {
122 fn cmp(&self, cursor_location: &D, cx: S::Context<'_>) -> Ordering;
123}
124
125impl<'a, S: Summary, D: Dimension<'a, S> + Ord> SeekTarget<'a, S, D> for D {
126 fn cmp(&self, cursor_location: &Self, _: S::Context<'_>) -> Ordering {
127 Ord::cmp(self, cursor_location)
128 }
129}
130
131impl<'a, T: Summary> Dimension<'a, T> for () {
132 fn zero(_: T::Context<'_>) -> Self {}
133
134 fn add_summary(&mut self, _: &'a T, _: T::Context<'_>) {}
135}
136
137#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
138pub struct Dimensions<D1, D2, D3 = ()>(pub D1, pub D2, pub D3);
139
140impl<'a, T: Summary, D1: Dimension<'a, T>, D2: Dimension<'a, T>, D3: Dimension<'a, T>>
141 Dimension<'a, T> for Dimensions<D1, D2, D3>
142{
143 fn zero(cx: T::Context<'_>) -> Self {
144 Dimensions(D1::zero(cx), D2::zero(cx), D3::zero(cx))
145 }
146
147 fn add_summary(&mut self, summary: &'a T, cx: T::Context<'_>) {
148 self.0.add_summary(summary, cx);
149 self.1.add_summary(summary, cx);
150 self.2.add_summary(summary, cx);
151 }
152}
153
154impl<'a, S, D1, D2, D3> SeekTarget<'a, S, Dimensions<D1, D2, D3>> for D1
155where
156 S: Summary,
157 D1: SeekTarget<'a, S, D1> + Dimension<'a, S>,
158 D2: Dimension<'a, S>,
159 D3: Dimension<'a, S>,
160{
161 fn cmp(&self, cursor_location: &Dimensions<D1, D2, D3>, cx: S::Context<'_>) -> Ordering {
162 self.cmp(&cursor_location.0, cx)
163 }
164}
165
166/// Bias is used to settle ambiguities when determining positions in an ordered sequence.
167///
168/// The primary use case is for text, where Bias influences
169/// which character an offset or anchor is associated with.
170///
171/// # Examples
172/// Given the buffer `AˇBCD`:
173/// - The offset of the cursor is 1
174/// - [Bias::Left] would attach the cursor to the character `A`
175/// - [Bias::Right] would attach the cursor to the character `B`
176///
177/// Given the buffer `A«BCˇ»D`:
178/// - The offset of the cursor is 3, and the selection is from 1 to 3
179/// - The left anchor of the selection has [Bias::Right], attaching it to the character `B`
180/// - The right anchor of the selection has [Bias::Left], attaching it to the character `C`
181///
182/// Given the buffer `{ˇ<...>`, where `<...>` is a folded region:
183/// - The display offset of the cursor is 1, but the offset in the buffer is determined by the bias
184/// - [Bias::Left] would attach the cursor to the character `{`, with a buffer offset of 1
185/// - [Bias::Right] would attach the cursor to the first character of the folded region,
186/// and the buffer offset would be the offset of the first character of the folded region
187#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Debug, Hash, Default)]
188pub enum Bias {
189 /// Attach to the character on the left
190 #[default]
191 Left,
192 /// Attach to the character on the right
193 Right,
194}
195
196impl Bias {
197 pub fn invert(self) -> Self {
198 match self {
199 Self::Left => Self::Right,
200 Self::Right => Self::Left,
201 }
202 }
203}
204
205/// A B+ tree in which each leaf node contains `Item`s of type `T` and a `Summary`s for each `Item`.
206/// Each internal node contains a `Summary` of the items in its subtree.
207///
208/// The maximum number of items per node is `TREE_BASE * 2`.
209///
210/// Any [`Dimension`] supported by the [`Summary`] type can be used to seek to a specific location in the tree.
211#[derive(Clone)]
212pub struct SumTree<T: Item>(Arc<Node<T>>);
213
214impl<T> fmt::Debug for SumTree<T>
215where
216 T: fmt::Debug + Item,
217 T::Summary: fmt::Debug,
218{
219 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
220 f.debug_tuple("SumTree").field(&self.0).finish()
221 }
222}
223
224impl<T: Item> SumTree<T> {
225 pub fn new(cx: <T::Summary as Summary>::Context<'_>) -> Self {
226 SumTree(Arc::new(Node::Leaf {
227 summary: <T::Summary as Summary>::zero(cx),
228 items: ArrayVec::new(),
229 item_summaries: ArrayVec::new(),
230 }))
231 }
232
233 /// Useful in cases where the item type has a non-trivial context type, but the zero value of the summary type doesn't depend on that context.
234 pub fn from_summary(summary: T::Summary) -> Self {
235 SumTree(Arc::new(Node::Leaf {
236 summary,
237 items: ArrayVec::new(),
238 item_summaries: ArrayVec::new(),
239 }))
240 }
241
242 pub fn from_item(item: T, cx: <T::Summary as Summary>::Context<'_>) -> Self {
243 let mut tree = Self::new(cx);
244 tree.push(item, cx);
245 tree
246 }
247
248 pub fn from_iter<I: IntoIterator<Item = T>>(
249 iter: I,
250 cx: <T::Summary as Summary>::Context<'_>,
251 ) -> Self {
252 let mut nodes = Vec::new();
253
254 let mut iter = iter.into_iter().fuse().peekable();
255 while iter.peek().is_some() {
256 let items: ArrayVec<T, { 2 * TREE_BASE }> = iter.by_ref().take(2 * TREE_BASE).collect();
257 let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> =
258 items.iter().map(|item| item.summary(cx)).collect();
259
260 let mut summary = item_summaries[0].clone();
261 for item_summary in &item_summaries[1..] {
262 <T::Summary as Summary>::add_summary(&mut summary, item_summary, cx);
263 }
264
265 nodes.push(Node::Leaf {
266 summary,
267 items,
268 item_summaries,
269 });
270 }
271
272 let mut parent_nodes = Vec::new();
273 let mut height = 0;
274 while nodes.len() > 1 {
275 height += 1;
276 let mut current_parent_node = None;
277 for child_node in nodes.drain(..) {
278 let parent_node = current_parent_node.get_or_insert_with(|| Node::Internal {
279 summary: <T::Summary as Summary>::zero(cx),
280 height,
281 child_summaries: ArrayVec::new(),
282 child_trees: ArrayVec::new(),
283 });
284 let Node::Internal {
285 summary,
286 child_summaries,
287 child_trees,
288 ..
289 } = parent_node
290 else {
291 unreachable!()
292 };
293 let child_summary = child_node.summary();
294 <T::Summary as Summary>::add_summary(summary, child_summary, cx);
295 child_summaries.push(child_summary.clone());
296 child_trees.push(Self(Arc::new(child_node)));
297
298 if child_trees.len() == 2 * TREE_BASE {
299 parent_nodes.extend(current_parent_node.take());
300 }
301 }
302 parent_nodes.extend(current_parent_node.take());
303 mem::swap(&mut nodes, &mut parent_nodes);
304 }
305
306 if nodes.is_empty() {
307 Self::new(cx)
308 } else {
309 debug_assert_eq!(nodes.len(), 1);
310 Self(Arc::new(nodes.pop().unwrap()))
311 }
312 }
313
314 pub async fn from_iter_async<I, S>(iter: I, spawn: S) -> Self
315 where
316 T: 'static + Send + Sync,
317 for<'a> T::Summary: Summary<Context<'a> = ()> + Send + Sync,
318 S: BackgroundSpawn,
319 I: IntoIterator<Item = T>,
320 {
321 let mut futures = vec![];
322 let chunks = iter.into_iter().chunks(2 * TREE_BASE);
323 for chunk in chunks.into_iter() {
324 let items: ArrayVec<T, { 2 * TREE_BASE }> = chunk.into_iter().collect();
325 futures.push(async move {
326 let item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> =
327 items.iter().map(|item| item.summary(())).collect();
328 let mut summary = item_summaries[0].clone();
329 for item_summary in &item_summaries[1..] {
330 <T::Summary as Summary>::add_summary(&mut summary, item_summary, ());
331 }
332 SumTree(Arc::new(Node::Leaf {
333 summary,
334 items,
335 item_summaries,
336 }))
337 });
338 }
339
340 let mut nodes = futures::stream::iter(futures)
341 .map(|future| spawn.background_spawn(future))
342 .buffered(4)
343 .collect::<Vec<_>>()
344 .await;
345
346 let mut height = 0;
347 while nodes.len() > 1 {
348 height += 1;
349 let current_nodes = mem::take(&mut nodes);
350 nodes = stream::iter(current_nodes)
351 .chunks(2 * TREE_BASE)
352 .map(|chunk| {
353 spawn.background_spawn(async move {
354 let child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }> =
355 chunk.into_iter().collect();
356 let child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }> = child_trees
357 .iter()
358 .map(|child_tree| child_tree.summary().clone())
359 .collect();
360 let mut summary = child_summaries[0].clone();
361 for child_summary in &child_summaries[1..] {
362 <T::Summary as Summary>::add_summary(&mut summary, child_summary, ());
363 }
364 SumTree(Arc::new(Node::Internal {
365 height,
366 summary,
367 child_summaries,
368 child_trees,
369 }))
370 })
371 })
372 .buffered(4)
373 .collect::<Vec<_>>()
374 .await;
375 }
376
377 if nodes.is_empty() {
378 Self::new(())
379 } else {
380 debug_assert_eq!(nodes.len(), 1);
381 nodes.pop().unwrap()
382 }
383 }
384
385 #[allow(unused)]
386 pub fn items<'a>(&'a self, cx: <T::Summary as Summary>::Context<'a>) -> Vec<T> {
387 let mut items = Vec::new();
388 let mut cursor = self.cursor::<()>(cx);
389 cursor.next();
390 while let Some(item) = cursor.item() {
391 items.push(item.clone());
392 cursor.next();
393 }
394 items
395 }
396
397 pub fn iter(&self) -> Iter<'_, T> {
398 Iter::new(self)
399 }
400
401 /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`.
402 ///
403 /// Only returns the item that exactly has the target match.
404 pub fn find_exact<'a, 'slf, D, Target>(
405 &'slf self,
406 cx: <T::Summary as Summary>::Context<'a>,
407 target: &Target,
408 bias: Bias,
409 ) -> (D, D, Option<&'slf T>)
410 where
411 D: Dimension<'slf, T::Summary>,
412 Target: SeekTarget<'slf, T::Summary, D>,
413 {
414 let tree_end = D::zero(cx).with_added_summary(self.summary(), cx);
415 let comparison = target.cmp(&tree_end, cx);
416 if comparison == Ordering::Greater || (comparison == Ordering::Equal && bias == Bias::Right)
417 {
418 return (tree_end.clone(), tree_end, None);
419 }
420
421 let mut pos = D::zero(cx);
422 return match Self::find_recurse::<_, _, true>(cx, target, bias, &mut pos, self) {
423 Some((item, end)) => (pos, end, Some(item)),
424 None => (pos.clone(), pos, None),
425 };
426 }
427
428 /// A more efficient version of `Cursor::new()` + `Cursor::seek()` + `Cursor::item()`
429 pub fn find<'a, 'slf, D, Target>(
430 &'slf self,
431 cx: <T::Summary as Summary>::Context<'a>,
432 target: &Target,
433 bias: Bias,
434 ) -> (D, D, Option<&'slf T>)
435 where
436 D: Dimension<'slf, T::Summary>,
437 Target: SeekTarget<'slf, T::Summary, D>,
438 {
439 let tree_end = D::zero(cx).with_added_summary(self.summary(), cx);
440 let comparison = target.cmp(&tree_end, cx);
441 if comparison == Ordering::Greater || (comparison == Ordering::Equal && bias == Bias::Right)
442 {
443 return (tree_end.clone(), tree_end, None);
444 }
445
446 let mut pos = D::zero(cx);
447 return match Self::find_recurse::<_, _, false>(cx, target, bias, &mut pos, self) {
448 Some((item, end)) => (pos, end, Some(item)),
449 None => (pos.clone(), pos, None),
450 };
451 }
452
453 fn find_recurse<'tree, 'a, D, Target, const EXACT: bool>(
454 cx: <T::Summary as Summary>::Context<'a>,
455 target: &Target,
456 bias: Bias,
457 position: &mut D,
458 this: &'tree SumTree<T>,
459 ) -> Option<(&'tree T, D)>
460 where
461 D: Dimension<'tree, T::Summary>,
462 Target: SeekTarget<'tree, T::Summary, D>,
463 {
464 match &*this.0 {
465 Node::Internal {
466 child_summaries,
467 child_trees,
468 ..
469 } => {
470 for (child_tree, child_summary) in child_trees.iter().zip(child_summaries) {
471 let child_end = position.clone().with_added_summary(child_summary, cx);
472
473 let comparison = target.cmp(&child_end, cx);
474 let target_in_child = comparison == Ordering::Less
475 || (comparison == Ordering::Equal && bias == Bias::Left);
476 if target_in_child {
477 return Self::find_recurse::<D, Target, EXACT>(
478 cx, target, bias, position, child_tree,
479 );
480 }
481 *position = child_end;
482 }
483 }
484 Node::Leaf {
485 items,
486 item_summaries,
487 ..
488 } => {
489 for (item, item_summary) in items.iter().zip(item_summaries) {
490 let mut child_end = position.clone();
491 child_end.add_summary(item_summary, cx);
492
493 let comparison = target.cmp(&child_end, cx);
494 let entry_found = if EXACT {
495 comparison == Ordering::Equal
496 } else {
497 comparison == Ordering::Less
498 || (comparison == Ordering::Equal && bias == Bias::Left)
499 };
500 if entry_found {
501 return Some((item, child_end));
502 }
503
504 *position = child_end;
505 }
506 }
507 }
508 None
509 }
510
511 pub fn cursor<'a, 'b, D>(
512 &'a self,
513 cx: <T::Summary as Summary>::Context<'b>,
514 ) -> Cursor<'a, 'b, T, D>
515 where
516 D: Dimension<'a, T::Summary>,
517 {
518 Cursor::new(self, cx)
519 }
520
521 /// Note: If the summary type requires a non `()` context, then the filter cursor
522 /// that is returned cannot be used with Rust's iterators.
523 pub fn filter<'a, 'b, F, U>(
524 &'a self,
525 cx: <T::Summary as Summary>::Context<'b>,
526 filter_node: F,
527 ) -> FilterCursor<'a, 'b, F, T, U>
528 where
529 F: FnMut(&T::Summary) -> bool,
530 U: Dimension<'a, T::Summary>,
531 {
532 FilterCursor::new(self, cx, filter_node)
533 }
534
535 #[allow(dead_code)]
536 pub fn first(&self) -> Option<&T> {
537 self.leftmost_leaf().0.items().first()
538 }
539
540 pub fn last(&self) -> Option<&T> {
541 self.rightmost_leaf().0.items().last()
542 }
543
544 pub fn update_last(
545 &mut self,
546 f: impl FnOnce(&mut T),
547 cx: <T::Summary as Summary>::Context<'_>,
548 ) {
549 self.update_last_recursive(f, cx);
550 }
551
552 fn update_last_recursive(
553 &mut self,
554 f: impl FnOnce(&mut T),
555 cx: <T::Summary as Summary>::Context<'_>,
556 ) -> Option<T::Summary> {
557 match Arc::make_mut(&mut self.0) {
558 Node::Internal {
559 summary,
560 child_summaries,
561 child_trees,
562 ..
563 } => {
564 let last_summary = child_summaries.last_mut().unwrap();
565 let last_child = child_trees.last_mut().unwrap();
566 *last_summary = last_child.update_last_recursive(f, cx).unwrap();
567 *summary = sum(child_summaries.iter(), cx);
568 Some(summary.clone())
569 }
570 Node::Leaf {
571 summary,
572 items,
573 item_summaries,
574 } => {
575 if let Some((item, item_summary)) = items.last_mut().zip(item_summaries.last_mut())
576 {
577 (f)(item);
578 *item_summary = item.summary(cx);
579 *summary = sum(item_summaries.iter(), cx);
580 Some(summary.clone())
581 } else {
582 None
583 }
584 }
585 }
586 }
587
588 pub fn extent<'a, D: Dimension<'a, T::Summary>>(
589 &'a self,
590 cx: <T::Summary as Summary>::Context<'_>,
591 ) -> D {
592 let mut extent = D::zero(cx);
593 match self.0.as_ref() {
594 Node::Internal { summary, .. } | Node::Leaf { summary, .. } => {
595 extent.add_summary(summary, cx);
596 }
597 }
598 extent
599 }
600
601 pub fn summary(&self) -> &T::Summary {
602 match self.0.as_ref() {
603 Node::Internal { summary, .. } => summary,
604 Node::Leaf { summary, .. } => summary,
605 }
606 }
607
608 pub fn is_empty(&self) -> bool {
609 match self.0.as_ref() {
610 Node::Internal { .. } => false,
611 Node::Leaf { items, .. } => items.is_empty(),
612 }
613 }
614
615 pub fn extend<I>(&mut self, iter: I, cx: <T::Summary as Summary>::Context<'_>)
616 where
617 I: IntoIterator<Item = T>,
618 {
619 self.append(Self::from_iter(iter, cx), cx);
620 }
621
622 pub async fn async_extend<S, I>(&mut self, iter: I, spawn: S)
623 where
624 S: BackgroundSpawn,
625 I: IntoIterator<Item = T> + 'static,
626 T: 'static + Send + Sync,
627 for<'b> T::Summary: Summary<Context<'b> = ()> + Send + Sync,
628 {
629 let other = Self::from_iter_async(iter, spawn);
630 self.append(other.await, ());
631 }
632
633 pub fn push(&mut self, item: T, cx: <T::Summary as Summary>::Context<'_>) {
634 let summary = item.summary(cx);
635 self.append(
636 SumTree(Arc::new(Node::Leaf {
637 summary: summary.clone(),
638 items: ArrayVec::from_iter(Some(item)),
639 item_summaries: ArrayVec::from_iter(Some(summary)),
640 })),
641 cx,
642 );
643 }
644
645 pub fn append(&mut self, other: Self, cx: <T::Summary as Summary>::Context<'_>) {
646 if self.is_empty() {
647 *self = other;
648 } else if !other.0.is_leaf() || !other.0.items().is_empty() {
649 if self.0.height() < other.0.height() {
650 for tree in other.0.child_trees() {
651 self.append(tree.clone(), cx);
652 }
653 } else if let Some(split_tree) = self.push_tree_recursive(other, cx) {
654 *self = Self::from_child_trees(self.clone(), split_tree, cx);
655 }
656 }
657 }
658
659 fn push_tree_recursive(
660 &mut self,
661 other: SumTree<T>,
662 cx: <T::Summary as Summary>::Context<'_>,
663 ) -> Option<SumTree<T>> {
664 match Arc::make_mut(&mut self.0) {
665 Node::Internal {
666 height,
667 summary,
668 child_summaries,
669 child_trees,
670 ..
671 } => {
672 let other_node = other.0.clone();
673 <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
674
675 let height_delta = *height - other_node.height();
676 let mut summaries_to_append = ArrayVec::<T::Summary, { 2 * TREE_BASE }>::new();
677 let mut trees_to_append = ArrayVec::<SumTree<T>, { 2 * TREE_BASE }>::new();
678 if height_delta == 0 {
679 summaries_to_append.extend(other_node.child_summaries().iter().cloned());
680 trees_to_append.extend(other_node.child_trees().iter().cloned());
681 } else if height_delta == 1 && !other_node.is_underflowing() {
682 summaries_to_append.push(other_node.summary().clone());
683 trees_to_append.push(other)
684 } else {
685 let tree_to_append = child_trees
686 .last_mut()
687 .unwrap()
688 .push_tree_recursive(other, cx);
689 *child_summaries.last_mut().unwrap() =
690 child_trees.last().unwrap().0.summary().clone();
691
692 if let Some(split_tree) = tree_to_append {
693 summaries_to_append.push(split_tree.0.summary().clone());
694 trees_to_append.push(split_tree);
695 }
696 }
697
698 let child_count = child_trees.len() + trees_to_append.len();
699 if child_count > 2 * TREE_BASE {
700 let left_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
701 let right_summaries: ArrayVec<_, { 2 * TREE_BASE }>;
702 let left_trees;
703 let right_trees;
704
705 let midpoint = (child_count + child_count % 2) / 2;
706 {
707 let mut all_summaries = child_summaries
708 .iter()
709 .chain(summaries_to_append.iter())
710 .cloned();
711 left_summaries = all_summaries.by_ref().take(midpoint).collect();
712 right_summaries = all_summaries.collect();
713 let mut all_trees =
714 child_trees.iter().chain(trees_to_append.iter()).cloned();
715 left_trees = all_trees.by_ref().take(midpoint).collect();
716 right_trees = all_trees.collect();
717 }
718 *summary = sum(left_summaries.iter(), cx);
719 *child_summaries = left_summaries;
720 *child_trees = left_trees;
721
722 Some(SumTree(Arc::new(Node::Internal {
723 height: *height,
724 summary: sum(right_summaries.iter(), cx),
725 child_summaries: right_summaries,
726 child_trees: right_trees,
727 })))
728 } else {
729 child_summaries.extend(summaries_to_append);
730 child_trees.extend(trees_to_append);
731 None
732 }
733 }
734 Node::Leaf {
735 summary,
736 items,
737 item_summaries,
738 } => {
739 let other_node = other.0;
740
741 let child_count = items.len() + other_node.items().len();
742 if child_count > 2 * TREE_BASE {
743 let left_items;
744 let right_items;
745 let left_summaries;
746 let right_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>;
747
748 let midpoint = (child_count + child_count % 2) / 2;
749 {
750 let mut all_items = items.iter().chain(other_node.items().iter()).cloned();
751 left_items = all_items.by_ref().take(midpoint).collect();
752 right_items = all_items.collect();
753
754 let mut all_summaries = item_summaries
755 .iter()
756 .chain(other_node.child_summaries())
757 .cloned();
758 left_summaries = all_summaries.by_ref().take(midpoint).collect();
759 right_summaries = all_summaries.collect();
760 }
761 *items = left_items;
762 *item_summaries = left_summaries;
763 *summary = sum(item_summaries.iter(), cx);
764 Some(SumTree(Arc::new(Node::Leaf {
765 items: right_items,
766 summary: sum(right_summaries.iter(), cx),
767 item_summaries: right_summaries,
768 })))
769 } else {
770 <T::Summary as Summary>::add_summary(summary, other_node.summary(), cx);
771 items.extend(other_node.items().iter().cloned());
772 item_summaries.extend(other_node.child_summaries().iter().cloned());
773 None
774 }
775 }
776 }
777 }
778
779 fn from_child_trees(
780 left: SumTree<T>,
781 right: SumTree<T>,
782 cx: <T::Summary as Summary>::Context<'_>,
783 ) -> Self {
784 let height = left.0.height() + 1;
785 let mut child_summaries = ArrayVec::new();
786 child_summaries.push(left.0.summary().clone());
787 child_summaries.push(right.0.summary().clone());
788 let mut child_trees = ArrayVec::new();
789 child_trees.push(left);
790 child_trees.push(right);
791 SumTree(Arc::new(Node::Internal {
792 height,
793 summary: sum(child_summaries.iter(), cx),
794 child_summaries,
795 child_trees,
796 }))
797 }
798
799 fn leftmost_leaf(&self) -> &Self {
800 match *self.0 {
801 Node::Leaf { .. } => self,
802 Node::Internal {
803 ref child_trees, ..
804 } => child_trees.first().unwrap().leftmost_leaf(),
805 }
806 }
807
808 fn rightmost_leaf(&self) -> &Self {
809 match *self.0 {
810 Node::Leaf { .. } => self,
811 Node::Internal {
812 ref child_trees, ..
813 } => child_trees.last().unwrap().rightmost_leaf(),
814 }
815 }
816}
817
818impl<T: Item + PartialEq> PartialEq for SumTree<T> {
819 fn eq(&self, other: &Self) -> bool {
820 self.iter().eq(other.iter())
821 }
822}
823
824impl<T: Item + Eq> Eq for SumTree<T> {}
825
826impl<T: KeyedItem> SumTree<T> {
827 pub fn insert_or_replace<'a, 'b>(
828 &'a mut self,
829 item: T,
830 cx: <T::Summary as Summary>::Context<'b>,
831 ) -> Option<T> {
832 let mut replaced = None;
833 {
834 let mut cursor = self.cursor::<T::Key>(cx);
835 let mut new_tree = cursor.slice(&item.key(), Bias::Left);
836 if let Some(cursor_item) = cursor.item()
837 && cursor_item.key() == item.key()
838 {
839 replaced = Some(cursor_item.clone());
840 cursor.next();
841 }
842 new_tree.push(item, cx);
843 new_tree.append(cursor.suffix(), cx);
844 drop(cursor);
845 *self = new_tree
846 };
847 replaced
848 }
849
850 pub fn remove(&mut self, key: &T::Key, cx: <T::Summary as Summary>::Context<'_>) -> Option<T> {
851 let mut removed = None;
852 *self = {
853 let mut cursor = self.cursor::<T::Key>(cx);
854 let mut new_tree = cursor.slice(key, Bias::Left);
855 if let Some(item) = cursor.item()
856 && item.key() == *key
857 {
858 removed = Some(item.clone());
859 cursor.next();
860 }
861 new_tree.append(cursor.suffix(), cx);
862 new_tree
863 };
864 removed
865 }
866
867 pub fn edit(
868 &mut self,
869 mut edits: Vec<Edit<T>>,
870 cx: <T::Summary as Summary>::Context<'_>,
871 ) -> Vec<T> {
872 if edits.is_empty() {
873 return Vec::new();
874 }
875
876 let mut removed = Vec::new();
877 edits.sort_unstable_by_key(|item| item.key());
878
879 *self = {
880 let mut cursor = self.cursor::<T::Key>(cx);
881 let mut new_tree = SumTree::new(cx);
882 let mut buffered_items = Vec::new();
883
884 cursor.seek(&T::Key::zero(cx), Bias::Left);
885 for edit in edits {
886 let new_key = edit.key();
887 let mut old_item = cursor.item();
888
889 if old_item
890 .as_ref()
891 .is_some_and(|old_item| old_item.key() < new_key)
892 {
893 new_tree.extend(buffered_items.drain(..), cx);
894 let slice = cursor.slice(&new_key, Bias::Left);
895 new_tree.append(slice, cx);
896 old_item = cursor.item();
897 }
898
899 if let Some(old_item) = old_item
900 && old_item.key() == new_key
901 {
902 removed.push(old_item.clone());
903 cursor.next();
904 }
905
906 match edit {
907 Edit::Insert(item) => {
908 buffered_items.push(item);
909 }
910 Edit::Remove(_) => {}
911 }
912 }
913
914 new_tree.extend(buffered_items, cx);
915 new_tree.append(cursor.suffix(), cx);
916 new_tree
917 };
918
919 removed
920 }
921
922 pub fn get<'a>(
923 &'a self,
924 key: &T::Key,
925 cx: <T::Summary as Summary>::Context<'a>,
926 ) -> Option<&'a T> {
927 if let (_, _, Some(item)) = self.find_exact::<T::Key, _>(cx, key, Bias::Left) {
928 Some(item)
929 } else {
930 None
931 }
932 }
933}
934
935impl<T, S> Default for SumTree<T>
936where
937 T: Item<Summary = S>,
938 S: for<'a> Summary<Context<'a> = ()>,
939{
940 fn default() -> Self {
941 Self::new(())
942 }
943}
944
945#[derive(Clone)]
946pub enum Node<T: Item> {
947 Internal {
948 height: u8,
949 summary: T::Summary,
950 child_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
951 child_trees: ArrayVec<SumTree<T>, { 2 * TREE_BASE }>,
952 },
953 Leaf {
954 summary: T::Summary,
955 items: ArrayVec<T, { 2 * TREE_BASE }>,
956 item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
957 },
958}
959
960impl<T> fmt::Debug for Node<T>
961where
962 T: Item + fmt::Debug,
963 T::Summary: fmt::Debug,
964{
965 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
966 match self {
967 Node::Internal {
968 height,
969 summary,
970 child_summaries,
971 child_trees,
972 } => f
973 .debug_struct("Internal")
974 .field("height", height)
975 .field("summary", summary)
976 .field("child_summaries", child_summaries)
977 .field("child_trees", child_trees)
978 .finish(),
979 Node::Leaf {
980 summary,
981 items,
982 item_summaries,
983 } => f
984 .debug_struct("Leaf")
985 .field("summary", summary)
986 .field("items", items)
987 .field("item_summaries", item_summaries)
988 .finish(),
989 }
990 }
991}
992
993impl<T: Item> Node<T> {
994 fn is_leaf(&self) -> bool {
995 matches!(self, Node::Leaf { .. })
996 }
997
998 fn height(&self) -> u8 {
999 match self {
1000 Node::Internal { height, .. } => *height,
1001 Node::Leaf { .. } => 0,
1002 }
1003 }
1004
1005 fn summary(&self) -> &T::Summary {
1006 match self {
1007 Node::Internal { summary, .. } => summary,
1008 Node::Leaf { summary, .. } => summary,
1009 }
1010 }
1011
1012 fn child_summaries(&self) -> &[T::Summary] {
1013 match self {
1014 Node::Internal {
1015 child_summaries, ..
1016 } => child_summaries.as_slice(),
1017 Node::Leaf { item_summaries, .. } => item_summaries.as_slice(),
1018 }
1019 }
1020
1021 fn child_trees(&self) -> &ArrayVec<SumTree<T>, { 2 * TREE_BASE }> {
1022 match self {
1023 Node::Internal { child_trees, .. } => child_trees,
1024 Node::Leaf { .. } => panic!("Leaf nodes have no child trees"),
1025 }
1026 }
1027
1028 fn items(&self) -> &ArrayVec<T, { 2 * TREE_BASE }> {
1029 match self {
1030 Node::Leaf { items, .. } => items,
1031 Node::Internal { .. } => panic!("Internal nodes have no items"),
1032 }
1033 }
1034
1035 fn is_underflowing(&self) -> bool {
1036 match self {
1037 Node::Internal { child_trees, .. } => child_trees.len() < TREE_BASE,
1038 Node::Leaf { items, .. } => items.len() < TREE_BASE,
1039 }
1040 }
1041}
1042
1043#[derive(Debug)]
1044pub enum Edit<T: KeyedItem> {
1045 Insert(T),
1046 Remove(T::Key),
1047}
1048
1049impl<T: KeyedItem> Edit<T> {
1050 fn key(&self) -> T::Key {
1051 match self {
1052 Edit::Insert(item) => item.key(),
1053 Edit::Remove(key) => key.clone(),
1054 }
1055 }
1056}
1057
1058fn sum<'a, T, I>(iter: I, cx: T::Context<'_>) -> T
1059where
1060 T: 'a + Summary,
1061 I: Iterator<Item = &'a T>,
1062{
1063 let mut sum = T::zero(cx);
1064 for value in iter {
1065 sum.add_summary(value, cx);
1066 }
1067 sum
1068}
1069
1070#[cfg(test)]
1071mod tests {
1072 use super::*;
1073 use rand::{distr::StandardUniform, prelude::*};
1074 use std::cmp;
1075
1076 #[ctor::ctor]
1077 fn init_logger() {
1078 zlog::init_test();
1079 }
1080
1081 #[test]
1082 fn test_extend_and_push_tree() {
1083 let mut tree1 = SumTree::default();
1084 tree1.extend(0..20, ());
1085
1086 let mut tree2 = SumTree::default();
1087 tree2.extend(50..100, ());
1088
1089 tree1.append(tree2, ());
1090 assert_eq!(tree1.items(()), (0..20).chain(50..100).collect::<Vec<u8>>());
1091 }
1092
1093 #[test]
1094 fn test_random() {
1095 struct NoSpawn;
1096 impl BackgroundSpawn for NoSpawn {
1097 type Task<R>
1098 = std::pin::Pin<Box<dyn Future<Output = R> + Sync + Send>>
1099 where
1100 R: Send + Sync;
1101 fn background_spawn<R>(
1102 &self,
1103 future: impl Future<Output = R> + Send + Sync + 'static,
1104 ) -> Self::Task<R>
1105 where
1106 R: Send + Sync + 'static,
1107 {
1108 Box::pin(future)
1109 }
1110 }
1111
1112 let mut starting_seed = 0;
1113 if let Ok(value) = std::env::var("SEED") {
1114 starting_seed = value.parse().expect("invalid SEED variable");
1115 }
1116 let mut num_iterations = 100;
1117 if let Ok(value) = std::env::var("ITERATIONS") {
1118 num_iterations = value.parse().expect("invalid ITERATIONS variable");
1119 }
1120 let num_operations = std::env::var("OPERATIONS")
1121 .map_or(5, |o| o.parse().expect("invalid OPERATIONS variable"));
1122
1123 for seed in starting_seed..(starting_seed + num_iterations) {
1124 eprintln!("seed = {}", seed);
1125 let mut rng = StdRng::seed_from_u64(seed);
1126
1127 let rng = &mut rng;
1128 let mut tree = SumTree::<u8>::default();
1129 let count = rng.random_range(0..10);
1130 if rng.random() {
1131 tree.extend(rng.sample_iter(StandardUniform).take(count), ());
1132 } else {
1133 let items = rng
1134 .sample_iter(StandardUniform)
1135 .take(count)
1136 .collect::<Vec<_>>();
1137 pollster::block_on(tree.async_extend(items, NoSpawn));
1138 }
1139
1140 for _ in 0..num_operations {
1141 let splice_end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
1142 let splice_start = rng.random_range(0..splice_end + 1);
1143 let count = rng.random_range(0..10);
1144 let tree_end = tree.extent::<Count>(());
1145 let new_items = rng
1146 .sample_iter(StandardUniform)
1147 .take(count)
1148 .collect::<Vec<u8>>();
1149
1150 let mut reference_items = tree.items(());
1151 reference_items.splice(splice_start..splice_end, new_items.clone());
1152
1153 tree = {
1154 let mut cursor = tree.cursor::<Count>(());
1155 let mut new_tree = cursor.slice(&Count(splice_start), Bias::Right);
1156 if rng.random() {
1157 new_tree.extend(new_items, ());
1158 } else {
1159 pollster::block_on(new_tree.async_extend(new_items, NoSpawn));
1160 }
1161 cursor.seek(&Count(splice_end), Bias::Right);
1162 new_tree.append(cursor.slice(&tree_end, Bias::Right), ());
1163 new_tree
1164 };
1165
1166 assert_eq!(tree.items(()), reference_items);
1167 assert_eq!(
1168 tree.iter().collect::<Vec<_>>(),
1169 tree.cursor::<()>(()).collect::<Vec<_>>()
1170 );
1171
1172 log::info!("tree items: {:?}", tree.items(()));
1173
1174 let mut filter_cursor =
1175 tree.filter::<_, Count>((), |summary| summary.contains_even);
1176 let expected_filtered_items = tree
1177 .items(())
1178 .into_iter()
1179 .enumerate()
1180 .filter(|(_, item)| (item & 1) == 0)
1181 .collect::<Vec<_>>();
1182
1183 let mut item_ix = if rng.random() {
1184 filter_cursor.next();
1185 0
1186 } else {
1187 filter_cursor.prev();
1188 expected_filtered_items.len().saturating_sub(1)
1189 };
1190 while item_ix < expected_filtered_items.len() {
1191 log::info!("filter_cursor, item_ix: {}", item_ix);
1192 let actual_item = filter_cursor.item().unwrap();
1193 let (reference_index, reference_item) = expected_filtered_items[item_ix];
1194 assert_eq!(actual_item, &reference_item);
1195 assert_eq!(filter_cursor.start().0, reference_index);
1196 log::info!("next");
1197 filter_cursor.next();
1198 item_ix += 1;
1199
1200 while item_ix > 0 && rng.random_bool(0.2) {
1201 log::info!("prev");
1202 filter_cursor.prev();
1203 item_ix -= 1;
1204
1205 if item_ix == 0 && rng.random_bool(0.2) {
1206 filter_cursor.prev();
1207 assert_eq!(filter_cursor.item(), None);
1208 assert_eq!(filter_cursor.start().0, 0);
1209 filter_cursor.next();
1210 }
1211 }
1212 }
1213 assert_eq!(filter_cursor.item(), None);
1214
1215 let mut before_start = false;
1216 let mut cursor = tree.cursor::<Count>(());
1217 let start_pos = rng.random_range(0..=reference_items.len());
1218 cursor.seek(&Count(start_pos), Bias::Right);
1219 let mut pos = rng.random_range(start_pos..=reference_items.len());
1220 cursor.seek_forward(&Count(pos), Bias::Right);
1221
1222 for i in 0..10 {
1223 assert_eq!(cursor.start().0, pos);
1224
1225 if pos > 0 {
1226 assert_eq!(cursor.prev_item().unwrap(), &reference_items[pos - 1]);
1227 } else {
1228 assert_eq!(cursor.prev_item(), None);
1229 }
1230
1231 if pos < reference_items.len() && !before_start {
1232 assert_eq!(cursor.item().unwrap(), &reference_items[pos]);
1233 } else {
1234 assert_eq!(cursor.item(), None);
1235 }
1236
1237 if before_start {
1238 assert_eq!(cursor.next_item(), reference_items.first());
1239 } else if pos + 1 < reference_items.len() {
1240 assert_eq!(cursor.next_item().unwrap(), &reference_items[pos + 1]);
1241 } else {
1242 assert_eq!(cursor.next_item(), None);
1243 }
1244
1245 if i < 5 {
1246 cursor.next();
1247 if pos < reference_items.len() {
1248 pos += 1;
1249 before_start = false;
1250 }
1251 } else {
1252 cursor.prev();
1253 if pos == 0 {
1254 before_start = true;
1255 }
1256 pos = pos.saturating_sub(1);
1257 }
1258 }
1259 }
1260
1261 for _ in 0..10 {
1262 let end = rng.random_range(0..tree.extent::<Count>(()).0 + 1);
1263 let start = rng.random_range(0..end + 1);
1264 let start_bias = if rng.random() {
1265 Bias::Left
1266 } else {
1267 Bias::Right
1268 };
1269 let end_bias = if rng.random() {
1270 Bias::Left
1271 } else {
1272 Bias::Right
1273 };
1274
1275 let mut cursor = tree.cursor::<Count>(());
1276 cursor.seek(&Count(start), start_bias);
1277 let slice = cursor.slice(&Count(end), end_bias);
1278
1279 cursor.seek(&Count(start), start_bias);
1280 let summary = cursor.summary::<_, Sum>(&Count(end), end_bias);
1281
1282 assert_eq!(summary.0, slice.summary().sum);
1283 }
1284 }
1285 }
1286
1287 #[test]
1288 fn test_cursor() {
1289 // Empty tree
1290 let tree = SumTree::<u8>::default();
1291 let mut cursor = tree.cursor::<IntegersSummary>(());
1292 assert_eq!(
1293 cursor.slice(&Count(0), Bias::Right).items(()),
1294 Vec::<u8>::new()
1295 );
1296 assert_eq!(cursor.item(), None);
1297 assert_eq!(cursor.prev_item(), None);
1298 assert_eq!(cursor.next_item(), None);
1299 assert_eq!(cursor.start().sum, 0);
1300 cursor.prev();
1301 assert_eq!(cursor.item(), None);
1302 assert_eq!(cursor.prev_item(), None);
1303 assert_eq!(cursor.next_item(), None);
1304 assert_eq!(cursor.start().sum, 0);
1305 cursor.next();
1306 assert_eq!(cursor.item(), None);
1307 assert_eq!(cursor.prev_item(), None);
1308 assert_eq!(cursor.next_item(), None);
1309 assert_eq!(cursor.start().sum, 0);
1310
1311 // Single-element tree
1312 let mut tree = SumTree::<u8>::default();
1313 tree.extend(vec![1], ());
1314 let mut cursor = tree.cursor::<IntegersSummary>(());
1315 assert_eq!(
1316 cursor.slice(&Count(0), Bias::Right).items(()),
1317 Vec::<u8>::new()
1318 );
1319 assert_eq!(cursor.item(), Some(&1));
1320 assert_eq!(cursor.prev_item(), None);
1321 assert_eq!(cursor.next_item(), None);
1322 assert_eq!(cursor.start().sum, 0);
1323
1324 cursor.next();
1325 assert_eq!(cursor.item(), None);
1326 assert_eq!(cursor.prev_item(), Some(&1));
1327 assert_eq!(cursor.next_item(), None);
1328 assert_eq!(cursor.start().sum, 1);
1329
1330 cursor.prev();
1331 assert_eq!(cursor.item(), Some(&1));
1332 assert_eq!(cursor.prev_item(), None);
1333 assert_eq!(cursor.next_item(), None);
1334 assert_eq!(cursor.start().sum, 0);
1335
1336 let mut cursor = tree.cursor::<IntegersSummary>(());
1337 assert_eq!(cursor.slice(&Count(1), Bias::Right).items(()), [1]);
1338 assert_eq!(cursor.item(), None);
1339 assert_eq!(cursor.prev_item(), Some(&1));
1340 assert_eq!(cursor.next_item(), None);
1341 assert_eq!(cursor.start().sum, 1);
1342
1343 cursor.seek(&Count(0), Bias::Right);
1344 assert_eq!(
1345 cursor
1346 .slice(&tree.extent::<Count>(()), Bias::Right)
1347 .items(()),
1348 [1]
1349 );
1350 assert_eq!(cursor.item(), None);
1351 assert_eq!(cursor.prev_item(), Some(&1));
1352 assert_eq!(cursor.next_item(), None);
1353 assert_eq!(cursor.start().sum, 1);
1354
1355 // Multiple-element tree
1356 let mut tree = SumTree::default();
1357 tree.extend(vec![1, 2, 3, 4, 5, 6], ());
1358 let mut cursor = tree.cursor::<IntegersSummary>(());
1359
1360 assert_eq!(cursor.slice(&Count(2), Bias::Right).items(()), [1, 2]);
1361 assert_eq!(cursor.item(), Some(&3));
1362 assert_eq!(cursor.prev_item(), Some(&2));
1363 assert_eq!(cursor.next_item(), Some(&4));
1364 assert_eq!(cursor.start().sum, 3);
1365
1366 cursor.next();
1367 assert_eq!(cursor.item(), Some(&4));
1368 assert_eq!(cursor.prev_item(), Some(&3));
1369 assert_eq!(cursor.next_item(), Some(&5));
1370 assert_eq!(cursor.start().sum, 6);
1371
1372 cursor.next();
1373 assert_eq!(cursor.item(), Some(&5));
1374 assert_eq!(cursor.prev_item(), Some(&4));
1375 assert_eq!(cursor.next_item(), Some(&6));
1376 assert_eq!(cursor.start().sum, 10);
1377
1378 cursor.next();
1379 assert_eq!(cursor.item(), Some(&6));
1380 assert_eq!(cursor.prev_item(), Some(&5));
1381 assert_eq!(cursor.next_item(), None);
1382 assert_eq!(cursor.start().sum, 15);
1383
1384 cursor.next();
1385 cursor.next();
1386 assert_eq!(cursor.item(), None);
1387 assert_eq!(cursor.prev_item(), Some(&6));
1388 assert_eq!(cursor.next_item(), None);
1389 assert_eq!(cursor.start().sum, 21);
1390
1391 cursor.prev();
1392 assert_eq!(cursor.item(), Some(&6));
1393 assert_eq!(cursor.prev_item(), Some(&5));
1394 assert_eq!(cursor.next_item(), None);
1395 assert_eq!(cursor.start().sum, 15);
1396
1397 cursor.prev();
1398 assert_eq!(cursor.item(), Some(&5));
1399 assert_eq!(cursor.prev_item(), Some(&4));
1400 assert_eq!(cursor.next_item(), Some(&6));
1401 assert_eq!(cursor.start().sum, 10);
1402
1403 cursor.prev();
1404 assert_eq!(cursor.item(), Some(&4));
1405 assert_eq!(cursor.prev_item(), Some(&3));
1406 assert_eq!(cursor.next_item(), Some(&5));
1407 assert_eq!(cursor.start().sum, 6);
1408
1409 cursor.prev();
1410 assert_eq!(cursor.item(), Some(&3));
1411 assert_eq!(cursor.prev_item(), Some(&2));
1412 assert_eq!(cursor.next_item(), Some(&4));
1413 assert_eq!(cursor.start().sum, 3);
1414
1415 cursor.prev();
1416 assert_eq!(cursor.item(), Some(&2));
1417 assert_eq!(cursor.prev_item(), Some(&1));
1418 assert_eq!(cursor.next_item(), Some(&3));
1419 assert_eq!(cursor.start().sum, 1);
1420
1421 cursor.prev();
1422 assert_eq!(cursor.item(), Some(&1));
1423 assert_eq!(cursor.prev_item(), None);
1424 assert_eq!(cursor.next_item(), Some(&2));
1425 assert_eq!(cursor.start().sum, 0);
1426
1427 cursor.prev();
1428 assert_eq!(cursor.item(), None);
1429 assert_eq!(cursor.prev_item(), None);
1430 assert_eq!(cursor.next_item(), Some(&1));
1431 assert_eq!(cursor.start().sum, 0);
1432
1433 cursor.next();
1434 assert_eq!(cursor.item(), Some(&1));
1435 assert_eq!(cursor.prev_item(), None);
1436 assert_eq!(cursor.next_item(), Some(&2));
1437 assert_eq!(cursor.start().sum, 0);
1438
1439 let mut cursor = tree.cursor::<IntegersSummary>(());
1440 assert_eq!(
1441 cursor
1442 .slice(&tree.extent::<Count>(()), Bias::Right)
1443 .items(()),
1444 tree.items(())
1445 );
1446 assert_eq!(cursor.item(), None);
1447 assert_eq!(cursor.prev_item(), Some(&6));
1448 assert_eq!(cursor.next_item(), None);
1449 assert_eq!(cursor.start().sum, 21);
1450
1451 cursor.seek(&Count(3), Bias::Right);
1452 assert_eq!(
1453 cursor
1454 .slice(&tree.extent::<Count>(()), Bias::Right)
1455 .items(()),
1456 [4, 5, 6]
1457 );
1458 assert_eq!(cursor.item(), None);
1459 assert_eq!(cursor.prev_item(), Some(&6));
1460 assert_eq!(cursor.next_item(), None);
1461 assert_eq!(cursor.start().sum, 21);
1462
1463 // Seeking can bias left or right
1464 cursor.seek(&Count(1), Bias::Left);
1465 assert_eq!(cursor.item(), Some(&1));
1466 cursor.seek(&Count(1), Bias::Right);
1467 assert_eq!(cursor.item(), Some(&2));
1468
1469 // Slicing without resetting starts from where the cursor is parked at.
1470 cursor.seek(&Count(1), Bias::Right);
1471 assert_eq!(cursor.slice(&Count(3), Bias::Right).items(()), vec![2, 3]);
1472 assert_eq!(cursor.slice(&Count(6), Bias::Left).items(()), vec![4, 5]);
1473 assert_eq!(cursor.slice(&Count(6), Bias::Right).items(()), vec![6]);
1474 }
1475
1476 #[test]
1477 fn test_edit() {
1478 let mut tree = SumTree::<u8>::default();
1479
1480 let removed = tree.edit(vec![Edit::Insert(1), Edit::Insert(2), Edit::Insert(0)], ());
1481 assert_eq!(tree.items(()), vec![0, 1, 2]);
1482 assert_eq!(removed, Vec::<u8>::new());
1483 assert_eq!(tree.get(&0, ()), Some(&0));
1484 assert_eq!(tree.get(&1, ()), Some(&1));
1485 assert_eq!(tree.get(&2, ()), Some(&2));
1486 assert_eq!(tree.get(&4, ()), None);
1487
1488 let removed = tree.edit(vec![Edit::Insert(2), Edit::Insert(4), Edit::Remove(0)], ());
1489 assert_eq!(tree.items(()), vec![1, 2, 4]);
1490 assert_eq!(removed, vec![0, 2]);
1491 assert_eq!(tree.get(&0, ()), None);
1492 assert_eq!(tree.get(&1, ()), Some(&1));
1493 assert_eq!(tree.get(&2, ()), Some(&2));
1494 assert_eq!(tree.get(&4, ()), Some(&4));
1495 }
1496
1497 #[test]
1498 fn test_from_iter() {
1499 assert_eq!(
1500 SumTree::from_iter(0..100, ()).items(()),
1501 (0..100).collect::<Vec<_>>()
1502 );
1503
1504 // Ensure `from_iter` works correctly when the given iterator restarts
1505 // after calling `next` if `None` was already returned.
1506 let mut ix = 0;
1507 let iterator = std::iter::from_fn(|| {
1508 ix = (ix + 1) % 2;
1509 if ix == 1 { Some(1) } else { None }
1510 });
1511 assert_eq!(SumTree::from_iter(iterator, ()).items(()), vec![1]);
1512 }
1513
1514 #[derive(Clone, Default, Debug)]
1515 pub struct IntegersSummary {
1516 count: usize,
1517 sum: usize,
1518 contains_even: bool,
1519 max: u8,
1520 }
1521
1522 #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1523 struct Count(usize);
1524
1525 #[derive(Ord, PartialOrd, Default, Eq, PartialEq, Clone, Debug)]
1526 struct Sum(usize);
1527
1528 impl Item for u8 {
1529 type Summary = IntegersSummary;
1530
1531 fn summary(&self, _cx: ()) -> Self::Summary {
1532 IntegersSummary {
1533 count: 1,
1534 sum: *self as usize,
1535 contains_even: (*self & 1) == 0,
1536 max: *self,
1537 }
1538 }
1539 }
1540
1541 impl KeyedItem for u8 {
1542 type Key = u8;
1543
1544 fn key(&self) -> Self::Key {
1545 *self
1546 }
1547 }
1548
1549 impl ContextLessSummary for IntegersSummary {
1550 fn zero() -> Self {
1551 Default::default()
1552 }
1553
1554 fn add_summary(&mut self, other: &Self) {
1555 self.count += other.count;
1556 self.sum += other.sum;
1557 self.contains_even |= other.contains_even;
1558 self.max = cmp::max(self.max, other.max);
1559 }
1560 }
1561
1562 impl Dimension<'_, IntegersSummary> for u8 {
1563 fn zero(_cx: ()) -> Self {
1564 Default::default()
1565 }
1566
1567 fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1568 *self = summary.max;
1569 }
1570 }
1571
1572 impl Dimension<'_, IntegersSummary> for Count {
1573 fn zero(_cx: ()) -> Self {
1574 Default::default()
1575 }
1576
1577 fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1578 self.0 += summary.count;
1579 }
1580 }
1581
1582 impl SeekTarget<'_, IntegersSummary, IntegersSummary> for Count {
1583 fn cmp(&self, cursor_location: &IntegersSummary, _: ()) -> Ordering {
1584 self.0.cmp(&cursor_location.count)
1585 }
1586 }
1587
1588 impl Dimension<'_, IntegersSummary> for Sum {
1589 fn zero(_cx: ()) -> Self {
1590 Default::default()
1591 }
1592
1593 fn add_summary(&mut self, summary: &IntegersSummary, _: ()) {
1594 self.0 += summary.sum;
1595 }
1596 }
1597}