interval_rbtree/
lib.rs

1use std::cmp::Ordering;
2use std::fmt::{Arguments, Debug, Write};
3pub mod range;
4pub use range::TextRange;
5
6fn safe_mut<'a, T>(n: *mut T) -> &'a mut T {
7    unsafe { n.as_mut().unwrap() }
8}
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
11pub enum Color {
12    Red = 0,
13    Black,
14}
15
16impl Color {
17    pub fn flip(&self) -> Color {
18        match self {
19            Color::Red => Color::Black,
20            Color::Black => Color::Red,
21        }
22    }
23}
24
25#[derive(Clone)]
26pub struct Node<T: Clone> {
27    pub key: TextRange,
28    pub val: T,
29    left: MaybeNode<T>,
30    right: MaybeNode<T>,
31    color: Color,
32    parent: *mut Node<T>,
33    is_right_child: bool,
34    n: usize,
35}
36pub type BoxedNode<T> = Box<Node<T>>;
37pub type MaybeNode<T> = Option<BoxedNode<T>>;
38
39impl<T: Clone> Node<T> {
40    #[inline]
41    pub fn red(node: &MaybeNode<T>) -> bool {
42        node.as_ref()
43            .map(|n| n.color == Color::Red)
44            .unwrap_or(false)
45    }
46
47    #[inline]
48    pub fn new(key: TextRange, val: T) -> Self {
49        Self::new_with_parent(key, val, std::ptr::null_mut(), false)
50    }
51
52    #[inline]
53    pub fn new_with_parent(
54        key: TextRange,
55        val: T,
56        parent: *mut Node<T>,
57        is_right_child: bool,
58    ) -> Node<T> {
59        Self {
60            key,
61            val,
62            left: None,
63            right: None,
64            color: Color::Red,
65            n: 1,
66            parent,
67            is_right_child,
68        }
69    }
70
71    #[inline]
72    pub fn new_boxed(key: TextRange, val: T) -> BoxedNode<T> {
73        Box::new(Self::new(key, val))
74    }
75
76    #[inline]
77    pub fn new_boxed_with_parent(
78        key: TextRange,
79        val: T,
80        parent: *mut Node<T>,
81        is_right_child: bool,
82    ) -> BoxedNode<T> {
83        Box::new(Self::new_with_parent(key, val, parent, is_right_child))
84    }
85
86    #[inline]
87    pub fn set_parent(&mut self, parent: &mut Node<T>) {
88        self.parent = parent
89    }
90
91    /// perform the following operation, \\ is the red link:
92    ///      |            |
93    ///      n            x
94    ///     / \\        // \
95    ///        x    =>  n
96    ///       / \      / \
97    ///      c            c
98    pub fn rotate_left<'a>(node: &'a mut BoxedNode<T>) -> Option<&'a mut BoxedNode<T>> {
99        let mut x = node.right.take()?;
100        let mut c = x.left.take();
101        if let Some(ref mut c) = c {
102            c.parent = node.as_mut();
103            c.is_right_child = true;
104        }
105        node.right = c;
106        x.color = node.color;
107        x.parent = node.parent;
108        x.is_right_child = node.is_right_child;
109        node.parent = x.as_mut();
110        node.is_right_child = false;
111        x.n = node.n;
112        node.color = Color::Red;
113        node.n = node.n();
114        // node and x swapped content
115        std::mem::swap(node, &mut x);
116        node.left.replace(x);
117        Some(node)
118    }
119
120    /// perform the following operation, \\ is the red link:
121    ///      |            |
122    ///      n            x
123    ///    // \          / \\
124    ///    x       =>       n
125    ///   / \              / \
126    ///      c            c
127    pub fn rotate_right<'a>(node: &'a mut BoxedNode<T>) -> Option<&'a mut BoxedNode<T>> {
128        let mut x = node.left.take()?;
129        let mut c = x.right.take();
130        if let Some(ref mut c) = c {
131            c.parent = node.as_mut();
132            c.is_right_child = false;
133        }
134        node.left = c;
135        x.color = node.color;
136        x.parent = node.parent;
137        x.is_right_child = node.is_right_child;
138        node.parent = x.as_mut();
139        node.is_right_child = true;
140        node.color = Color::Red;
141        x.n = node.n;
142        node.n = node.n();
143
144        std::mem::swap(node, &mut x);
145        node.right.replace(x);
146        Some(node)
147    }
148
149    /// total number of items in this sub-tree
150    #[inline]
151    pub fn n(&self) -> usize {
152        let l = match self.left {
153            Some(ref n) => n.n,
154            None => 0,
155        };
156        let r = match self.right {
157            Some(ref n) => n.n,
158            None => 0,
159        };
160        1 + l + r
161    }
162
163    fn min(&self) -> &Node<T> {
164        if let Some(ref l) = self.left {
165            l.min()
166        } else {
167            self
168        }
169    }
170
171    fn min_mut<'a>(&'a mut self) -> &'a mut Node<T> {
172        if let Some(ref mut l) = self.left {
173            l.min_mut()
174        } else {
175            self
176        }
177    }
178
179    #[inline]
180    fn parent(&self) -> Option<&Node<T>> {
181        unsafe { self.parent.as_ref() }
182    }
183
184    #[inline]
185    fn parent_mut(&self) -> Option<&mut Node<T>> {
186        unsafe { self.parent.as_mut() }
187    }
188
189    fn get(&self, key: TextRange) -> Option<T> {
190        match key.cmp(&self.key) {
191            std::cmp::Ordering::Equal => Some(self.val.clone()), // TODO determine clone here
192            std::cmp::Ordering::Less => self.left.as_ref().and_then(|n| n.get(key)),
193            std::cmp::Ordering::Greater => self.right.as_ref().and_then(|n| n.get(key)),
194        }
195    }
196
197    /* ------------------------------------------------------------ */
198    /*                       insertion                              */
199    /* ------------------------------------------------------------ */
200
201    pub fn insert_at<'a, F: Fn(T, T) -> T>(
202        node: &'a mut MaybeNode<T>,
203        key: TextRange,
204        val: T,
205        parent: *mut Node<T>,
206        is_right_child: bool,
207        merge_fn: &F,
208    ) -> Option<&'a mut BoxedNode<T>> {
209        match node {
210            Some(ref mut n) => Node::insert_at_inner(n, key, val, merge_fn),
211            None => {
212                node.replace(Node::new_boxed_with_parent(
213                    key,
214                    val.clone(),
215                    parent,
216                    is_right_child,
217                ));
218                node.as_mut()
219            }
220        }
221    }
222
223    fn insert_at_inner<'a, F: Fn(T, T) -> T>(
224        node: &'a mut BoxedNode<T>,
225        mut key: TextRange,
226        val: T,
227        merge_fn: &F,
228    ) -> Option<&'a mut BoxedNode<T>> {
229        let intersect = key.intersects(node.key);
230        // TODO too taunting, also, plist should be cloned?
231        let ptr: *mut Node<T> = node.as_mut();
232        if intersect {
233            if key.start < node.key.start {
234                let key_left = key.split_at(node.key.start, true);
235                Node::insert_at(&mut node.left, key_left, val.clone(), ptr, false, merge_fn);
236
237                if key.end < node.key.end {
238                    let key_right = node.key.split_at(key.end, false);
239                    Node::insert_at(
240                        &mut node.right,
241                        key_right,
242                        node.val.clone(),
243                        ptr,
244                        true,
245                        merge_fn,
246                    );
247                } else if key.end > node.key.end {
248                    let key_right = key.split_at(node.key.end, false);
249                    Node::insert_at(&mut node.right, key_right, val.clone(), ptr, true, merge_fn);
250                }
251            } else {
252                let key_left = node.key.split_at(key.start, true);
253                Node::insert_at(
254                    &mut node.left,
255                    key_left,
256                    node.val.clone(),
257                    ptr,
258                    false,
259                    merge_fn,
260                );
261
262                if key.end < node.key.end {
263                    let key_right = node.key.split_at(key.end, false);
264                    Node::insert_at(
265                        &mut node.right,
266                        key_right,
267                        node.val.clone(),
268                        ptr,
269                        true,
270                        merge_fn,
271                    );
272                } else if key.end > node.key.end {
273                    let key_right = key.split_at(node.key.end, false);
274                    Node::insert_at(&mut node.right, key_right, val.clone(), ptr, true, merge_fn);
275                }
276            }
277        }
278        let cmp = key.cmp(&node.key);
279        match cmp {
280            Ordering::Less => {
281                Node::insert_at(&mut node.left, key, val, ptr, false, merge_fn)?;
282            }
283            Ordering::Equal => {
284                node.val = merge_fn(val, node.val.clone());
285            }
286            Ordering::Greater => {
287                Node::insert_at(&mut node.right, key, val, ptr, true, merge_fn)?;
288            }
289        };
290
291        // cond1: r_red && !l_red
292        if Node::red(&node.right) && !Node::red(&node.left) {
293            Node::rotate_left(node)?;
294        }
295
296        // cond2: l_red && ll_red
297        let cond2 = match node.left {
298            Some(ref nl) => nl.color == Color::Red && Node::red(&nl.left),
299            None => false,
300        };
301        if cond2 {
302            Node::rotate_right(node)?;
303        }
304
305        // cond3: l_red && r_red
306        if let (Some(l), Some(r)) = (&mut node.left, &mut node.right) {
307            if l.color == Color::Red && r.color == Color::Red {
308                l.color = Color::Black;
309                r.color = Color::Black;
310                node.color = Color::Red;
311            }
312        }
313        // update node's size
314        node.n = node.n();
315        Some(node)
316    }
317
318    /* ------------------------------------------------------------ */
319    /*                        deletion                              */
320    /* ------------------------------------------------------------ */
321
322    /// if node.left and node.right are both red, mark them black turn node to red.
323    fn flip_colors(n: &mut BoxedNode<T>) {
324        if let Some(ref mut l) = n.left {
325            l.color = l.color.flip();
326        }
327        if let Some(ref mut r) = n.right {
328            r.color = r.color.flip();
329        }
330        n.color = n.color.flip();
331    }
332
333    /// rotate left if node.right is red
334    fn balance(node: &mut BoxedNode<T>) -> Option<()> {
335        if Node::red(&node.right) {
336            Node::rotate_left(node)?;
337        }
338        Some(())
339    }
340
341    fn move_red_left(node: &mut BoxedNode<T>) -> Option<()> {
342        Node::flip_colors(node);
343        let nr = node.right.as_mut()?;
344        if Node::red(&nr.left) {
345            Node::rotate_right(nr)?;
346            Node::rotate_left(node)?;
347            Node::flip_colors(node);
348        }
349        Some(())
350    }
351
352    fn move_red_right(node: &mut BoxedNode<T>) -> Option<()> {
353        Node::flip_colors(node);
354        // h.left.left == Red
355        let cond = match node.left {
356            Some(ref l) => Node::red(&l.left),
357            None => false,
358        };
359        if cond {
360            Node::rotate_right(node)?;
361            Node::flip_colors(node);
362        }
363        Some(())
364    }
365
366    fn delete_min(node: &mut MaybeNode<T>) -> MaybeNode<T> {
367        let n = node.as_mut()?;
368        match n.left {
369            Some(ref mut l) => {
370                if l.color == Color::Black && !Node::red(&l.left) {
371                    Node::move_red_left(n)?;
372                }
373            }
374            None => {
375                return node.take();
376            }
377        }
378        let result = Node::delete_min(&mut n.left);
379        Node::balance(n)?;
380        result
381    }
382
383    fn delete_max(node: &mut MaybeNode<T>) -> MaybeNode<T> {
384        let n = node.as_mut()?;
385        if Node::red(&n.left) {
386            Node::rotate_right(n);
387        }
388        let n = node.as_mut()?;
389        match n.right {
390            Some(ref mut r) => {
391                if r.color == Color::Black && !Node::red(&r.left) {
392                    Node::move_red_right(n)?;
393                }
394            }
395            None => {
396                return node.take();
397            }
398        }
399        let result = Node::delete_max(&mut n.right);
400        Node::balance(n)?;
401        result
402    }
403
404    fn delete(node: &mut MaybeNode<T>, key: TextRange) -> MaybeNode<T> {
405        let n = node.as_mut()?;
406        let result = if key < n.key {
407            // n.left != red && n.left.left != red
408            if let Some(ref mut l) = n.left {
409                if l.color == Color::Black && !Node::red(&l.left) {
410                    Node::move_red_left(n).unwrap();
411                }
412            }
413            Node::delete(&mut n.left, key)
414        } else {
415            if Node::red(&n.left) {
416                Node::rotate_right(n).unwrap();
417            }
418            if key == n.key && n.right.is_none() {
419                return node.take();
420                // return None;
421            }
422
423            let cond = if let Some(ref r) = n.right {
424                r.color == Color::Black && !Node::red(&r.left)
425            } else {
426                true
427            };
428
429            if cond {
430                Node::move_red_right(n).unwrap();
431            }
432
433            // if let Some(ref mut r) = n.right {
434            //     if r.color == Color::Black && !Node::red(&r.left) {
435            //         Node::move_red_right(n).unwrap();
436            //     }
437            // }
438
439            if key == n.key {
440                let mut result = Node::delete_min(&mut n.right);
441                let r_min = result.as_mut().unwrap();
442                std::mem::swap(&mut n.val, &mut r_min.val);
443                std::mem::swap(&mut n.key, &mut r_min.key);
444                result
445            } else {
446                Node::delete(&mut n.right, key)
447            }
448        };
449
450        Node::balance(n)?;
451        result
452    }
453
454    /* ------------------------------------------------------------ */
455    /*                        intersection                          */
456    /* ------------------------------------------------------------ */
457
458    pub fn next(&self) -> Option<&Node<T>> {
459        let mut n = self;
460        if let Some(ref r) = self.right {
461            n = r;
462            while let Some(ref l) = n.left {
463                n = l;
464            }
465            return Some(n);
466        }
467        while let Some(parent) = n.parent() {
468            if !n.is_right_child {
469                return Some(parent);
470            }
471            n = parent;
472        }
473        None
474    }
475
476    pub fn next_mut(&mut self) -> Option<&mut Node<T>> {
477        let mut n: *mut Node<T> = self;
478        if let Some(r) = self.right.as_mut() {
479            n = r.as_mut();
480            while let Some(ref mut l) = safe_mut(n).left {
481                n = l.as_mut();
482            }
483            return Some(safe_mut(n));
484        }
485        while let Some(parent) = (safe_mut(n)).parent_mut() {
486            if !(safe_mut(n)).is_right_child {
487                return Some(parent);
488            }
489            n = parent;
490        }
491        None
492    }
493
494    pub fn next_raw_box(node: *mut Node<T>) -> Option<*mut Node<T>> {
495        let mut n = node;
496        if let Some(r) = safe_mut(node).right.as_mut() {
497            n = r.as_mut();
498            while let Some(ref mut l) = safe_mut(n).left {
499                n = l.as_mut();
500            }
501            return Some(n);
502        }
503        while let Some(parent) = (safe_mut(n)).parent_mut() {
504            if !(safe_mut(n)).is_right_child {
505                return Some(parent);
506            }
507            n = parent;
508        }
509
510        None
511    }
512    pub fn next_raw(&mut self) -> Option<*mut Node<T>> {
513        let mut n: *mut Node<T> = self;
514        if let Some(r) = self.right.as_mut() {
515            n = r.as_mut();
516            while let Some(ref mut l) = safe_mut(n).left {
517                n = l.as_mut();
518            }
519            return Some(n);
520        }
521        while let Some(parent) = (safe_mut(n)).parent_mut() {
522            if !(safe_mut(n)).is_right_child {
523                return Some(parent);
524            }
525            n = parent;
526        }
527
528        None
529    }
530
531    pub fn find_intersects<'a, 'b>(&'a self, range: TextRange, results: &'b mut Vec<&'a Node<T>>) {
532        let ord = range.strict_order(&self.key);
533        match ord {
534            Some(Ordering::Less) => {
535                if let Some(ref l) = self.left {
536                    l.find_intersects(range, results);
537                }
538            }
539            Some(Ordering::Greater) => {
540                if let Some(ref r) = self.right {
541                    r.find_intersects(range, results);
542                }
543            }
544            _ => {
545                if let Some(ref l) = self.left {
546                    l.find_intersects(range, results);
547                }
548                results.push(self);
549                if let Some(ref r) = self.right {
550                    r.find_intersects(range, results);
551                }
552            }
553        }
554    }
555
556    /// Recursively applies a function to each node in the tree in order.
557    /// f is mutable and has type FnMut because it may modify its parameters
558    fn apply<F>(&self, f: &mut F)
559    where
560        F: FnMut(&Node<T>),
561    {
562        if let Some(ref l) = self.left {
563            l.apply(f);
564        }
565        f(self);
566        if let Some(ref r) = self.right {
567            r.apply(f);
568        }
569    }
570
571    /// Recursively applies a function to each node in the tree in order.
572    /// The function may modify `Node`.
573    fn apply_mut<F>(&mut self, f: &mut F)
574    where
575        F: FnMut(&mut Node<T>),
576    {
577        if let Some(ref mut l) = self.left {
578            l.apply_mut(f);
579        }
580        f(self);
581        if let Some(ref mut r) = self.right {
582            r.apply_mut(f);
583        }
584    }
585
586    pub fn advance(&mut self, position: usize, length: usize) {
587        let cmp = self.key.start > position;
588        if cmp {
589            if let Some(ref mut l) = self.left {
590                l.advance(position, length);
591            }
592            self.key.advance(length);
593            if let Some(ref mut r) = self.right {
594                r.apply_mut(&mut |n| n.key.advance(length));
595            }
596        } else {
597            if self.key.end > position {
598                // position is inside this interval
599                self.key.end += length;
600            }
601            if let Some(ref mut l) = self.left {
602                l.advance(position, length);
603            }
604            if let Some(ref mut r) = self.right {
605                r.advance(position, length)
606            }
607        }
608    }
609}
610
611#[derive(Default)]
612/// A interval tree using red-black tree, whereas keys are intervals, and values are
613/// plists in elisp.
614///
615/// All intervals are half-opened intervals, i.e. `I = [start, end)`.These intervals
616/// are sorted by their starting point, then their ending point.
617///
618/// NOTE When inserting an interval I, if I intersects with some existing interval
619/// J but I != J, then we split & merge I and J into sub-intervals. Thus all intervals
620/// inside a interval tree will not overlap. Adjacant intervals with identical props
621/// should be merged afterwards, maybe during redisplay.
622pub struct IntervalTree<T: Clone> {
623    root: MaybeNode<T>,
624}
625
626impl<T: Clone> IntervalTree<T> {
627    /// Creates an empty interval tree.
628    pub fn new() -> Self {
629        Self { root: None }
630    }
631
632    /// Inserts a new interval with the specified `key` and `val` into the interval tree.
633    ///
634    /// If the interval `key` is degenerate (i.e., its start equals its end), the function
635    /// returns `None` as such intervals are not allowed in the tree. Otherwise, it delegates
636    /// the insertion to the underlying node structure.
637    ///
638    /// # Arguments
639    ///
640    /// * `key` - The text range representing the interval to insert.
641    /// * `val` - The value associated with the interval.
642    /// * `merge` - A closure that specifies how to merge intervals if they overlap
643    ///
644    /// # Returns
645    ///
646    /// An optional mutable reference to the newly inserted node, or `None` if the interval is
647    /// degenerate.
648    pub fn insert<'a, F: Fn(T, T) -> T>(
649        &'a mut self,
650        key: impl Into<TextRange>,
651        val: T,
652        merge_fn: F,
653    ) -> Option<&'a mut Box<Node<T>>> {
654        let key = key.into();
655        if key.start == key.end {
656            return None;
657        }
658        let mut result = Node::insert_at(
659            &mut self.root,
660            key,
661            val,
662            std::ptr::null_mut(),
663            false,
664            &merge_fn,
665        );
666        result.as_mut().unwrap().color = Color::Black;
667        result
668    }
669
670    /// Finds the node with key `key` in the tree and returns its value if found.
671    ///
672    /// # Arguments
673    ///
674    /// * `key` - The text range representing the interval to search for.
675    ///
676    /// # Returns
677    ///
678    /// An optional value associated with the node if it exists, `None` otherwise.
679    pub fn get(&self, key: impl Into<TextRange>) -> Option<T> {
680        match self.root {
681            Some(ref r) => r.get(key.into()),
682            None => None,
683        }
684    }
685
686    /// Delete the node with key `key` from the tree.
687    ///
688    /// If the root node is the only black node, then we have to make it red
689    /// before deleting. Otherwise, the tree would become unbalanced.
690    ///
691    /// After deleting, we make sure the root node is black again.
692    pub fn delete(&mut self, key: impl Into<TextRange>) -> MaybeNode<T> {
693        let key = key.into();
694        let result = match self.root {
695            Some(ref mut root) => {
696                if !Node::red(&root.left) && !Node::red(&root.right) {
697                    root.color = Color::Red;
698                }
699
700                Node::delete(&mut self.root, key)
701            }
702            None => None,
703        };
704        if let Some(ref mut root) = self.root {
705            root.color = Color::Black;
706        }
707        result
708    }
709
710    /// Deletes the node with the minimum key from the interval tree.
711    ///
712    /// If the root node is the only black node, it is temporarily colored red
713    /// to maintain tree balance during deletion. After deletion, the root node
714    /// is recolored black to ensure the red-black tree properties are preserved.
715    ///
716    /// # Returns
717    ///
718    /// An optional `Node<T>` representing the removed node, or `None` if
719    /// the tree is empty.
720    pub fn delete_min(&mut self) -> MaybeNode<T> {
721        let root = self.root.as_mut()?;
722        if !Node::red(&root.left) && !Node::red(&root.right) {
723            root.color = Color::Red;
724        }
725        let result = Node::delete_min(&mut self.root);
726        if let Some(ref mut root) = self.root {
727            root.color = Color::Black;
728        }
729        result
730    }
731
732    pub fn delete_max(&mut self) -> MaybeNode<T> {
733        let root = self.root.as_mut()?;
734        if !Node::red(&root.left) && !Node::red(&root.right) {
735            root.color = Color::Red;
736        }
737        let result = Node::delete_max(&mut self.root);
738        if let Some(ref mut root) = self.root {
739            root.color = Color::Black;
740        }
741        result
742    }
743
744    /// Advances all intervals in the tree by `length`, starting at
745    /// `position`. This is typically used to implement operations that insert
746    /// or delete text in a buffer.
747    pub fn advance(&mut self, position: usize, length: usize) {
748        if let Some(ref mut node) = self.root {
749            node.advance(position, length);
750        }
751    }
752
753    /// Find the node whose interval contains the given `position`. If no interval
754    /// contains the position, returns `None`.
755    pub fn find(&self, position: usize) -> Option<&Node<T>> {
756        let range = TextRange::new(position, position + 1);
757        let res = self.find_intersects(range);
758        res.first().copied()
759    }
760
761    /// Find all nodes in the tree whose intervals intersect the given
762    /// `range`. The result is a vector of references to the found nodes.
763    pub fn find_intersects(&self, range: TextRange) -> Vec<&Node<T>> {
764        let mut result = Vec::new();
765        if let Some(ref r) = self.root {
766            r.find_intersects(range, &mut result);
767        }
768        result
769    }
770
771    /// Return the minimum node in the tree, or `None` if the tree is empty.
772    pub fn min(&self) -> Option<&Node<T>> {
773        self.root.as_ref().map(|n| n.min())
774    }
775
776    fn min_mut(&mut self) -> Option<*mut Node<T>> {
777        self.root.as_mut().map(|n| n.min_mut() as *mut Node<T>)
778    }
779
780    /// Merges adjacent intervals in the tree that have equal properties.
781    ///
782    /// This function iterates over the nodes in the interval tree, starting from
783    /// the minimum node. It checks if the current node's end equals the next node's
784    /// start and if their values are considered equal by the provided `equal`
785    /// function. If both conditions are met, it merges the intervals by extending
786    /// the current node's end to the next node's end and deletes the next node.
787    ///
788    /// # Arguments
789    ///
790    /// * `equal` - A closure that takes references to two values and returns `true`
791    ///   if they are considered equal, `false` otherwise.
792    pub fn merge<F: Fn(&T, &T) -> bool>(&mut self, equal: F) {
793        if let Some(node_ptr) = self.min_mut() {
794            let mut node = safe_mut(node_ptr);
795            while let Some(next_ptr) = node.next_raw() {
796                let next = safe_mut(next_ptr);
797                if node.key.end == next.key.start && equal(&node.val, &next.val) {
798                    let del = self.delete(next.key).unwrap();
799                    node.key.end = del.key.end;
800                } else {
801                    node = next;
802                }
803            }
804        }
805    }
806
807    pub fn apply<F: FnMut(&T)>(&self, f: &mut F) {
808        if let Some(r) = self.root.as_ref() {
809            r.apply(&mut |n: &Node<T>| f(&n.val));
810        }
811    }
812
813    pub fn apply_mut<F: FnMut(&mut Node<T>)>(&mut self, f: &mut F) {
814        if let Some(r) = self.root.as_mut() {
815            r.apply_mut(&mut |n| f(n));
816        }
817    }
818}
819impl<T: Clone + Debug> IntervalTree<T> {
820    /// Recursively print out the tree, for debugging purposes. The output format
821    /// is not guaranteed to be stable.
822    pub fn print(&self) {
823        println!("{self:?}");
824    }
825
826    fn print_inner(node: &Node<T>, f: &mut std::fmt::Formatter, level: usize) -> std::fmt::Result {
827        write_fmt_with_level(
828            f,
829            level,
830            format_args!(
831                "[key: {:?}, val: {:?}, color: {:?}]\n",
832                node.key, node.val, node.color
833            ),
834        )?;
835        if let Some(parent) = unsafe { node.parent.as_ref() } {
836            let direction = if node.is_right_child { "right" } else { "left" };
837            write_fmt_with_level(
838                f,
839                level,
840                format_args!("parent({} child): {:?}", direction, parent.key),
841            )?;
842        } else {
843            write_fmt_with_level(f, level, format_args!("parent: not found"))?;
844        }
845        f.write_char('\n')?;
846        if let Some(ref l) = node.left {
847            write_fmt_with_level(f, level, format_args!("left: \n"))?;
848            IntervalTree::print_inner(l, f, level + 1)?;
849            write_fmt_with_level(f, level, format_args!("left end for {:?}\n", node.key))?;
850        }
851        if let Some(ref r) = node.right {
852            write_fmt_with_level(f, level, format_args!("right: \n"))?;
853            IntervalTree::print_inner(r, f, level + 1)?;
854            write_fmt_with_level(f, level, format_args!("right end for {:?}\n", node.key))?;
855        }
856        Ok(())
857    }
858}
859
860fn write_fmt_with_level(
861    f: &mut std::fmt::Formatter,
862    level: usize,
863    fmt: Arguments<'_>,
864) -> std::fmt::Result {
865    for _ in 0..level {
866        f.write_char('\t')?;
867    }
868    f.write_fmt(fmt)
869}
870
871impl<T: Clone + Debug> Debug for IntervalTree<T> {
872    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
873        f.write_str("Interval Tree:\n")?;
874        if let Some(root) = self.root.as_ref() {
875            IntervalTree::print_inner(root, f, 0)?
876        }
877        Ok(())
878    }
879}
880
881#[cfg(test)]
882mod tests {
883
884    use super::*;
885
886    fn merge<T: Clone + Debug>(a: T, _b: T) -> T {
887        a
888    }
889
890    fn build_tree<T: Clone + Debug>(val: T) -> IntervalTree<T> {
891        let mut tree = IntervalTree::new();
892        tree.insert(TextRange::new(0, 1), val.clone(), merge);
893        tree.insert(TextRange::new(1, 2), val.clone(), merge);
894        tree.insert(TextRange::new(2, 3), val.clone(), merge);
895        tree.insert(TextRange::new(3, 4), val.clone(), merge);
896        tree.insert(TextRange::new(4, 5), val.clone(), merge);
897        tree.insert(TextRange::new(5, 6), val.clone(), merge);
898        tree.insert(TextRange::new(6, 7), val.clone(), merge);
899        tree.insert(TextRange::new(7, 8), val.clone(), merge);
900        tree.insert(TextRange::new(8, 9), val.clone(), merge);
901        tree.insert(TextRange::new(9, 10), val.clone(), merge);
902        tree.print();
903        tree
904    }
905
906    #[test]
907    fn rotate_left() {
908        let val = 1;
909        let mut node1 = Node::new_boxed(TextRange::new(0, 3), val);
910        node1.color = Color::Black;
911        let mut node2 = Node::new_boxed(TextRange::new(3, 6), val);
912        let mut node3 = Node::new_boxed(TextRange::new(6, 9), val);
913        node3.color = Color::Black;
914        let mut node4 = Node::new_boxed(TextRange::new(9, 12), val);
915        node4.color = Color::Black;
916        let mut node5 = Node::new_boxed(TextRange::new(12, 15), val);
917        node5.color = Color::Black;
918
919        node2.left = Some(node3);
920        node2.right = Some(node4);
921        node2.color = Color::Red;
922        node1.right = Some(node2);
923        node1.left = Some(node5);
924        Node::rotate_left(&mut node1);
925        assert_eq!(node1.key.start, 3);
926        let n2 = node1.left.unwrap();
927        assert_eq!(n2.key.start, 0);
928        let n3 = n2.right.unwrap();
929        assert_eq!(n3.key.start, 6);
930        let n4 = node1.right.unwrap();
931        assert_eq!(n4.key.start, 9);
932        let n5 = n2.left.unwrap();
933        assert_eq!(n5.key.start, 12);
934
935        assert_eq!(node1.color, Color::Black);
936        assert_eq!(n2.color, Color::Red);
937        assert_eq!(n2.n, 3);
938    }
939
940    #[test]
941    fn insert() {
942        let val = 1;
943        let tree = build_tree(val);
944        let root = tree.root.as_ref().unwrap();
945        assert_eq!(root.key.start, 3);
946        let n1 = root.left.as_ref().unwrap();
947        assert_eq!(n1.key.start, 1);
948        let n2 = root.right.as_ref().unwrap();
949        assert_eq!(n2.key.start, 7);
950        let n3 = n2.right.as_ref().unwrap();
951        assert_eq!(n3.key.start, 9);
952        let n4 = n3.left.as_ref().unwrap();
953        assert_eq!(n4.key.start, 8);
954        assert!(n3.right.is_none())
955    }
956
957    #[test]
958    fn delete() {
959        let val = 1;
960        let mut tree = build_tree(val);
961        // let mut tree = dbg!(tree);
962        for k in vec![8, 4, 5, 7, 3, 6].into_iter() {
963            let i = TextRange::new(k, k + 1);
964            let a = tree.delete(i).unwrap();
965            assert_eq!(a.key, i);
966        }
967    }
968
969    #[test]
970    fn delete_min() {
971        let val = 1;
972        let mut tree = build_tree(val);
973        // let mut tree = dbg!(tree);
974        let a = tree.delete_min().unwrap();
975        assert_eq!(a.key.start, 0);
976    }
977
978    #[test]
979    fn find_intersect() {
980        let val = 1;
981        let tree = build_tree(val);
982        let re = tree.find_intersects(TextRange::new(2, 4));
983        let k1 = re[0].key;
984        let k2 = re[1].key;
985        assert_eq!(k1, TextRange::new(2, 3));
986        assert_eq!(k2, TextRange::new(3, 4));
987        assert_eq!(re.len(), 2);
988    }
989
990    #[test]
991    fn advance() {
992        let val = 1;
993        let mut tree = build_tree(val);
994        tree.advance(7, 5);
995        // let mut tree = dbg!(tree);
996        tree.get(TextRange::new(6, 7)).unwrap();
997        tree.get(TextRange::new(7, 13)).unwrap();
998        tree.get(TextRange::new(13, 14)).unwrap();
999        tree.get(TextRange::new(14, 15)).unwrap();
1000    }
1001
1002    #[test]
1003    fn find_next() {
1004        let val = 1;
1005        let mut tree = build_tree(val);
1006        tree.delete(TextRange::new(5, 6));
1007        let mut n = tree.min().unwrap();
1008
1009        loop {
1010            match n.next() {
1011                Some(ne) => {
1012                    n = ne;
1013                }
1014                None => break,
1015            }
1016        }
1017        assert_eq!(n.key.start, 9)
1018    }
1019
1020    #[test]
1021    fn test_merge_adjacent_intervals_with_equal_properties() {
1022        let mut tree = IntervalTree::new();
1023        tree.insert(TextRange::new(1, 5), 1, merge);
1024        tree.insert(TextRange::new(5, 10), 1, merge);
1025        tree.merge(|a, b| *a == *b);
1026        assert_eq!(tree.find_intersects(TextRange::new(1, 10)).len(), 1);
1027    }
1028
1029    #[test]
1030    fn test_not_merge_adjacent_intervals_with_different_properties() {
1031        let mut tree = IntervalTree::new();
1032        tree.insert(TextRange::new(1, 5), 1, merge);
1033        tree.insert(TextRange::new(5, 10), 2, merge);
1034        tree.merge(|a, b| *a == *b);
1035        assert_eq!(tree.find_intersects(TextRange::new(1, 10)).len(), 2);
1036    }
1037
1038    #[test]
1039    fn test_not_merge_non_adjacent_intervals() {
1040        let mut tree = IntervalTree::new();
1041        tree.insert(TextRange::new(1, 5), 1, merge);
1042        tree.insert(TextRange::new(10, 15), 1, merge);
1043        tree.merge(|a, b| *a == *b);
1044        assert_eq!(tree.find_intersects(TextRange::new(1, 15)).len(), 2);
1045    }
1046
1047    #[test]
1048    fn test_merge_multiple_adjacent_intervals_with_equal_properties() {
1049        let mut tree = IntervalTree::new();
1050        tree.insert(TextRange::new(5, 10), 1, merge);
1051        tree.insert(TextRange::new(1, 5), 1, merge);
1052        tree.insert(TextRange::new(10, 15), 1, merge);
1053        tree.merge(|a, b| *a == *b);
1054        assert_eq!(tree.find_intersects(TextRange::new(1, 15)).len(), 1);
1055    }
1056
1057    #[test]
1058    fn test_handle_empty_tree() {
1059        let mut tree: IntervalTree<i32> = IntervalTree::new();
1060        tree.merge(|a, b| *a == *b);
1061        assert_eq!(tree.find_intersects(TextRange::new(1, 10)).len(), 0);
1062    }
1063
1064    #[test]
1065    fn test_handle_tree_with_single_node() {
1066        let mut tree = IntervalTree::new();
1067        tree.insert(TextRange::new(1, 5), 1, merge);
1068        tree.merge(|a, b| *a == *b);
1069        assert_eq!(tree.find_intersects(TextRange::new(1, 5)).len(), 1);
1070    }
1071}
interval_rbtree/lib.rs

interval_rbtree/
lib.rs
