1use std::cmp::Ordering;
2use std::fmt::{Arguments, Debug, Write};
3pub mod range;
4pub use range::TextRange;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
7pub enum Color {
8 Red = 0,
9 Black,
10}
11
12impl Color {
13 pub fn flip(&self) -> Color {
14 match self {
15 Color::Red => Color::Black,
16 Color::Black => Color::Red,
17 }
18 }
19}
20
21#[derive(Clone)]
22pub struct Node<T: Clone> {
23 pub key: TextRange,
24 pub val: T,
25 left: MaybeNode<T>,
26 right: MaybeNode<T>,
27 color: Color,
28 is_right_child: bool,
29 n: usize,
30}
31pub type BoxedNode<T> = Box<Node<T>>;
32pub type MaybeNode<T> = Option<BoxedNode<T>>;
33
34impl<T: Clone> Node<T> {
35 #[inline]
36 pub fn red(node: &MaybeNode<T>) -> bool {
37 node.as_ref().is_some_and(|n| n.color == Color::Red)
38 }
39
40 #[inline]
41 pub fn new(key: TextRange, val: T, is_right_child: bool) -> Node<T> {
42 Self { key, val, left: None, right: None, color: Color::Red, n: 1, is_right_child }
43 }
44
45 #[inline]
46 pub fn new_boxed(key: TextRange, val: T, is_right_child: bool) -> BoxedNode<T> {
47 Box::new(Self::new(key, val, is_right_child))
48 }
49
50 pub fn rotate_left(node: &mut BoxedNode<T>) -> Option<&mut BoxedNode<T>> {
58 let mut x = node.right.take()?;
59 let mut c = x.left.take();
60 if let Some(ref mut c) = c {
61 c.is_right_child = true;
62 }
63 node.right = c;
64 x.color = node.color;
65 x.is_right_child = node.is_right_child;
66 node.is_right_child = false;
67 x.n = node.n;
68 node.color = Color::Red;
69 node.n = node.n();
70 std::mem::swap(node, &mut x);
72 node.left.replace(x);
73 Some(node)
74 }
75
76 pub fn rotate_right(node: &mut BoxedNode<T>) -> Option<&mut BoxedNode<T>> {
84 let mut x = node.left.take()?;
85 let mut c = x.right.take();
86 if let Some(ref mut c) = c {
87 c.is_right_child = false;
88 }
89 node.left = c;
90 x.color = node.color;
91 x.is_right_child = node.is_right_child;
92 node.is_right_child = true;
93 node.color = Color::Red;
94 x.n = node.n;
95 node.n = node.n();
96
97 std::mem::swap(node, &mut x);
98 node.right.replace(x);
99 Some(node)
100 }
101
102 #[inline]
104 pub fn n(&self) -> usize {
105 let l = match self.left {
106 Some(ref n) => n.n,
107 None => 0,
108 };
109 let r = match self.right {
110 Some(ref n) => n.n,
111 None => 0,
112 };
113 1 + l + r
114 }
115
116 fn min(&self) -> &Node<T> {
117 if let Some(ref l) = self.left { l.min() } else { self }
118 }
119
120 pub fn get_node(&self, key: TextRange) -> Option<&Node<T>> {
121 match key.cmp(&self.key) {
122 Ordering::Equal => Some(self),
123 Ordering::Less => self.left.as_ref().and_then(|n| n.get_node(key)),
124 Ordering::Greater => self.right.as_ref().and_then(|n| n.get_node(key)),
125 }
126 }
127
128 fn get_node_mut(&mut self, key: TextRange) -> Option<&mut Node<T>> {
129 match key.cmp(&self.key) {
130 Ordering::Equal => Some(self),
131 Ordering::Less => self.left.as_mut().and_then(|n| n.get_node_mut(key)),
132 Ordering::Greater => self.right.as_mut().and_then(|n| n.get_node_mut(key)),
133 }
134 }
135
136 fn get(&self, key: TextRange) -> Option<T> {
137 match key.cmp(&self.key) {
138 Ordering::Equal => Some(self.val.clone()),
139 Ordering::Less => self.left.as_ref().and_then(|n| n.get(key)),
140 Ordering::Greater => self.right.as_ref().and_then(|n| n.get(key)),
141 }
142 }
143
144 pub fn insert_at<'a, F: Fn(T, T) -> T>(
149 node: &'a mut MaybeNode<T>,
150 key: TextRange,
151 val: T,
152 is_right_child: bool,
153 merge_fn: &F,
154 ) -> Option<&'a mut BoxedNode<T>> {
155 if key.start == key.end {
156 return None;
157 }
158 match node {
159 Some(n) => Node::insert_at_inner(n, key, val, merge_fn),
160 None => {
161 node.replace(Node::new_boxed(key, val.clone(), is_right_child));
162 node.as_mut()
163 }
164 }
165 }
166
167 fn insert_at_inner<'a, F: Fn(T, T) -> T>(
168 node: &'a mut BoxedNode<T>,
169 mut key: TextRange,
170 val: T,
171 merge_fn: &F,
172 ) -> Option<&'a mut BoxedNode<T>> {
173 let intersect = key.intersects(node.key);
174 if intersect {
176 if key.start < node.key.start {
177 let key_left = key.split_at(node.key.start, true);
178 Node::insert_at(&mut node.left, key_left, val.clone(), false, merge_fn);
179
180 if key.end < node.key.end {
181 let key_right = node.key.split_at(key.end, false);
182 Node::insert_at(&mut node.right, key_right, node.val.clone(), true, merge_fn);
183 } else if key.end > node.key.end {
184 let key_right = key.split_at(node.key.end, false);
185 Node::insert_at(&mut node.right, key_right, val.clone(), true, merge_fn);
186 }
187 } else {
188 let key_left = node.key.split_at(key.start, true);
189 Node::insert_at(&mut node.left, key_left, node.val.clone(), false, merge_fn);
190
191 if key.end < node.key.end {
192 let key_right = node.key.split_at(key.end, false);
193 Node::insert_at(&mut node.right, key_right, node.val.clone(), true, merge_fn);
194 } else if key.end > node.key.end {
195 let key_right = key.split_at(node.key.end, false);
196 Node::insert_at(&mut node.right, key_right, val.clone(), true, merge_fn);
197 }
198 }
199 }
200 if key.start == key.end {
201 return None;
202 }
203 let cmp = key.cmp(&node.key);
204 match cmp {
205 Ordering::Less => {
206 Node::insert_at(&mut node.left, key, val, false, merge_fn)?;
207 }
208 Ordering::Equal => {
209 node.val = merge_fn(val, node.val.clone());
210 }
211 Ordering::Greater => {
212 Node::insert_at(&mut node.right, key, val, true, merge_fn)?;
213 }
214 }
215
216 if Node::red(&node.right) && !Node::red(&node.left) {
218 Node::rotate_left(node)?;
219 }
220
221 let cond2 = match node.left {
223 Some(ref nl) => nl.color == Color::Red && Node::red(&nl.left),
224 None => false,
225 };
226 if cond2 {
227 Node::rotate_right(node)?;
228 }
229
230 if let (Some(l), Some(r)) = (&mut node.left, &mut node.right) {
232 if l.color == Color::Red && r.color == Color::Red {
233 l.color = Color::Black;
234 r.color = Color::Black;
235 node.color = Color::Red;
236 }
237 }
238 node.n = node.n();
240 Some(node)
241 }
242
243 fn flip_colors(n: &mut BoxedNode<T>) {
249 if let Some(ref mut l) = n.left {
250 l.color = l.color.flip();
251 }
252 if let Some(ref mut r) = n.right {
253 r.color = r.color.flip();
254 }
255 n.color = n.color.flip();
256 }
257
258 fn balance(node: &mut BoxedNode<T>) -> Option<()> {
260 if Node::red(&node.right) {
261 Node::rotate_left(node)?;
262 }
263 Some(())
264 }
265
266 fn move_red_left(node: &mut BoxedNode<T>) -> Option<()> {
269 Node::flip_colors(node);
270 if let Some(ref mut nr) = node.right.as_mut() {
272 if Node::red(&nr.left) {
273 Node::rotate_right(nr)?;
274 Node::rotate_left(node)?;
275 Node::flip_colors(node);
276 }
277 }
278 Some(())
279 }
280
281 fn move_red_right(node: &mut BoxedNode<T>) -> Option<()> {
282 Node::flip_colors(node);
283 let cond = match node.left {
285 Some(ref l) => Node::red(&l.left),
286 None => false,
287 };
288 if cond {
289 Node::rotate_right(node)?;
290 Node::flip_colors(node);
291 }
292 Some(())
293 }
294
295 fn delete_min(node: &mut MaybeNode<T>) -> MaybeNode<T> {
296 let n = node.as_mut()?;
297 match n.left {
298 Some(ref mut l) => {
299 if l.color == Color::Black && !Node::red(&l.left) {
300 Node::move_red_left(n)?;
301 }
302 }
303 None => {
304 return node.take();
305 }
306 }
307 let result = Node::delete_min(&mut n.left);
308 Node::balance(n)?;
309 result
310 }
311
312 fn delete_max(node: &mut MaybeNode<T>) -> MaybeNode<T> {
313 let n = node.as_mut()?;
314 if Node::red(&n.left) {
315 Node::rotate_right(n);
316 }
317 let n = node.as_mut()?;
318 match n.right {
319 Some(ref mut r) => {
320 if r.color == Color::Black && !Node::red(&r.left) {
321 Node::move_red_right(n)?;
322 }
323 }
324 None => {
325 return node.take();
326 }
327 }
328 let result = Node::delete_max(&mut n.right);
329 Node::balance(n)?;
330 result
331 }
332
333 fn delete(node: &mut MaybeNode<T>, key: TextRange) -> MaybeNode<T> {
334 let n = node.as_mut()?;
335 let result = if key < n.key {
336 if let Some(ref mut l) = n.left {
338 if l.color == Color::Black && !Node::red(&l.left) {
339 Node::move_red_left(n).unwrap();
340 }
341 }
342 Node::delete(&mut n.left, key)
343 } else {
344 if Node::red(&n.left) {
345 Node::rotate_right(n).unwrap();
346 }
347 if key == n.key && n.right.is_none() {
348 return node.take();
349 }
351
352 let cond = if let Some(ref r) = n.right {
353 r.color == Color::Black && !Node::red(&r.left)
354 } else {
355 true
356 };
357
358 if cond {
359 Node::move_red_right(n).unwrap();
360 }
361
362 if key == n.key {
363 let mut result = Node::delete_min(&mut n.right);
364 let r_min = result.as_mut().unwrap();
365 r_min.left = n.left.take();
366 r_min.right = n.right.take();
367 r_min.color = n.color;
368 r_min.is_right_child = n.is_right_child;
369 r_min.n = n.n;
370 std::mem::swap(r_min, n);
371 result
372 } else {
373 Node::delete(&mut n.right, key)
374 }
375 };
376
377 Node::balance(n)?;
378 result
379 }
380
381 pub fn find_intersects<'a>(&'a self, range: TextRange, results: &mut Vec<&'a Node<T>>) {
386 let ord = range.strict_order(&self.key);
387 match ord {
388 Some(Ordering::Less) => {
389 if let Some(ref l) = self.left {
390 l.find_intersects(range, results);
391 }
392 }
393 Some(Ordering::Greater) => {
394 if let Some(ref r) = self.right {
395 r.find_intersects(range, results);
396 }
397 }
398 _ => {
399 if let Some(ref l) = self.left {
400 l.find_intersects(range, results);
401 }
402 results.push(self);
403 if let Some(ref r) = self.right {
404 r.find_intersects(range, results);
405 }
406 }
407 }
408 }
409
410 fn find_intersect_min(&self, range: TextRange) -> Option<&Node<T>> {
411 if range.start == range.end {
412 return None;
413 }
414 let ord = range.strict_order(&self.key);
415 match ord {
416 Some(Ordering::Less) => self.left.as_ref().and_then(|l| l.find_intersect_min(range)),
417 Some(Ordering::Equal) => Some(self),
418 Some(Ordering::Greater) => {
419 self.right.as_ref().and_then(|r| r.find_intersect_min(range))
420 }
421 _ => self.left.as_ref().and_then(|l| l.find_intersect_min(range)).or(Some(self)),
422 }
423 }
424
425 fn find_intersect_max(&self, range: TextRange) -> Option<&Node<T>> {
426 if range.start == range.end {
427 return None;
428 }
429 let ord = range.strict_order(&self.key);
430 match ord {
431 Some(Ordering::Less) => self.left.as_ref().and_then(|l| l.find_intersect_max(range)),
432 Some(Ordering::Equal) => Some(self),
433 Some(Ordering::Greater) => {
434 self.right.as_ref().and_then(|l| l.find_intersect_max(range))
435 }
436 _ => self.right.as_ref().and_then(|l| l.find_intersect_max(range)).or(Some(self)),
437 }
438 }
439
440 fn apply<F>(&self, f: &mut F)
443 where
444 F: FnMut(&Node<T>),
445 {
446 if let Some(ref l) = self.left {
447 l.apply(f);
448 }
449 f(self);
450 if let Some(ref r) = self.right {
451 r.apply(f);
452 }
453 }
454
455 fn apply_mut<F>(&mut self, f: &mut F)
458 where
459 F: FnMut(&mut Node<T>),
460 {
461 if let Some(ref mut l) = self.left {
462 l.apply_mut(f);
463 }
464 f(self);
465 if let Some(ref mut r) = self.right {
466 r.apply_mut(f);
467 }
468 }
469
470 pub fn advance(&mut self, position: usize, length: usize) {
471 let cmp = self.key.start > position;
472 if cmp {
473 if let Some(ref mut l) = self.left {
474 l.advance(position, length);
475 }
476 self.key.advance(length);
477 if let Some(ref mut r) = self.right {
478 r.apply_mut(&mut |n| n.key.advance(length));
479 }
480 } else {
481 if self.key.end > position {
482 self.key.end += length;
484 }
485 if let Some(ref mut l) = self.left {
486 l.advance(position, length);
487 }
488 if let Some(ref mut r) = self.right {
489 r.advance(position, length);
490 }
491 }
492 }
493}
494
495#[derive(Default, Clone)]
496pub struct IntervalTree<T: Clone> {
507 root: MaybeNode<T>,
508}
509
510impl<T: Clone> IntervalTree<T> {
511 pub fn new() -> Self {
513 Self { root: None }
514 }
515
516 pub fn size(&self) -> usize {
517 self.root.as_ref().map_or(0, |n| n.n)
518 }
519
520 pub fn insert<F: Fn(T, T) -> T>(
537 &mut self,
538 key: impl Into<TextRange>,
539 val: T,
540 merge_fn: F,
541 ) -> Option<&mut Box<Node<T>>> {
542 let key = key.into();
543 if key.start == key.end {
544 return None;
545 }
546 let mut result = Node::insert_at(&mut self.root, key, val, false, &merge_fn);
547 result.as_mut().unwrap().color = Color::Black;
548 result
549 }
550
551 pub fn get(&self, key: impl Into<TextRange>) -> Option<T> {
561 match self.root {
562 Some(ref r) => r.get(key.into()),
563 None => None,
564 }
565 }
566
567 pub fn get_node_mut(&mut self, key: impl Into<TextRange>) -> Option<&mut Node<T>> {
568 match self.root {
569 Some(ref mut r) => r.get_node_mut(key.into()),
570 None => None,
571 }
572 }
573
574 pub fn delete_exact(&mut self, key: impl Into<TextRange>) -> MaybeNode<T> {
582 let key = key.into();
583 let result = match self.root {
584 Some(ref mut root) => {
585 if !Node::red(&root.left) && !Node::red(&root.right) {
586 root.color = Color::Red;
587 }
588
589 Node::delete(&mut self.root, key)
590 }
591 None => None,
592 };
593 if let Some(ref mut root) = self.root {
594 root.color = Color::Black;
595 }
596 result
597 }
598
599 pub fn delete_min(&mut self) -> MaybeNode<T> {
610 let root = self.root.as_mut()?;
611 if !Node::red(&root.left) && !Node::red(&root.right) {
612 root.color = Color::Red;
613 }
614 let result = Node::delete_min(&mut self.root);
615 if let Some(ref mut root) = self.root {
616 root.color = Color::Black;
617 }
618 result
619 }
620
621 pub fn delete_max(&mut self) -> MaybeNode<T> {
622 let root = self.root.as_mut()?;
623 if !Node::red(&root.left) && !Node::red(&root.right) {
624 root.color = Color::Red;
625 }
626 let result = Node::delete_max(&mut self.root);
627 if let Some(ref mut root) = self.root {
628 root.color = Color::Black;
629 }
630 result
631 }
632
633 pub fn delete(&mut self, range: impl Into<TextRange>, del_extend: bool) {
665 let range: TextRange = range.into();
666 for key in self.find_intersects(range).map(|n| n.key).collect::<Vec<_>>() {
667 if del_extend {
668 self.delete_exact(key);
669 continue;
670 }
671 if key.start >= range.start && key.end <= range.end {
674 self.delete_exact(key);
675 }
676
677 if key.start < range.start {
680 let n = self.get_node_mut(key).unwrap();
681 n.key.end = range.start;
682 if key.end > range.end {
683 let val = n.val.clone();
684 let f = |_, _| unreachable!(); self.insert(TextRange::new(range.start, key.end), val, f);
686 }
687 }
688
689 if key.end > range.end {
692 let n = self.get_node_mut(key).unwrap();
693 n.key.start = range.end;
694 }
695 }
697 }
698
699 pub fn advance(&mut self, position: usize, length: usize) {
703 if let Some(ref mut node) = self.root {
704 node.advance(position, length);
705 }
706 }
707
708 pub fn find(&self, position: usize) -> Option<&Node<T>> {
711 let range = TextRange::new(position, position + 1);
712 self.find_intersects(range).next()
713 }
714
715 pub fn find_intersects(&self, range: impl Into<TextRange>) -> impl Iterator<Item = &Node<T>> {
718 let range = range.into();
719 let node = self.find_intersect_min(range);
720 let key = node.map(|n| n.key);
721
722 StackIterator::new(self, key, false).take_while(move |n| n.key.intersects(range))
723 }
734
735 pub fn find_intersect_min(&self, range: impl Into<TextRange>) -> Option<&Node<T>> {
763 self.root.as_ref().and_then(|r| r.find_intersect_min(range.into()))
764 }
765
766 pub fn find_intersect_max(&self, range: impl Into<TextRange>) -> Option<&Node<T>> {
768 self.root.as_ref().and_then(|r| r.find_intersect_max(range.into()))
769 }
770
771 pub fn min(&self) -> Option<&Node<T>> {
773 self.root.as_ref().map(|n| n.min())
774 }
775
776 pub fn clean<F: Fn(&T, &T) -> bool, G: Fn(&T) -> bool>(&mut self, eq: F, empty: G) {
806 let min = self.min().map(|n| n.key);
808 self.clean_from_node(min, eq, empty);
809 }
810
811 pub fn clean_from<F: Fn(&T, &T) -> bool, G: Fn(&T) -> bool>(
812 &mut self,
813 start: TextRange,
814 eq: F,
815 empty: G,
816 ) {
817 let start_key = self.find_intersect_min(start).map(|n| n.key);
818 self.clean_from_node(start_key, eq, empty);
819 }
820
821 fn clean_from_node<F: Fn(&T, &T) -> bool, G: Fn(&T) -> bool>(
822 &mut self,
823 start_key: Option<TextRange>,
824 eq: F,
825 empty: G,
826 ) {
827 let mut operations: Vec<Operation<T>> = Vec::new();
831 let iter = StackIterator::new(self, start_key, false);
832
833 let mut prev_node: Option<&Node<T>> = None;
834 let mut current_merge: Option<(TextRange, Vec<TextRange>)> = None;
835
836 for node in iter {
837 if node.key.empty() || empty(&node.val) {
839 operations.push(Operation::Delete(node.key));
840 continue;
841 }
842
843 if let Some(prev) = prev_node {
845 if prev.key.end == node.key.start && eq(&prev.val, &node.val) {
846 match &mut current_merge {
848 Some((_start_key, keys)) => {
849 keys.push(node.key);
850 }
851 None => {
852 current_merge = Some((prev.key, vec![node.key]));
853 }
854 }
855 } else if let Some((start_key, keys)) = current_merge.take() {
856 operations.push(Operation::Merge(start_key, keys));
858 }
859 }
860
861 prev_node = Some(node);
862 }
863
864 if let Some((start_key, keys)) = current_merge {
866 operations.push(Operation::Merge(start_key, keys));
867 }
868
869 self.resolve_ops(operations, &|_, _| unreachable!());
870 }
871
872 pub fn merge<F: Fn(&T, &T) -> bool>(&mut self, eq: F) {
885 self.clean(eq, |_| false);
886 }
887
888 pub fn apply<F: FnMut(&T)>(&self, f: &mut F) {
889 if let Some(r) = self.root.as_ref() {
890 r.apply(&mut |n: &Node<T>| f(&n.val));
891 }
892 }
893
894 pub fn apply_mut<F: FnMut(&mut Node<T>)>(&mut self, f: &mut F) {
895 if let Some(r) = self.root.as_mut() {
896 r.apply_mut(&mut |n| f(n));
897 }
898 }
899
900 pub fn apply_with_split<F: Fn(T) -> Option<T>>(&mut self, f: F, range: impl Into<TextRange>) {
930 let range = range.into();
931 let merge = |new_val, _old_val| new_val;
932
933 let mut operations = Vec::new();
935 let intersects = self.find_intersects(range);
936
937 for node in intersects {
938 if let Some(overlap) = node.key.intersection(range) {
939 if node.key.start < overlap.start {
941 operations.push(Operation::Insert(
942 TextRange::new(node.key.start, overlap.start),
943 node.val.clone(),
944 ));
945 }
946
947 if node.key.end > overlap.end {
949 operations.push(Operation::Insert(
950 TextRange::new(overlap.end, node.key.end),
951 node.val.clone(),
952 ));
953 }
954
955 if let Some(new_val) = f(node.val.clone()) {
957 operations.push(Operation::Insert(overlap, new_val));
958 } else {
959 operations.push(Operation::Delete(overlap));
960 }
961 }
962 }
963
964 self.resolve_ops(operations, &merge);
966 }
967
968 fn resolve_ops<F: Fn(T, T) -> T>(&mut self, operations: Vec<Operation<T>>, merge_fn: &F) {
969 for op in operations {
970 match op {
971 Operation::Insert(key, val) => {
972 self.insert(key, val, merge_fn);
973 }
974 Operation::Delete(key) => {
975 self.delete(key, false);
977 }
978 Operation::Merge(start_key, keys) => {
979 if let Some(node) = self.get_node_mut(start_key) {
980 let last_key = *keys.last().unwrap();
982 node.key.end = last_key.end;
983 for key in keys {
984 self.delete_exact(key);
985 }
986 }
987 }
988 }
989 }
990 }
991}
992
993impl<T: Clone + Debug> Node<T> {
996 fn print_inner(&self, f: &mut std::fmt::Formatter, level: usize) -> std::fmt::Result {
997 write_fmt_with_level(
998 f,
999 level,
1000 format_args!("[key: {:?}, val: {:?}, color: {:?}]\n", self.key, self.val, self.color),
1001 )?;
1002 f.write_char('\n')?;
1003 if let Some(ref l) = self.left {
1004 write_fmt_with_level(f, level, format_args!("left: \n"))?;
1005 l.print_inner(f, level + 1)?;
1006 write_fmt_with_level(f, level, format_args!("left end for {:?}\n", self.key))?;
1007 } else {
1008 write_fmt_with_level(f, level, format_args!("no left child \n"))?;
1009 }
1010 if let Some(ref r) = self.right {
1011 write_fmt_with_level(f, level, format_args!("right: \n"))?;
1012 r.print_inner(f, level + 1)?;
1013 write_fmt_with_level(f, level, format_args!("right end for {:?}\n", self.key))?;
1014 } else {
1015 write_fmt_with_level(f, level, format_args!("no right child \n"))?;
1016 }
1017 Ok(())
1018 }
1019}
1020
1021impl<T: Clone + Debug> IntervalTree<T> {
1022 pub fn print(&self) {
1025 println!("{self:?}");
1026 }
1027}
1028
1029fn write_fmt_with_level(
1030 f: &mut std::fmt::Formatter,
1031 level: usize,
1032 fmt: Arguments<'_>,
1033) -> std::fmt::Result {
1034 for _ in 0..level {
1035 f.write_char('\t')?;
1036 }
1037 f.write_fmt(fmt)
1038}
1039
1040impl<T: Clone + Debug> Debug for Node<T> {
1041 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1042 f.write_str("Node:\n")?;
1043 self.print_inner(f, 0)
1044 }
1045}
1046
1047impl<T: Clone + Debug> Debug for IntervalTree<T> {
1048 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1049 f.write_str("Interval Tree:\n")?;
1050 if let Some(root) = self.root.as_ref() {
1051 root.print_inner(f, 0)?;
1052 }
1053 Ok(())
1054 }
1055}
1056
1057#[derive(Debug)]
1058enum Operation<T> {
1059 Delete(TextRange),
1060 Merge(TextRange, Vec<TextRange>), Insert(TextRange, T),
1062}
1063
1064pub struct StackIterator<'tree, T: Clone> {
1065 stack: Vec<&'tree Node<T>>,
1066 reverse_order: bool,
1067}
1068
1069impl<'tree, T: Clone> StackIterator<'tree, T> {
1070 pub fn new(tree: &'tree IntervalTree<T>, key: Option<TextRange>, reverse_order: bool) -> Self {
1071 let Some(key) = key else { return Self { stack: Vec::new(), reverse_order } };
1072 let mut stack = Vec::new();
1073 let mut current = tree.root.as_ref();
1074
1075 while let Some(node) = current {
1077 let strict_order = key.strict_order(&node.key);
1078 let push_to_stack = strict_order.is_none()
1079 || (strict_order == Some(Ordering::Less) && !reverse_order)
1080 || (strict_order == Some(Ordering::Greater) && reverse_order);
1081 if push_to_stack {
1082 stack.push(node.as_ref());
1083 }
1084 current = match key.cmp(&node.key) {
1085 Ordering::Less => node.left.as_ref(),
1086 Ordering::Greater => node.right.as_ref(),
1087 Ordering::Equal => None,
1088 };
1089 }
1090
1091 Self { stack, reverse_order }
1092 }
1093}
1094
1095impl<'tree, T: Clone> Iterator for StackIterator<'tree, T> {
1096 type Item = &'tree Node<T>;
1097
1098 fn next(&mut self) -> Option<Self::Item> {
1099 let node = self.stack.pop()?;
1101
1102 let next_branch = if self.reverse_order { node.left.as_ref() } else { node.right.as_ref() };
1104 if let Some(mut current) = next_branch {
1105 loop {
1107 self.stack.push(current);
1108 let next_branch =
1109 if self.reverse_order { current.right.as_ref() } else { current.left.as_ref() };
1110 current = match next_branch {
1111 Some(n) => n,
1112 None => break,
1113 };
1114 }
1115 }
1116
1117 Some(node)
1118 }
1119}
1120
1121#[cfg(test)]
1122mod tests {
1123 use std::ptr;
1124
1125 use super::*;
1126
1127 fn merge<T: Clone + Debug>(a: T, _b: T) -> T {
1128 a
1129 }
1130
1131 #[test]
1132 fn test_stack_iterator_empty_tree() {
1133 let tree: IntervalTree<i32> = IntervalTree::new();
1134 let mut iter = StackIterator::new(&tree, None, false);
1135 assert_eq!(iter.next().map(|n| n.val), None);
1136 }
1137
1138 #[test]
1139 fn test_stack_iterator_single_node() {
1140 let mut tree = IntervalTree::new();
1141 tree.insert(TextRange::new(0, 1), 1, merge);
1142
1143 let min_key = tree.min().map(|n| n.key);
1144 let mut iter = StackIterator::new(&tree, min_key, false);
1145 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(0, 1)));
1146 assert_eq!(iter.next().map(|n| n.key), None);
1147 }
1148
1149 #[test]
1150 fn test_stack_iterator_multiple_nodes() {
1151 let mut tree = IntervalTree::new();
1152 tree.insert(TextRange::new(2, 3), 2, merge);
1153 tree.insert(TextRange::new(1, 2), 1, merge);
1154 tree.insert(TextRange::new(3, 4), 3, merge);
1155
1156 let min_key = tree.min().map(|n| n.key);
1157 let mut iter = StackIterator::new(&tree, min_key, false);
1158 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(1, 2)));
1159 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(2, 3)));
1160 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(3, 4)));
1161 assert_eq!(iter.next().map(|n| n.key), None);
1162 }
1163
1164 #[test]
1165 fn test_stack_iterator_complex_tree() {
1166 let mut tree = IntervalTree::new();
1167 tree.insert(TextRange::new(4, 5), 4, merge);
1174 tree.insert(TextRange::new(2, 3), 2, merge);
1175 tree.insert(TextRange::new(6, 7), 6, merge);
1176 tree.insert(TextRange::new(1, 2), 1, merge);
1177 tree.insert(TextRange::new(3, 4), 3, merge);
1178 tree.insert(TextRange::new(5, 6), 5, merge);
1179 tree.insert(TextRange::new(7, 8), 7, merge);
1180
1181 let min_key = tree.min().map(|n| n.key);
1182 let mut iter = StackIterator::new(&tree, min_key, false);
1183 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(1, 2)));
1184 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(2, 3)));
1185 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(3, 4)));
1186 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(4, 5)));
1187 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(5, 6)));
1188 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(6, 7)));
1189 assert_eq!(iter.next().map(|n| n.key), Some(TextRange::new(7, 8)));
1190 assert_eq!(iter.next().map(|n| n.key), None);
1191 }
1192
1193 fn build_tree<T: Clone + Debug>(val: T) -> IntervalTree<T> {
1194 let mut tree = IntervalTree::new();
1195 tree.insert(TextRange::new(0, 1), val.clone(), merge);
1196 tree.insert(TextRange::new(1, 2), val.clone(), merge);
1197 tree.insert(TextRange::new(2, 3), val.clone(), merge);
1198 tree.insert(TextRange::new(3, 4), val.clone(), merge);
1199 tree.insert(TextRange::new(4, 5), val.clone(), merge);
1200 tree.insert(TextRange::new(5, 6), val.clone(), merge);
1201 tree.insert(TextRange::new(6, 7), val.clone(), merge);
1202 tree.insert(TextRange::new(7, 8), val.clone(), merge);
1203 tree.insert(TextRange::new(8, 9), val.clone(), merge);
1204 tree.insert(TextRange::new(9, 10), val.clone(), merge);
1205 tree.print();
1206 tree
1207 }
1208
1209 #[test]
1210 fn rotate_left() {
1211 let val = 1;
1212 let mut node1 = Node::new_boxed(TextRange::new(0, 3), val, false);
1213 node1.color = Color::Black;
1214 let mut node2 = Node::new_boxed(TextRange::new(3, 6), val, false);
1215 let mut node3 = Node::new_boxed(TextRange::new(6, 9), val, false);
1216 node3.color = Color::Black;
1217 let mut node4 = Node::new_boxed(TextRange::new(9, 12), val, false);
1218 node4.color = Color::Black;
1219 let mut node5 = Node::new_boxed(TextRange::new(12, 15), val, false);
1220 node5.color = Color::Black;
1221
1222 node2.left = Some(node3);
1223 node2.right = Some(node4);
1224 node2.color = Color::Red;
1225 node1.right = Some(node2);
1226 node1.left = Some(node5);
1227 Node::rotate_left(&mut node1);
1228 assert_eq!(node1.key.start, 3);
1229 let n2 = node1.left.unwrap();
1230 assert_eq!(n2.key.start, 0);
1231 let n3 = n2.right.unwrap();
1232 assert_eq!(n3.key.start, 6);
1233 let n4 = node1.right.unwrap();
1234 assert_eq!(n4.key.start, 9);
1235 let n5 = n2.left.unwrap();
1236 assert_eq!(n5.key.start, 12);
1237
1238 assert_eq!(node1.color, Color::Black);
1239 assert_eq!(n2.color, Color::Red);
1240 assert_eq!(n2.n, 3);
1241 }
1242
1243 #[test]
1244 fn insert() {
1245 let val = 1;
1246 let tree = build_tree(val);
1247 let root = tree.root.as_ref().unwrap();
1248 assert_eq!(root.key.start, 3);
1249 let n1 = root.left.as_ref().unwrap();
1250 assert_eq!(n1.key.start, 1);
1251 let n2 = root.right.as_ref().unwrap();
1252 assert_eq!(n2.key.start, 7);
1253 let n3 = n2.right.as_ref().unwrap();
1254 assert_eq!(n3.key.start, 9);
1255 let n4 = n3.left.as_ref().unwrap();
1256 assert_eq!(n4.key.start, 8);
1257 assert!(n3.right.is_none());
1258 }
1259
1260 #[test]
1261 fn delete() {
1262 let val = 1;
1263 let mut tree = build_tree(val);
1264 for k in [8, 4, 5, 7, 3, 6] {
1266 let i = TextRange::new(k, k + 1);
1267 let a = tree.delete_exact(i).unwrap();
1268 assert_eq!(a.key, i);
1269 }
1270 }
1271
1272 #[test]
1273 fn delete_ptr() {
1274 let val = 1;
1275 let mut tree = build_tree(val);
1276 let a = 3;
1277 let key = TextRange::new(a, a + 1);
1278 let n: *mut Node<i32> = tree.get_node_mut(key).unwrap();
1279 let mut a = tree.delete_exact(key).unwrap();
1280 assert_eq!(n, ptr::from_mut(&mut *a));
1281 }
1282
1283 #[test]
1284 fn delete_min() {
1285 let val = 1;
1286 let mut tree = build_tree(val);
1287 let a = tree.delete_min().unwrap();
1289 assert_eq!(a.key.start, 0);
1290 }
1291
1292 #[test]
1293 fn test_find_intersects() {
1294 let mut tree = IntervalTree::new();
1295 tree.insert(TextRange::new(0, 5), 1, merge);
1296 tree.insert(TextRange::new(5, 10), 2, merge);
1297 tree.insert(TextRange::new(10, 15), 3, merge);
1298 tree.insert(TextRange::new(15, 20), 4, merge);
1299
1300 let results = tree.find_intersects(TextRange::new(5, 10)).collect::<Vec<_>>();
1302 assert_eq!(results.len(), 1);
1303 assert_eq!(results[0].key, TextRange::new(5, 10));
1304 assert_eq!(results[0].val, 2);
1305
1306 let results = tree.find_intersects(TextRange::new(3, 7)).collect::<Vec<_>>();
1308 assert_eq!(results.len(), 2);
1309 assert_eq!(results[0].key, TextRange::new(0, 5));
1310 assert_eq!(results[1].key, TextRange::new(5, 10));
1311
1312 let results = tree.find_intersects(TextRange::new(12, 18)).collect::<Vec<_>>();
1314 assert_eq!(results.len(), 2);
1315 assert_eq!(results[0].key, TextRange::new(10, 15));
1316 assert_eq!(results[1].key, TextRange::new(15, 20));
1317
1318 let results = tree.find_intersects(TextRange::new(8, 17)).collect::<Vec<_>>();
1320 assert_eq!(results.len(), 3);
1321 assert_eq!(results[0].key, TextRange::new(5, 10));
1322 assert_eq!(results[1].key, TextRange::new(10, 15));
1323 assert_eq!(results[2].key, TextRange::new(15, 20));
1324
1325 let results = tree.find_intersects(TextRange::new(6, 8)).collect::<Vec<_>>();
1327 assert_eq!(results.len(), 1);
1328 assert_eq!(results[0].key, TextRange::new(5, 10));
1329
1330 let results = tree.find_intersects(TextRange::new(25, 30)).collect::<Vec<_>>();
1332 assert!(results.is_empty());
1333
1334 let results = tree.find_intersects(TextRange::new(3, 3)).collect::<Vec<_>>();
1336 assert!(results.is_empty());
1337
1338 let results = tree.find_intersects(TextRange::new(0, 25)).collect::<Vec<_>>();
1340 assert_eq!(results.len(), 4);
1341
1342 let results = tree.find_intersects(TextRange::new(10, 11)).collect::<Vec<_>>();
1344 assert_eq!(results.len(), 1);
1345 assert_eq!(results[0].key, TextRange::new(10, 15));
1346 }
1347
1348 #[test]
1349 fn advance() {
1350 let val = 1;
1351 let mut tree = build_tree(val);
1352 tree.advance(7, 5);
1353 tree.get(TextRange::new(6, 7)).unwrap();
1355 tree.get(TextRange::new(7, 13)).unwrap();
1356 tree.get(TextRange::new(13, 14)).unwrap();
1357 tree.get(TextRange::new(14, 15)).unwrap();
1358 }
1359
1360 #[test]
1361 fn test_merge_adjacent_intervals_with_equal_properties() {
1362 let mut tree = IntervalTree::new();
1363 tree.insert(TextRange::new(1, 5), 1, merge);
1364 tree.insert(TextRange::new(5, 10), 1, merge);
1365 tree.merge(|a, b| *a == *b);
1366 assert_eq!(tree.find_intersects(TextRange::new(1, 10)).collect::<Vec<_>>().len(), 1);
1367 }
1368
1369 #[test]
1370 fn test_not_merge_adjacent_intervals_with_different_properties() {
1371 let mut tree = IntervalTree::new();
1372 tree.insert(TextRange::new(1, 5), 1, merge);
1373 tree.insert(TextRange::new(5, 10), 2, merge);
1374 tree.merge(|a, b| *a == *b);
1375 assert_eq!(tree.find_intersects(TextRange::new(1, 10)).collect::<Vec<_>>().len(), 2);
1376 }
1377
1378 #[test]
1379 fn test_not_merge_non_adjacent_intervals() {
1380 let mut tree = IntervalTree::new();
1381 tree.insert(TextRange::new(1, 5), 1, merge);
1382 tree.insert(TextRange::new(10, 15), 1, merge);
1383 tree.merge(|a, b| *a == *b);
1384 assert_eq!(tree.find_intersects(TextRange::new(1, 15)).collect::<Vec<_>>().len(), 2);
1385 }
1386
1387 #[test]
1388 fn test_merge_multiple_adjacent_intervals_with_equal_properties() {
1389 let mut tree = IntervalTree::new();
1390 tree.insert(TextRange::new(5, 10), 1, merge);
1391 tree.insert(TextRange::new(1, 5), 1, merge);
1392 tree.insert(TextRange::new(10, 15), 1, merge);
1393 tree.merge(|a, b| *a == *b);
1394 assert_eq!(tree.find_intersects(TextRange::new(1, 15)).collect::<Vec<_>>().len(), 1);
1395 }
1396
1397 #[test]
1398 fn test_handle_empty_tree() {
1399 let mut tree: IntervalTree<i32> = IntervalTree::new();
1400 tree.merge(|a, b| *a == *b);
1401 assert_eq!(tree.find_intersects(TextRange::new(1, 10)).collect::<Vec<_>>().len(), 0);
1402 }
1403
1404 #[test]
1405 fn test_handle_tree_with_single_node() {
1406 let mut tree = IntervalTree::new();
1407 tree.insert(TextRange::new(1, 5), 1, merge);
1408 tree.merge(|a, b| *a == *b);
1409 assert_eq!(tree.find_intersects(TextRange::new(1, 5)).collect::<Vec<_>>().len(), 1);
1410 }
1411
1412 #[test]
1413 fn test_apply_with_split() {
1414 let mut tree = IntervalTree::new();
1415 tree.insert(TextRange::new(0, 10), 1, merge);
1416
1417 tree.apply_with_split(|val| Some(val * 2), TextRange::new(5, 15));
1419 let nodes = tree.find_intersects(TextRange::new(0, 15)).collect::<Vec<_>>();
1420 assert_eq!(nodes.len(), 2);
1421 assert_eq!(nodes[0].val, 1);
1422 assert_eq!(nodes[1].val, 2);
1423
1424 tree.apply_with_split(|_| None, TextRange::new(7, 8));
1426 let nodes = tree.find_intersects(TextRange::new(0, 15)).collect::<Vec<_>>();
1427 assert_eq!(nodes.len(), 3);
1428 assert_eq!(nodes[1].key, TextRange::new(5, 7));
1429 assert_eq!(nodes[1].val, 2); tree.apply_with_split(|val| Some(val + 1), TextRange::new(2, 4));
1433 let node = tree.find_intersects(TextRange::new(2, 4)).collect::<Vec<_>>()[0];
1434 assert_eq!(node.val, 2);
1435 }
1436
1437 #[test]
1438 fn test_clone() {
1439 let mut tree = IntervalTree::new();
1440 tree.insert(TextRange::new(0, 5), 1, merge);
1441 tree.insert(TextRange::new(5, 10), 3, merge);
1442 tree.insert(TextRange::new(10, 15), 2, merge);
1443 let mut tree_cloned = tree.clone();
1444 let n1 = tree.get_node_mut(TextRange::new(0, 5)).unwrap();
1445 let n1c = tree_cloned.get_node_mut(TextRange::new(0, 5)).unwrap();
1446 assert!(!ptr::eq(n1, n1c));
1447 let n2 = tree.get_node_mut(TextRange::new(5, 10)).unwrap();
1448 let n2c = tree_cloned.get_node_mut(TextRange::new(5, 10)).unwrap();
1449 assert!(!ptr::eq(n2, n2c));
1450 let n3 = tree.get_node_mut(TextRange::new(10, 15)).unwrap();
1451 let n3c = tree_cloned.get_node_mut(TextRange::new(10, 15)).unwrap();
1452 assert!(!ptr::eq(n3, n3c));
1453 }
1454
1455 #[test]
1456 fn test_clean() {
1457 let mut tree = IntervalTree::new();
1458 tree.insert(TextRange::new(0, 5), 1, merge);
1459 tree.insert(TextRange::new(5, 10), 1, merge);
1460 tree.insert(TextRange::new(10, 15), 2, merge);
1461 tree.insert(TextRange::new(15, 20), 0, merge); tree.insert(TextRange::new(20, 25), 2, merge);
1463 tree.insert(TextRange::new(25, 30), 2, merge);
1464
1465 tree.clean(|a, b| a == b, |v| *v == 0);
1467
1468 let nodes = tree.find_intersects(TextRange::new(0, 30)).collect::<Vec<_>>();
1469 assert_eq!(nodes.len(), 3);
1470 assert_eq!(nodes[0].key, TextRange::new(0, 10));
1471 assert_eq!(nodes[2].key, TextRange::new(20, 30));
1472 }
1473
1474 #[test]
1475 fn test_clean_empty_tree() {
1476 let mut tree: IntervalTree<i32> = IntervalTree::new();
1477 tree.clean(|a, b| a == b, |v| *v == 0);
1478 assert!(tree.find_intersects(TextRange::new(0, 1)).next().is_none());
1479 }
1480
1481 #[test]
1482 fn test_clean_with_empty_intervals() {
1483 let mut tree = IntervalTree::new();
1484 tree.insert(TextRange::new(0, 0), 1, merge); tree.insert(TextRange::new(0, 5), 1, merge);
1486 tree.insert(TextRange::new(5, 5), 2, merge); tree.clean(|a, b| a == b, |v| *v == 0);
1489
1490 assert_eq!(tree.find_intersects(TextRange::new(0, 5)).collect::<Vec<_>>().len(), 1);
1491 }
1492
1493 #[test]
1494 fn test_find_intersect_min() {
1495 let mut tree = IntervalTree::new();
1496 tree.insert(TextRange::new(0, 5), 1, merge);
1497 tree.insert(TextRange::new(5, 10), 2, merge);
1498 tree.insert(TextRange::new(14, 18), 3, merge);
1499
1500 assert_eq!(
1502 tree.find_intersect_min(TextRange::new(5, 10)).unwrap().key,
1503 TextRange::new(5, 10)
1504 );
1505
1506 assert_eq!(
1508 tree.find_intersect_min(TextRange::new(3, 7)).unwrap().key,
1509 TextRange::new(0, 5)
1510 );
1511
1512 assert_eq!(
1513 tree.find_intersect_min(TextRange::new(6, 15)).unwrap().key,
1514 TextRange::new(5, 10)
1515 );
1516
1517 assert!(tree.find_intersect_min(TextRange::new(19, 23)).is_none());
1519
1520 let empty_tree: IntervalTree<i32> = IntervalTree::new();
1522 assert!(empty_tree.find_intersect_min(TextRange::new(0, 1)).is_none());
1523 }
1524}