netstack3_tcp/
seq_ranges.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use alloc::collections::VecDeque;
6use core::ops::Range;
7
8use derivative::Derivative;
9use netstack3_base::SeqNum;
10
11/// A generic data structure that keeps track of ordered sequence number ranges.
12///
13/// Each kept range has associated metadata of type `M`.
14#[derive(Debug, Derivative)]
15#[derivative(Default(bound = ""))]
16#[cfg_attr(test, derive(PartialEq, Eq))]
17pub(crate) struct SeqRanges<M> {
18    blocks: VecDeque<SeqRange<M>>,
19}
20
21impl<M> SeqRanges<M> {
22    pub(crate) fn is_empty(&self) -> bool {
23        self.blocks.is_empty()
24    }
25
26    pub(crate) fn pop_front_if<F: FnOnce(&SeqRange<M>) -> bool>(
27        &mut self,
28        f: F,
29    ) -> Option<SeqRange<M>> {
30        let front = self.blocks.front()?;
31        if f(front) { self.blocks.pop_front() } else { None }
32    }
33
34    fn find_first_after(blocks: &VecDeque<SeqRange<M>>, start: SeqNum) -> usize {
35        match blocks.binary_search_by(|block| {
36            if block.start() == start {
37                return core::cmp::Ordering::Equal;
38            }
39            if block.start().before(start) {
40                core::cmp::Ordering::Less
41            } else {
42                core::cmp::Ordering::Greater
43            }
44        }) {
45            Ok(r) => {
46                // We found the exact same start point, so the first segment
47                // whose start is greater must be the next one.
48                r + 1
49            }
50            Err(e) => {
51                // When binary search doesn't find the exact place it returns
52                // the index where this block should be in, which should be the
53                // next greater range.
54                e
55            }
56        }
57    }
58
59    /// Inserts `range` into this tracking structure.
60    ///
61    /// No-op if the range is empty.
62    ///
63    /// `meta` is attached to the newly created range and all the ranges it
64    /// touches, including if `range` is a subset of a currently tracked range.
65    ///
66    /// Returns `true` iff `range` insertion increases the total number of
67    /// tracked bytes contained in `SeqRanges`.
68    pub(crate) fn insert(&mut self, range: Range<SeqNum>, meta: M) -> bool
69    where
70        M: Clone,
71    {
72        let Some(range) = SeqRange::new(range, meta) else {
73            return false;
74        };
75        self.insert_seq_range(range)
76    }
77
78    fn insert_seq_range(&mut self, mut range: SeqRange<M>) -> bool
79    where
80        M: Clone,
81    {
82        let Self { blocks } = self;
83
84        if blocks.is_empty() {
85            blocks.push_back(range);
86            return true;
87        }
88
89        // Search for the first segment whose `start` is greater.
90        let first_after = Self::find_first_after(blocks, range.start());
91
92        let mut merge_right = 0;
93        for block in blocks.range(first_after..blocks.len()) {
94            match range.merge_right(block) {
95                MergeRightResult::Before => break,
96                MergeRightResult::After { merged } => {
97                    merge_right += 1;
98                    if merged {
99                        break;
100                    }
101                }
102            }
103        }
104
105        // Given we're always sorted and we know the first range strictly after
106        // the inserting one, we only need to check to the left once.
107        let merge_left = match first_after
108            .checked_sub(1)
109            .and_then(|first_before| blocks.get_mut(first_before))
110        {
111            Some(block) => {
112                match block.merge_right(&range) {
113                    MergeRightResult::Before => 0,
114                    MergeRightResult::After { merged } => {
115                        if merged {
116                            range.clone_range_from(&block);
117                            1
118                        } else {
119                            // The new range fits entirely within an existing
120                            // block. Update the metadata and return early.
121                            block.set_meta(range.into_meta());
122                            return false;
123                        }
124                    }
125                }
126            }
127            None => 0,
128        };
129
130        if merge_left == 0 && merge_right == 0 {
131            // If the new segment cannot merge with any of its neighbors, we
132            // add a new entry for it.
133            blocks.insert(first_after, range);
134        } else {
135            // Otherwise, we put the new segment at the left edge of the merge
136            // window and remove all other existing segments.
137            let left_edge = first_after - merge_left;
138            let right_edge = first_after + merge_right;
139            blocks[left_edge] = range;
140            for i in right_edge..blocks.len() {
141                blocks[i - merge_left - merge_right + 1] = blocks[i].clone();
142            }
143            blocks.truncate(blocks.len() - merge_left - merge_right + 1);
144        }
145
146        true
147    }
148
149    pub(crate) fn iter(&self) -> impl Iterator<Item = &SeqRange<M>> + '_ {
150        self.blocks.iter()
151    }
152
153    /// Provides an iterator that allows modifying the metadata for each stored
154    /// range.
155    pub(crate) fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut SeqRange<M>> + '_ {
156        self.blocks.iter_mut()
157    }
158
159    /// Trims the ranges discarding any values at or before `value`.
160    pub(crate) fn discard_starting_at_or_before(&mut self, value: SeqNum) {
161        let Self { blocks } = self;
162        let first_after = Self::find_first_after(blocks, value);
163        // All the blocks starting at or before `value` can be discarded.
164        let _drain = blocks.drain(0..first_after);
165    }
166
167    /// Finds the first hole containing or after the sequence number `marker`.
168    ///
169    /// Returns [`FirstHoleResult::None`] if no hole can be identified, i.e.,
170    /// `marker`` is at or above the highest tracked range.
171    ///
172    /// Returns [`FirstHoleResult::Right`] if the hole is bounded to its
173    /// right side by the first stored sequence range,
174    /// [`FirstHoleResult::Both`] otherwise.
175    pub(crate) fn first_hole_on_or_after(&self, marker: SeqNum) -> FirstHoleResult<'_, M> {
176        let Self { blocks } = self;
177        let first_after = Self::find_first_after(blocks, marker);
178        blocks
179            .get(first_after)
180            .map(|right| {
181                first_after
182                    .checked_sub(1)
183                    .map(|left| FirstHoleResult::Both(&blocks[left], right))
184                    .unwrap_or_else(|| FirstHoleResult::Right(right))
185            })
186            .unwrap_or(FirstHoleResult::None)
187    }
188
189    /// Returns the highest sequence number range stored, if there are any.
190    pub(crate) fn last(&self) -> Option<&SeqRange<M>> {
191        let Self { blocks } = self;
192        blocks.back()
193    }
194
195    pub(crate) fn clear(&mut self) {
196        let Self { blocks } = self;
197        blocks.clear();
198    }
199}
200
201#[derive(Debug)]
202pub(crate) enum FirstHoleResult<'a, M> {
203    None,
204    Right(&'a SeqRange<M>),
205    Both(&'a SeqRange<M>, &'a SeqRange<M>),
206}
207
208impl<M: Clone> FromIterator<SeqRange<M>> for SeqRanges<M> {
209    fn from_iter<T: IntoIterator<Item = SeqRange<M>>>(iter: T) -> Self {
210        let mut ranges = SeqRanges::default();
211        for range in iter {
212            let _: bool = ranges.insert_seq_range(range);
213        }
214        ranges
215    }
216}
217
218mod range {
219    use netstack3_base::SackBlock;
220
221    use super::*;
222
223    /// A range kept in [`SeqRanges`].
224    #[derive(Debug, Clone)]
225    #[cfg_attr(test, derive(PartialEq, Eq))]
226    pub(crate) struct SeqRange<M> {
227        range: Range<SeqNum>,
228        meta: M,
229    }
230
231    impl<M> SeqRange<M> {
232        pub(crate) fn new(range: Range<SeqNum>, meta: M) -> Option<Self> {
233            range.end.after(range.start).then(|| Self { range, meta })
234        }
235
236        pub(crate) fn start(&self) -> SeqNum {
237            self.range.start
238        }
239
240        pub(crate) fn end(&self) -> SeqNum {
241            self.range.end
242        }
243
244        pub(crate) fn set_meta(&mut self, meta: M) {
245            self.meta = meta;
246        }
247
248        pub(crate) fn meta(&self) -> &M {
249            &self.meta
250        }
251
252        pub(crate) fn into_meta(self) -> M {
253            self.meta
254        }
255
256        pub(super) fn clone_range_from(&mut self, other: &Self) {
257            let Self { range, meta: _ } = self;
258            *range = other.range.clone();
259        }
260
261        pub(crate) fn len(&self) -> u32 {
262            let Self { range: Range { start, end }, meta: _ } = self;
263            let len = *end - *start;
264            // Assert on SeqRange invariant in debug only.
265            debug_assert!(len >= 0);
266            len as u32
267        }
268
269        pub(crate) fn cap_right(self, seq: SeqNum) -> Option<Self> {
270            let Self { range: Range { start, end }, meta } = self;
271            seq.after(start).then(|| Self { range: Range { start, end: end.earliest(seq) }, meta })
272        }
273
274        pub(crate) fn to_sack_block(&self) -> SackBlock {
275            let Self { range: Range { start, end }, meta: _ } = self;
276            // SAFETY: SackBlock requires that end is after start, which is the
277            // same invariant held by SeqRange.
278            unsafe { SackBlock::new_unchecked(*start, *end) }
279        }
280
281        pub(super) fn merge_right(&mut self, other: &Self) -> MergeRightResult {
282            if self.range.end.before(other.range.start) {
283                return MergeRightResult::Before;
284            }
285
286            let merged = self.range.end.before(other.range.end);
287            if merged {
288                self.range.end = other.range.end;
289            }
290
291            MergeRightResult::After { merged }
292        }
293    }
294
295    pub(super) enum MergeRightResult {
296        Before,
297        After { merged: bool },
298    }
299}
300use range::MergeRightResult;
301pub(crate) use range::SeqRange;
302
303#[cfg(test)]
304mod test {
305    use super::*;
306
307    use alloc::vec::Vec;
308    use alloc::{format, vec};
309
310    use netstack3_base::{SackBlock, WindowSize};
311    use proptest::strategy::{Just, Strategy};
312    use proptest::test_runner::Config;
313    use proptest::{prop_assert, prop_assert_eq, proptest};
314    use proptest_support::failed_seeds_no_std;
315    use test_case::test_case;
316
317    impl SeqRanges<()> {
318        fn insert_u32(&mut self, range: Range<u32>) -> bool {
319            let Range { start, end } = range;
320            self.insert(SeqNum::new(start)..SeqNum::new(end), ())
321        }
322    }
323
324    impl FromIterator<Range<u32>> for SeqRanges<()> {
325        fn from_iter<T: IntoIterator<Item = Range<u32>>>(iter: T) -> Self {
326            let mut ranges = SeqRanges::default();
327            for range in iter {
328                let _: bool = ranges.insert_u32(range);
329            }
330            ranges
331        }
332    }
333
334    proptest! {
335        #![proptest_config(Config {
336            // Add all failed seeds here.
337            failure_persistence: failed_seeds_no_std!(
338                "cc f621ca7d3a2b108e0dc41f7169ad028f4329b79e90e73d5f68042519a9f63999",
339                "cc c449aebed201b4ec4f137f3c224f20325f4cfee0b7fd596d9285176b6d811aa9"
340            ),
341            ..Config::default()
342        })]
343
344        #[test]
345        fn seq_ranges_insert(insertions in proptest::collection::vec(insertions(), 200)) {
346            let mut seq_ranges = SeqRanges::<()>::default();
347            let mut num_insertions_performed = 0;
348            let mut min_seq = SeqNum::new(WindowSize::MAX.into());
349            let mut max_seq = SeqNum::new(0);
350            for Range { start, end } in insertions {
351                if min_seq.after(start) {
352                    min_seq = start;
353                }
354                if max_seq.before(end) {
355                    max_seq = end;
356                }
357                // assert that it's impossible to have more entries than the
358                // number of insertions performed.
359                prop_assert!(seq_ranges.blocks.len() <= num_insertions_performed);
360                let _: bool = seq_ranges.insert(start..end, ());
361                num_insertions_performed += 1;
362
363                // assert that the ranges are sorted and don't overlap with
364                // each other.
365                for i in 1..seq_ranges.blocks.len() {
366                    prop_assert!(
367                        seq_ranges.blocks[i-1].end().before(seq_ranges.blocks[i].start())
368                    );
369                }
370            }
371            prop_assert_eq!(seq_ranges.blocks.front().unwrap().start(), min_seq);
372            prop_assert_eq!(seq_ranges.blocks.back().unwrap().end(), max_seq);
373        }
374
375        // Test that the invariants between SackBlock and SeqRange creation
376        // match. Supports unsafe block in SeqRange::to_sack_block.
377        #[test]
378        fn seq_range_to_sack_block((start, end) in sequence_numbers()) {
379            prop_assert_eq!(
380                SeqRange::new(start..end, ()).map(|sr| sr.to_sack_block()),
381                SackBlock::try_new(start, end).ok()
382            );
383        }
384
385    }
386
387    fn insertions() -> impl Strategy<Value = Range<SeqNum>> {
388        (0..u32::from(WindowSize::MAX)).prop_flat_map(|start| {
389            (start + 1..=u32::from(WindowSize::MAX)).prop_flat_map(move |end| {
390                Just(Range { start: SeqNum::new(start), end: SeqNum::new(end) })
391            })
392        })
393    }
394
395    fn sequence_numbers() -> impl Strategy<Value = (SeqNum, SeqNum)> {
396        (0u32..5).prop_flat_map(|start| {
397            (0u32..5).prop_flat_map(move |end| Just((SeqNum::new(start), SeqNum::new(end))))
398        })
399    }
400
401    #[test]
402    fn insert_return() {
403        let mut sr = SeqRanges::default();
404        assert!(sr.insert_u32(10..20));
405
406        assert!(!sr.insert_u32(10..20));
407        assert!(!sr.insert_u32(11..20));
408        assert!(!sr.insert_u32(11..12));
409        assert!(!sr.insert_u32(19..20));
410
411        assert!(sr.insert_u32(0..5));
412        assert!(sr.insert_u32(25..35));
413        assert!(sr.insert_u32(5..7));
414        assert!(sr.insert_u32(22..25));
415
416        assert!(sr.insert_u32(7..22));
417        assert!(!sr.insert_u32(0..35));
418    }
419
420    #[test_case(&[], 0 => Vec::<Range<u32>>::new(); "empty")]
421    #[test_case(&[10..20], 0 => vec![10..20]; "before 1")]
422    #[test_case(&[10..20, 30..40], 0 => vec![10..20, 30..40]; "before 2")]
423    #[test_case(&[10..20], 10 =>  Vec::<Range<u32>>::new(); "same 1")]
424    #[test_case(&[10..20, 30..40], 10 => vec![30..40]; "same 2")]
425    #[test_case(&[10..20], 20 =>  Vec::<Range<u32>>::new(); "after 1")]
426    #[test_case(&[10..20, 30..40], 20 => vec![30..40]; "after 2")]
427    #[test_case(&[10..20, 30..40], 30 =>  Vec::<Range<u32>>::new(); "after 3")]
428    #[test_case(&[10..20, 30..40], 15 =>  vec![30..40]; "mid 1")]
429    #[test_case(&[10..20, 30..40], 35 =>  Vec::<Range<u32>>::new(); "mid 2")]
430    fn discard_starting_at_or_before(ranges: &[Range<u32>], discard: u32) -> Vec<Range<u32>> {
431        let mut sr = ranges.into_iter().cloned().collect::<SeqRanges<()>>();
432        sr.discard_starting_at_or_before(SeqNum::new(discard));
433        sr.iter()
434            .map(|seq_range| u32::from(seq_range.start())..u32::from(seq_range.end()))
435            .collect()
436    }
437
438    #[test_case(&[], 0 => (None, None))]
439    #[test_case(&[10..20, 30..40], 0 => (None, Some(10)))]
440    #[test_case(&[10..20, 30..40], 10 => (Some(20), Some(30)))]
441    #[test_case(&[10..20, 30..40], 20 => (Some(20), Some(30)))]
442    #[test_case(&[10..20, 30..40], 30 => (None, None))]
443    #[test_case(&[10..20, 30..40], 40 => (None, None))]
444    #[test_case(&[10..20, 30..40, 50..60], 30 => (Some(40), Some(50)))]
445    #[test_case(&[10..20, 30..40, 50..60], 35 => (Some(40), Some(50)))]
446    #[test_case(&[10..20, 30..40, 50..60], 40 => (Some(40), Some(50)))]
447    fn first_hole_on_or_after(ranges: &[Range<u32>], marker: u32) -> (Option<u32>, Option<u32>) {
448        match ranges
449            .into_iter()
450            .cloned()
451            .collect::<SeqRanges<()>>()
452            .first_hole_on_or_after(SeqNum::new(marker))
453        {
454            FirstHoleResult::None => (None, None),
455            FirstHoleResult::Right(right) => (None, Some(u32::from(right.start()))),
456            FirstHoleResult::Both(left, right) => {
457                (Some(u32::from(left.end())), Some(u32::from(right.start())))
458            }
459        }
460    }
461}