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 #[inline(always)]
86 fn fixup(&mut self) {
87 match self {
88 RangeSet::Inline(set) if set.inner.len() == MAX_INLINE_CAPACITY => {
89 let old_inner = std::mem::take(&mut set.inner);
90 *self = RangeSet::BTree(BTreeRangeSet {
91 inner: old_inner.into_inner().expect("At capacity").into(),
92 capacity: set.capacity,
93 });
94 },
95
96 RangeSet::BTree(set) if set.inner.len() <= MIN_TO_INLINE => {
97 let old_inner = std::mem::take(&mut set.inner);
98 *self = RangeSet::Inline(InlineRangeSet {
99 inner: SmallVec::from_iter(old_inner),
100 capacity: set.capacity,
101 })
102 },
103
104 _ => {},
105 }
106 }
107
108 #[inline]
114 pub fn insert(&mut self, item: Range<u64>) {
115 match self {
116 RangeSet::Inline(set) => set.insert(item),
117 RangeSet::BTree(set) => set.insert(item),
118 }
119
120 self.fixup();
121 }
122
123 pub fn iter(
125 &self,
126 ) -> impl DoubleEndedIterator<Item = Range<u64>> + ExactSizeIterator + '_
127 {
128 match self {
129 RangeSet::BTree(set) =>
130 Either::Left(set.inner.iter().map(|(k, v)| *k..*v)),
131
132 RangeSet::Inline(set) =>
133 Either::Right(set.inner.iter().map(|(s, e)| *s..*e)),
134 }
135 }
136
137 pub fn flatten(&self) -> impl DoubleEndedIterator<Item = u64> + '_ {
140 match self {
141 RangeSet::BTree(set) =>
142 Either::Left(set.inner.iter().flat_map(|(k, v)| *k..*v)),
143
144 RangeSet::Inline(set) =>
145 Either::Right(set.inner.iter().flat_map(|(s, e)| *s..*e)),
146 }
147 }
148
149 pub fn first(&self) -> Option<u64> {
151 match self {
152 RangeSet::Inline(set) => set.inner.first().map(|(s, _)| *s),
153
154 RangeSet::BTree(set) => set.inner.first_key_value().map(|(k, _)| *k),
155 }
156 }
157
158 pub fn last(&self) -> Option<u64> {
160 match self {
161 RangeSet::Inline(set) => set.inner.last().map(|(_, e)| *e - 1),
162
163 RangeSet::BTree(set) =>
164 set.inner.last_key_value().map(|(_, v)| *v - 1),
165 }
166 }
167
168 #[inline]
169 pub fn remove_until(&mut self, largest: u64) {
170 match self {
171 RangeSet::Inline(set) => set.remove_until(largest),
172 RangeSet::BTree(set) => set.remove_until(largest),
173 }
174
175 self.fixup();
176 }
177
178 pub fn push_item(&mut self, item: u64) {
179 self.insert(item..item + 1)
180 }
181}
182
183impl InlineRangeSet {
184 fn insert(&mut self, item: Range<u64>) {
185 let start = item.start;
186 let mut end = item.end;
187 let mut pos = 0;
188
189 loop {
190 match self.inner.get_mut(pos) {
191 Some((s, e)) => {
192 if start > *e {
193 pos += 1;
195 continue;
196 }
197
198 if end < *s {
199 if self.inner.len() == self.capacity {
202 self.inner.remove(0);
203 pos -= 1;
204 }
205
206 self.inner.insert(pos, (start, end));
207 return;
208 }
209
210 if start < *s {
212 *s = start;
215 }
216
217 if end > *e {
218 *e = end;
221 break;
222 } else {
223 return;
224 }
225 },
226
227 None => {
228 if self.inner.len() == self.capacity {
229 self.inner.remove(0);
230 }
231
232 self.inner.push((start, end));
233 return;
234 },
235 }
236 }
237
238 while let Some((s, e)) = self.inner.get(pos + 1).copied() {
240 if end < s {
241 break;
243 }
244
245 let new_e = e.max(end);
246 self.inner[pos].1 = new_e;
247 end = new_e;
248 self.inner.remove(pos + 1);
249 }
250 }
251
252 fn remove_until(&mut self, largest: u64) {
253 while let Some((s, e)) = self.inner.first_mut() {
254 if largest >= *e {
255 self.inner.remove(0);
256 continue;
257 }
258
259 *s = (largest + 1).max(*s);
260 if *s == *e {
261 self.inner.remove(0);
262 }
263
264 break;
265 }
266 }
267}
268
269impl BTreeRangeSet {
270 fn insert(&mut self, item: Range<u64>) {
272 let mut start = item.start;
273 let mut end = item.end;
274
275 if let Some(r) = self.prev_to(start) {
277 if range_overlaps(&r, &item) {
279 self.inner.remove(&r.start);
280
281 start = cmp::min(start, r.start);
282 end = cmp::max(end, r.end);
283 }
284 }
285
286 while let Some(r) = self.next_to(start) {
288 if item.contains(&r.start) && item.contains(&r.end) {
290 self.inner.remove(&r.start);
291 continue;
292 }
293
294 if !range_overlaps(&r, &item) {
296 break;
297 }
298
299 self.inner.remove(&r.start);
301
302 start = cmp::min(start, r.start);
303 end = cmp::max(end, r.end);
304 }
305
306 if self.inner.len() >= self.capacity {
307 self.inner.pop_first();
308 }
309
310 self.inner.insert(start, end);
311 }
312
313 fn remove_until(&mut self, largest: u64) {
314 let ranges: Vec<Range<u64>> = self
315 .inner
316 .range((Bound::Unbounded, Bound::Included(&largest)))
317 .map(|(&s, &e)| (s..e))
318 .collect();
319
320 for r in ranges {
321 self.inner.remove(&r.start);
322
323 if r.end > largest + 1 {
324 let start = largest + 1;
325 self.insert(start..r.end);
326 }
327 }
328 }
329
330 fn prev_to(&self, item: u64) -> Option<Range<u64>> {
331 self.inner
332 .range((Bound::Unbounded, Bound::Included(item)))
333 .map(|(&s, &e)| (s..e))
334 .next_back()
335 }
336
337 fn next_to(&self, item: u64) -> Option<Range<u64>> {
338 self.inner
339 .range((Bound::Included(item), Bound::Unbounded))
340 .map(|(&s, &e)| (s..e))
341 .next()
342 }
343}
344
345impl Default for RangeSet {
346 fn default() -> Self {
347 RangeSet::Inline(InlineRangeSet {
348 inner: Default::default(),
349 capacity: usize::MAX,
350 })
351 }
352}
353
354impl PartialEq<Range<u64>> for RangeSet {
359 fn eq(&self, other: &Range<u64>) -> bool {
360 if self.len() != 1 {
363 return false;
364 }
365
366 let range = self.iter().next().unwrap();
368 range == *other
369 }
370}
371
372impl std::fmt::Debug for RangeSet {
373 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
374 let ranges: Vec<Range<u64>> = self
375 .iter()
376 .map(|mut r| {
377 r.end -= 1;
378 r
379 })
380 .collect();
381
382 write!(f, "{ranges:?}")
383 }
384}
385
386fn range_overlaps(r: &Range<u64>, other: &Range<u64>) -> bool {
387 other.start >= r.start && other.start <= r.end ||
388 other.end >= r.start && other.end <= r.end
389}
390
391#[cfg(test)]
392mod tests {
393 use super::*;
394
395 #[test]
396 fn insert_non_overlapping() {
397 let mut r = RangeSet::default();
398 assert_eq!(r.len(), 0);
399 let empty: &[u64] = &[];
400 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
401
402 r.insert(4..7);
403 assert_eq!(r.len(), 1);
404 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
405
406 r.insert(9..12);
407 assert_eq!(r.len(), 2);
408 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
409 }
410
411 #[test]
412 fn insert_contained() {
413 let mut r = RangeSet::default();
414
415 r.insert(4..7);
416 r.insert(9..12);
417 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
418
419 r.insert(4..7);
420 assert_eq!(r.len(), 2);
421 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
422
423 r.insert(4..6);
424 assert_eq!(r.len(), 2);
425 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
426
427 r.insert(5..6);
428 assert_eq!(r.len(), 2);
429 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
430
431 r.insert(10..11);
432 assert_eq!(r.len(), 2);
433 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
434
435 r.insert(9..11);
436 assert_eq!(r.len(), 2);
437 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
438 }
439
440 #[test]
441 fn insert_overlapping() {
442 let mut r = RangeSet::default();
443
444 r.insert(3..6);
445 r.insert(9..12);
446 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 9, 10, 11]);
447
448 r.insert(5..7);
449 assert_eq!(r.len(), 2);
450 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 6, 9, 10, 11]);
451
452 r.insert(10..15);
453 assert_eq!(r.len(), 2);
454 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
455 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
456 ]);
457
458 r.insert(2..5);
459 assert_eq!(r.len(), 2);
460 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
461 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
462 ]);
463
464 r.insert(8..10);
465 assert_eq!(r.len(), 2);
466 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
467 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14
468 ]);
469
470 r.insert(6..10);
471 assert_eq!(r.len(), 1);
472 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
473 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14
474 ]);
475 }
476
477 #[test]
478 fn insert_overlapping_multi() {
479 let mut r = RangeSet::default();
480
481 r.insert(3..6);
482 r.insert(16..20);
483 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
484 3, 4, 5, 16, 17, 18, 19
485 ]);
486
487 r.insert(10..11);
488 assert_eq!(r.len(), 3);
489 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
490 3, 4, 5, 10, 16, 17, 18, 19
491 ]);
492
493 assert!(matches!(r, RangeSet::Inline(_)));
494
495 r.insert(13..14);
496 assert_eq!(r.len(), 4);
497 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
498 3, 4, 5, 10, 13, 16, 17, 18, 19
499 ]);
500
501 assert!(matches!(r, RangeSet::BTree(_)));
503
504 r.insert(4..17);
505 assert_eq!(r.len(), 1);
506 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
507 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
508 ]);
509
510 assert!(matches!(r, RangeSet::Inline(_)));
512 }
513
514 #[test]
515 fn prev_to() {
516 let mut r = BTreeRangeSet {
517 inner: Default::default(),
518 capacity: usize::MAX,
519 };
520
521 r.insert(4..7);
522 r.insert(9..12);
523
524 assert_eq!(r.prev_to(2), None);
525 assert_eq!(r.prev_to(4), Some(4..7));
526 assert_eq!(r.prev_to(15), Some(9..12));
527 assert_eq!(r.prev_to(5), Some(4..7));
528 assert_eq!(r.prev_to(8), Some(4..7));
529 }
530
531 #[test]
532 fn next_to() {
533 let mut r = BTreeRangeSet {
534 inner: Default::default(),
535 capacity: usize::MAX,
536 };
537
538 r.insert(4..7);
539 r.insert(9..12);
540
541 assert_eq!(r.next_to(2), Some(4..7));
542 assert_eq!(r.next_to(12), None);
543 assert_eq!(r.next_to(15), None);
544 assert_eq!(r.next_to(5), Some(9..12));
545 assert_eq!(r.next_to(8), Some(9..12));
546 }
547
548 #[test]
549 fn push_item() {
550 let mut r = RangeSet::default();
551
552 r.insert(4..7);
553 r.insert(9..12);
554 assert_eq!(r.len(), 2);
555 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
556
557 r.push_item(15);
558 assert_eq!(r.len(), 3);
559 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
560 4, 5, 6, 9, 10, 11, 15
561 ]);
562
563 r.push_item(15);
564 assert_eq!(r.len(), 3);
565 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
566 4, 5, 6, 9, 10, 11, 15
567 ]);
568
569 r.push_item(1);
570 assert_eq!(r.len(), 4);
571 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
572 1, 4, 5, 6, 9, 10, 11, 15
573 ]);
574
575 r.push_item(12);
576 r.push_item(13);
577 r.push_item(14);
578
579 assert_eq!(r.len(), 3);
580 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
581 1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
582 ]);
583
584 r.push_item(2);
585 r.push_item(3);
586 assert_eq!(r.len(), 2);
587 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
588 1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
589 ]);
590
591 r.push_item(8);
592 r.push_item(7);
593 assert_eq!(r.len(), 1);
594 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
595 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
596 ]);
597 }
598
599 #[test]
600 fn flatten_rev() {
601 let mut r = RangeSet::default();
602 assert_eq!(r.len(), 0);
603
604 let empty: &[u64] = &[];
605 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
606
607 r.insert(4..7);
608 assert_eq!(r.len(), 1);
609 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
610 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[6, 5, 4]);
611
612 r.insert(9..12);
613 assert_eq!(r.len(), 2);
614 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
615 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[
616 11, 10, 9, 6, 5, 4
617 ]);
618 }
619
620 #[test]
621 fn flatten_one() {
622 let mut r = RangeSet::default();
623 assert_eq!(r.len(), 0);
624
625 let empty: &[u64] = &[];
626 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
627
628 r.insert(0..1);
629 assert_eq!(r.len(), 1);
630 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[0]);
631 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[0]);
632 }
633
634 #[test]
635 fn remove_largest() {
636 let mut r = RangeSet::default();
637
638 r.insert(3..6);
639 r.insert(9..11);
640 r.insert(13..14);
641 r.insert(16..20);
642 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
643 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
644 ]);
645
646 r.remove_until(2);
647 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
648 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
649 ]);
650
651 r.remove_until(4);
652 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
653 5, 9, 10, 13, 16, 17, 18, 19
654 ]);
655
656 r.remove_until(6);
657 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
658 9, 10, 13, 16, 17, 18, 19
659 ]);
660
661 r.remove_until(10);
662 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[13, 16, 17, 18, 19]);
663
664 r.remove_until(17);
665 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[18, 19]);
666
667 r.remove_until(18);
668 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[19]);
669
670 r.remove_until(20);
671
672 let empty: &[u64] = &[];
673 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
674 }
675
676 #[test]
677 fn eq_range() {
678 let mut r = RangeSet::default();
679 assert_ne!(r, 0..0);
680
681 let expected = 3..20;
682
683 r.insert(3..6);
684 assert_ne!(r, expected);
685
686 r.insert(16..20);
687 assert_ne!(r, expected);
688
689 r.insert(10..11);
690 assert_ne!(r, expected);
691
692 r.insert(13..14);
693 assert_ne!(r, expected);
694
695 r.insert(4..17);
696
697 assert_eq!(r, expected);
698 }
699
700 #[test]
701 fn first_last() {
702 let mut r = RangeSet::default();
703 assert_eq!(r.first(), None);
704 assert_eq!(r.last(), None);
705
706 r.insert(10..11);
707 assert_eq!(r.first(), Some(10));
708 assert_eq!(r.last(), Some(10));
709
710 r.insert(13..14);
711 assert_eq!(r.first(), Some(10));
712 assert_eq!(r.last(), Some(13));
713
714 r.insert(3..6);
715 assert_eq!(r.first(), Some(3));
716 assert_eq!(r.last(), Some(13));
717
718 r.insert(16..20);
719 assert_eq!(r.first(), Some(3));
720 assert_eq!(r.last(), Some(19));
721
722 r.insert(4..17);
723 assert_eq!(r.first(), Some(3));
724 assert_eq!(r.last(), Some(19));
725 }
726
727 #[test]
728 fn capacity() {
729 let mut r = RangeSet::new(3);
730 assert_eq!(r.first(), None);
731 assert_eq!(r.last(), None);
732
733 r.insert(10..11);
734 assert_eq!(r.first(), Some(10));
735 assert_eq!(r.last(), Some(10));
736
737 r.insert(13..14);
738 assert_eq!(r.first(), Some(10));
739 assert_eq!(r.last(), Some(13));
740
741 r.insert(3..6);
742 assert_eq!(r.first(), Some(3));
743 assert_eq!(r.last(), Some(13));
744
745 r.insert(16..20);
746 assert_eq!(r.first(), Some(10));
747 assert_eq!(r.last(), Some(19));
748
749 r.insert(4..17);
750 assert_eq!(r.first(), Some(4));
751 assert_eq!(r.last(), Some(19));
752 }
753}