noq_proto/range_set/
array_range_set.rs

1use std::cmp::Ordering;
2use std::fmt::{self, Write};
3use std::iter::Sum;
4use std::ops::{Add, Range, Sub};
5
6use tinyvec::TinyVec;
7
8/// A set of u64 values optimized for long runs and random insert/delete/contains
9///
10/// `ArrayRangeSet` uses an array representation, where each array entry represents
11/// a range.
12///
13/// The array-based RangeSet provides 2 benefits:
14/// - There exists an inline representation, which avoids the need of heap
15///   allocating ACK ranges for SentFrames for small ranges.
16/// - Iterating over ranges should usually be faster since there is only
17///   a single cache-friendly contiguous range.
18///
19/// `ArrayRangeSet` is especially useful for tracking ACK ranges where the amount
20/// of ranges is usually very low (since ACK numbers are in consecutive fashion
21/// unless reordering or packet loss occur).
22#[derive(Default, PartialEq, Eq)]
23pub(crate) struct ArrayRangeSet<const N: usize = ARRAY_RANGE_SET_INLINE_CAPACITY, T: Default = u64>(
24    TinyVec<[Range<T>; N]>,
25);
26
27impl<const N: usize, T: fmt::Debug + Default> fmt::Debug for ArrayRangeSet<N, T> {
28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29        f.write_char('[')?;
30        let mut first = true;
31        for range in self.0.iter() {
32            if !first {
33                f.write_char(',')?;
34            }
35            write!(f, "{range:?}")?;
36            first = false;
37        }
38        f.write_char(']')?;
39        Ok(())
40    }
41}
42
43/// The capacity of elements directly stored in [`ArrayRangeSet`]
44///
45/// An inline capacity of 2 is chosen to keep `SentFrame` below 128 bytes.
46pub(crate) const ARRAY_RANGE_SET_INLINE_CAPACITY: usize = 2;
47
48impl<const N: usize> Clone for ArrayRangeSet<N> {
49    fn clone(&self) -> Self {
50        // tinyvec keeps the heap representation after clones.
51        // We rather prefer the inline representation for clones if possible,
52        // since clones (e.g. for storage in `SentFrames`) are rarely mutated
53        if self.0.is_inline() || self.0.len() > ARRAY_RANGE_SET_INLINE_CAPACITY {
54            return Self(self.0.clone());
55        }
56
57        let mut vec = TinyVec::new();
58        vec.extend_from_slice(self.0.as_slice());
59        Self(vec)
60    }
61}
62
63impl<const N: usize, T> ArrayRangeSet<N, T>
64where
65    T: Default
66        + Clone
67        + Copy
68        + PartialOrd
69        + Ord
70        + From<u32>
71        + Add<T, Output = T>
72        + Sub<T, Output = T>
73        + Sum,
74{
75    pub(crate) fn new() -> Self {
76        Default::default()
77    }
78
79    pub(crate) fn iter(&self) -> impl DoubleEndedIterator<Item = Range<T>> + '_ {
80        self.0.iter().cloned()
81    }
82
83    pub(crate) fn range_count(&self) -> usize {
84        self.0.len()
85    }
86
87    pub(crate) fn elts_count(&self) -> T {
88        self.0.iter().map(|r| r.end - r.start).sum()
89    }
90
91    pub(crate) fn contains(&self, x: T) -> bool {
92        self.0
93            .binary_search_by(|range| {
94                if range.end <= x {
95                    Ordering::Less
96                } else if x < range.start {
97                    Ordering::Greater
98                } else {
99                    // range.start <= x < range.end
100                    Ordering::Equal
101                }
102            })
103            .is_ok()
104    }
105
106    pub(crate) fn iter_range(&self, range: Range<T>) -> impl Iterator<Item = Range<T>> + '_ {
107        self.iter().filter_map(move |r| {
108            if r.end > range.start && r.start < range.end {
109                Some(r.start.max(range.start)..r.end.min(range.end))
110            } else {
111                None
112            }
113        })
114    }
115
116    pub(crate) fn insert_one(&mut self, x: T) -> bool {
117        self.insert(x..x + T::from(1u32))
118    }
119
120    pub(crate) fn insert(&mut self, x: Range<T>) -> bool {
121        let mut result = false;
122
123        if x.is_empty() {
124            // Don't try to deal with ranges where x.end <= x.start
125            return false;
126        }
127
128        // Find the first range that might interact with `x`.
129        // Unlike removal, we use a strict comparison so that ranges
130        // adjacent to `x` are included in the right-hand partition and can be merged.
131        let idx = self.0.partition_point(|r| r.end < x.start);
132
133        if idx == self.0.len() {
134            self.0.push(x);
135            return true;
136        }
137
138        let range = &mut self.0[idx];
139
140        if x.end < range.start {
141            // The range is fully before this range and therefore not extensible.
142            // Add a new range to the left
143            self.0.insert(idx, x);
144            return true;
145        } else if range.start > x.start {
146            // The new range starts before this range but overlaps.
147            // Extend the current range to the left
148            // Note that we don't have to merge a potential left range, since
149            // this case would have been captured by merging the right range
150            // in the previous loop iteration
151            result = true;
152            range.start = x.start;
153        }
154
155        // At this point we have handled all parts of the new range which
156        // are in front of the current range. Now we handle everything from
157        // the start of the current range
158
159        if x.end <= range.end {
160            // Fully contained
161            return result;
162        }
163
164        // Extend the current range to the end of the new range.
165        // Since it's not contained it must be bigger
166        range.end = x.end;
167
168        // Merge all follow-up ranges which overlap
169        while idx != self.0.len() - 1 {
170            let curr = self.0[idx].clone();
171            let next = self.0[idx + 1].clone();
172            if curr.end >= next.start {
173                self.0[idx].end = next.end.max(curr.end);
174                self.0.remove(idx + 1);
175            } else {
176                break;
177            }
178        }
179
180        true
181    }
182
183    pub(crate) fn remove(&mut self, x: Range<T>) -> bool {
184        let mut result = false;
185
186        if x.is_empty() {
187            // Don't try to deal with ranges where x.end <= x.start
188            return false;
189        }
190
191        // Find the first range that might overlap with `x`.
192        // Unlike insertion, we use an inclusive comparison since removal does not
193        // affect the adjacent range on the left.
194        let mut idx = self.0.partition_point(|r| r.end <= x.start);
195
196        while idx != self.0.len() {
197            let range = self.0[idx].clone();
198
199            if x.end <= range.start {
200                // The range is not in the set
201                break;
202            }
203
204            // The range overlaps with this range
205            result = true;
206
207            let left = range.start..x.start;
208            let right = x.end..range.end;
209
210            if left.is_empty() && right.is_empty() {
211                self.0.remove(idx);
212            } else if left.is_empty() {
213                self.0[idx] = right;
214                idx += 1;
215            } else if right.is_empty() {
216                self.0[idx] = left;
217                idx += 1;
218            } else {
219                self.0[idx] = right;
220                self.0.insert(idx, left);
221                idx += 2;
222            }
223        }
224
225        result
226    }
227
228    pub(crate) fn is_empty(&self) -> bool {
229        self.0.is_empty()
230    }
231
232    pub(crate) fn pop_min(&mut self) -> Option<Range<T>> {
233        if !self.0.is_empty() {
234            Some(self.0.remove(0))
235        } else {
236            None
237        }
238    }
239
240    pub(crate) fn min(&self) -> Option<T> {
241        self.iter().next().map(|x| x.start)
242    }
243
244    pub(crate) fn max(&self) -> Option<T> {
245        self.iter().next_back().map(|x| x.end - T::from(1))
246    }
247}
248
249/// Functions which need `Range<T>` to impl IntoIterator for u64.
250///
251/// `Range<T>` only implements [`IntoIterator`] for types implementing `std::iter::Step`,
252/// but that trait is unstable. We can work around this by duplicating these functions for
253/// [`u32`] and [`u64`]. Only we don't currently use the u32 version so it is u64-only for
254/// now.
255impl<const N: usize> ArrayRangeSet<N, u64> {
256    pub(crate) fn elts(&self) -> impl Iterator<Item = u64> + '_ {
257        self.iter().flatten()
258    }
259}
260
261#[cfg(test)]
262impl proptest::arbitrary::Arbitrary for ArrayRangeSet {
263    type Parameters = ();
264    type Strategy = proptest::strategy::BoxedStrategy<Self>;
265
266    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
267        use proptest::prelude::*;
268        // Generate 1-8 ranges. Each range is defined by a gap from the previous and a size.
269        // We use small values to keep encoding reasonable.
270        prop::collection::vec((1u64..100, 1u64..50), 1..8)
271            .prop_map(|gaps_and_sizes| {
272                let mut ranges = Self::new();
273                let mut pos = 0u64;
274                for (gap, size) in gaps_and_sizes {
275                    let start = pos + gap;
276                    let end = start + size;
277                    ranges.insert(start..end);
278                    pos = end;
279                }
280                ranges
281            })
282            .boxed()
283    }
284}
noq_proto/range_set/array_range_set.rs

noq_proto/range_set/
array_range_set.rs
