Skip to main content

simoom_quic/
ranges.rs

1//! A set of disjoint half-open `u64` ranges, kept sorted. Used for ACK ranges, received
2//! packet numbers and stream retransmission bookkeeping.
3
4#[derive(Clone, Debug, Default, PartialEq, Eq)]
5pub struct RangeSet {
6    r: Vec<(u64, u64)>,
7}
8
9impl RangeSet {
10    pub fn new() -> Self {
11        RangeSet { r: Vec::new() }
12    }
13    pub fn is_empty(&self) -> bool {
14        self.r.is_empty()
15    }
16    /// Number of disjoint ranges.
17    pub fn len(&self) -> usize {
18        self.r.len()
19    }
20    pub fn min(&self) -> Option<u64> {
21        self.r.first().map(|x| x.0)
22    }
23    /// The largest contained value.
24    pub fn max(&self) -> Option<u64> {
25        self.r.last().map(|x| x.1 - 1)
26    }
27    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (u64, u64)> + '_ {
28        self.r.iter().copied()
29    }
30    pub fn first(&self) -> Option<(u64, u64)> {
31        self.r.first().copied()
32    }
33    pub fn contains(&self, v: u64) -> bool {
34        match self.r.binary_search_by(|&(s, _)| s.cmp(&v)) {
35            Ok(_) => true,
36            Err(0) => false,
37            Err(i) => v < self.r[i - 1].1,
38        }
39    }
40    /// True if all of `[s, e)` is in the set.
41    pub fn contains_range(&self, s: u64, e: u64) -> bool {
42        if s >= e {
43            return true;
44        }
45        match self.r.binary_search_by(|&(x, _)| x.cmp(&s)) {
46            Ok(i) => e <= self.r[i].1,
47            Err(0) => false,
48            Err(i) => e <= self.r[i - 1].1,
49        }
50    }
51    /// Inserts `[s, e)`; returns true if anything new was added.
52    pub fn insert(&mut self, s: u64, e: u64) -> bool {
53        if s >= e {
54            return false;
55        }
56        // Fast path: append at the end (the common in-order case).
57        if let Some(last) = self.r.last_mut() {
58            if s > last.1 {
59                self.r.push((s, e));
60                return true;
61            }
62            if s >= last.0 {
63                if e <= last.1 {
64                    return false;
65                }
66                last.1 = e;
67                return true;
68            }
69        } else {
70            self.r.push((s, e));
71            return true;
72        }
73        // First range whose end >= s.
74        let i = self.r.partition_point(|&(_, end)| end < s);
75        let mut j = i;
76        let (mut ns, mut ne) = (s, e);
77        while j < self.r.len() && self.r[j].0 <= e {
78            ns = ns.min(self.r[j].0);
79            ne = ne.max(self.r[j].1);
80            j += 1;
81        }
82        if j == i + 1 && self.r[i] == (ns, ne) {
83            return false;
84        }
85        self.r.splice(i..j, std::iter::once((ns, ne)));
86        true
87    }
88    /// Removes `[s, e)`.
89    pub fn remove(&mut self, s: u64, e: u64) {
90        if s >= e || self.r.is_empty() {
91            return;
92        }
93        let mut out = Vec::with_capacity(self.r.len() + 1);
94        for &(a, b) in &self.r {
95            if b <= s || a >= e {
96                out.push((a, b));
97            } else {
98                if a < s {
99                    out.push((a, s));
100                }
101                if b > e {
102                    out.push((e, b));
103                }
104            }
105        }
106        self.r = out;
107    }
108    /// Removes everything below `v`.
109    pub fn remove_below(&mut self, v: u64) {
110        let i = self.r.partition_point(|&(_, e)| e <= v);
111        self.r.drain(..i);
112        if let Some(f) = self.r.first_mut() {
113            if f.0 < v {
114                f.0 = v;
115            }
116        }
117    }
118    /// Drops the lowest ranges so at most `n` remain.
119    pub fn truncate_low(&mut self, n: usize) {
120        if self.r.len() > n {
121            let k = self.r.len() - n;
122            self.r.drain(..k);
123        }
124    }
125    /// Removes and returns the lowest range, clipped to at most `max_len` values.
126    pub fn pop_min(&mut self, max_len: u64) -> Option<(u64, u64)> {
127        let f = self.r.first_mut()?;
128        let (s, e) = *f;
129        if e - s > max_len {
130            f.0 = s + max_len;
131            Some((s, s + max_len))
132        } else {
133            self.r.remove(0);
134            Some((s, e))
135        }
136    }
137    pub fn clear(&mut self) {
138        self.r.clear();
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn insert_merge_remove() {
148        let mut s = RangeSet::new();
149        assert!(s.insert(10, 20));
150        assert!(s.insert(30, 40));
151        assert!(!s.insert(12, 15));
152        assert!(s.insert(5, 10));
153        assert_eq!(s.iter().collect::<Vec<_>>(), vec![(5, 20), (30, 40)]);
154        assert!(s.insert(20, 30));
155        assert_eq!(s.iter().collect::<Vec<_>>(), vec![(5, 40)]);
156        s.remove(10, 12);
157        assert_eq!(s.iter().collect::<Vec<_>>(), vec![(5, 10), (12, 40)]);
158        assert!(s.contains(5) && !s.contains(10) && s.contains(39) && !s.contains(40));
159        assert!(s.contains_range(12, 40) && !s.contains_range(9, 13));
160        s.remove_below(13);
161        assert_eq!(s.iter().collect::<Vec<_>>(), vec![(13, 40)]);
162        assert_eq!(s.pop_min(5), Some((13, 18)));
163        assert_eq!(s.max(), Some(39));
164        let mut t = RangeSet::new();
165        for i in (0..100).rev().step_by(2) {
166            t.insert(i, i + 1);
167        }
168        assert_eq!(t.len(), 50);
169        t.insert(0, 100);
170        assert_eq!(t.len(), 1);
171    }
172}