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