1 // Copyright (C) 2018-2019, Cloudflare, Inc.
2 // All rights reserved.
3 //
4 // Redistribution and use in source and binary forms, with or without
5 // modification, are permitted provided that the following conditions are
6 // met:
7 //
8 // * Redistributions of source code must retain the above copyright notice,
9 // this list of conditions and the following disclaimer.
10 //
11 // * Redistributions in binary form must reproduce the above copyright
12 // notice, this list of conditions and the following disclaimer in the
13 // documentation and/or other materials provided with the distribution.
14 //
15 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
16 // IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
17 // THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
19 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
23 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
24 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
27 use std::ops::Range;
28
29 use std::collections::btree_map;
30 use std::collections::BTreeMap;
31 use std::collections::Bound;
32
33 #[derive(Clone, PartialEq, Eq, PartialOrd)]
34 pub struct RangeSet {
35 inner: BTreeMap<u64, u64>,
36
37 capacity: usize,
38 }
39
40 impl RangeSet {
new(capacity: usize) -> Self41 pub fn new(capacity: usize) -> Self {
42 RangeSet {
43 inner: BTreeMap::default(),
44 capacity,
45 }
46 }
47
48 // TODO: use RangeInclusive
insert(&mut self, item: Range<u64>)49 pub fn insert(&mut self, item: Range<u64>) {
50 let mut start = item.start;
51 let mut end = item.end;
52
53 // Check if preceding existing range overlaps with the new one.
54 if let Some(r) = self.prev_to(start) {
55 // New range overlaps with existing range in the set, merge them.
56 if range_overlaps(&r, &item) {
57 self.inner.remove(&r.start);
58
59 start = std::cmp::min(start, r.start);
60 end = std::cmp::max(end, r.end);
61 }
62 }
63
64 // Check if following existing ranges overlap with the new one.
65 while let Some(r) = self.next_to(start) {
66 // Existing range is fully contained in the new range, remove it.
67 if item.contains(&r.start) && item.contains(&r.end) {
68 self.inner.remove(&r.start);
69 continue;
70 }
71
72 // New range doesn't overlap anymore, we are done.
73 if !range_overlaps(&r, &item) {
74 break;
75 }
76
77 // New range overlaps with existing range in the set, merge them.
78 self.inner.remove(&r.start);
79
80 start = std::cmp::min(start, r.start);
81 end = std::cmp::max(end, r.end);
82 }
83
84 if self.inner.len() >= self.capacity {
85 self.inner.pop_first();
86 }
87
88 self.inner.insert(start, end);
89 }
90
remove_until(&mut self, largest: u64)91 pub fn remove_until(&mut self, largest: u64) {
92 let ranges: Vec<Range<u64>> = self
93 .inner
94 .range((Bound::Unbounded, Bound::Included(&largest)))
95 .map(|(&s, &e)| (s..e))
96 .collect();
97
98 for r in ranges {
99 self.inner.remove(&r.start);
100
101 if r.end > largest + 1 {
102 let start = largest + 1;
103 self.insert(start..r.end);
104 }
105 }
106 }
107
push_item(&mut self, item: u64)108 pub fn push_item(&mut self, item: u64) {
109 self.insert(item..item + 1);
110 }
111
first(&self) -> Option<u64>112 pub fn first(&self) -> Option<u64> {
113 self.flatten().next()
114 }
115
last(&self) -> Option<u64>116 pub fn last(&self) -> Option<u64> {
117 self.flatten().next_back()
118 }
119
len(&self) -> usize120 pub fn len(&self) -> usize {
121 self.inner.len()
122 }
123
iter(&self) -> Iter124 pub fn iter(&self) -> Iter {
125 Iter {
126 inner: self.inner.iter(),
127 }
128 }
129
flatten(&self) -> Flatten130 pub fn flatten(&self) -> Flatten {
131 Flatten {
132 inner: self.inner.iter(),
133 next: 0,
134 end: 0,
135 }
136 }
137
prev_to(&self, item: u64) -> Option<Range<u64>>138 fn prev_to(&self, item: u64) -> Option<Range<u64>> {
139 self.inner
140 .range((Bound::Unbounded, Bound::Included(item)))
141 .map(|(&s, &e)| (s..e))
142 .next_back()
143 }
144
next_to(&self, item: u64) -> Option<Range<u64>>145 fn next_to(&self, item: u64) -> Option<Range<u64>> {
146 self.inner
147 .range((Bound::Included(item), Bound::Unbounded))
148 .map(|(&s, &e)| (s..e))
149 .next()
150 }
151 }
152
153 impl Default for RangeSet {
default() -> Self154 fn default() -> Self {
155 Self::new(usize::MAX)
156 }
157 }
158
159 // This implements comparison between `RangeSet` and standard `Range`. The idea
160 // is that a `RangeSet` with no gaps (i.e. that only contains a single range)
161 // is basically equvalent to a normal `Range` so they should be comparable.
162 impl PartialEq<Range<u64>> for RangeSet {
eq(&self, other: &Range<u64>) -> bool163 fn eq(&self, other: &Range<u64>) -> bool {
164 // If there is more than one range it means that the range set is not
165 // contiguous, so can't be equal to a single range.
166 if self.inner.len() != 1 {
167 return false;
168 }
169
170 // Get the first and only range in the set.
171 let (first_start, first_end) = self.inner.iter().next().unwrap();
172
173 if (*first_start..*first_end) != *other {
174 return false;
175 }
176
177 true
178 }
179 }
180
181 impl std::fmt::Debug for RangeSet {
fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result182 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
183 let ranges: Vec<Range<u64>> = self
184 .iter()
185 .map(|mut r| {
186 r.end -= 1;
187 r
188 })
189 .collect();
190
191 write!(f, "{ranges:?}")
192 }
193 }
194
195 pub struct Iter<'a> {
196 inner: btree_map::Iter<'a, u64, u64>,
197 }
198
199 impl<'a> Iterator for Iter<'a> {
200 type Item = Range<u64>;
201
next(&mut self) -> Option<Range<u64>>202 fn next(&mut self) -> Option<Range<u64>> {
203 let (&start, &end) = self.inner.next()?;
204 Some(start..end)
205 }
206 }
207
208 impl<'a> DoubleEndedIterator for Iter<'a> {
next_back(&mut self) -> Option<Range<u64>>209 fn next_back(&mut self) -> Option<Range<u64>> {
210 let (&start, &end) = self.inner.next_back()?;
211 Some(start..end)
212 }
213 }
214
215 impl<'a> ExactSizeIterator for Iter<'a> {
len(&self) -> usize216 fn len(&self) -> usize {
217 self.inner.len()
218 }
219 }
220
221 pub struct Flatten<'a> {
222 inner: btree_map::Iter<'a, u64, u64>,
223 next: u64,
224 end: u64,
225 }
226
227 impl<'a> Iterator for Flatten<'a> {
228 type Item = u64;
229
next(&mut self) -> Option<u64>230 fn next(&mut self) -> Option<u64> {
231 if self.next == self.end {
232 let (&start, &end) = self.inner.next()?;
233
234 self.next = start;
235 self.end = end;
236 }
237
238 let next = self.next;
239 self.next += 1;
240
241 Some(next)
242 }
243 }
244
245 impl<'a> DoubleEndedIterator for Flatten<'a> {
next_back(&mut self) -> Option<u64>246 fn next_back(&mut self) -> Option<u64> {
247 if self.next == self.end {
248 let (&start, &end) = self.inner.next_back()?;
249
250 self.next = start;
251 self.end = end;
252 }
253
254 self.end -= 1;
255
256 Some(self.end)
257 }
258 }
259
range_overlaps(r: &Range<u64>, other: &Range<u64>) -> bool260 fn range_overlaps(r: &Range<u64>, other: &Range<u64>) -> bool {
261 other.start >= r.start && other.start <= r.end ||
262 other.end >= r.start && other.end <= r.end
263 }
264
265 #[cfg(test)]
266 mod tests {
267 use super::*;
268
269 #[test]
insert_non_overlapping()270 fn insert_non_overlapping() {
271 let mut r = RangeSet::default();
272 assert_eq!(r.inner.len(), 0);
273 let empty: &[u64] = &[];
274 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
275
276 r.insert(4..7);
277 assert_eq!(r.inner.len(), 1);
278 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
279
280 r.insert(9..12);
281 assert_eq!(r.inner.len(), 2);
282 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
283 }
284
285 #[test]
insert_contained()286 fn insert_contained() {
287 let mut r = RangeSet::default();
288
289 r.insert(4..7);
290 r.insert(9..12);
291 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
292
293 r.insert(4..7);
294 assert_eq!(r.inner.len(), 2);
295 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
296
297 r.insert(4..6);
298 assert_eq!(r.inner.len(), 2);
299 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
300
301 r.insert(5..6);
302 assert_eq!(r.inner.len(), 2);
303 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
304
305 r.insert(10..11);
306 assert_eq!(r.inner.len(), 2);
307 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
308
309 r.insert(9..11);
310 assert_eq!(r.inner.len(), 2);
311 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
312 }
313
314 #[test]
insert_overlapping()315 fn insert_overlapping() {
316 let mut r = RangeSet::default();
317
318 r.insert(3..6);
319 r.insert(9..12);
320 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 9, 10, 11]);
321
322 r.insert(5..7);
323 assert_eq!(r.inner.len(), 2);
324 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[3, 4, 5, 6, 9, 10, 11]);
325
326 r.insert(10..15);
327 assert_eq!(r.inner.len(), 2);
328 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
329 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
330 ]);
331
332 r.insert(2..5);
333 assert_eq!(r.inner.len(), 2);
334 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
335 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14
336 ]);
337
338 r.insert(8..10);
339 assert_eq!(r.inner.len(), 2);
340 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
341 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14
342 ]);
343
344 r.insert(6..10);
345 assert_eq!(r.inner.len(), 1);
346 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
347 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14
348 ]);
349 }
350
351 #[test]
insert_overlapping_multi()352 fn insert_overlapping_multi() {
353 let mut r = RangeSet::default();
354
355 r.insert(3..6);
356 r.insert(16..20);
357 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
358 3, 4, 5, 16, 17, 18, 19
359 ]);
360
361 r.insert(10..11);
362 assert_eq!(r.inner.len(), 3);
363 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
364 3, 4, 5, 10, 16, 17, 18, 19
365 ]);
366
367 r.insert(13..14);
368 assert_eq!(r.inner.len(), 4);
369 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
370 3, 4, 5, 10, 13, 16, 17, 18, 19
371 ]);
372
373 r.insert(4..17);
374 assert_eq!(r.inner.len(), 1);
375 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
376 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
377 ]);
378 }
379
380 #[test]
prev_to()381 fn prev_to() {
382 let mut r = RangeSet::default();
383
384 r.insert(4..7);
385 r.insert(9..12);
386
387 assert_eq!(r.prev_to(2), None);
388 assert_eq!(r.prev_to(4), Some(4..7));
389 assert_eq!(r.prev_to(15), Some(9..12));
390 assert_eq!(r.prev_to(5), Some(4..7));
391 assert_eq!(r.prev_to(8), Some(4..7));
392 }
393
394 #[test]
next_to()395 fn next_to() {
396 let mut r = RangeSet::default();
397
398 r.insert(4..7);
399 r.insert(9..12);
400
401 assert_eq!(r.next_to(2), Some(4..7));
402 assert_eq!(r.next_to(12), None);
403 assert_eq!(r.next_to(15), None);
404 assert_eq!(r.next_to(5), Some(9..12));
405 assert_eq!(r.next_to(8), Some(9..12));
406 }
407
408 #[test]
push_item()409 fn push_item() {
410 let mut r = RangeSet::default();
411
412 r.insert(4..7);
413 r.insert(9..12);
414 assert_eq!(r.inner.len(), 2);
415 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
416
417 r.push_item(15);
418 assert_eq!(r.inner.len(), 3);
419 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
420 4, 5, 6, 9, 10, 11, 15
421 ]);
422
423 r.push_item(15);
424 assert_eq!(r.inner.len(), 3);
425 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
426 4, 5, 6, 9, 10, 11, 15
427 ]);
428
429 r.push_item(1);
430 assert_eq!(r.inner.len(), 4);
431 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
432 1, 4, 5, 6, 9, 10, 11, 15
433 ]);
434
435 r.push_item(12);
436 r.push_item(13);
437 r.push_item(14);
438 assert_eq!(r.inner.len(), 3);
439 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
440 1, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
441 ]);
442
443 r.push_item(2);
444 r.push_item(3);
445 assert_eq!(r.inner.len(), 2);
446 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
447 1, 2, 3, 4, 5, 6, 9, 10, 11, 12, 13, 14, 15
448 ]);
449
450 r.push_item(8);
451 r.push_item(7);
452 assert_eq!(r.inner.len(), 1);
453 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
454 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15
455 ]);
456 }
457
458 #[test]
flatten_rev()459 fn flatten_rev() {
460 let mut r = RangeSet::default();
461 assert_eq!(r.inner.len(), 0);
462
463 let empty: &[u64] = &[];
464 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
465
466 r.insert(4..7);
467 assert_eq!(r.inner.len(), 1);
468 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6]);
469 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[6, 5, 4]);
470
471 r.insert(9..12);
472 assert_eq!(r.inner.len(), 2);
473 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[4, 5, 6, 9, 10, 11]);
474 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[
475 11, 10, 9, 6, 5, 4
476 ]);
477 }
478
479 #[test]
flatten_one()480 fn flatten_one() {
481 let mut r = RangeSet::default();
482 assert_eq!(r.inner.len(), 0);
483
484 let empty: &[u64] = &[];
485 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
486
487 r.insert(0..1);
488 assert_eq!(r.inner.len(), 1);
489 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[0]);
490 assert_eq!(&r.flatten().rev().collect::<Vec<u64>>(), &[0]);
491 }
492
493 #[test]
remove_largest()494 fn remove_largest() {
495 let mut r = RangeSet::default();
496
497 r.insert(3..6);
498 r.insert(9..11);
499 r.insert(13..14);
500 r.insert(16..20);
501 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
502 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
503 ]);
504
505 r.remove_until(2);
506 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
507 3, 4, 5, 9, 10, 13, 16, 17, 18, 19
508 ]);
509
510 r.remove_until(4);
511 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
512 5, 9, 10, 13, 16, 17, 18, 19
513 ]);
514
515 r.remove_until(6);
516 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[
517 9, 10, 13, 16, 17, 18, 19
518 ]);
519
520 r.remove_until(10);
521 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[13, 16, 17, 18, 19]);
522
523 r.remove_until(17);
524 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[18, 19]);
525
526 r.remove_until(18);
527 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &[19]);
528
529 r.remove_until(20);
530
531 let empty: &[u64] = &[];
532 assert_eq!(&r.flatten().collect::<Vec<u64>>(), &empty);
533 }
534
535 #[test]
eq_range()536 fn eq_range() {
537 let mut r = RangeSet::default();
538 assert_ne!(r, 0..0);
539
540 let expected = 3..20;
541
542 r.insert(3..6);
543 assert_ne!(r, expected);
544
545 r.insert(16..20);
546 assert_ne!(r, expected);
547
548 r.insert(10..11);
549 assert_ne!(r, expected);
550
551 r.insert(13..14);
552 assert_ne!(r, expected);
553
554 r.insert(4..17);
555 assert_eq!(r, expected);
556 }
557
558 #[test]
first_last()559 fn first_last() {
560 let mut r = RangeSet::default();
561 assert_eq!(r.first(), None);
562 assert_eq!(r.last(), None);
563
564 r.insert(10..11);
565 assert_eq!(r.first(), Some(10));
566 assert_eq!(r.last(), Some(10));
567
568 r.insert(13..14);
569 assert_eq!(r.first(), Some(10));
570 assert_eq!(r.last(), Some(13));
571
572 r.insert(3..6);
573 assert_eq!(r.first(), Some(3));
574 assert_eq!(r.last(), Some(13));
575
576 r.insert(16..20);
577 assert_eq!(r.first(), Some(3));
578 assert_eq!(r.last(), Some(19));
579
580 r.insert(4..17);
581 assert_eq!(r.first(), Some(3));
582 assert_eq!(r.last(), Some(19));
583 }
584
585 #[test]
capacity()586 fn capacity() {
587 let mut r = RangeSet::new(3);
588 assert_eq!(r.first(), None);
589 assert_eq!(r.last(), None);
590
591 r.insert(10..11);
592 assert_eq!(r.first(), Some(10));
593 assert_eq!(r.last(), Some(10));
594
595 r.insert(13..14);
596 assert_eq!(r.first(), Some(10));
597 assert_eq!(r.last(), Some(13));
598
599 r.insert(3..6);
600 assert_eq!(r.first(), Some(3));
601 assert_eq!(r.last(), Some(13));
602
603 r.insert(16..20);
604 assert_eq!(r.first(), Some(10));
605 assert_eq!(r.last(), Some(19));
606
607 r.insert(4..17);
608 assert_eq!(r.first(), Some(4));
609 assert_eq!(r.last(), Some(19));
610 }
611 }
612