1use std::ops::Range;
2
3pub type Interval = Range<usize>;
4
5#[derive(PartialEq, Eq, Debug)]
7pub enum Mutation<V> {
8 ModifiedBegin(Interval, usize),
17 ModifiedEnd(Interval, usize),
26 Split(Interval, Interval, Interval),
35 Removed(Interval, V),
43}
44
45pub struct ItemIter<'a, V> {
46 map: &'a IntervalMap<V>,
47 i: usize,
48}
49
50impl<'a, V> Iterator for ItemIter<'a, V> {
51 type Item = (Interval, &'a V);
52
53 fn next(&mut self) -> Option<Self::Item> {
54 let i = self.i;
55 let m = self.map;
56 if i >= m.starts.len() {
57 return None;
58 }
59 let rv = Some(((m.starts[i]..m.ends[i]), &m.vals[i]));
60 self.i += 1;
61 rv
62 }
63}
64
65pub struct KeyIter<'a, V> {
66 map: &'a IntervalMap<V>,
67 i: usize,
68}
69
70impl<V> Iterator for KeyIter<'_, V> {
71 type Item = Interval;
72
73 fn next(&mut self) -> Option<Self::Item> {
74 let i = self.i;
75 let m = self.map;
76 if i >= m.starts.len() {
77 return None;
78 }
79 let rv = Some(m.starts[i]..m.ends[i]);
80 self.i += 1;
81 rv
82 }
83}
84
85#[derive(Clone, Debug)]
86pub struct IntervalMap<V> {
87 starts: Vec<usize>,
88 ends: Vec<usize>,
89 vals: Vec<V>,
90}
91
92impl<V: Clone> IntervalMap<V> {
94 pub fn new() -> IntervalMap<V> {
95 IntervalMap {
96 starts: Vec::new(),
97 ends: Vec::new(),
98 vals: Vec::new(),
99 }
100 }
101
102 pub fn keys(&self) -> KeyIter<V> {
104 KeyIter { map: self, i: 0 }
105 }
106
107 pub fn iter(&self) -> ItemIter<V> {
109 ItemIter { map: self, i: 0 }
110 }
111
112 pub fn iter_from(&self, begin: usize) -> ItemIter<V> {
115 let idx = match self.starts.binary_search(&begin) {
116 Ok(i) => i,
117 Err(i) if (i > 0 && begin < self.ends[i - 1]) => i - 1,
118 Err(i) => i,
119 };
120 ItemIter { map: self, i: idx }
121 }
122
123 pub fn clear(&mut self, interval: Interval) -> Vec<Mutation<V>> {
130 self.splice(interval, None)
131 }
132
133 pub fn insert(&mut self, interval: Interval, val: V) -> Vec<Mutation<V>> {
136 self.splice(interval, Some(val))
137 }
138
139 fn splice(&mut self, interval: Interval, val: Option<V>) -> Vec<Mutation<V>> {
141 let start = interval.start;
142 let end = interval.end;
143
144 assert!(start < end);
146
147 let mut mutations = Vec::new();
149
150 let mut starts_insertions = Vec::new();
152 let mut ends_insertions = Vec::new();
153 let mut vals_insertions = Vec::new();
154
155 if let Some(v) = val {
157 starts_insertions.push(start);
158 ends_insertions.push(end);
159 vals_insertions.push(v);
160 };
161
162 let splice_start = match self.starts.binary_search(&start) {
165 Ok(i) | Err(i) => i,
166 };
167
168 if splice_start > 0 && self.ends[splice_start - 1] > start {
171 let overlapping_idx = splice_start - 1;
172 let overlapping_int = self.starts[overlapping_idx]..self.ends[overlapping_idx];
173
174 if overlapping_int.end <= end {
175 self.ends[overlapping_idx] = start;
179 mutations.push(Mutation::ModifiedEnd(
180 overlapping_int,
181 self.ends[overlapping_idx],
182 ));
183 } else {
184 let new1 = overlapping_int.start..start;
189 let new2 = end..overlapping_int.end;
190
191 self.ends[overlapping_idx] = new1.end;
193
194 starts_insertions.push(new2.start);
196 ends_insertions.push(new2.end);
197 vals_insertions.push(self.vals[overlapping_idx].clone());
198 mutations.push(Mutation::Split(overlapping_int, new1, new2));
199 }
200 }
201
202 let splice_end = match self.ends.binary_search(&end) {
206 Ok(i) => i + 1,
207 Err(i) => i,
208 };
209
210 let mut modified_start: Option<Mutation<V>> = None;
212 if splice_end < self.starts.len()
213 && self.starts[splice_end] < end
214 && self.ends[splice_end] > end
215 {
216 let overlapping_idx = splice_end;
217 let overlapping_int = self.starts[overlapping_idx]..self.ends[overlapping_idx];
218 self.starts[overlapping_idx] = end;
222 modified_start = Some(Mutation::ModifiedBegin(
225 overlapping_int,
226 self.starts[overlapping_idx],
227 ));
228 }
229
230 let dropped_starts = self
235 .starts
236 .splice(splice_start..splice_end, starts_insertions);
237 let dropped_ends = self.ends.splice(splice_start..splice_end, ends_insertions);
238 let mut dropped_vals = self.vals.splice(splice_start..splice_end, vals_insertions);
239 for (dropped_start, dropped_end) in dropped_starts.zip(dropped_ends) {
240 mutations.push(Mutation::Removed(
241 dropped_start..dropped_end,
242 dropped_vals.next().unwrap(),
243 ));
244 }
245
246 if let Some(m) = modified_start {
248 mutations.push(m)
249 }
250
251 mutations
252 }
253
254 fn get_index(&self, x: usize) -> Option<usize> {
256 match self.starts.binary_search(&x) {
257 Ok(i) => Some(i),
258 Err(i) => {
259 if i == 0 {
260 None
261 } else if x < self.ends[i - 1] {
262 Some(i - 1)
263 } else {
264 None
265 }
266 }
267 }
268 }
269
270 pub fn get(&self, x: usize) -> Option<(Interval, &V)> {
272 self.get_index(x)
273 .map(|i| (self.starts[i]..self.ends[i], &self.vals[i]))
274 }
275
276 pub fn get_mut(&mut self, x: usize) -> Option<(Interval, &mut V)> {
278 match self.get_index(x) {
279 None => None,
280 Some(i) => Some((self.starts[i]..self.ends[i], &mut self.vals[i])),
281 }
282 }
283}
284
285impl<V: Clone> Default for IntervalMap<V> {
286 fn default() -> Self {
287 Self::new()
288 }
289}
290
291#[cfg(test)]
292mod tests {
293 use super::*;
294
295 fn interval_sum<I>(i: I) -> usize
296 where
297 I: Iterator<Item = Interval>,
298 {
299 i.map(|x| x.end - x.start).sum()
300 }
301
302 fn validate_map<V: Clone>(m: &IntervalMap<V>) -> Result<(), String> {
303 for i in m.keys() {
305 if i.start >= i.end {
307 return Err(format!("Invalid interval {i:?}"));
308 }
309 }
310
311 for (i1, i2) in m.keys().zip(m.keys().skip(1)) {
313 if i1.end > i2.start {
314 return Err(format!("Overlapping intervals {i1:?} and {i2:?}"));
315 }
316 }
317
318 Ok(())
319 }
320
321 fn insert_and_sanity_check(
322 m: &mut IntervalMap<String>,
323 interval: Interval,
324 val: &str,
325 ) -> Result<Vec<Mutation<String>>, String> {
326 let old_len_sum = interval_sum(m.keys());
327
328 let mutations = m.insert(interval.clone(), val.to_string());
330
331 validate_map(m)?;
333
334 let new_len_sum = interval_sum(m.keys());
335 let new_len = m.keys().count();
336 if new_len_sum < old_len_sum {
337 return Err(format!(
338 "length-sum shrunk from {old_len_sum} to {new_len_sum}"
339 ));
340 }
341 if new_len_sum < (interval.end - interval.start) {
342 return Err(format!(
343 "length-sum {} is smaller than inserted interval length {}",
344 new_len_sum,
345 interval.end - interval.start
346 ));
347 }
348 if new_len == 0 {
349 return Err("new length is zero".to_string());
350 }
351
352 Ok(mutations)
353 }
354
355 #[test]
356 #[cfg_attr(miri, ignore)]
358 fn test_insert_random() {
359 use rand::Rng;
360 use rand_core::SeedableRng;
361 let dist = rand::distr::Uniform::new_inclusive(0, 10).unwrap();
362 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(10);
363 for i in 0..1000 {
364 let mut m = IntervalMap::new();
365 for j in 0..10 {
366 let start: usize = rng.sample(dist);
367 let end: usize = start + rng.sample(dist) + 1;
368 let mut m_clone = m.clone();
369 let res = std::panic::catch_unwind(move || {
371 insert_and_sanity_check(&mut m_clone, start..end, &format!("{i}.{j}")).unwrap();
372 m_clone
373 });
374 if res.is_err() {
375 println!(
376 "Failed inserting {} -> {} into {:?}",
377 start,
378 end,
379 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>()
380 );
381 }
382 m = res.unwrap();
383 }
384 }
385 }
386
387 fn insert_and_validate(
388 m: &mut IntervalMap<String>,
389 interval: Interval,
390 val: &str,
391 expected_mutations: &[Mutation<String>],
392 expected_val: &[(Interval, &str)],
393 ) {
394 let mutations = insert_and_sanity_check(m, interval, val).unwrap();
395
396 assert_eq!(mutations, expected_mutations);
398
399 assert_eq!(
401 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>(),
402 expected_val
403 .iter()
404 .map(|(i, s)| (i.clone(), s.to_string()))
405 .collect::<Vec<_>>()
406 );
407 }
408
409 #[test]
410 fn test_insert_into_empty() {
411 let mut m = IntervalMap::new();
412 insert_and_validate(&mut m, 10..20, "x", &[], &[(10..20, "x")]);
413 }
414
415 #[test]
416 fn test_insert_after() {
417 let mut m = IntervalMap::new();
418 m.insert(1..3, "i1".to_string());
419 insert_and_validate(&mut m, 4..6, "i2", &[], &[(1..3, "i1"), (4..6, "i2")]);
420 }
421
422 #[test]
423 fn test_insert_before() {
424 let mut m = IntervalMap::new();
425 m.insert(4..6, "i1".to_string());
426 insert_and_validate(&mut m, 1..3, "i2", &[], &[(1..3, "i2"), (4..6, "i1")]);
427 }
428
429 #[test]
430 fn test_insert_just_before() {
431 let mut m = IntervalMap::new();
432 m.insert(20..30, "first".to_string());
433 insert_and_validate(
434 &mut m,
435 10..20,
436 "second",
437 &[],
438 &[(10..20, "second"), (20..30, "first")],
439 );
440 }
441
442 #[test]
443 fn test_insert_over_start() {
444 let mut m = IntervalMap::new();
445 m.insert(20..30, "first".to_string());
446 insert_and_validate(
447 &mut m,
448 10..21,
449 "second",
450 &[Mutation::ModifiedBegin(20..30, 21)],
451 &[(10..21, "second"), (21..30, "first")],
452 );
453 }
454
455 #[test]
456 fn test_insert_on_start() {
457 let mut m = IntervalMap::new();
458 m.insert(20..30, "first".to_string());
459 insert_and_validate(
460 &mut m,
461 20..21,
462 "second",
463 &[Mutation::ModifiedBegin(20..30, 21)],
464 &[(20..21, "second"), (21..30, "first")],
465 );
466 }
467
468 #[test]
469 fn test_insert_just_after() {
470 let mut m = IntervalMap::new();
471 m.insert(20..30, "first".to_string());
472 insert_and_validate(
473 &mut m,
474 30..40,
475 "second",
476 &[],
477 &[(20..30, "first"), (30..40, "second")],
478 );
479 }
480
481 #[test]
482 fn test_insert_over_end() {
483 let mut m = IntervalMap::new();
484 m.insert(20..30, "first".to_string());
485 insert_and_validate(
486 &mut m,
487 29..31,
488 "second",
489 &[Mutation::ModifiedEnd(20..30, 29)],
490 &[(20..29, "first"), (29..31, "second")],
491 );
492 }
493
494 #[test]
495 fn test_insert_on_end() {
496 let mut m = IntervalMap::new();
497 m.insert(20..30, "first".to_string());
498 insert_and_validate(
499 &mut m,
500 29..30,
501 "second",
502 &[Mutation::ModifiedEnd(20..30, 29)],
503 &[(20..29, "first"), (29..30, "second")],
504 );
505 }
506
507 #[test]
508 fn test_insert_removing() {
509 let mut m = IntervalMap::new();
510 m.insert(20..30, "first".to_string());
511 insert_and_validate(
512 &mut m,
513 10..40,
514 "second",
515 &[Mutation::Removed(20..30, "first".to_string())],
516 &[(10..40, "second")],
517 );
518 }
519
520 #[test]
521 fn test_insert_forcing_split() {
522 let mut m = IntervalMap::new();
523 m.insert(20..30, "first".to_string());
524 insert_and_validate(
525 &mut m,
526 24..25,
527 "second",
528 &[Mutation::Split(20..30, 20..24, 25..30)],
529 &[(20..24, "first"), (24..25, "second"), (25..30, "first")],
530 );
531 }
532
533 #[test]
534 fn test_insert_over_exact() {
535 let mut m = IntervalMap::new();
536 m.insert(1..2, "old".to_string());
537 insert_and_validate(
538 &mut m,
539 1..2,
540 "new",
541 &[Mutation::Removed(1..2, "old".to_string())],
542 &[(1..2, "new")],
543 );
544 }
545
546 #[test]
547 fn test_insert_all_mutations() {
548 let mut m = IntervalMap::new();
549 m.insert(0..10, "first".to_string());
550 m.insert(20..30, "second".to_string());
551 m.insert(40..50, "third".to_string());
552 insert_and_validate(
553 &mut m,
554 5..45,
555 "clobbering",
556 &[
557 Mutation::ModifiedEnd(0..10, 5),
558 Mutation::Removed(20..30, "second".to_string()),
559 Mutation::ModifiedBegin(40..50, 45),
560 ],
561 &[(0..5, "first"), (5..45, "clobbering"), (45..50, "third")],
562 );
563 }
564
565 fn clear_and_sanity_check(
566 m: &mut IntervalMap<String>,
567 start: usize,
568 end: usize,
569 ) -> Result<Vec<Mutation<String>>, String> {
570 let old_len = interval_sum(m.keys());
571
572 let mutations = m.clear(start..end);
574
575 validate_map(m)?;
577
578 let new_len = interval_sum(m.keys());
579 if new_len > old_len {
580 return Err(format!("new_len {new_len} > old_len {old_len}"));
581 }
582
583 Ok(mutations)
584 }
585
586 #[test]
587 #[cfg_attr(miri, ignore)]
589 fn test_clear_random() {
590 use rand::Rng;
591 use rand_core::SeedableRng;
592 let dist = rand::distr::Uniform::new_inclusive(0, 10).unwrap();
593 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(10);
594 for i in 0..1000 {
595 let mut m = IntervalMap::new();
596 for j in 0..10 {
597 let insert_start = rng.sample(dist);
598 let insert_end = insert_start + 1 + rng.sample(dist);
599 let clear_start = rng.sample(dist);
600 let clear_end = clear_start + 1 + rng.sample(dist);
601 let mut m_clone = m.clone();
602 let res = std::panic::catch_unwind(move || {
604 clear_and_sanity_check(&mut m_clone, clear_start, clear_end).unwrap();
605 insert_and_sanity_check(
606 &mut m_clone,
607 insert_start..insert_end,
608 &format!("{i}.{j}"),
609 )
610 .unwrap();
611 m_clone
612 });
613 if res.is_err() {
614 println!(
615 "Failed after inserting {} -> {} and clearing {} -> {} in {:?}",
616 insert_start,
617 insert_end,
618 clear_start,
619 clear_end,
620 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>()
621 );
622 }
623 m = res.unwrap();
624 }
625 }
626 }
627
628 fn clear_and_validate(
629 m: &mut IntervalMap<String>,
630 interval: Interval,
631 expected_mutations: &[Mutation<String>],
632 expected_val: &[(Interval, &str)],
633 ) {
634 let start = interval.start;
635 let end = interval.end;
636 let mutations = clear_and_sanity_check(m, start, end).unwrap();
637
638 assert_eq!(mutations, expected_mutations);
640
641 assert_eq!(
643 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>(),
644 expected_val
645 .iter()
646 .map(|(i, s)| (i.clone(), s.to_string()))
647 .collect::<Vec<_>>()
648 );
649 }
650
651 #[test]
652 fn test_clear_over_start() {
653 let mut m = IntervalMap::new();
654 m.insert(20..30, "first".to_string());
655 clear_and_validate(
656 &mut m,
657 10..21,
658 &[Mutation::ModifiedBegin(20..30, 21)],
659 &[(21..30, "first")],
660 );
661 }
662
663 #[test]
664 fn test_clear_over_end() {
665 let mut m = IntervalMap::new();
666 m.insert(20..30, "first".to_string());
667 clear_and_validate(
668 &mut m,
669 29..31,
670 &[Mutation::ModifiedEnd(20..30, 29)],
671 &[(20..29, "first")],
672 );
673 }
674
675 #[test]
676 fn test_clear_forcing_split() {
677 let mut m = IntervalMap::new();
678 m.insert(20..30, "first".to_string());
679 clear_and_validate(
680 &mut m,
681 24..25,
682 &[Mutation::Split(20..30, 20..24, 25..30)],
683 &[(20..24, "first"), (25..30, "first")],
684 );
685 }
686
687 #[test]
688 fn test_clear_removing() {
689 let mut m = IntervalMap::new();
690 m.insert(20..30, "first".to_string());
691 clear_and_validate(
692 &mut m,
693 10..40,
694 &[Mutation::Removed(20..30, "first".to_string())],
695 &[],
696 );
697 }
698
699 #[test]
700 fn test_get_empty() {
701 let m = IntervalMap::<String>::new();
702 assert_eq!(m.get(10), None);
703 }
704
705 #[test]
706 fn test_get_single_interval() {
707 let mut m = IntervalMap::<String>::new();
708 m.insert(1..4, "interval".to_string());
709 assert_eq!(m.get(0), None);
710 assert_eq!(m.get(1), Some((1..4, &"interval".to_string())));
711 assert_eq!(m.get(2), Some((1..4, &"interval".to_string())));
712 assert_eq!(m.get(3), Some((1..4, &"interval".to_string())));
713 assert_eq!(m.get(4), None);
714 assert_eq!(m.get(5), None);
715 }
716
717 #[test]
718 fn test_get_two_intervals_with_gap() {
719 let mut m = IntervalMap::<String>::new();
720 m.insert(1..4, "i1".to_string());
721 m.insert(5..9, "i2".to_string());
722 assert_eq!(m.get(0), None);
723 assert_eq!(m.get(1), Some((1..4, &"i1".to_string())));
724 assert_eq!(m.get(2), Some((1..4, &"i1".to_string())));
725 assert_eq!(m.get(3), Some((1..4, &"i1".to_string())));
726 assert_eq!(m.get(4), None);
727 assert_eq!(m.get(5), Some((5..9, &"i2".to_string())));
728 assert_eq!(m.get(6), Some((5..9, &"i2".to_string())));
729 assert_eq!(m.get(7), Some((5..9, &"i2".to_string())));
730 assert_eq!(m.get(8), Some((5..9, &"i2".to_string())));
731 assert_eq!(m.get(9), None);
732 }
733
734 #[test]
735 fn test_get_two_intervals_without_gap() {
736 let mut m = IntervalMap::<String>::new();
737 m.insert(1..3, "i1".to_string());
738 m.insert(3..6, "i2".to_string());
739 assert_eq!(m.get(0), None);
740 assert_eq!(m.get(1), Some((1..3, &"i1".to_string())));
741 assert_eq!(m.get(2), Some((1..3, &"i1".to_string())));
742 assert_eq!(m.get(3), Some((3..6, &"i2".to_string())));
743 assert_eq!(m.get(4), Some((3..6, &"i2".to_string())));
744 assert_eq!(m.get(5), Some((3..6, &"i2".to_string())));
745 assert_eq!(m.get(6), None);
746 }
747
748 #[test]
749 fn test_iter_from() {
750 let mut m = IntervalMap::<&str>::new();
751 m.insert(1..3, "i1");
752 m.insert(4..6, "i2");
753 assert_eq!(
754 m.iter_from(0).collect::<Vec<_>>(),
755 vec![(1..3, &"i1"), (4..6, &"i2")]
756 );
757 assert_eq!(
758 m.iter_from(1).collect::<Vec<_>>(),
759 vec![(1..3, &"i1"), (4..6, &"i2")]
760 );
761 assert_eq!(
762 m.iter_from(2).collect::<Vec<_>>(),
763 vec![(1..3, &"i1"), (4..6, &"i2")]
764 );
765 assert_eq!(m.iter_from(3).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
766 assert_eq!(m.iter_from(4).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
767 assert_eq!(m.iter_from(5).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
768 assert_eq!(m.iter_from(6).collect::<Vec<_>>(), vec![]);
769 assert_eq!(m.iter_from(7).collect::<Vec<_>>(), vec![]);
770 }
771}