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 pub fn insert_with_type(&mut self, start: u64, end: u64, marker_type: MarkerType) -> MarkerId {
171 let id = self.next_id;
172 self.next_id += 1;
173 let marker = Marker {
174 id,
175 interval: Interval { start, end },
176 marker_type,
177 };
178
179 let new_node = Node::new(marker.clone(), Weak::new());
180 self.root = Self::insert_recursive(self.root.take(), new_node.clone());
181
182 self.marker_map.insert(id, new_node);
183 id
184 }
185
186 pub fn insert_line_anchor(
188 &mut self,
189 start: u64,
190 end: u64,
191 estimated_line: usize,
192 confidence: AnchorConfidence,
193 ) -> MarkerId {
194 self.insert_with_type(
195 start,
196 end,
197 MarkerType::LineAnchor {
198 estimated_line,
199 confidence,
200 },
201 )
202 }
203
204 pub fn get_position(&self, id: MarkerId) -> Option<(u64, u64)> {
206 let node_rc = self.marker_map.get(&id)?;
207 let mut node_opt = Some(Rc::clone(node_rc));
208 let mut current_delta: i64 = 0;
209
210 while let Some(current_rc) = node_opt {
212 let current = current_rc.borrow();
213
214 current_delta += current.lazy_delta;
216
217 node_opt = current.parent.upgrade();
219 }
220
221 let raw_marker = node_rc.borrow().marker.interval.clone();
222
223 let start = (raw_marker.start as i64 + current_delta) as u64;
224 let end = (raw_marker.end as i64 + current_delta) as u64;
225
226 Some((start, end))
227 }
228
229 pub fn delete(&mut self, id: MarkerId) -> bool {
231 let (start, _) = match self.get_position(id) {
232 Some(pos) => pos,
233 None => return false,
234 };
235
236 if !self.marker_map.contains_key(&id) {
237 return false;
238 }
239
240 self.root = Self::delete_recursive(self.root.take(), start, id);
241
242 self.marker_map.remove(&id).is_some()
243 }
244
245 pub fn adjust_for_edit(&mut self, pos: u64, delta: i64) {
248 Self::adjust_recursive(&mut self.root, pos, delta);
249 }
250
251 pub fn query(&self, query_start: u64, query_end: u64) -> Vec<Marker> {
254 let mut results = Vec::new();
255 Self::query_recursive(&self.root, query_start, query_end, &mut results);
256 results
257 }
258
259 pub fn get_marker(&self, id: MarkerId) -> Option<Marker> {
261 let node_rc = self.marker_map.get(&id)?;
262 Some(node_rc.borrow().marker.clone())
263 }
264
265 pub fn update_line_anchor(
267 &mut self,
268 id: MarkerId,
269 estimated_line: usize,
270 confidence: AnchorConfidence,
271 ) -> bool {
272 if let Some(node_rc) = self.marker_map.get(&id) {
273 let mut node = node_rc.borrow_mut();
274 node.marker.marker_type = MarkerType::LineAnchor {
275 estimated_line,
276 confidence,
277 };
278 true
279 } else {
280 false
281 }
282 }
283
284 pub fn query_line_anchors(&self, query_start: u64, query_end: u64) -> Vec<Marker> {
286 self.query(query_start, query_end)
287 .into_iter()
288 .filter(|m| matches!(m.marker_type, MarkerType::LineAnchor { .. }))
289 .collect()
290 }
291}
292
293impl IntervalTree {
298 fn insert_recursive(root: NodePtr, new_node: Rc<RefCell<Node>>) -> NodePtr {
300 let root = match root {
302 Some(r) => r,
303 None => return Some(new_node),
304 };
305
306 Node::push_delta(&root);
307
308 let (start, id) = (
309 new_node.borrow().marker.interval.start,
310 new_node.borrow().marker.id,
311 );
312
313 let mut root_mut = root.borrow_mut();
314 let (root_start, root_id) = (root_mut.marker.interval.start, root_mut.marker.id);
315
316 if start < root_start || (start == root_start && id < root_id) {
317 root_mut.left = Self::insert_recursive(root_mut.left.take(), Rc::clone(&new_node));
318 root_mut.left.as_ref().unwrap().borrow_mut().parent = Rc::downgrade(&root);
319 } else {
320 root_mut.right = Self::insert_recursive(root_mut.right.take(), Rc::clone(&new_node));
321 root_mut.right.as_ref().unwrap().borrow_mut().parent = Rc::downgrade(&root);
322 }
323
324 drop(root_mut);
325 Node::update_stats(&root);
326 Self::balance(root)
327 }
328
329 fn delete_recursive(root: NodePtr, start: u64, id: MarkerId) -> NodePtr {
331 let root = root?;
333
334 Node::push_delta(&root);
335
336 let mut root_mut = root.borrow_mut();
337 let (root_start, root_id) = (root_mut.marker.interval.start, root_mut.marker.id);
338
339 match start.cmp(&root_start) {
340 Ordering::Less => {
341 root_mut.left = Self::delete_recursive(root_mut.left.take(), start, id);
342 }
343 Ordering::Greater => {
344 root_mut.right = Self::delete_recursive(root_mut.right.take(), start, id);
345 }
346 Ordering::Equal => match id.cmp(&root_id) {
347 Ordering::Less => {
348 root_mut.left = Self::delete_recursive(root_mut.left.take(), start, id);
349 }
350 Ordering::Greater => {
351 root_mut.right = Self::delete_recursive(root_mut.right.take(), start, id);
352 }
353 Ordering::Equal => {
354 return Self::perform_node_deletion(root_mut, Rc::clone(&root));
355 }
356 },
357 }
358
359 drop(root_mut);
360 Node::update_stats(&root);
361 Self::balance(root)
362 }
363
364 fn perform_node_deletion(mut node: RefMut<Node>, node_rc: Rc<RefCell<Node>>) -> NodePtr {
366 if node.left.is_none() {
367 let right = node.right.take();
368 if let Some(ref r) = right {
369 r.borrow_mut().parent = node.parent.clone();
370 }
371 right
372 } else if node.right.is_none() {
373 let left = node.left.take();
374 if let Some(ref l) = left {
375 l.borrow_mut().parent = node.parent.clone();
376 }
377 left
378 } else {
379 let successor_rc = Self::min_node(node.right.as_ref().unwrap());
380
381 let (successor_start, successor_id) = {
382 let s = successor_rc.borrow();
383 (s.marker.interval.start, s.marker.id)
384 };
385
386 mem::swap(&mut node.marker, &mut successor_rc.borrow_mut().marker);
387
388 node.right = Self::delete_recursive(node.right.take(), successor_start, successor_id);
389
390 drop(node);
391 Node::update_stats(&node_rc);
392 Self::balance(node_rc)
393 }
394 }
395
396 fn min_node(node_rc: &Rc<RefCell<Node>>) -> Rc<RefCell<Node>> {
398 let mut current = Rc::clone(node_rc);
399 loop {
400 Node::push_delta(¤t);
401
402 let next_left_opt = current.borrow().left.clone();
405
406 if let Some(next) = next_left_opt {
407 current = next;
408 } else {
409 break current;
410 }
411 }
412 }
413
414 fn adjust_recursive(node_opt: &mut NodePtr, pos: u64, delta: i64) {
416 let node_rc = match node_opt {
417 Some(n) => n,
418 None => return,
419 };
420
421 Node::push_delta(node_rc);
422
423 let mut node = node_rc.borrow_mut();
424 let start = node.marker.interval.start;
425
426 if pos <= start {
427 if delta < 0 {
432 node.marker.interval.start = (start as i64 + delta).max(pos as i64) as u64;
433 } else {
434 node.marker.interval.start = (start as i64 + delta) as u64;
435 }
436
437 if delta < 0 {
441 Self::adjust_recursive(&mut node.right, pos, delta);
443 } else {
444 if let Some(ref right) = node.right {
446 right.borrow_mut().lazy_delta += delta;
447 }
448 }
449
450 Self::adjust_recursive(&mut node.left, pos, delta);
452 } else {
453 Self::adjust_recursive(&mut node.right, pos, delta);
458 }
459
460 if node.marker.interval.end >= pos {
462 node.marker.interval.end = (node.marker.interval.end as i64 + delta)
463 .max(node.marker.interval.start as i64)
464 as u64;
465 }
466
467 drop(node);
468 Node::update_stats(node_rc);
469 }
470
471 fn query_recursive(
473 node_opt: &NodePtr,
474 query_start: u64,
475 query_end: u64,
476 results: &mut Vec<Marker>,
477 ) {
478 let node_rc = match node_opt {
479 Some(n) => n,
480 None => return,
481 };
482
483 Node::push_delta(node_rc);
484 let node = node_rc.borrow();
485
486 let i = &node.marker.interval;
487 if i.start <= query_end && i.end >= query_start {
488 results.push(node.marker.clone());
489 }
490
491 if node.left.is_some() && node.left.as_ref().unwrap().borrow().max_end >= query_start {
492 Self::query_recursive(&node.left, query_start, query_end, results);
493 }
494
495 if node.right.is_some() && node.marker.interval.start <= query_end {
496 Self::query_recursive(&node.right, query_start, query_end, results);
497 }
498 }
499
500 fn balance(node: Rc<RefCell<Node>>) -> NodePtr {
503 let bf = Node::balance_factor(&node);
504
505 if bf > 1 {
506 let left_rc = node.borrow().left.as_ref().unwrap().clone();
507 if Node::balance_factor(&left_rc) < 0 {
508 let left_child = node.borrow_mut().left.take().unwrap();
510 node.borrow_mut().left = Self::rotate_left(left_child);
511 }
512 Self::rotate_right(node)
513 } else if bf < -1 {
514 let right_rc = node.borrow().right.as_ref().unwrap().clone();
515 if Node::balance_factor(&right_rc) > 0 {
516 let right_child = node.borrow_mut().right.take().unwrap();
518 node.borrow_mut().right = Self::rotate_right(right_child);
519 }
520 Self::rotate_left(node)
521 } else {
522 Some(node)
523 }
524 }
525
526 fn rotate_left(node_rc: Rc<RefCell<Node>>) -> NodePtr {
527 Node::push_delta(&node_rc);
528 let x_rc = node_rc.borrow_mut().right.take().unwrap();
529 Node::push_delta(&x_rc);
530
531 let mut y = node_rc.borrow_mut();
532 let mut x = x_rc.borrow_mut();
533
534 y.right = x.left.take();
535 if let Some(ref r) = y.right {
536 r.borrow_mut().parent = Rc::downgrade(&node_rc);
537 }
538 x.left = Some(Rc::clone(&node_rc));
539 x.parent = y.parent.clone();
540 y.parent = Rc::downgrade(&x_rc);
541
542 drop(x);
543 drop(y);
544
545 Node::update_stats(&node_rc);
546 Node::update_stats(&x_rc);
547 Some(x_rc)
548 }
549
550 fn rotate_right(node_rc: Rc<RefCell<Node>>) -> NodePtr {
551 Node::push_delta(&node_rc);
552 let x_rc = node_rc.borrow_mut().left.take().unwrap();
553 Node::push_delta(&x_rc);
554
555 let mut y = node_rc.borrow_mut();
556 let mut x = x_rc.borrow_mut();
557
558 y.left = x.right.take();
559 if let Some(ref l) = y.left {
560 l.borrow_mut().parent = Rc::downgrade(&node_rc);
561 }
562 x.right = Some(Rc::clone(&node_rc));
563 x.parent = y.parent.clone();
564 y.parent = Rc::downgrade(&x_rc);
565
566 drop(x);
567 drop(y);
568
569 Node::update_stats(&node_rc);
570 Node::update_stats(&x_rc);
571 Some(x_rc)
572 }
573}
574
575#[cfg(test)]
576mod tests {
577 use super::*;
578
579 fn insert_marker(tree: &mut IntervalTree, start: u64, end: u64) -> MarkerId {
581 tree.insert(start, end)
582 }
583
584 fn get_pos(tree: &IntervalTree, id: MarkerId) -> (u64, u64) {
586 tree.get_position(id)
587 .unwrap_or_else(|| panic!("Marker ID {} not found.", id))
588 }
589
590 #[test]
591 fn test_initial_insert_and_delete() {
592 let mut tree = IntervalTree::new();
593 let id1 = insert_marker(&mut tree, 10, 20);
594 let id2 = insert_marker(&mut tree, 30, 40);
595
596 assert_eq!(get_pos(&tree, id1), (10, 20));
597 assert_eq!(get_pos(&tree, id2), (30, 40));
598
599 assert!(tree.delete(id1));
600 assert_eq!(tree.get_position(id1), None);
601 assert_eq!(get_pos(&tree, id2), (30, 40));
602 }
603
604 #[test]
605 fn test_basic_edit_adjustment() {
606 let mut tree = IntervalTree::new();
607 let id1 = insert_marker(&mut tree, 10, 20); let id2 = insert_marker(&mut tree, 30, 40); tree.adjust_for_edit(30, 5);
612
613 assert_eq!(
615 get_pos(&tree, id1),
616 (10, 20),
617 "Marker before edit should not move."
618 );
619
620 assert_eq!(
622 get_pos(&tree, id2),
623 (35, 45),
624 "Marker at/after edit should move."
625 );
626
627 tree.adjust_for_edit(5, -10); assert_eq!(
632 get_pos(&tree, id1),
633 (5, 10),
634 "Marker moved back by deletion."
635 );
636
637 assert_eq!(
639 get_pos(&tree, id2),
640 (25, 35),
641 "Marker moved back by deletion."
642 );
643 }
644
645 #[test]
646 fn test_problematic_lazy_delta_scenario() {
647 let mut tree = IntervalTree::new();
651
652 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!(
661 get_pos(&tree, id_l),
662 (100, 150),
663 "L initial position incorrect."
664 );
665 assert_eq!(
666 get_pos(&tree, id_p),
667 (200, 250),
668 "P initial position incorrect."
669 );
670 assert_eq!(
671 get_pos(&tree, id_r),
672 (300, 350),
673 "R initial position incorrect."
674 );
675
676 tree.adjust_for_edit(150, 50);
681
682 assert_eq!(
686 get_pos(&tree, id_l),
687 (100, 200),
688 "L(100) should expand to (100, 200)."
689 );
690
691 assert_eq!(
693 get_pos(&tree, id_p),
694 (250, 300),
695 "P(200) did not shift correctly. Should be 250."
696 );
697
698 assert_eq!(
700 get_pos(&tree, id_r),
701 (350, 400),
702 "R(300) did not shift correctly. Should be 350."
703 );
704 }
705
706 #[test]
707 fn test_interval_spanning_edit() {
708 let mut tree = IntervalTree::new();
709 let id_s = insert_marker(&mut tree, 50, 200);
711
712 tree.adjust_for_edit(100, 10);
714
715 assert_eq!(
718 get_pos(&tree, id_s),
719 (50, 210),
720 "Spanning marker end did not move correctly."
721 );
722 }
723
724 #[test]
725 fn test_deletion_engulfing_marker_start() {
726 let mut tree = IntervalTree::new();
727 let id1 = insert_marker(&mut tree, 8, 20);
728
729 tree.adjust_for_edit(5, -10);
735
736 assert_eq!(
737 get_pos(&tree, id1),
738 (5, 10),
739 "Marker should be clamped and shrunk."
740 );
741 }
742
743 #[test]
744 fn test_zero_length_marker() {
745 let mut tree = IntervalTree::new();
746 let id1 = insert_marker(&mut tree, 10, 10);
747
748 tree.adjust_for_edit(10, 5);
750 assert_eq!(
751 get_pos(&tree, id1),
752 (15, 15),
753 "Insertion at zero-length marker."
754 );
755
756 tree.adjust_for_edit(5, 5);
758 assert_eq!(
759 get_pos(&tree, id1),
760 (20, 20),
761 "Insertion before zero-length marker."
762 );
763
764 tree.adjust_for_edit(25, -5);
766 assert_eq!(
767 get_pos(&tree, id1),
768 (20, 20),
769 "Deletion after zero-length marker."
770 );
771
772 tree.adjust_for_edit(15, -10);
774 assert_eq!(
778 get_pos(&tree, id1),
779 (15, 15),
780 "Deletion containing zero-length marker."
781 );
782 }
783
784 #[test]
785 fn test_edit_at_pos_zero() {
786 let mut tree = IntervalTree::new();
787 let id1 = insert_marker(&mut tree, 10, 20);
788
789 tree.adjust_for_edit(0, 5);
791 assert_eq!(get_pos(&tree, id1), (15, 25), "Insertion at pos 0.");
792
793 tree.adjust_for_edit(0, -5);
795 assert_eq!(get_pos(&tree, id1), (10, 20), "Deletion at pos 0.");
796
797 tree.adjust_for_edit(0, -15);
799 assert_eq!(get_pos(&tree, id1), (0, 5), "Engulfing deletion at pos 0.");
803 }
804
805 #[test]
806 fn test_deletion_preserves_marker_ordering() {
807 let mut tree = IntervalTree::new();
810
811 let id0 = insert_marker(&mut tree, 0, 0);
813 let id1 = insert_marker(&mut tree, 10, 10);
814 let id2 = insert_marker(&mut tree, 20, 20);
815 let id3 = insert_marker(&mut tree, 30, 30);
816 let id4 = insert_marker(&mut tree, 40, 40);
817
818 assert_eq!(get_pos(&tree, id0), (0, 0));
820 assert_eq!(get_pos(&tree, id1), (10, 10));
821 assert_eq!(get_pos(&tree, id2), (20, 20));
822 assert_eq!(get_pos(&tree, id3), (30, 30));
823 assert_eq!(get_pos(&tree, id4), (40, 40));
824
825 tree.adjust_for_edit(5, -16);
829
830 let positions = vec![
832 get_pos(&tree, id0).0,
833 get_pos(&tree, id1).0,
834 get_pos(&tree, id2).0,
835 get_pos(&tree, id3).0,
836 get_pos(&tree, id4).0,
837 ];
838
839 for i in 0..positions.len() - 1 {
841 assert!(
842 positions[i] <= positions[i + 1],
843 "Ordering violated at index {}: {:?}[{}]={} > {:?}[{}]={}",
844 i,
845 positions,
846 i,
847 positions[i],
848 positions,
849 i + 1,
850 positions[i + 1]
851 );
852 }
853
854 assert_eq!(get_pos(&tree, id0), (0, 0), "Marker at 0 should stay at 0");
856 assert_eq!(
857 get_pos(&tree, id1),
858 (5, 5),
859 "Marker at 10 should clamp to 5"
860 );
861 assert_eq!(
862 get_pos(&tree, id2),
863 (5, 5),
864 "Marker at 20 should clamp to 5"
865 );
866 assert_eq!(
867 get_pos(&tree, id3),
868 (14, 14),
869 "Marker at 30 should shift to 14"
870 );
871 assert_eq!(
872 get_pos(&tree, id4),
873 (24, 24),
874 "Marker at 40 should shift to 24"
875 );
876 }
877}