1use std::cell::{RefCell, RefMut};
2use std::cmp::{max, Ordering};
3use std::collections::HashMap;
4use std::mem;
5use std::rc::{Rc, Weak};
6
7pub type MarkerId = u64;
9
10#[derive(Debug, Clone, PartialEq)]
15pub struct Interval {
16 pub start: u64,
17 pub end: u64,
18}
19
20#[derive(Debug, Clone, PartialEq)]
22pub enum MarkerType {
23 Position,
25 LineAnchor {
27 estimated_line: usize,
28 confidence: AnchorConfidence,
29 },
30}
31
32#[derive(Debug, Clone, PartialEq)]
34pub enum AnchorConfidence {
35 Exact,
37 Estimated,
39 Relative(MarkerId),
41}
42
43#[derive(Debug, Clone, PartialEq)]
44pub struct Marker {
45 pub id: MarkerId,
46 pub interval: Interval,
47 pub marker_type: MarkerType,
48}
49
50type NodePtr = Option<Rc<RefCell<Node>>>;
52type WeakNodePtr = Weak<RefCell<Node>>;
54
55#[derive(Debug)]
57struct Node {
58 pub marker: Marker,
59
60 pub height: i32,
62 pub max_end: u64,
64 pub lazy_delta: i64,
66
67 pub parent: WeakNodePtr,
68 pub left: NodePtr,
69 pub right: NodePtr,
70}
71
72#[derive(Debug, Default)]
74pub struct IntervalTree {
75 root: NodePtr,
76 next_id: u64,
77 marker_map: HashMap<MarkerId, Rc<RefCell<Node>>>,
79}
80
81impl Node {
86 fn new(marker: Marker, parent: WeakNodePtr) -> Rc<RefCell<Self>> {
87 let max_end_val = marker.interval.end;
89
90 Rc::new(RefCell::new(Node {
91 marker,
92 height: 1,
93 max_end: max_end_val,
94 lazy_delta: 0,
95 parent,
96 left: None,
97 right: None,
98 }))
99 }
100
101 fn height(node: &NodePtr) -> i32 {
103 node.as_ref().map_or(0, |n| n.borrow().height)
104 }
105
106 fn balance_factor(node: &Rc<RefCell<Self>>) -> i32 {
108 let n = node.borrow();
109 Self::height(&n.left) - Self::height(&n.right)
110 }
111
112 fn push_delta(node_rc: &Rc<RefCell<Self>>) {
114 let mut node = node_rc.borrow_mut();
115 if node.lazy_delta == 0 {
116 return;
117 }
118
119 let delta = node.lazy_delta;
120
121 node.marker.interval.start = (node.marker.interval.start as i64 + delta) as u64;
123 node.marker.interval.end = (node.marker.interval.end as i64 + delta) as u64;
124
125 if let Some(ref left) = node.left {
127 left.borrow_mut().lazy_delta += delta;
128 }
129 if let Some(ref right) = node.right {
130 right.borrow_mut().lazy_delta += delta;
131 }
132
133 node.lazy_delta = 0;
134
135 let max_l = node.left.as_ref().map_or(0, |l| l.borrow().max_end);
137 let max_r = node.right.as_ref().map_or(0, |r| r.borrow().max_end);
138 node.max_end = max(node.marker.interval.end, max(max_l, max_r));
139 }
140
141 fn update_stats(node: &Rc<RefCell<Self>>) {
143 let mut n = node.borrow_mut();
144 let height_l = Self::height(&n.left);
145 let height_r = Self::height(&n.right);
146
147 n.height = 1 + max(height_l, height_r);
148
149 let max_l = n.left.as_ref().map_or(0, |l| l.borrow().max_end);
150 let max_r = n.right.as_ref().map_or(0, |r| r.borrow().max_end);
151 n.max_end = max(n.marker.interval.end, max(max_l, max_r));
152 }
153}
154
155impl IntervalTree {
160 pub fn new() -> Self {
161 Self::default()
162 }
163
164 pub fn insert(&mut self, start: u64, end: u64) -> MarkerId {
166 self.insert_with_type(start, end, MarkerType::Position)
167 }
168
169 fn insert_with_id(&mut self, id: MarkerId, start: u64, end: u64, marker_type: MarkerType) {
172 debug_assert!(
173 id < self.next_id,
174 "insert_with_id: id {} must be < next_id {}",
175 id,
176 self.next_id
177 );
178 debug_assert!(
179 !self.marker_map.contains_key(&id),
180 "insert_with_id: id {} already in use",
181 id
182 );
183 let marker = Marker {
184 id,
185 interval: Interval { start, end },
186 marker_type,
187 };
188
189 let new_node = Node::new(marker.clone(), Weak::new());
190 self.root = Self::insert_recursive(self.root.take(), new_node.clone());
191 self.marker_map.insert(id, new_node);
192 }
193
194 pub fn insert_with_type(&mut self, start: u64, end: u64, marker_type: MarkerType) -> MarkerId {
196 let id = self.next_id;
197 self.next_id += 1;
198 let marker = Marker {
199 id,
200 interval: Interval { start, end },
201 marker_type,
202 };
203
204 let new_node = Node::new(marker.clone(), Weak::new());
205 self.root = Self::insert_recursive(self.root.take(), new_node.clone());
206
207 self.marker_map.insert(id, new_node);
208 id
209 }
210
211 pub fn insert_line_anchor(
213 &mut self,
214 start: u64,
215 end: u64,
216 estimated_line: usize,
217 confidence: AnchorConfidence,
218 ) -> MarkerId {
219 self.insert_with_type(
220 start,
221 end,
222 MarkerType::LineAnchor {
223 estimated_line,
224 confidence,
225 },
226 )
227 }
228
229 pub fn get_position(&self, id: MarkerId) -> Option<(u64, u64)> {
231 let node_rc = self.marker_map.get(&id)?;
232 let mut node_opt = Some(Rc::clone(node_rc));
233 let mut current_delta: i64 = 0;
234
235 while let Some(current_rc) = node_opt {
237 let current = current_rc.borrow();
238
239 current_delta += current.lazy_delta;
241
242 node_opt = current.parent.upgrade();
244 }
245
246 let raw_marker = node_rc.borrow().marker.interval.clone();
247
248 let start = (raw_marker.start as i64 + current_delta) as u64;
249 let end = (raw_marker.end as i64 + current_delta) as u64;
250
251 Some((start, end))
252 }
253
254 pub fn delete(&mut self, id: MarkerId) -> bool {
256 let (start, _) = match self.get_position(id) {
257 Some(pos) => pos,
258 None => return false,
259 };
260
261 if !self.marker_map.contains_key(&id) {
262 return false;
263 }
264
265 self.root = Self::delete_recursive(self.root.take(), start, id);
266
267 self.marker_map.remove(&id).is_some()
268 }
269
270 pub fn set_position(&mut self, id: MarkerId, new_start: u64, new_end: u64) -> bool {
275 let marker_type = match self.get_marker(id) {
277 Some(m) => m.marker_type,
278 None => return false,
279 };
280
281 if !self.delete(id) {
283 return false;
284 }
285
286 self.insert_with_id(id, new_start, new_end, marker_type);
288 true
289 }
290
291 pub fn adjust_for_edit(&mut self, pos: u64, delta: i64) {
294 Self::adjust_recursive(&mut self.root, pos, delta);
295 }
296
297 pub fn query(&self, query_start: u64, query_end: u64) -> Vec<Marker> {
300 let mut results = Vec::new();
301 Self::query_recursive(&self.root, query_start, query_end, &mut results);
302 results
303 }
304
305 pub fn get_marker(&self, id: MarkerId) -> Option<Marker> {
307 let node_rc = self.marker_map.get(&id)?;
308 Some(node_rc.borrow().marker.clone())
309 }
310
311 pub fn update_line_anchor(
313 &mut self,
314 id: MarkerId,
315 estimated_line: usize,
316 confidence: AnchorConfidence,
317 ) -> bool {
318 if let Some(node_rc) = self.marker_map.get(&id) {
319 let mut node = node_rc.borrow_mut();
320 node.marker.marker_type = MarkerType::LineAnchor {
321 estimated_line,
322 confidence,
323 };
324 true
325 } else {
326 false
327 }
328 }
329
330 pub fn query_line_anchors(&self, query_start: u64, query_end: u64) -> Vec<Marker> {
332 self.query(query_start, query_end)
333 .into_iter()
334 .filter(|m| matches!(m.marker_type, MarkerType::LineAnchor { .. }))
335 .collect()
336 }
337}
338
339impl IntervalTree {
344 fn insert_recursive(root: NodePtr, new_node: Rc<RefCell<Node>>) -> NodePtr {
346 let root = match root {
348 Some(r) => r,
349 None => return Some(new_node),
350 };
351
352 Node::push_delta(&root);
353
354 let (start, id) = (
355 new_node.borrow().marker.interval.start,
356 new_node.borrow().marker.id,
357 );
358
359 let mut root_mut = root.borrow_mut();
360 let (root_start, root_id) = (root_mut.marker.interval.start, root_mut.marker.id);
361
362 if start < root_start || (start == root_start && id < root_id) {
363 root_mut.left = Self::insert_recursive(root_mut.left.take(), Rc::clone(&new_node));
364 root_mut.left.as_ref().unwrap().borrow_mut().parent = Rc::downgrade(&root);
365 } else {
366 root_mut.right = Self::insert_recursive(root_mut.right.take(), Rc::clone(&new_node));
367 root_mut.right.as_ref().unwrap().borrow_mut().parent = Rc::downgrade(&root);
368 }
369
370 drop(root_mut);
371 Node::update_stats(&root);
372 Self::balance(root)
373 }
374
375 fn delete_recursive(root: NodePtr, start: u64, id: MarkerId) -> NodePtr {
377 let root = root?;
379
380 Node::push_delta(&root);
381
382 let mut root_mut = root.borrow_mut();
383 let (root_start, root_id) = (root_mut.marker.interval.start, root_mut.marker.id);
384
385 match start.cmp(&root_start) {
386 Ordering::Less => {
387 root_mut.left = Self::delete_recursive(root_mut.left.take(), start, id);
388 }
389 Ordering::Greater => {
390 root_mut.right = Self::delete_recursive(root_mut.right.take(), start, id);
391 }
392 Ordering::Equal => match id.cmp(&root_id) {
393 Ordering::Less => {
394 root_mut.left = Self::delete_recursive(root_mut.left.take(), start, id);
395 }
396 Ordering::Greater => {
397 root_mut.right = Self::delete_recursive(root_mut.right.take(), start, id);
398 }
399 Ordering::Equal => {
400 return Self::perform_node_deletion(root_mut, Rc::clone(&root));
401 }
402 },
403 }
404
405 drop(root_mut);
406 Node::update_stats(&root);
407 Self::balance(root)
408 }
409
410 fn perform_node_deletion(mut node: RefMut<Node>, node_rc: Rc<RefCell<Node>>) -> NodePtr {
412 if node.left.is_none() {
413 let right = node.right.take();
414 if let Some(ref r) = right {
415 r.borrow_mut().parent = node.parent.clone();
416 }
417 right
418 } else if node.right.is_none() {
419 let left = node.left.take();
420 if let Some(ref l) = left {
421 l.borrow_mut().parent = node.parent.clone();
422 }
423 left
424 } else {
425 let successor_rc = Self::min_node(node.right.as_ref().unwrap());
426
427 let (successor_start, successor_id) = {
428 let s = successor_rc.borrow();
429 (s.marker.interval.start, s.marker.id)
430 };
431
432 mem::swap(&mut node.marker, &mut successor_rc.borrow_mut().marker);
433
434 node.right = Self::delete_recursive(node.right.take(), successor_start, successor_id);
435
436 drop(node);
437 Node::update_stats(&node_rc);
438 Self::balance(node_rc)
439 }
440 }
441
442 fn min_node(node_rc: &Rc<RefCell<Node>>) -> Rc<RefCell<Node>> {
444 let mut current = Rc::clone(node_rc);
445 loop {
446 Node::push_delta(¤t);
447
448 let next_left_opt = current.borrow().left.clone();
451
452 if let Some(next) = next_left_opt {
453 current = next;
454 } else {
455 break current;
456 }
457 }
458 }
459
460 fn adjust_recursive(node_opt: &mut NodePtr, pos: u64, delta: i64) {
462 let node_rc = match node_opt {
463 Some(n) => n,
464 None => return,
465 };
466
467 Node::push_delta(node_rc);
468
469 let mut node = node_rc.borrow_mut();
470 let start = node.marker.interval.start;
471
472 if pos <= start {
473 if delta < 0 {
478 node.marker.interval.start = (start as i64 + delta).max(pos as i64) as u64;
479 } else {
480 node.marker.interval.start = (start as i64 + delta) as u64;
481 }
482
483 if delta < 0 {
487 Self::adjust_recursive(&mut node.right, pos, delta);
489 } else {
490 if let Some(ref right) = node.right {
492 right.borrow_mut().lazy_delta += delta;
493 }
494 }
495
496 Self::adjust_recursive(&mut node.left, pos, delta);
498 } else {
499 Self::adjust_recursive(&mut node.right, pos, delta);
504 }
505
506 if node.marker.interval.end >= pos {
508 node.marker.interval.end = (node.marker.interval.end as i64 + delta)
509 .max(node.marker.interval.start as i64)
510 as u64;
511 }
512
513 drop(node);
514 Node::update_stats(node_rc);
515 }
516
517 fn query_recursive(
519 node_opt: &NodePtr,
520 query_start: u64,
521 query_end: u64,
522 results: &mut Vec<Marker>,
523 ) {
524 let node_rc = match node_opt {
525 Some(n) => n,
526 None => return,
527 };
528
529 Node::push_delta(node_rc);
530 let node = node_rc.borrow();
531
532 let i = &node.marker.interval;
533 if i.start <= query_end && i.end >= query_start {
534 results.push(node.marker.clone());
535 }
536
537 if node.left.is_some() && node.left.as_ref().unwrap().borrow().max_end >= query_start {
538 Self::query_recursive(&node.left, query_start, query_end, results);
539 }
540
541 if node.right.is_some() && node.marker.interval.start <= query_end {
542 Self::query_recursive(&node.right, query_start, query_end, results);
543 }
544 }
545
546 fn balance(node: Rc<RefCell<Node>>) -> NodePtr {
549 let bf = Node::balance_factor(&node);
550
551 if bf > 1 {
552 let left_rc = node.borrow().left.as_ref().unwrap().clone();
553 if Node::balance_factor(&left_rc) < 0 {
554 let left_child = node.borrow_mut().left.take().unwrap();
556 node.borrow_mut().left = Self::rotate_left(left_child);
557 }
558 Self::rotate_right(node)
559 } else if bf < -1 {
560 let right_rc = node.borrow().right.as_ref().unwrap().clone();
561 if Node::balance_factor(&right_rc) > 0 {
562 let right_child = node.borrow_mut().right.take().unwrap();
564 node.borrow_mut().right = Self::rotate_right(right_child);
565 }
566 Self::rotate_left(node)
567 } else {
568 Some(node)
569 }
570 }
571
572 fn rotate_left(node_rc: Rc<RefCell<Node>>) -> NodePtr {
573 Node::push_delta(&node_rc);
574 let x_rc = node_rc.borrow_mut().right.take().unwrap();
575 Node::push_delta(&x_rc);
576
577 let mut y = node_rc.borrow_mut();
578 let mut x = x_rc.borrow_mut();
579
580 y.right = x.left.take();
581 if let Some(ref r) = y.right {
582 r.borrow_mut().parent = Rc::downgrade(&node_rc);
583 }
584 x.left = Some(Rc::clone(&node_rc));
585 x.parent = y.parent.clone();
586 y.parent = Rc::downgrade(&x_rc);
587
588 drop(x);
589 drop(y);
590
591 Node::update_stats(&node_rc);
592 Node::update_stats(&x_rc);
593 Some(x_rc)
594 }
595
596 fn rotate_right(node_rc: Rc<RefCell<Node>>) -> NodePtr {
597 Node::push_delta(&node_rc);
598 let x_rc = node_rc.borrow_mut().left.take().unwrap();
599 Node::push_delta(&x_rc);
600
601 let mut y = node_rc.borrow_mut();
602 let mut x = x_rc.borrow_mut();
603
604 y.left = x.right.take();
605 if let Some(ref l) = y.left {
606 l.borrow_mut().parent = Rc::downgrade(&node_rc);
607 }
608 x.right = Some(Rc::clone(&node_rc));
609 x.parent = y.parent.clone();
610 y.parent = Rc::downgrade(&x_rc);
611
612 drop(x);
613 drop(y);
614
615 Node::update_stats(&node_rc);
616 Node::update_stats(&x_rc);
617 Some(x_rc)
618 }
619}
620
621#[cfg(test)]
622mod tests {
623 use super::*;
624
625 fn insert_marker(tree: &mut IntervalTree, start: u64, end: u64) -> MarkerId {
627 tree.insert(start, end)
628 }
629
630 fn get_pos(tree: &IntervalTree, id: MarkerId) -> (u64, u64) {
632 tree.get_position(id)
633 .unwrap_or_else(|| panic!("Marker ID {} not found.", id))
634 }
635
636 #[test]
637 fn test_initial_insert_and_delete() {
638 let mut tree = IntervalTree::new();
639 let id1 = insert_marker(&mut tree, 10, 20);
640 let id2 = insert_marker(&mut tree, 30, 40);
641
642 assert_eq!(get_pos(&tree, id1), (10, 20));
643 assert_eq!(get_pos(&tree, id2), (30, 40));
644
645 assert!(tree.delete(id1));
646 assert_eq!(tree.get_position(id1), None);
647 assert_eq!(get_pos(&tree, id2), (30, 40));
648 }
649
650 #[test]
651 fn test_basic_edit_adjustment() {
652 let mut tree = IntervalTree::new();
653 let id1 = insert_marker(&mut tree, 10, 20); let id2 = insert_marker(&mut tree, 30, 40); tree.adjust_for_edit(30, 5);
658
659 assert_eq!(
661 get_pos(&tree, id1),
662 (10, 20),
663 "Marker before edit should not move."
664 );
665
666 assert_eq!(
668 get_pos(&tree, id2),
669 (35, 45),
670 "Marker at/after edit should move."
671 );
672
673 tree.adjust_for_edit(5, -10); assert_eq!(
678 get_pos(&tree, id1),
679 (5, 10),
680 "Marker moved back by deletion."
681 );
682
683 assert_eq!(
685 get_pos(&tree, id2),
686 (25, 35),
687 "Marker moved back by deletion."
688 );
689 }
690
691 #[test]
692 fn test_problematic_lazy_delta_scenario() {
693 let mut tree = IntervalTree::new();
697
698 let id_p = insert_marker(&mut tree, 200, 250); let id_r = insert_marker(&mut tree, 300, 350); let id_l = insert_marker(&mut tree, 100, 150); assert_eq!(
707 get_pos(&tree, id_l),
708 (100, 150),
709 "L initial position incorrect."
710 );
711 assert_eq!(
712 get_pos(&tree, id_p),
713 (200, 250),
714 "P initial position incorrect."
715 );
716 assert_eq!(
717 get_pos(&tree, id_r),
718 (300, 350),
719 "R initial position incorrect."
720 );
721
722 tree.adjust_for_edit(150, 50);
727
728 assert_eq!(
732 get_pos(&tree, id_l),
733 (100, 200),
734 "L(100) should expand to (100, 200)."
735 );
736
737 assert_eq!(
739 get_pos(&tree, id_p),
740 (250, 300),
741 "P(200) did not shift correctly. Should be 250."
742 );
743
744 assert_eq!(
746 get_pos(&tree, id_r),
747 (350, 400),
748 "R(300) did not shift correctly. Should be 350."
749 );
750 }
751
752 #[test]
753 fn test_interval_spanning_edit() {
754 let mut tree = IntervalTree::new();
755 let id_s = insert_marker(&mut tree, 50, 200);
757
758 tree.adjust_for_edit(100, 10);
760
761 assert_eq!(
764 get_pos(&tree, id_s),
765 (50, 210),
766 "Spanning marker end did not move correctly."
767 );
768 }
769
770 #[test]
771 fn test_deletion_engulfing_marker_start() {
772 let mut tree = IntervalTree::new();
773 let id1 = insert_marker(&mut tree, 8, 20);
774
775 tree.adjust_for_edit(5, -10);
781
782 assert_eq!(
783 get_pos(&tree, id1),
784 (5, 10),
785 "Marker should be clamped and shrunk."
786 );
787 }
788
789 #[test]
790 fn test_zero_length_marker() {
791 let mut tree = IntervalTree::new();
792 let id1 = insert_marker(&mut tree, 10, 10);
793
794 tree.adjust_for_edit(10, 5);
796 assert_eq!(
797 get_pos(&tree, id1),
798 (15, 15),
799 "Insertion at zero-length marker."
800 );
801
802 tree.adjust_for_edit(5, 5);
804 assert_eq!(
805 get_pos(&tree, id1),
806 (20, 20),
807 "Insertion before zero-length marker."
808 );
809
810 tree.adjust_for_edit(25, -5);
812 assert_eq!(
813 get_pos(&tree, id1),
814 (20, 20),
815 "Deletion after zero-length marker."
816 );
817
818 tree.adjust_for_edit(15, -10);
820 assert_eq!(
824 get_pos(&tree, id1),
825 (15, 15),
826 "Deletion containing zero-length marker."
827 );
828 }
829
830 #[test]
831 fn test_edit_at_pos_zero() {
832 let mut tree = IntervalTree::new();
833 let id1 = insert_marker(&mut tree, 10, 20);
834
835 tree.adjust_for_edit(0, 5);
837 assert_eq!(get_pos(&tree, id1), (15, 25), "Insertion at pos 0.");
838
839 tree.adjust_for_edit(0, -5);
841 assert_eq!(get_pos(&tree, id1), (10, 20), "Deletion at pos 0.");
842
843 tree.adjust_for_edit(0, -15);
845 assert_eq!(get_pos(&tree, id1), (0, 5), "Engulfing deletion at pos 0.");
849 }
850
851 #[test]
852 fn test_deletion_preserves_marker_ordering() {
853 let mut tree = IntervalTree::new();
856
857 let id0 = insert_marker(&mut tree, 0, 0);
859 let id1 = insert_marker(&mut tree, 10, 10);
860 let id2 = insert_marker(&mut tree, 20, 20);
861 let id3 = insert_marker(&mut tree, 30, 30);
862 let id4 = insert_marker(&mut tree, 40, 40);
863
864 assert_eq!(get_pos(&tree, id0), (0, 0));
866 assert_eq!(get_pos(&tree, id1), (10, 10));
867 assert_eq!(get_pos(&tree, id2), (20, 20));
868 assert_eq!(get_pos(&tree, id3), (30, 30));
869 assert_eq!(get_pos(&tree, id4), (40, 40));
870
871 tree.adjust_for_edit(5, -16);
875
876 let positions = vec![
878 get_pos(&tree, id0).0,
879 get_pos(&tree, id1).0,
880 get_pos(&tree, id2).0,
881 get_pos(&tree, id3).0,
882 get_pos(&tree, id4).0,
883 ];
884
885 for i in 0..positions.len() - 1 {
887 assert!(
888 positions[i] <= positions[i + 1],
889 "Ordering violated at index {}: {:?}[{}]={} > {:?}[{}]={}",
890 i,
891 positions,
892 i,
893 positions[i],
894 positions,
895 i + 1,
896 positions[i + 1]
897 );
898 }
899
900 assert_eq!(get_pos(&tree, id0), (0, 0), "Marker at 0 should stay at 0");
902 assert_eq!(
903 get_pos(&tree, id1),
904 (5, 5),
905 "Marker at 10 should clamp to 5"
906 );
907 assert_eq!(
908 get_pos(&tree, id2),
909 (5, 5),
910 "Marker at 20 should clamp to 5"
911 );
912 assert_eq!(
913 get_pos(&tree, id3),
914 (14, 14),
915 "Marker at 30 should shift to 14"
916 );
917 assert_eq!(
918 get_pos(&tree, id4),
919 (24, 24),
920 "Marker at 40 should shift to 24"
921 );
922 }
923}