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, &mut self.marker_map);
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(
377 root: NodePtr,
378 start: u64,
379 id: MarkerId,
380 marker_map: &mut HashMap<MarkerId, Rc<RefCell<Node>>>,
381 ) -> NodePtr {
382 let root = root?;
384
385 Node::push_delta(&root);
386
387 let mut root_mut = root.borrow_mut();
388 let (root_start, root_id) = (root_mut.marker.interval.start, root_mut.marker.id);
389
390 match start.cmp(&root_start) {
391 Ordering::Less => {
392 root_mut.left = Self::delete_recursive(root_mut.left.take(), start, id, marker_map);
393 }
394 Ordering::Greater => {
395 root_mut.right =
396 Self::delete_recursive(root_mut.right.take(), start, id, marker_map);
397 }
398 Ordering::Equal => match id.cmp(&root_id) {
399 Ordering::Less => {
400 root_mut.left =
401 Self::delete_recursive(root_mut.left.take(), start, id, marker_map);
402 }
403 Ordering::Greater => {
404 root_mut.right =
405 Self::delete_recursive(root_mut.right.take(), start, id, marker_map);
406 }
407 Ordering::Equal => {
408 return Self::perform_node_deletion(root_mut, Rc::clone(&root), marker_map);
409 }
410 },
411 }
412
413 drop(root_mut);
414 Node::update_stats(&root);
415 Self::balance(root)
416 }
417
418 fn perform_node_deletion(
420 mut node: RefMut<Node>,
421 node_rc: Rc<RefCell<Node>>,
422 marker_map: &mut HashMap<MarkerId, Rc<RefCell<Node>>>,
423 ) -> NodePtr {
424 if node.left.is_none() {
425 let right = node.right.take();
426 if let Some(ref r) = right {
427 r.borrow_mut().parent = node.parent.clone();
428 }
429 right
430 } else if node.right.is_none() {
431 let left = node.left.take();
432 if let Some(ref l) = left {
433 l.borrow_mut().parent = node.parent.clone();
434 }
435 left
436 } else {
437 let successor_rc = Self::min_node(node.right.as_ref().unwrap());
438
439 let (_successor_start, successor_id) = {
440 let s = successor_rc.borrow();
441 (s.marker.interval.start, s.marker.id)
442 };
443
444 let (orig_start, orig_id) = (node.marker.interval.start, node.marker.id);
449
450 mem::swap(&mut node.marker, &mut successor_rc.borrow_mut().marker);
451
452 marker_map.insert(successor_id, Rc::clone(&node_rc));
462
463 node.right = Self::delete_recursive(node.right.take(), orig_start, orig_id, marker_map);
468
469 drop(node);
470 Node::update_stats(&node_rc);
471 Self::balance(node_rc)
472 }
473 }
474
475 fn min_node(node_rc: &Rc<RefCell<Node>>) -> Rc<RefCell<Node>> {
477 let mut current = Rc::clone(node_rc);
478 loop {
479 Node::push_delta(¤t);
480
481 let next_left_opt = current.borrow().left.clone();
484
485 if let Some(next) = next_left_opt {
486 current = next;
487 } else {
488 break current;
489 }
490 }
491 }
492
493 fn adjust_recursive(node_opt: &mut NodePtr, pos: u64, delta: i64) {
495 let node_rc = match node_opt {
496 Some(n) => n,
497 None => return,
498 };
499
500 Node::push_delta(node_rc);
501
502 let mut node = node_rc.borrow_mut();
503 let start = node.marker.interval.start;
504
505 if pos <= start {
506 if delta < 0 {
511 node.marker.interval.start = (start as i64 + delta).max(pos as i64) as u64;
512 } else {
513 node.marker.interval.start = (start as i64 + delta) as u64;
514 }
515
516 if delta < 0 {
520 Self::adjust_recursive(&mut node.right, pos, delta);
522 } else {
523 if let Some(ref right) = node.right {
525 right.borrow_mut().lazy_delta += delta;
526 }
527 }
528
529 Self::adjust_recursive(&mut node.left, pos, delta);
531 } else {
532 Self::adjust_recursive(&mut node.right, pos, delta);
537 }
538
539 if node.marker.interval.end >= pos {
541 node.marker.interval.end = (node.marker.interval.end as i64 + delta)
542 .max(node.marker.interval.start as i64)
543 as u64;
544 }
545
546 drop(node);
547 Node::update_stats(node_rc);
548 }
549
550 fn query_recursive(
552 node_opt: &NodePtr,
553 query_start: u64,
554 query_end: u64,
555 results: &mut Vec<Marker>,
556 ) {
557 let node_rc = match node_opt {
558 Some(n) => n,
559 None => return,
560 };
561
562 Node::push_delta(node_rc);
563 let node = node_rc.borrow();
564
565 let i = &node.marker.interval;
566 if i.start <= query_end && i.end >= query_start {
567 results.push(node.marker.clone());
568 }
569
570 if node.left.is_some() && node.left.as_ref().unwrap().borrow().max_end >= query_start {
571 Self::query_recursive(&node.left, query_start, query_end, results);
572 }
573
574 if node.right.is_some() && node.marker.interval.start <= query_end {
575 Self::query_recursive(&node.right, query_start, query_end, results);
576 }
577 }
578
579 fn balance(node: Rc<RefCell<Node>>) -> NodePtr {
582 let bf = Node::balance_factor(&node);
583
584 if bf > 1 {
585 let left_rc = node.borrow().left.as_ref().unwrap().clone();
586 if Node::balance_factor(&left_rc) < 0 {
587 let left_child = node.borrow_mut().left.take().unwrap();
589 node.borrow_mut().left = Self::rotate_left(left_child);
590 }
591 Self::rotate_right(node)
592 } else if bf < -1 {
593 let right_rc = node.borrow().right.as_ref().unwrap().clone();
594 if Node::balance_factor(&right_rc) > 0 {
595 let right_child = node.borrow_mut().right.take().unwrap();
597 node.borrow_mut().right = Self::rotate_right(right_child);
598 }
599 Self::rotate_left(node)
600 } else {
601 Some(node)
602 }
603 }
604
605 fn rotate_left(node_rc: Rc<RefCell<Node>>) -> NodePtr {
606 Node::push_delta(&node_rc);
607 let x_rc = node_rc.borrow_mut().right.take().unwrap();
608 Node::push_delta(&x_rc);
609
610 let mut y = node_rc.borrow_mut();
611 let mut x = x_rc.borrow_mut();
612
613 y.right = x.left.take();
614 if let Some(ref r) = y.right {
615 r.borrow_mut().parent = Rc::downgrade(&node_rc);
616 }
617 x.left = Some(Rc::clone(&node_rc));
618 x.parent = y.parent.clone();
619 y.parent = Rc::downgrade(&x_rc);
620
621 drop(x);
622 drop(y);
623
624 Node::update_stats(&node_rc);
625 Node::update_stats(&x_rc);
626 Some(x_rc)
627 }
628
629 fn rotate_right(node_rc: Rc<RefCell<Node>>) -> NodePtr {
630 Node::push_delta(&node_rc);
631 let x_rc = node_rc.borrow_mut().left.take().unwrap();
632 Node::push_delta(&x_rc);
633
634 let mut y = node_rc.borrow_mut();
635 let mut x = x_rc.borrow_mut();
636
637 y.left = x.right.take();
638 if let Some(ref l) = y.left {
639 l.borrow_mut().parent = Rc::downgrade(&node_rc);
640 }
641 x.right = Some(Rc::clone(&node_rc));
642 x.parent = y.parent.clone();
643 y.parent = Rc::downgrade(&x_rc);
644
645 drop(x);
646 drop(y);
647
648 Node::update_stats(&node_rc);
649 Node::update_stats(&x_rc);
650 Some(x_rc)
651 }
652}
653
654#[cfg(test)]
655mod tests {
656 use super::*;
657
658 fn insert_marker(tree: &mut IntervalTree, start: u64, end: u64) -> MarkerId {
660 tree.insert(start, end)
661 }
662
663 fn get_pos(tree: &IntervalTree, id: MarkerId) -> (u64, u64) {
665 tree.get_position(id)
666 .unwrap_or_else(|| panic!("Marker ID {} not found.", id))
667 }
668
669 #[test]
670 fn test_initial_insert_and_delete() {
671 let mut tree = IntervalTree::new();
672 let id1 = insert_marker(&mut tree, 10, 20);
673 let id2 = insert_marker(&mut tree, 30, 40);
674
675 assert_eq!(get_pos(&tree, id1), (10, 20));
676 assert_eq!(get_pos(&tree, id2), (30, 40));
677
678 assert!(tree.delete(id1));
679 assert_eq!(tree.get_position(id1), None);
680 assert_eq!(get_pos(&tree, id2), (30, 40));
681 }
682
683 #[test]
684 fn test_basic_edit_adjustment() {
685 let mut tree = IntervalTree::new();
686 let id1 = insert_marker(&mut tree, 10, 20); let id2 = insert_marker(&mut tree, 30, 40); tree.adjust_for_edit(30, 5);
691
692 assert_eq!(
694 get_pos(&tree, id1),
695 (10, 20),
696 "Marker before edit should not move."
697 );
698
699 assert_eq!(
701 get_pos(&tree, id2),
702 (35, 45),
703 "Marker at/after edit should move."
704 );
705
706 tree.adjust_for_edit(5, -10); assert_eq!(
711 get_pos(&tree, id1),
712 (5, 10),
713 "Marker moved back by deletion."
714 );
715
716 assert_eq!(
718 get_pos(&tree, id2),
719 (25, 35),
720 "Marker moved back by deletion."
721 );
722 }
723
724 #[test]
725 fn test_problematic_lazy_delta_scenario() {
726 let mut tree = IntervalTree::new();
730
731 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!(
740 get_pos(&tree, id_l),
741 (100, 150),
742 "L initial position incorrect."
743 );
744 assert_eq!(
745 get_pos(&tree, id_p),
746 (200, 250),
747 "P initial position incorrect."
748 );
749 assert_eq!(
750 get_pos(&tree, id_r),
751 (300, 350),
752 "R initial position incorrect."
753 );
754
755 tree.adjust_for_edit(150, 50);
760
761 assert_eq!(
765 get_pos(&tree, id_l),
766 (100, 200),
767 "L(100) should expand to (100, 200)."
768 );
769
770 assert_eq!(
772 get_pos(&tree, id_p),
773 (250, 300),
774 "P(200) did not shift correctly. Should be 250."
775 );
776
777 assert_eq!(
779 get_pos(&tree, id_r),
780 (350, 400),
781 "R(300) did not shift correctly. Should be 350."
782 );
783 }
784
785 #[test]
786 fn test_interval_spanning_edit() {
787 let mut tree = IntervalTree::new();
788 let id_s = insert_marker(&mut tree, 50, 200);
790
791 tree.adjust_for_edit(100, 10);
793
794 assert_eq!(
797 get_pos(&tree, id_s),
798 (50, 210),
799 "Spanning marker end did not move correctly."
800 );
801 }
802
803 #[test]
804 fn test_deletion_engulfing_marker_start() {
805 let mut tree = IntervalTree::new();
806 let id1 = insert_marker(&mut tree, 8, 20);
807
808 tree.adjust_for_edit(5, -10);
814
815 assert_eq!(
816 get_pos(&tree, id1),
817 (5, 10),
818 "Marker should be clamped and shrunk."
819 );
820 }
821
822 #[test]
823 fn test_zero_length_marker() {
824 let mut tree = IntervalTree::new();
825 let id1 = insert_marker(&mut tree, 10, 10);
826
827 tree.adjust_for_edit(10, 5);
829 assert_eq!(
830 get_pos(&tree, id1),
831 (15, 15),
832 "Insertion at zero-length marker."
833 );
834
835 tree.adjust_for_edit(5, 5);
837 assert_eq!(
838 get_pos(&tree, id1),
839 (20, 20),
840 "Insertion before zero-length marker."
841 );
842
843 tree.adjust_for_edit(25, -5);
845 assert_eq!(
846 get_pos(&tree, id1),
847 (20, 20),
848 "Deletion after zero-length marker."
849 );
850
851 tree.adjust_for_edit(15, -10);
853 assert_eq!(
857 get_pos(&tree, id1),
858 (15, 15),
859 "Deletion containing zero-length marker."
860 );
861 }
862
863 #[test]
864 fn test_edit_at_pos_zero() {
865 let mut tree = IntervalTree::new();
866 let id1 = insert_marker(&mut tree, 10, 20);
867
868 tree.adjust_for_edit(0, 5);
870 assert_eq!(get_pos(&tree, id1), (15, 25), "Insertion at pos 0.");
871
872 tree.adjust_for_edit(0, -5);
874 assert_eq!(get_pos(&tree, id1), (10, 20), "Deletion at pos 0.");
875
876 tree.adjust_for_edit(0, -15);
878 assert_eq!(get_pos(&tree, id1), (0, 5), "Engulfing deletion at pos 0.");
882 }
883
884 #[test]
885 fn test_deletion_preserves_marker_ordering() {
886 let mut tree = IntervalTree::new();
889
890 let id0 = insert_marker(&mut tree, 0, 0);
892 let id1 = insert_marker(&mut tree, 10, 10);
893 let id2 = insert_marker(&mut tree, 20, 20);
894 let id3 = insert_marker(&mut tree, 30, 30);
895 let id4 = insert_marker(&mut tree, 40, 40);
896
897 assert_eq!(get_pos(&tree, id0), (0, 0));
899 assert_eq!(get_pos(&tree, id1), (10, 10));
900 assert_eq!(get_pos(&tree, id2), (20, 20));
901 assert_eq!(get_pos(&tree, id3), (30, 30));
902 assert_eq!(get_pos(&tree, id4), (40, 40));
903
904 tree.adjust_for_edit(5, -16);
908
909 let positions = vec![
911 get_pos(&tree, id0).0,
912 get_pos(&tree, id1).0,
913 get_pos(&tree, id2).0,
914 get_pos(&tree, id3).0,
915 get_pos(&tree, id4).0,
916 ];
917
918 for i in 0..positions.len() - 1 {
920 assert!(
921 positions[i] <= positions[i + 1],
922 "Ordering violated at index {}: {:?}[{}]={} > {:?}[{}]={}",
923 i,
924 positions,
925 i,
926 positions[i],
927 positions,
928 i + 1,
929 positions[i + 1]
930 );
931 }
932
933 assert_eq!(get_pos(&tree, id0), (0, 0), "Marker at 0 should stay at 0");
935 assert_eq!(
936 get_pos(&tree, id1),
937 (5, 5),
938 "Marker at 10 should clamp to 5"
939 );
940 assert_eq!(
941 get_pos(&tree, id2),
942 (5, 5),
943 "Marker at 20 should clamp to 5"
944 );
945 assert_eq!(
946 get_pos(&tree, id3),
947 (14, 14),
948 "Marker at 30 should shift to 14"
949 );
950 assert_eq!(
951 get_pos(&tree, id4),
952 (24, 24),
953 "Marker at 40 should shift to 24"
954 );
955 }
956}