1use super::*;
2use arrayvec::ArrayVec;
3use std::{cmp::Ordering, mem, sync::Arc};
4
5#[derive(Clone)]
6struct StackEntry<'a, T: Item, D> {
7 tree: &'a SumTree<T>,
8 index: usize,
9 position: D,
10}
11
12#[derive(Clone)]
13pub struct Cursor<'a, T: Item, D> {
14 tree: &'a SumTree<T>,
15 stack: ArrayVec<StackEntry<'a, T, D>, 16>,
16 position: D,
17 did_seek: bool,
18 at_end: bool,
19}
20
21impl<'a, T, D> Cursor<'a, T, D>
22where
23 T: Item,
24 D: Dimension<'a, T::Summary>,
25{
26 pub fn new(tree: &'a SumTree<T>) -> Self {
27 Self {
28 tree,
29 stack: ArrayVec::new(),
30 position: D::default(),
31 did_seek: false,
32 at_end: false,
33 }
34 }
35
36 fn reset(&mut self) {
37 self.did_seek = false;
38 self.at_end = false;
39 self.stack.truncate(0);
40 self.position = D::default();
41 }
42
43 pub fn start(&self) -> &D {
44 &self.position
45 }
46
47 pub fn end(&self, cx: &<T::Summary as Summary>::Context) -> D {
48 if let Some(item_summary) = self.item_summary() {
49 let mut end = self.start().clone();
50 end.add_summary(item_summary, cx);
51 end
52 } else {
53 self.start().clone()
54 }
55 }
56
57 pub fn item(&self) -> Option<&'a T> {
58 assert!(self.did_seek, "Must seek before calling this method");
59 if let Some(entry) = self.stack.last() {
60 match *entry.tree.0 {
61 Node::Leaf { ref items, .. } => {
62 if entry.index == items.len() {
63 None
64 } else {
65 Some(&items[entry.index])
66 }
67 }
68 _ => unreachable!(),
69 }
70 } else {
71 None
72 }
73 }
74
75 pub fn item_summary(&self) -> Option<&'a T::Summary> {
76 assert!(self.did_seek, "Must seek before calling this method");
77 if let Some(entry) = self.stack.last() {
78 match *entry.tree.0 {
79 Node::Leaf {
80 ref item_summaries, ..
81 } => {
82 if entry.index == item_summaries.len() {
83 None
84 } else {
85 Some(&item_summaries[entry.index])
86 }
87 }
88 _ => unreachable!(),
89 }
90 } else {
91 None
92 }
93 }
94
95 pub fn prev_item(&self) -> Option<&'a T> {
96 assert!(self.did_seek, "Must seek before calling this method");
97 if let Some(entry) = self.stack.last() {
98 if entry.index == 0 {
99 if let Some(prev_leaf) = self.prev_leaf() {
100 Some(prev_leaf.0.items().last().unwrap())
101 } else {
102 None
103 }
104 } else {
105 match *entry.tree.0 {
106 Node::Leaf { ref items, .. } => Some(&items[entry.index - 1]),
107 _ => unreachable!(),
108 }
109 }
110 } else if self.at_end {
111 self.tree.last()
112 } else {
113 None
114 }
115 }
116
117 fn prev_leaf(&self) -> Option<&'a SumTree<T>> {
118 for entry in self.stack.iter().rev().skip(1) {
119 if entry.index != 0 {
120 match *entry.tree.0 {
121 Node::Internal {
122 ref child_trees, ..
123 } => return Some(child_trees[entry.index - 1].rightmost_leaf()),
124 Node::Leaf { .. } => unreachable!(),
125 };
126 }
127 }
128 None
129 }
130
131 pub fn prev(&mut self, cx: &<T::Summary as Summary>::Context) {
132 assert!(self.did_seek, "Must seek before calling this method");
133
134 if self.at_end {
135 self.position = D::default();
136 self.descend_to_last_item(self.tree, cx);
137 self.at_end = false;
138 } else {
139 while let Some(entry) = self.stack.pop() {
140 if entry.index > 0 {
141 let new_index = entry.index - 1;
142
143 if let Some(StackEntry { position, .. }) = self.stack.last() {
144 self.position = position.clone();
145 } else {
146 self.position = D::default();
147 }
148
149 match entry.tree.0.as_ref() {
150 Node::Internal {
151 child_trees,
152 child_summaries,
153 ..
154 } => {
155 for summary in &child_summaries[0..new_index] {
156 self.position.add_summary(summary, cx);
157 }
158 self.stack.push(StackEntry {
159 tree: entry.tree,
160 index: new_index,
161 position: self.position.clone(),
162 });
163 self.descend_to_last_item(&child_trees[new_index], cx);
164 }
165 Node::Leaf { item_summaries, .. } => {
166 for item_summary in &item_summaries[0..new_index] {
167 self.position.add_summary(item_summary, cx);
168 }
169 self.stack.push(StackEntry {
170 tree: entry.tree,
171 index: new_index,
172 position: self.position.clone(),
173 });
174 }
175 }
176
177 break;
178 }
179 }
180 }
181 }
182
183 pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
184 self.next_internal(|_| true, cx)
185 }
186
187 fn next_internal<F>(&mut self, mut filter_node: F, cx: &<T::Summary as Summary>::Context)
188 where
189 F: FnMut(&T::Summary) -> bool,
190 {
191 let mut descend = false;
192
193 if self.stack.is_empty() && !self.at_end {
194 self.stack.push(StackEntry {
195 tree: self.tree,
196 index: 0,
197 position: D::default(),
198 });
199 descend = true;
200 self.did_seek = true;
201 }
202
203 while self.stack.len() > 0 {
204 let new_subtree = {
205 let entry = self.stack.last_mut().unwrap();
206 match entry.tree.0.as_ref() {
207 Node::Internal {
208 child_trees,
209 child_summaries,
210 ..
211 } => {
212 if !descend {
213 entry.position = self.position.clone();
214 entry.index += 1;
215 }
216
217 while entry.index < child_summaries.len() {
218 let next_summary = &child_summaries[entry.index];
219 if filter_node(next_summary) {
220 break;
221 } else {
222 self.position.add_summary(next_summary, cx);
223 }
224 entry.index += 1;
225 }
226
227 child_trees.get(entry.index)
228 }
229 Node::Leaf { item_summaries, .. } => {
230 if !descend {
231 let item_summary = &item_summaries[entry.index];
232 self.position.add_summary(item_summary, cx);
233 entry.position.add_summary(item_summary, cx);
234 entry.index += 1;
235 }
236
237 loop {
238 if let Some(next_item_summary) = item_summaries.get(entry.index) {
239 if filter_node(next_item_summary) {
240 return;
241 } else {
242 self.position.add_summary(next_item_summary, cx);
243 entry.position.add_summary(next_item_summary, cx);
244 entry.index += 1;
245 }
246 } else {
247 break None;
248 }
249 }
250 }
251 }
252 };
253
254 if let Some(subtree) = new_subtree {
255 descend = true;
256 self.stack.push(StackEntry {
257 tree: subtree,
258 index: 0,
259 position: self.position.clone(),
260 });
261 } else {
262 descend = false;
263 self.stack.pop();
264 }
265 }
266
267 self.at_end = self.stack.is_empty();
268 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
269 }
270
271 fn descend_to_last_item(
272 &mut self,
273 mut subtree: &'a SumTree<T>,
274 cx: &<T::Summary as Summary>::Context,
275 ) {
276 self.did_seek = true;
277 loop {
278 match subtree.0.as_ref() {
279 Node::Internal {
280 child_trees,
281 child_summaries,
282 ..
283 } => {
284 for summary in &child_summaries[0..child_summaries.len() - 1] {
285 self.position.add_summary(summary, cx);
286 }
287
288 self.stack.push(StackEntry {
289 tree: subtree,
290 index: child_trees.len() - 1,
291 position: self.position.clone(),
292 });
293 subtree = child_trees.last().unwrap();
294 }
295 Node::Leaf { item_summaries, .. } => {
296 let last_index = item_summaries.len().saturating_sub(1);
297 for item_summary in &item_summaries[0..last_index] {
298 self.position.add_summary(item_summary, cx);
299 }
300 self.stack.push(StackEntry {
301 tree: subtree,
302 index: last_index,
303 position: self.position.clone(),
304 });
305 break;
306 }
307 }
308 }
309 }
310}
311
312impl<'a, T, D> Cursor<'a, T, D>
313where
314 T: Item,
315 D: Dimension<'a, T::Summary>,
316{
317 pub fn seek<Target>(
318 &mut self,
319 pos: &Target,
320 bias: Bias,
321 cx: &<T::Summary as Summary>::Context,
322 ) -> bool
323 where
324 Target: SeekTarget<'a, T::Summary, D>,
325 {
326 self.reset();
327 self.seek_internal(pos, bias, &mut (), cx)
328 }
329
330 pub fn seek_forward<Target>(
331 &mut self,
332 pos: &Target,
333 bias: Bias,
334 cx: &<T::Summary as Summary>::Context,
335 ) -> bool
336 where
337 Target: SeekTarget<'a, T::Summary, D>,
338 {
339 self.seek_internal(pos, bias, &mut (), cx)
340 }
341
342 pub fn slice<Target>(
343 &mut self,
344 end: &Target,
345 bias: Bias,
346 cx: &<T::Summary as Summary>::Context,
347 ) -> SumTree<T>
348 where
349 Target: SeekTarget<'a, T::Summary, D>,
350 {
351 let mut slice = SliceSeekAggregate {
352 tree: SumTree::new(),
353 leaf_items: ArrayVec::new(),
354 leaf_item_summaries: ArrayVec::new(),
355 leaf_summary: T::Summary::default(),
356 };
357 self.seek_internal(end, bias, &mut slice, cx);
358 slice.tree
359 }
360
361 pub fn suffix(&mut self, cx: &<T::Summary as Summary>::Context) -> SumTree<T> {
362 self.slice(&End::new(), Bias::Right, cx)
363 }
364
365 pub fn summary<Target, Output>(
366 &mut self,
367 end: &Target,
368 bias: Bias,
369 cx: &<T::Summary as Summary>::Context,
370 ) -> Output
371 where
372 Target: SeekTarget<'a, T::Summary, D>,
373 Output: Dimension<'a, T::Summary>,
374 {
375 let mut summary = SummarySeekAggregate(Output::default());
376 self.seek_internal(end, bias, &mut summary, cx);
377 summary.0
378 }
379
380 fn seek_internal(
381 &mut self,
382 target: &dyn SeekTarget<'a, T::Summary, D>,
383 bias: Bias,
384 aggregate: &mut dyn SeekAggregate<'a, T>,
385 cx: &<T::Summary as Summary>::Context,
386 ) -> bool {
387 debug_assert!(
388 target.cmp(&self.position, cx) >= Ordering::Equal,
389 "cannot seek backward from {:?} to {:?}",
390 self.position,
391 target
392 );
393
394 if !self.did_seek {
395 self.did_seek = true;
396 self.stack.push(StackEntry {
397 tree: self.tree,
398 index: 0,
399 position: Default::default(),
400 });
401 }
402
403 let mut ascending = false;
404 'outer: while let Some(entry) = self.stack.last_mut() {
405 match *entry.tree.0 {
406 Node::Internal {
407 ref child_summaries,
408 ref child_trees,
409 ..
410 } => {
411 if ascending {
412 entry.index += 1;
413 }
414
415 for (child_tree, child_summary) in child_trees[entry.index..]
416 .iter()
417 .zip(&child_summaries[entry.index..])
418 {
419 let mut child_end = self.position.clone();
420 child_end.add_summary(&child_summary, cx);
421
422 let comparison = target.cmp(&child_end, cx);
423 if comparison == Ordering::Greater
424 || (comparison == Ordering::Equal && bias == Bias::Right)
425 {
426 self.position = child_end;
427 aggregate.push_tree(child_tree, child_summary, cx);
428 entry.index += 1;
429 entry.position = self.position.clone();
430 } else {
431 self.stack.push(StackEntry {
432 tree: child_tree,
433 index: 0,
434 position: self.position.clone(),
435 });
436 ascending = false;
437 continue 'outer;
438 }
439 }
440 }
441 Node::Leaf {
442 ref items,
443 ref item_summaries,
444 ..
445 } => {
446 aggregate.begin_leaf();
447
448 for (item, item_summary) in items[entry.index..]
449 .iter()
450 .zip(&item_summaries[entry.index..])
451 {
452 let mut child_end = self.position.clone();
453 child_end.add_summary(item_summary, cx);
454
455 let comparison = target.cmp(&child_end, cx);
456 if comparison == Ordering::Greater
457 || (comparison == Ordering::Equal && bias == Bias::Right)
458 {
459 self.position = child_end;
460 aggregate.push_item(item, item_summary, cx);
461 entry.index += 1;
462 } else {
463 aggregate.end_leaf(cx);
464 break 'outer;
465 }
466 }
467
468 aggregate.end_leaf(cx);
469 }
470 }
471
472 self.stack.pop();
473 ascending = true;
474 }
475
476 self.at_end = self.stack.is_empty();
477 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
478
479 let mut end = self.position.clone();
480 if bias == Bias::Left {
481 if let Some(summary) = self.item_summary() {
482 end.add_summary(summary, cx);
483 }
484 }
485
486 target.cmp(&end, cx) == Ordering::Equal
487 }
488}
489
490impl<'a, T, S, D> Iterator for Cursor<'a, T, D>
491where
492 T: Item<Summary = S>,
493 S: Summary<Context = ()>,
494 D: Dimension<'a, T::Summary>,
495{
496 type Item = &'a T;
497
498 fn next(&mut self) -> Option<Self::Item> {
499 if !self.did_seek {
500 self.next(&());
501 }
502
503 if let Some(item) = self.item() {
504 self.next(&());
505 Some(item)
506 } else {
507 None
508 }
509 }
510}
511
512pub struct FilterCursor<'a, F, T: Item, D> {
513 cursor: Cursor<'a, T, D>,
514 filter_node: F,
515}
516
517impl<'a, F, T, D> FilterCursor<'a, F, T, D>
518where
519 F: FnMut(&T::Summary) -> bool,
520 T: Item,
521 D: Dimension<'a, T::Summary>,
522{
523 pub fn new(
524 tree: &'a SumTree<T>,
525 mut filter_node: F,
526 cx: &<T::Summary as Summary>::Context,
527 ) -> Self {
528 let mut cursor = tree.cursor::<D>();
529 cursor.next_internal(&mut filter_node, cx);
530 Self {
531 cursor,
532 filter_node,
533 }
534 }
535
536 pub fn start(&self) -> &D {
537 self.cursor.start()
538 }
539
540 pub fn item(&self) -> Option<&'a T> {
541 self.cursor.item()
542 }
543
544 pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
545 self.cursor.next_internal(&mut self.filter_node, cx);
546 }
547}
548
549impl<'a, F, T, S, U> Iterator for FilterCursor<'a, F, T, U>
550where
551 F: Fn(&T::Summary) -> bool,
552 T: Item<Summary = S>,
553 S: Summary<Context = ()>,
554 U: Dimension<'a, T::Summary>,
555{
556 type Item = &'a T;
557
558 fn next(&mut self) -> Option<Self::Item> {
559 if let Some(item) = self.item() {
560 self.cursor.next_internal(&self.filter_node, &());
561 Some(item)
562 } else {
563 None
564 }
565 }
566}
567
568trait SeekAggregate<'a, T: Item> {
569 fn begin_leaf(&mut self);
570 fn end_leaf(&mut self, cx: &<T::Summary as Summary>::Context);
571 fn push_item(
572 &mut self,
573 item: &'a T,
574 summary: &'a T::Summary,
575 cx: &<T::Summary as Summary>::Context,
576 );
577 fn push_tree(
578 &mut self,
579 tree: &'a SumTree<T>,
580 summary: &'a T::Summary,
581 cx: &<T::Summary as Summary>::Context,
582 );
583}
584
585struct SliceSeekAggregate<T: Item> {
586 tree: SumTree<T>,
587 leaf_items: ArrayVec<T, { 2 * TREE_BASE }>,
588 leaf_item_summaries: ArrayVec<T::Summary, { 2 * TREE_BASE }>,
589 leaf_summary: T::Summary,
590}
591
592struct SummarySeekAggregate<D>(D);
593
594impl<'a, T: Item> SeekAggregate<'a, T> for () {
595 fn begin_leaf(&mut self) {}
596 fn end_leaf(&mut self, _: &<T::Summary as Summary>::Context) {}
597 fn push_item(&mut self, _: &T, _: &T::Summary, _: &<T::Summary as Summary>::Context) {}
598 fn push_tree(&mut self, _: &SumTree<T>, _: &T::Summary, _: &<T::Summary as Summary>::Context) {}
599}
600
601impl<'a, T: Item> SeekAggregate<'a, T> for SliceSeekAggregate<T> {
602 fn begin_leaf(&mut self) {}
603 fn end_leaf(&mut self, cx: &<T::Summary as Summary>::Context) {
604 self.tree.push_tree(
605 SumTree(Arc::new(Node::Leaf {
606 summary: mem::take(&mut self.leaf_summary),
607 items: mem::take(&mut self.leaf_items),
608 item_summaries: mem::take(&mut self.leaf_item_summaries),
609 })),
610 cx,
611 );
612 }
613 fn push_item(&mut self, item: &T, summary: &T::Summary, cx: &<T::Summary as Summary>::Context) {
614 self.leaf_items.push(item.clone());
615 self.leaf_item_summaries.push(summary.clone());
616 Summary::add_summary(&mut self.leaf_summary, summary, cx);
617 }
618 fn push_tree(
619 &mut self,
620 tree: &SumTree<T>,
621 _: &T::Summary,
622 cx: &<T::Summary as Summary>::Context,
623 ) {
624 self.tree.push_tree(tree.clone(), cx);
625 }
626}
627
628impl<'a, T: Item, D> SeekAggregate<'a, T> for SummarySeekAggregate<D>
629where
630 D: Dimension<'a, T::Summary>,
631{
632 fn begin_leaf(&mut self) {}
633 fn end_leaf(&mut self, _: &<T::Summary as Summary>::Context) {}
634 fn push_item(&mut self, _: &T, summary: &'a T::Summary, cx: &<T::Summary as Summary>::Context) {
635 self.0.add_summary(summary, cx);
636 }
637 fn push_tree(
638 &mut self,
639 _: &SumTree<T>,
640 summary: &'a T::Summary,
641 cx: &<T::Summary as Summary>::Context,
642 ) {
643 self.0.add_summary(summary, cx);
644 }
645}