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: &mut 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
172impl<M: Clone> FromIterator<SeqRange<M>> for SeqRanges<M> {
173    fn from_iter<T: IntoIterator<Item = SeqRange<M>>>(iter: T) -> Self {
174        let mut ranges = SeqRanges::default();
175        for range in iter {
176            let _: bool = ranges.insert_seq_range(range);
177        }
178        ranges
179    }
180}
181
182mod range {
183    use netstack3_base::SackBlock;
184
185    use super::*;
186
187    /// A range kept in [`SeqRanges`].
188    #[derive(Debug, Clone)]
189    #[cfg_attr(test, derive(PartialEq, Eq))]
190    pub(crate) struct SeqRange<M> {
191        range: Range<SeqNum>,
192        meta: M,
193    }
194
195    impl<M> SeqRange<M> {
196        pub(crate) fn new(range: Range<SeqNum>, meta: M) -> Option<Self> {
197            range.end.after(range.start).then(|| Self { range, meta })
198        }
199
200        pub(crate) fn start(&self) -> SeqNum {
201            self.range.start
202        }
203
204        pub(crate) fn end(&self) -> SeqNum {
205            self.range.end
206        }
207
208        pub(crate) fn set_meta(&mut self, meta: M) {
209            self.meta = meta;
210        }
211
212        pub(crate) fn meta(&self) -> &M {
213            &self.meta
214        }
215
216        pub(crate) fn into_meta(self) -> M {
217            self.meta
218        }
219
220        pub(super) fn clone_range_from(&mut self, other: &Self) {
221            let Self { range, meta: _ } = self;
222            *range = other.range.clone();
223        }
224
225        pub(crate) fn len(&self) -> u32 {
226            let Self { range: Range { start, end }, meta: _ } = self;
227            let len = *end - *start;
228            // Assert on SeqRange invariant in debug only.
229            debug_assert!(len >= 0);
230            len as u32
231        }
232
233        pub(crate) fn cap_right(self, seq: SeqNum) -> Option<Self> {
234            let Self { range: Range { start, end }, meta } = self;
235            seq.after(start).then(|| Self { range: Range { start, end: end.earliest(seq) }, meta })
236        }
237
238        pub(crate) fn to_sack_block(&self) -> SackBlock {
239            let Self { range: Range { start, end }, meta: _ } = self;
240            // SAFETY: SackBlock requires that end is after start, which is the
241            // same invariant held by SeqRange.
242            unsafe { SackBlock::new_unchecked(*start, *end) }
243        }
244
245        pub(super) fn merge_right(&mut self, other: &Self) -> MergeRightResult {
246            if self.range.end.before(other.range.start) {
247                return MergeRightResult::Before;
248            }
249
250            let merged = self.range.end.before(other.range.end);
251            if merged {
252                self.range.end = other.range.end;
253            }
254
255            MergeRightResult::After { merged }
256        }
257    }
258
259    pub(super) enum MergeRightResult {
260        Before,
261        After { merged: bool },
262    }
263}
264use range::MergeRightResult;
265pub(crate) use range::SeqRange;
266
267#[cfg(test)]
268mod test {
269    use super::*;
270
271    use alloc::vec::Vec;
272    use alloc::{format, vec};
273
274    use netstack3_base::{SackBlock, WindowSize};
275    use proptest::strategy::{Just, Strategy};
276    use proptest::test_runner::Config;
277    use proptest::{prop_assert, prop_assert_eq, proptest};
278    use proptest_support::failed_seeds_no_std;
279    use test_case::test_case;
280
281    impl SeqRanges<()> {
282        fn insert_u32(&mut self, range: Range<u32>) -> bool {
283            let Range { start, end } = range;
284            self.insert(SeqNum::new(start)..SeqNum::new(end), ())
285        }
286    }
287
288    impl FromIterator<Range<u32>> for SeqRanges<()> {
289        fn from_iter<T: IntoIterator<Item = Range<u32>>>(iter: T) -> Self {
290            let mut ranges = SeqRanges::default();
291            for range in iter {
292                let _: bool = ranges.insert_u32(range);
293            }
294            ranges
295        }
296    }
297
298    proptest! {
299        #![proptest_config(Config {
300            // Add all failed seeds here.
301            failure_persistence: failed_seeds_no_std!(
302                "cc f621ca7d3a2b108e0dc41f7169ad028f4329b79e90e73d5f68042519a9f63999",
303                "cc c449aebed201b4ec4f137f3c224f20325f4cfee0b7fd596d9285176b6d811aa9"
304            ),
305            ..Config::default()
306        })]
307
308        #[test]
309        fn seq_ranges_insert(insertions in proptest::collection::vec(insertions(), 200)) {
310            let mut seq_ranges = SeqRanges::<()>::default();
311            let mut num_insertions_performed = 0;
312            let mut min_seq = SeqNum::new(WindowSize::MAX.into());
313            let mut max_seq = SeqNum::new(0);
314            for Range { start, end } in insertions {
315                if min_seq.after(start) {
316                    min_seq = start;
317                }
318                if max_seq.before(end) {
319                    max_seq = end;
320                }
321                // assert that it's impossible to have more entries than the
322                // number of insertions performed.
323                prop_assert!(seq_ranges.blocks.len() <= num_insertions_performed);
324                let _: bool = seq_ranges.insert(start..end, ());
325                num_insertions_performed += 1;
326
327                // assert that the ranges are sorted and don't overlap with
328                // each other.
329                for i in 1..seq_ranges.blocks.len() {
330                    prop_assert!(
331                        seq_ranges.blocks[i-1].end().before(seq_ranges.blocks[i].start())
332                    );
333                }
334            }
335            prop_assert_eq!(seq_ranges.blocks.front().unwrap().start(), min_seq);
336            prop_assert_eq!(seq_ranges.blocks.back().unwrap().end(), max_seq);
337        }
338
339        // Test that the invariants between SackBlock and SeqRange creation
340        // match. Supports unsafe block in SeqRange::to_sack_block.
341        #[test]
342        fn seq_range_to_sack_block((start, end) in sequence_numbers()) {
343            prop_assert_eq!(
344                SeqRange::new(start..end, ()).map(|sr| sr.to_sack_block()),
345                SackBlock::try_new(start, end).ok()
346            );
347        }
348
349    }
350
351    fn insertions() -> impl Strategy<Value = Range<SeqNum>> {
352        (0..u32::from(WindowSize::MAX)).prop_flat_map(|start| {
353            (start + 1..=u32::from(WindowSize::MAX)).prop_flat_map(move |end| {
354                Just(Range { start: SeqNum::new(start), end: SeqNum::new(end) })
355            })
356        })
357    }
358
359    fn sequence_numbers() -> impl Strategy<Value = (SeqNum, SeqNum)> {
360        (0u32..5).prop_flat_map(|start| {
361            (0u32..5).prop_flat_map(move |end| Just((SeqNum::new(start), SeqNum::new(end))))
362        })
363    }
364
365    #[test]
366    fn insert_return() {
367        let mut sr = SeqRanges::default();
368        assert!(sr.insert_u32(10..20));
369
370        assert!(!sr.insert_u32(10..20));
371        assert!(!sr.insert_u32(11..20));
372        assert!(!sr.insert_u32(11..12));
373        assert!(!sr.insert_u32(19..20));
374
375        assert!(sr.insert_u32(0..5));
376        assert!(sr.insert_u32(25..35));
377        assert!(sr.insert_u32(5..7));
378        assert!(sr.insert_u32(22..25));
379
380        assert!(sr.insert_u32(7..22));
381        assert!(!sr.insert_u32(0..35));
382    }
383
384    #[test_case(&[], 0 => Vec::<Range<u32>>::new(); "empty")]
385    #[test_case(&[10..20], 0 => vec![10..20]; "before 1")]
386    #[test_case(&[10..20, 30..40], 0 => vec![10..20, 30..40]; "before 2")]
387    #[test_case(&[10..20], 10 =>  Vec::<Range<u32>>::new(); "same 1")]
388    #[test_case(&[10..20, 30..40], 10 => vec![30..40]; "same 2")]
389    #[test_case(&[10..20], 20 =>  Vec::<Range<u32>>::new(); "after 1")]
390    #[test_case(&[10..20, 30..40], 20 => vec![30..40]; "after 2")]
391    #[test_case(&[10..20, 30..40], 30 =>  Vec::<Range<u32>>::new(); "after 3")]
392    #[test_case(&[10..20, 30..40], 15 =>  vec![30..40]; "mid 1")]
393    #[test_case(&[10..20, 30..40], 35 =>  Vec::<Range<u32>>::new(); "mid 2")]
394    fn discard_starting_at_or_before(ranges: &[Range<u32>], discard: u32) -> Vec<Range<u32>> {
395        let mut sr = ranges.into_iter().cloned().collect::<SeqRanges<()>>();
396        sr.discard_starting_at_or_before(SeqNum::new(discard));
397        sr.iter()
398            .map(|seq_range| u32::from(seq_range.start())..u32::from(seq_range.end()))
399            .collect()
400    }
401}
netstack3_tcp/seq_ranges.rs

netstack3_tcp/
seq_ranges.rs
