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 {:?} and {:?}", i1, 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 {} to {}",
339 old_len_sum, new_len_sum
340 ));
341 }
342 if new_len_sum < (interval.end - interval.start) {
343 return Err(format!(
344 "length-sum {} is smaller than inserted interval length {}",
345 new_len_sum,
346 interval.end - interval.start
347 ));
348 }
349 if new_len == 0 {
350 return Err("new length is zero".to_string());
351 }
352
353 Ok(mutations)
354 }
355
356 #[test]
357 #[cfg_attr(miri, ignore)]
359 fn test_insert_random() {
360 use rand::Rng;
361 use rand_core::SeedableRng;
362 let dist = rand::distr::Uniform::new_inclusive(0, 10).unwrap();
363 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(10);
364 for i in 0..1000 {
365 let mut m = IntervalMap::new();
366 for j in 0..10 {
367 let start: usize = rng.sample(dist);
368 let end: usize = start + rng.sample(dist) + 1;
369 let mut m_clone = m.clone();
370 let res = std::panic::catch_unwind(move || {
372 insert_and_sanity_check(&mut m_clone, start..end, &format!("{}.{}", i, j))
373 .unwrap();
374 m_clone
375 });
376 if res.is_err() {
377 println!(
378 "Failed inserting {} -> {} into {:?}",
379 start,
380 end,
381 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>()
382 );
383 }
384 m = res.unwrap();
385 }
386 }
387 }
388
389 fn insert_and_validate(
390 m: &mut IntervalMap<String>,
391 interval: Interval,
392 val: &str,
393 expected_mutations: &[Mutation<String>],
394 expected_val: &[(Interval, &str)],
395 ) {
396 let mutations = insert_and_sanity_check(m, interval, val).unwrap();
397
398 assert_eq!(mutations, expected_mutations);
400
401 assert_eq!(
403 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>(),
404 expected_val
405 .iter()
406 .map(|(i, s)| (i.clone(), s.to_string()))
407 .collect::<Vec<_>>()
408 );
409 }
410
411 #[test]
412 fn test_insert_into_empty() {
413 let mut m = IntervalMap::new();
414 insert_and_validate(&mut m, 10..20, "x", &[], &[(10..20, "x")]);
415 }
416
417 #[test]
418 fn test_insert_after() {
419 let mut m = IntervalMap::new();
420 m.insert(1..3, "i1".to_string());
421 insert_and_validate(&mut m, 4..6, "i2", &[], &[(1..3, "i1"), (4..6, "i2")]);
422 }
423
424 #[test]
425 fn test_insert_before() {
426 let mut m = IntervalMap::new();
427 m.insert(4..6, "i1".to_string());
428 insert_and_validate(&mut m, 1..3, "i2", &[], &[(1..3, "i2"), (4..6, "i1")]);
429 }
430
431 #[test]
432 fn test_insert_just_before() {
433 let mut m = IntervalMap::new();
434 m.insert(20..30, "first".to_string());
435 insert_and_validate(
436 &mut m,
437 10..20,
438 "second",
439 &[],
440 &[(10..20, "second"), (20..30, "first")],
441 );
442 }
443
444 #[test]
445 fn test_insert_over_start() {
446 let mut m = IntervalMap::new();
447 m.insert(20..30, "first".to_string());
448 insert_and_validate(
449 &mut m,
450 10..21,
451 "second",
452 &[Mutation::ModifiedBegin(20..30, 21)],
453 &[(10..21, "second"), (21..30, "first")],
454 );
455 }
456
457 #[test]
458 fn test_insert_on_start() {
459 let mut m = IntervalMap::new();
460 m.insert(20..30, "first".to_string());
461 insert_and_validate(
462 &mut m,
463 20..21,
464 "second",
465 &[Mutation::ModifiedBegin(20..30, 21)],
466 &[(20..21, "second"), (21..30, "first")],
467 );
468 }
469
470 #[test]
471 fn test_insert_just_after() {
472 let mut m = IntervalMap::new();
473 m.insert(20..30, "first".to_string());
474 insert_and_validate(
475 &mut m,
476 30..40,
477 "second",
478 &[],
479 &[(20..30, "first"), (30..40, "second")],
480 );
481 }
482
483 #[test]
484 fn test_insert_over_end() {
485 let mut m = IntervalMap::new();
486 m.insert(20..30, "first".to_string());
487 insert_and_validate(
488 &mut m,
489 29..31,
490 "second",
491 &[Mutation::ModifiedEnd(20..30, 29)],
492 &[(20..29, "first"), (29..31, "second")],
493 );
494 }
495
496 #[test]
497 fn test_insert_on_end() {
498 let mut m = IntervalMap::new();
499 m.insert(20..30, "first".to_string());
500 insert_and_validate(
501 &mut m,
502 29..30,
503 "second",
504 &[Mutation::ModifiedEnd(20..30, 29)],
505 &[(20..29, "first"), (29..30, "second")],
506 );
507 }
508
509 #[test]
510 fn test_insert_removing() {
511 let mut m = IntervalMap::new();
512 m.insert(20..30, "first".to_string());
513 insert_and_validate(
514 &mut m,
515 10..40,
516 "second",
517 &[Mutation::Removed(20..30, "first".to_string())],
518 &[(10..40, "second")],
519 );
520 }
521
522 #[test]
523 fn test_insert_forcing_split() {
524 let mut m = IntervalMap::new();
525 m.insert(20..30, "first".to_string());
526 insert_and_validate(
527 &mut m,
528 24..25,
529 "second",
530 &[Mutation::Split(20..30, 20..24, 25..30)],
531 &[(20..24, "first"), (24..25, "second"), (25..30, "first")],
532 );
533 }
534
535 #[test]
536 fn test_insert_over_exact() {
537 let mut m = IntervalMap::new();
538 m.insert(1..2, "old".to_string());
539 insert_and_validate(
540 &mut m,
541 1..2,
542 "new",
543 &[Mutation::Removed(1..2, "old".to_string())],
544 &[(1..2, "new")],
545 );
546 }
547
548 #[test]
549 fn test_insert_all_mutations() {
550 let mut m = IntervalMap::new();
551 m.insert(0..10, "first".to_string());
552 m.insert(20..30, "second".to_string());
553 m.insert(40..50, "third".to_string());
554 insert_and_validate(
555 &mut m,
556 5..45,
557 "clobbering",
558 &[
559 Mutation::ModifiedEnd(0..10, 5),
560 Mutation::Removed(20..30, "second".to_string()),
561 Mutation::ModifiedBegin(40..50, 45),
562 ],
563 &[(0..5, "first"), (5..45, "clobbering"), (45..50, "third")],
564 );
565 }
566
567 fn clear_and_sanity_check(
568 m: &mut IntervalMap<String>,
569 start: usize,
570 end: usize,
571 ) -> Result<Vec<Mutation<String>>, String> {
572 let old_len = interval_sum(m.keys());
573
574 let mutations = m.clear(start..end);
576
577 validate_map(m)?;
579
580 let new_len = interval_sum(m.keys());
581 if new_len > old_len {
582 return Err(format!("new_len {} > old_len {}", new_len, old_len));
583 }
584
585 Ok(mutations)
586 }
587
588 #[test]
589 #[cfg_attr(miri, ignore)]
591 fn test_clear_random() {
592 use rand::Rng;
593 use rand_core::SeedableRng;
594 let dist = rand::distr::Uniform::new_inclusive(0, 10).unwrap();
595 let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(10);
596 for i in 0..1000 {
597 let mut m = IntervalMap::new();
598 for j in 0..10 {
599 let insert_start = rng.sample(dist);
600 let insert_end = insert_start + 1 + rng.sample(dist);
601 let clear_start = rng.sample(dist);
602 let clear_end = clear_start + 1 + rng.sample(dist);
603 let mut m_clone = m.clone();
604 let res = std::panic::catch_unwind(move || {
606 clear_and_sanity_check(&mut m_clone, clear_start, clear_end).unwrap();
607 insert_and_sanity_check(
608 &mut m_clone,
609 insert_start..insert_end,
610 &format!("{}.{}", i, j),
611 )
612 .unwrap();
613 m_clone
614 });
615 if res.is_err() {
616 println!(
617 "Failed after inserting {} -> {} and clearing {} -> {} in {:?}",
618 insert_start,
619 insert_end,
620 clear_start,
621 clear_end,
622 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>()
623 );
624 }
625 m = res.unwrap();
626 }
627 }
628 }
629
630 fn clear_and_validate(
631 m: &mut IntervalMap<String>,
632 interval: Interval,
633 expected_mutations: &[Mutation<String>],
634 expected_val: &[(Interval, &str)],
635 ) {
636 let start = interval.start;
637 let end = interval.end;
638 let mutations = clear_and_sanity_check(m, start, end).unwrap();
639
640 assert_eq!(mutations, expected_mutations);
642
643 assert_eq!(
645 m.iter().map(|(i, s)| (i, s.clone())).collect::<Vec<_>>(),
646 expected_val
647 .iter()
648 .map(|(i, s)| (i.clone(), s.to_string()))
649 .collect::<Vec<_>>()
650 );
651 }
652
653 #[test]
654 fn test_clear_over_start() {
655 let mut m = IntervalMap::new();
656 m.insert(20..30, "first".to_string());
657 clear_and_validate(
658 &mut m,
659 10..21,
660 &[Mutation::ModifiedBegin(20..30, 21)],
661 &[(21..30, "first")],
662 );
663 }
664
665 #[test]
666 fn test_clear_over_end() {
667 let mut m = IntervalMap::new();
668 m.insert(20..30, "first".to_string());
669 clear_and_validate(
670 &mut m,
671 29..31,
672 &[Mutation::ModifiedEnd(20..30, 29)],
673 &[(20..29, "first")],
674 );
675 }
676
677 #[test]
678 fn test_clear_forcing_split() {
679 let mut m = IntervalMap::new();
680 m.insert(20..30, "first".to_string());
681 clear_and_validate(
682 &mut m,
683 24..25,
684 &[Mutation::Split(20..30, 20..24, 25..30)],
685 &[(20..24, "first"), (25..30, "first")],
686 );
687 }
688
689 #[test]
690 fn test_clear_removing() {
691 let mut m = IntervalMap::new();
692 m.insert(20..30, "first".to_string());
693 clear_and_validate(
694 &mut m,
695 10..40,
696 &[Mutation::Removed(20..30, "first".to_string())],
697 &[],
698 );
699 }
700
701 #[test]
702 fn test_get_empty() {
703 let m = IntervalMap::<String>::new();
704 assert_eq!(m.get(10), None);
705 }
706
707 #[test]
708 fn test_get_single_interval() {
709 let mut m = IntervalMap::<String>::new();
710 m.insert(1..4, "interval".to_string());
711 assert_eq!(m.get(0), None);
712 assert_eq!(m.get(1), Some((1..4, &"interval".to_string())));
713 assert_eq!(m.get(2), Some((1..4, &"interval".to_string())));
714 assert_eq!(m.get(3), Some((1..4, &"interval".to_string())));
715 assert_eq!(m.get(4), None);
716 assert_eq!(m.get(5), None);
717 }
718
719 #[test]
720 fn test_get_two_intervals_with_gap() {
721 let mut m = IntervalMap::<String>::new();
722 m.insert(1..4, "i1".to_string());
723 m.insert(5..9, "i2".to_string());
724 assert_eq!(m.get(0), None);
725 assert_eq!(m.get(1), Some((1..4, &"i1".to_string())));
726 assert_eq!(m.get(2), Some((1..4, &"i1".to_string())));
727 assert_eq!(m.get(3), Some((1..4, &"i1".to_string())));
728 assert_eq!(m.get(4), None);
729 assert_eq!(m.get(5), Some((5..9, &"i2".to_string())));
730 assert_eq!(m.get(6), Some((5..9, &"i2".to_string())));
731 assert_eq!(m.get(7), Some((5..9, &"i2".to_string())));
732 assert_eq!(m.get(8), Some((5..9, &"i2".to_string())));
733 assert_eq!(m.get(9), None);
734 }
735
736 #[test]
737 fn test_get_two_intervals_without_gap() {
738 let mut m = IntervalMap::<String>::new();
739 m.insert(1..3, "i1".to_string());
740 m.insert(3..6, "i2".to_string());
741 assert_eq!(m.get(0), None);
742 assert_eq!(m.get(1), Some((1..3, &"i1".to_string())));
743 assert_eq!(m.get(2), Some((1..3, &"i1".to_string())));
744 assert_eq!(m.get(3), Some((3..6, &"i2".to_string())));
745 assert_eq!(m.get(4), Some((3..6, &"i2".to_string())));
746 assert_eq!(m.get(5), Some((3..6, &"i2".to_string())));
747 assert_eq!(m.get(6), None);
748 }
749
750 #[test]
751 fn test_iter_from() {
752 let mut m = IntervalMap::<&str>::new();
753 m.insert(1..3, "i1");
754 m.insert(4..6, "i2");
755 assert_eq!(
756 m.iter_from(0).collect::<Vec<_>>(),
757 vec![(1..3, &"i1"), (4..6, &"i2")]
758 );
759 assert_eq!(
760 m.iter_from(1).collect::<Vec<_>>(),
761 vec![(1..3, &"i1"), (4..6, &"i2")]
762 );
763 assert_eq!(
764 m.iter_from(2).collect::<Vec<_>>(),
765 vec![(1..3, &"i1"), (4..6, &"i2")]
766 );
767 assert_eq!(m.iter_from(3).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
768 assert_eq!(m.iter_from(4).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
769 assert_eq!(m.iter_from(5).collect::<Vec<_>>(), vec![(4..6, &"i2")]);
770 assert_eq!(m.iter_from(6).collect::<Vec<_>>(), vec![]);
771 assert_eq!(m.iter_from(7).collect::<Vec<_>>(), vec![]);
772 }
773}