1use super::*;
2use arrayvec::ArrayVec;
3use std::{cmp::Ordering, 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, filter_node: F, cx: &<T::Summary as Summary>::Context)
188 where
189 F: Fn(&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 SeekAggregate::None, 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 SeekAggregate::None, 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 = SeekAggregate::Slice(SumTree::new());
352 self.seek_internal::<_, ()>(end, bias, &mut slice, cx);
353 if let SeekAggregate::Slice(slice) = slice {
354 slice
355 } else {
356 unreachable!()
357 }
358 }
359
360 pub fn suffix(&mut self, cx: &<T::Summary as Summary>::Context) -> SumTree<T> {
361 let mut slice = SeekAggregate::Slice(SumTree::new());
362 self.seek_internal::<_, ()>(&End::new(), Bias::Right, &mut slice, cx);
363 if let SeekAggregate::Slice(slice) = slice {
364 slice
365 } else {
366 unreachable!()
367 }
368 }
369
370 pub fn summary<Target, Output>(
371 &mut self,
372 end: &Target,
373 bias: Bias,
374 cx: &<T::Summary as Summary>::Context,
375 ) -> Output
376 where
377 Target: SeekTarget<'a, T::Summary, D>,
378 Output: Dimension<'a, T::Summary>,
379 {
380 let mut summary = SeekAggregate::Summary(Output::default());
381 self.seek_internal(end, bias, &mut summary, cx);
382 if let SeekAggregate::Summary(summary) = summary {
383 summary
384 } else {
385 unreachable!()
386 }
387 }
388
389 fn seek_internal<Target, Output>(
390 &mut self,
391 target: &Target,
392 bias: Bias,
393 aggregate: &mut SeekAggregate<T, Output>,
394 cx: &<T::Summary as Summary>::Context,
395 ) -> bool
396 where
397 Target: SeekTarget<'a, T::Summary, D>,
398 Output: Dimension<'a, T::Summary>,
399 {
400 debug_assert!(
401 target.cmp(&self.position, cx) >= Ordering::Equal,
402 "cannot seek backward from {:?} to {:?}",
403 self.position,
404 target
405 );
406
407 if !self.did_seek {
408 self.did_seek = true;
409 self.stack.push(StackEntry {
410 tree: self.tree,
411 index: 0,
412 position: Default::default(),
413 });
414 }
415
416 let mut ascending = false;
417 'outer: while let Some(entry) = self.stack.last_mut() {
418 match *entry.tree.0 {
419 Node::Internal {
420 ref child_summaries,
421 ref child_trees,
422 ..
423 } => {
424 if ascending {
425 entry.index += 1;
426 }
427
428 for (child_tree, child_summary) in child_trees[entry.index..]
429 .iter()
430 .zip(&child_summaries[entry.index..])
431 {
432 let mut child_end = self.position.clone();
433 child_end.add_summary(&child_summary, cx);
434
435 let comparison = target.cmp(&child_end, cx);
436 if comparison == Ordering::Greater
437 || (comparison == Ordering::Equal && bias == Bias::Right)
438 {
439 self.position = child_end;
440 match aggregate {
441 SeekAggregate::None => {}
442 SeekAggregate::Slice(slice) => {
443 slice.push_tree(child_tree.clone(), cx);
444 }
445 SeekAggregate::Summary(summary) => {
446 summary.add_summary(child_summary, cx);
447 }
448 }
449 entry.index += 1;
450 entry.position = self.position.clone();
451 } else {
452 self.stack.push(StackEntry {
453 tree: child_tree,
454 index: 0,
455 position: self.position.clone(),
456 });
457 ascending = false;
458 continue 'outer;
459 }
460 }
461 }
462 Node::Leaf {
463 ref items,
464 ref item_summaries,
465 ..
466 } => {
467 let mut slice_items = ArrayVec::<T, { 2 * TREE_BASE }>::new();
468 let mut slice_item_summaries = ArrayVec::<T::Summary, { 2 * TREE_BASE }>::new();
469 let mut slice_items_summary = match aggregate {
470 SeekAggregate::Slice(_) => Some(T::Summary::default()),
471 _ => None,
472 };
473
474 for (item, item_summary) in items[entry.index..]
475 .iter()
476 .zip(&item_summaries[entry.index..])
477 {
478 let mut child_end = self.position.clone();
479 child_end.add_summary(item_summary, cx);
480
481 let comparison = target.cmp(&child_end, cx);
482 if comparison == Ordering::Greater
483 || (comparison == Ordering::Equal && bias == Bias::Right)
484 {
485 self.position = child_end;
486 match aggregate {
487 SeekAggregate::None => {}
488 SeekAggregate::Slice(_) => {
489 slice_items.push(item.clone());
490 slice_item_summaries.push(item_summary.clone());
491 <T::Summary as Summary>::add_summary(
492 slice_items_summary.as_mut().unwrap(),
493 item_summary,
494 cx,
495 );
496 }
497 SeekAggregate::Summary(summary) => {
498 summary.add_summary(item_summary, cx);
499 }
500 }
501 entry.index += 1;
502 } else {
503 if let SeekAggregate::Slice(slice) = aggregate {
504 slice.push_tree(
505 SumTree(Arc::new(Node::Leaf {
506 summary: slice_items_summary.unwrap(),
507 items: slice_items,
508 item_summaries: slice_item_summaries,
509 })),
510 cx,
511 );
512 }
513 break 'outer;
514 }
515 }
516
517 if let SeekAggregate::Slice(slice) = aggregate {
518 if !slice_items.is_empty() {
519 slice.push_tree(
520 SumTree(Arc::new(Node::Leaf {
521 summary: slice_items_summary.unwrap(),
522 items: slice_items,
523 item_summaries: slice_item_summaries,
524 })),
525 cx,
526 );
527 }
528 }
529 }
530 }
531
532 self.stack.pop();
533 ascending = true;
534 }
535
536 self.at_end = self.stack.is_empty();
537 debug_assert!(self.stack.is_empty() || self.stack.last().unwrap().tree.0.is_leaf());
538
539 let mut end = self.position.clone();
540 if bias == Bias::Left {
541 if let Some(summary) = self.item_summary() {
542 end.add_summary(summary, cx);
543 }
544 }
545
546 target.cmp(&end, cx) == Ordering::Equal
547 }
548}
549
550impl<'a, T, S, D> Iterator for Cursor<'a, T, D>
551where
552 T: Item<Summary = S>,
553 S: Summary<Context = ()>,
554 D: Dimension<'a, T::Summary>,
555{
556 type Item = &'a T;
557
558 fn next(&mut self) -> Option<Self::Item> {
559 if !self.did_seek {
560 self.next(&());
561 }
562
563 if let Some(item) = self.item() {
564 self.next(&());
565 Some(item)
566 } else {
567 None
568 }
569 }
570}
571
572pub struct FilterCursor<'a, F: Fn(&T::Summary) -> bool, T: Item, D> {
573 cursor: Cursor<'a, T, D>,
574 filter_node: F,
575}
576
577impl<'a, F, T, D> FilterCursor<'a, F, T, D>
578where
579 F: Fn(&T::Summary) -> bool,
580 T: Item,
581 D: Dimension<'a, T::Summary>,
582{
583 pub fn new(
584 tree: &'a SumTree<T>,
585 filter_node: F,
586 cx: &<T::Summary as Summary>::Context,
587 ) -> Self {
588 let mut cursor = tree.cursor::<D>();
589 cursor.next_internal(&filter_node, cx);
590 Self {
591 cursor,
592 filter_node,
593 }
594 }
595
596 pub fn start(&self) -> &D {
597 self.cursor.start()
598 }
599
600 pub fn item(&self) -> Option<&'a T> {
601 self.cursor.item()
602 }
603
604 pub fn next(&mut self, cx: &<T::Summary as Summary>::Context) {
605 self.cursor.next_internal(&self.filter_node, cx);
606 }
607}
608
609impl<'a, F, T, S, U> Iterator for FilterCursor<'a, F, T, U>
610where
611 F: Fn(&T::Summary) -> bool,
612 T: Item<Summary = S>,
613 S: Summary<Context = ()>,
614 U: Dimension<'a, T::Summary>,
615{
616 type Item = &'a T;
617
618 fn next(&mut self) -> Option<Self::Item> {
619 if let Some(item) = self.item() {
620 self.cursor.next_internal(&self.filter_node, &());
621 Some(item)
622 } else {
623 None
624 }
625 }
626}
627
628enum SeekAggregate<T: Item, D> {
629 None,
630 Slice(SumTree<T>),
631 Summary(D),
632}