1use std::cmp;
28use std::iter::FromIterator;
29use std::ops::Range;
30
31use std::collections::BTreeMap;
32use std::collections::Bound;
33
34use either::Either;
35use smallvec::SmallVec;
36
37const MAX_INLINE_CAPACITY: usize = 4;
38const MIN_TO_INLINE: usize = 2;
39
40#[derive(Clone, PartialEq, Eq, PartialOrd)]
42pub enum RangeSet {
43 Inline(InlineRangeSet),
44 BTree(BTreeRangeSet),
45}
46
47#[derive(Clone, PartialEq, Eq, PartialOrd)]
50pub struct InlineRangeSet {
51 inner: SmallVec<[(u64, u64); MAX_INLINE_CAPACITY]>,
52 capacity: usize,
53}
54
55#[derive(Clone, PartialEq, Eq, PartialOrd)]
58pub struct BTreeRangeSet {
59 inner: BTreeMap<u64, u64>,
60 capacity: usize,
61}
62
63impl RangeSet {
64 pub fn new(capacity: usize) -> Self {
69 RangeSet::Inline(InlineRangeSet {
70 inner: Default::default(),
71 capacity,
72 })
73 }
74
75 pub fn len(&self) -> usize {
77 match self {
78 RangeSet::Inline(set) => set.inner.len(),
79 RangeSet::BTree(set) => set.inner.len(),
80 }
81 }
82
83 pub fn contains(&self, item: u64) -> bool {
85 match self {
86 RangeSet::Inline(set) => set
87 .inner
88 .iter()
89 .any(|&(start, end)| start <= item && item < end),
90
91 RangeSet::BTree(set) =>
92 set.prev_to(item).is_some_and(|range| item < range.end),
93 }
94 }
95
96 #[inline(always)]
99 fn fixup(&mut self) {
100 match self {
101 RangeSet::Inline(set) if set.inner.len() == MAX_INLINE_CAPACITY => {
102 let old_inner = std::mem::take(&mut set.inner);
103 *self = RangeSet::BTree(BTreeRangeSet {
104 inner: old_inner.into_inner().expect("At capacity").into(),
105 capacity: set.capacity,
106 });
107 },
108
109 RangeSet::BTree(set) if set.inner.len() <= MIN_TO_INLINE => {
110 let old_inner = std::mem::take(&mut set.inner);
111 *self = RangeSet::Inline(InlineRangeSet {
112 inner: SmallVec::from_iter(old_inner),
113 capacity: set.capacity,
114 })
115 },
116
117 _ => {},
118 }
119 }
120
121 #[inline]
127 pub fn insert(&mut self, item: Range<u64>) {
128 match self {
129 RangeSet::Inline(set) => set.insert(item),
130 RangeSet::BTree(set) => set.insert(item),
131 }
132
133 self.fixup();
134 }
135
136 pub fn iter(
138 &self,
139 ) -> impl DoubleEndedIterator<Item = Range<u64>> + ExactSizeIterator + '_
140 {
141 match self {
142 RangeSet::BTree(set) =>
143 Either::Left(set.inner.iter().map(|(k, v)| *k..*v)),
144
145 RangeSet::Inline(set) =>
146 Either::Right(set.inner.iter().map(|(s, e)| *s..*e)),
147 }
148 }
149
150 #[cfg(test)]
153 pub fn flatten(&self) -> impl DoubleEndedIterator<Item = u64> + '_ {
154 match self {
155 RangeSet::BTree(set) =>
156 Either::Left(set.inner.iter().flat_map(|(k, v)| *k..*v)),
157
158 RangeSet::Inline(set) =>
159 Either::Right(set.inner.iter().flat_map(|(s, e)| *s..*e)),
160 }
161 }
162
163 #[cfg(test)]
165 pub fn first(&self) -> Option<u64> {
166 match self {
167 RangeSet::Inline(set) => set.inner.first().map(|(s, _)| *s),
168
169 RangeSet::BTree(set) => set.inner.first_key_value().map(|(k, _)| *k),
170 }
171 }
172
173 pub fn last(&self) -> Option<u64> {
175 match self {
176 RangeSet::Inline(set) => set.inner.last().map(|(_, e)| *e - 1),
177
178 RangeSet::BTree(set) =>
179 set.inner.last_key_value().map(|(_, v)| *v - 1),
180 }
181 }
182
183 #[inline]
184 pub fn remove_until(&mut self, largest: u64) {
185 match self {
186 RangeSet::Inline(set) => set.remove_until(largest),
187 RangeSet::BTree(set) => set.remove_until(largest),
188 }
189
190 self.fixup();
191 }
192
193 pub fn push_item(&mut self, item: u64) {
194 self.insert(item..item + 1)
195 }
196}
197
198impl InlineRangeSet {
199 fn insert(&mut self, item: Range<u64>) {
200 let start = item.start;
201 let mut end = item.end;
202 let mut pos = 0;
203
204 loop {
205 match self.inner.get_mut(pos) {
206 Some((s, e)) => {
207 if start > *e {
208 pos += 1;
210 continue;
211 }
212
213 if end < *s {
214 if self.inner.len() == self.capacity {
217 self.inner.remove(0);
218 pos -= 1;
219 }
220
221 self.inner.insert(pos, (start, end));
222 return;
223 }
224
225 if start < *s {
227 *s = start;
230 }
231
232 if end > *e {
233 *e = end;
236 break;
237 } else {
238 return;
239 }
240 },
241
242 None => {
243 if self.inner.len() == self.capacity {
244 self.inner.remove(0);
245 }
246
247 self.inner.push((start, end));
248 return;
249 },
250 }
251 }
252
253 while let Some((s, e)) = self.inner.get(pos + 1).copied() {
255 if end < s {
256 break;
258 }
259
260 let new_e = e.max(end);
261 self.inner[pos].1 = new_e;
262 end = new_e;
263 self.inner.remove(pos + 1);
264 }
265 }
266
267 fn remove_until(&mut self, largest: u64) {
268 while let Some((s, e)) = self.inner.first_mut() {
269 if largest >= *e {
270 self.inner.remove(0);
271 continue;
272 }
273
274 *s = (largest + 1).max(*s);
275 if *s == *e {
276 self.inner.remove(0);
277 }
278
279 break;
280 }
281 }
282}
283
284impl BTreeRangeSet {
285 fn insert(&mut self, item: Range<u64>) {
287 let mut start = item.start;
288 let mut end = item.end;
289
290 if let Some(r) = self.prev_to(start) {
292 if range_overlaps(&r, &item) {
294 self.inner.remove(&r.start);
295
296 start = cmp::min(start, r.start);
297 end = cmp::max(end, r.end);
298 }
299 }
300
301 while let Some(r) = self.next_to(start) {
303 if item.contains(&r.start) && item.contains(&r.end) {
305 self.inner.remove(&r.start);
306 continue;
307 }
308
309 if !range_overlaps(&r, &item) {
311 break;
312 }
313
314 self.inner.remove(&r.start);
316
317 start = cmp::min(start, r.start);
318 end = cmp::max(end, r.end);
319 }
320
321 if self.inner.len() >= self.capacity {
322 self.inner.pop_first();
323 }
324
325 self.inner.insert(start, end);
326 }
327
328 fn remove_until(&mut self, largest: u64) {
329 let ranges: Vec<Range<u64>> = self
330 .inner
331 .range((Bound::Unbounded, Bound::Included(&largest)))
332 .map(|(&s, &e)| s..e)
333 .collect();
334
335 for r in ranges {
336 self.inner.remove(&r.start);
337
338 if r.end > largest + 1 {
339 let start = largest + 1;
340 self.insert(start..r.end);
341 }
342 }
343 }
344
345 fn prev_to(&self, item: u64) -> Option<Range<u64>> {
346 self.inner
347 .range((Bound::Unbounded, Bound::Included(item)))
348 .map(|(&s, &e)| s..e)
349 .next_back()
350 }
351
352 fn next_to(&self, item: u64) -> Option<Range<u64>> {
353 self.inner
354 .range((Bound::Included(item), Bound::Unbounded))
355 .map(|(&s, &e)| s..e)
356 .next()
357 }
358}
359
360impl Default for RangeSet {
361 fn default() -> Self {
362 RangeSet::Inline(InlineRangeSet {
363 inner: Default::default(),
364 capacity: usize::MAX,
365 })
366 }
367}
368
369impl PartialEq<Range<u64>> for RangeSet {
374 fn eq(&self, other: &Range<u64>) -> bool {
375 if self.len() != 1 {
378 return false;
379 }
380
381 let range = self.iter().next().unwrap();
383 range == *other
384 }
385}
386
387impl std::fmt::Debug for RangeSet {
388 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
389 let ranges: Vec<Range<u64>> = self
390 .iter()
391 .map(|mut r| {
392 r.end -= 1;
393 r
394 })
395 .collect();
396
397 write!(f, "{ranges:?}")
398 }
399}
400
401fn range_overlaps(r: &Range<u64>, other: &Range<u64>) -> bool {
402 other.start >= r.start && other.start <= r.end ||
403 other.end >= r.start && other.end <= r.end
404}
405
406#[cfg(test)]
407mod tests {
408 use super::*;
409
410 #[test]
411 fn contains_inline() {
412 let mut r = RangeSet::default();
413 assert!(!r.contains(0));
414 assert!(!r.contains(u64::MAX));
415
416 r.insert(4..7);
417 r.insert(9..12);
418 r.insert(u64::MAX - 2..u64::MAX);
419 assert!(matches!(r, RangeSet::Inline(_)));
420
421 for item in 0..15 {
422 assert_eq!(
423 r.contains(item),
424 (4..7).contains(&item) || (9..12).contains(&item)
425 );
426 }
427
428 assert!(!r.contains(u64::MAX - 3));
429 assert!(r.contains(u64::MAX - 2));
430 assert!(r.contains(u64::MAX - 1));
431 assert!(!r.contains(u64::MAX));
432 }
433
434 #[test]
435 fn contains_btree() {
436 let mut r = RangeSet::default();
437
438 for start in [4, 9, 14, u64::MAX - 3] {
439 r.insert(start..start + 3);
440 }
441
442 assert!(matches!(r, RangeSet::BTree(_)));
443
444 for item in 0..20 {
445 assert_eq!(
446 r.contains(item),
447 (4..7).contains(&item) ||
448 (9..12).contains(&item) ||
449 (14..17).contains(&item)
450 );
451 }
452
453 assert!(!r.contains(u64::MAX - 4));
454 assert!(r.contains(u64::MAX - 3));
455 assert!(r.contains(u64::MAX - 1));
456 assert!(!r.contains(u64::MAX));
457
458 r.remove_until(13);
459 assert!(matches!(r, RangeSet::Inline(_)));
460 assert!(!r.contains(12));
461 assert!(r.contains(14));
462 assert!(r.contains(16));
463 assert!(!r.contains(17));
464 }
465
466 #[test]
467 fn adjacent_singletons_merge() {
468 let mut r = RangeSet::default();
469
470 for item in [0, 2, 4, 6] {
471 r.push_item(item);
472 }
473
474 assert!(matches!(r, RangeSet::BTree(_)));
475
476 for item in [5, 1, 3] {
477 assert!(!r.contains(item));
478 r.push_item(item);
479 assert!(r.contains(item));
480 }
481
482 assert_eq!(r, 0..7);
483 assert!(matches!(r, RangeSet::Inline(_)));
484
485 for item in 0..7 {
486 assert!(r.contains(item));
487 }
488
489 assert!(!r.contains(7));
490 }
491
492 #[test]
493 fn insert_non_overlapping() {
494 let mut r = RangeSet::default();
495 assert_eq!(r.len(), 0);
496 let empty: &[u64] = &[];
497 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
498
499 r.insert(4..7);
500 assert_eq!(r.len(), 1);
501 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
502
503 r.insert(9..12);
504 assert_eq!(r.len(), 2);
505 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
506 }
507
508 #[test]
509 fn insert_contained() {
510 let mut r = RangeSet::default();
511
512 r.insert(4..7);
513 r.insert(9..12);
514 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
515
516 r.insert(4..7);
517 assert_eq!(r.len(), 2);
518 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
519
520 r.insert(4..6);
521 assert_eq!(r.len(), 2);
522 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
523
524 r.insert(5..6);
525 assert_eq!(r.len(), 2);
526 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
527
528 r.insert(10..11);
529 assert_eq!(r.len(), 2);
530 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
531
532 r.insert(9..11);
533 assert_eq!(r.len(), 2);
534 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
535 }
536
537 #[test]
538 fn insert_overlapping() {
539 let mut r = RangeSet::default();
540
541 r.insert(3..6);
542 r.insert(9..12);
543 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 9, 10, 11]);
544
545 r.insert(5..7);
546 assert_eq!(r.len(), 2);
547 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 6, 9, 10, 11]);
548
549 r.insert(10..15);
550 assert_eq!(r.len(), 2);
551 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
552 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
553 ]);
554
555 r.insert(2..5);
556 assert_eq!(r.len(), 2);
557 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
558 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
559 ]);
560
561 r.insert(8..10);
562 assert_eq!(r.len(), 2);
563 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
564 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14
565 ]);
566
567 r.insert(6..10);
568 assert_eq!(r.len(), 1);
569 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
570 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14
571 ]);
572 }
573
574 #[test]
575 fn insert_overlapping_multi() {
576 let mut r = RangeSet::default();
577
578 r.insert(3..6);
579 r.insert(16..20);
580 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
581 3, 4, 5, 16, 17, 18, 19
582 ]);
583
584 r.insert(10..11);
585 assert_eq!(r.len(), 3);
586 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
587 3, 4, 5, 10, 16, 17, 18, 19
588 ]);
589
590 assert!(matches!(r, RangeSet::Inline(_)));
591
592 r.insert(13..14);
593 assert_eq!(r.len(), 4);
594 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
595 3, 4, 5, 10, 13, 16, 17, 18, 19
596 ]);
597
598 assert!(matches!(r, RangeSet::BTree(_)));
600
601 r.insert(4..17);
602 assert_eq!(r.len(), 1);
603 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
604 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
605 ]);
606
607 assert!(matches!(r, RangeSet::Inline(_)));
609 }
610
611 #[test]
612 fn prev_to() {
613 let mut r = BTreeRangeSet {
614 inner: Default::default(),
615 capacity: usize::MAX,
616 };
617
618 r.insert(4..7);
619 r.insert(9..12);
620
621 assert_eq!(r.prev_to(2), None);
622 assert_eq!(r.prev_to(4), Some(4..7));
623 assert_eq!(r.prev_to(15), Some(9..12));
624 assert_eq!(r.prev_to(5), Some(4..7));
625 assert_eq!(r.prev_to(8), Some(4..7));
626 }
627
628 #[test]
629 fn next_to() {
630 let mut r = BTreeRangeSet {
631 inner: Default::default(),
632 capacity: usize::MAX,
633 };
634
635 r.insert(4..7);
636 r.insert(9..12);
637
638 assert_eq!(r.next_to(2), Some(4..7));
639 assert_eq!(r.next_to(12), None);
640 assert_eq!(r.next_to(15), None);
641 assert_eq!(r.next_to(5), Some(9..12));
642 assert_eq!(r.next_to(8), Some(9..12));
643 }
644
645 #[test]
646 fn push_item() {
647 let mut r = RangeSet::default();
648
649 r.insert(4..7);
650 r.insert(9..12);
651 assert_eq!(r.len(), 2);
652 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
653
654 r.push_item(15);
655 assert_eq!(r.len(), 3);
656 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
657 4, 5, 6, 9, 10, 11, 15
658 ]);
659
660 r.push_item(15);
661 assert_eq!(r.len(), 3);
662 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
663 4, 5, 6, 9, 10, 11, 15
664 ]);
665
666 r.push_item(1);
667 assert_eq!(r.len(), 4);
668 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
669 1, 4, 5, 6, 9, 10, 11, 15
670 ]);
671
672 r.push_item(12);
673 r.push_item(13);
674 r.push_item(14);
675
676 assert_eq!(r.len(), 3);
677 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
678 1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
679 ]);
680
681 r.push_item(2);
682 r.push_item(3);
683 assert_eq!(r.len(), 2);
684 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
685 1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
686 ]);
687
688 r.push_item(8);
689 r.push_item(7);
690 assert_eq!(r.len(), 1);
691 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
692 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
693 ]);
694 }
695
696 #[test]
697 fn flatten_rev() {
698 let mut r = RangeSet::default();
699 assert_eq!(r.len(), 0);
700
701 let empty: &[u64] = &[];
702 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
703
704 r.insert(4..7);
705 assert_eq!(r.len(), 1);
706 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
707 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[6, 5, 4]);
708
709 r.insert(9..12);
710 assert_eq!(r.len(), 2);
711 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
712 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[
713 11, 10, 9, 6, 5, 4
714 ]);
715 }
716
717 #[test]
718 fn flatten_one() {
719 let mut r = RangeSet::default();
720 assert_eq!(r.len(), 0);
721
722 let empty: &[u64] = &[];
723 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
724
725 r.insert(0..1);
726 assert_eq!(r.len(), 1);
727 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[0]);
728 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[0]);
729 }
730
731 #[test]
732 fn remove_largest() {
733 let mut r = RangeSet::default();
734
735 r.insert(3..6);
736 r.insert(9..11);
737 r.insert(13..14);
738 r.insert(16..20);
739 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
740 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
741 ]);
742
743 r.remove_until(2);
744 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
745 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
746 ]);
747
748 r.remove_until(4);
749 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
750 5, 9, 10, 13, 16, 17, 18, 19
751 ]);
752
753 r.remove_until(6);
754 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
755 9, 10, 13, 16, 17, 18, 19
756 ]);
757
758 r.remove_until(10);
759 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[13, 16, 17, 18, 19]);
760
761 r.remove_until(17);
762 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[18, 19]);
763
764 r.remove_until(18);
765 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[19]);
766
767 r.remove_until(20);
768
769 let empty: &[u64] = &[];
770 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
771 }
772
773 #[test]
774 fn eq_range() {
775 let mut r = RangeSet::default();
776 assert_ne!(r, 0..0);
777
778 let expected = 3..20;
779
780 r.insert(3..6);
781 assert_ne!(r, expected);
782
783 r.insert(16..20);
784 assert_ne!(r, expected);
785
786 r.insert(10..11);
787 assert_ne!(r, expected);
788
789 r.insert(13..14);
790 assert_ne!(r, expected);
791
792 r.insert(4..17);
793
794 assert_eq!(r, expected);
795 }
796
797 #[test]
798 fn first_last() {
799 let mut r = RangeSet::default();
800 assert_eq!(r.first(), None);
801 assert_eq!(r.last(), None);
802
803 r.insert(10..11);
804 assert_eq!(r.first(), Some(10));
805 assert_eq!(r.last(), Some(10));
806
807 r.insert(13..14);
808 assert_eq!(r.first(), Some(10));
809 assert_eq!(r.last(), Some(13));
810
811 r.insert(3..6);
812 assert_eq!(r.first(), Some(3));
813 assert_eq!(r.last(), Some(13));
814
815 r.insert(16..20);
816 assert_eq!(r.first(), Some(3));
817 assert_eq!(r.last(), Some(19));
818
819 r.insert(4..17);
820 assert_eq!(r.first(), Some(3));
821 assert_eq!(r.last(), Some(19));
822 }
823
824 #[test]
825 fn capacity() {
826 let mut r = RangeSet::new(3);
827 assert_eq!(r.first(), None);
828 assert_eq!(r.last(), None);
829
830 r.insert(10..11);
831 assert_eq!(r.first(), Some(10));
832 assert_eq!(r.last(), Some(10));
833
834 r.insert(13..14);
835 assert_eq!(r.first(), Some(10));
836 assert_eq!(r.last(), Some(13));
837
838 r.insert(3..6);
839 assert_eq!(r.first(), Some(3));
840 assert_eq!(r.last(), Some(13));
841
842 r.insert(16..20);
843 assert_eq!(r.first(), Some(10));
844 assert_eq!(r.last(), Some(19));
845
846 r.insert(4..17);
847 assert_eq!(r.first(), Some(4));
848 assert_eq!(r.last(), Some(19));
849 }
850}