smoltcp/storage/
assembler.rs

1use core::fmt;
2
3use crate::config::ASSEMBLER_MAX_SEGMENT_COUNT;
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub struct TooManyHolesError;
7
8impl fmt::Display for TooManyHolesError {
9    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
10        write!(f, "too many holes")
11    }
12}
13
14#[cfg(feature = "std")]
15impl std::error::Error for TooManyHolesError {}
16
17/// A contiguous chunk of absent data, followed by a contiguous chunk of present data.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19struct Contig {
20    hole_size: usize,
21    data_size: usize,
22}
23
24impl fmt::Display for Contig {
25    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
26        if self.has_hole() {
27            write!(f, "({})", self.hole_size)?;
28        }
29        if self.has_hole() && self.has_data() {
30            write!(f, " ")?;
31        }
32        if self.has_data() {
33            write!(f, "{}", self.data_size)?;
34        }
35        Ok(())
36    }
37}
38
39#[cfg(feature = "defmt")]
40impl defmt::Format for Contig {
41    fn format(&self, fmt: defmt::Formatter) {
42        if self.has_hole() {
43            defmt::write!(fmt, "({})", self.hole_size);
44        }
45        if self.has_hole() && self.has_data() {
46            defmt::write!(fmt, " ");
47        }
48        if self.has_data() {
49            defmt::write!(fmt, "{}", self.data_size);
50        }
51    }
52}
53
54impl Contig {
55    const fn empty() -> Contig {
56        Contig {
57            hole_size: 0,
58            data_size: 0,
59        }
60    }
61
62    fn hole_and_data(hole_size: usize, data_size: usize) -> Contig {
63        Contig {
64            hole_size,
65            data_size,
66        }
67    }
68
69    fn has_hole(&self) -> bool {
70        self.hole_size != 0
71    }
72
73    fn has_data(&self) -> bool {
74        self.data_size != 0
75    }
76
77    fn total_size(&self) -> usize {
78        self.hole_size + self.data_size
79    }
80
81    fn shrink_hole_by(&mut self, size: usize) {
82        self.hole_size -= size;
83    }
84
85    fn shrink_hole_to(&mut self, size: usize) {
86        debug_assert!(self.hole_size >= size);
87
88        let total_size = self.total_size();
89        self.hole_size = size;
90        self.data_size = total_size - size;
91    }
92}
93
94/// A buffer (re)assembler.
95///
96/// Currently, up to a hardcoded limit of 4 or 32 holes can be tracked in the buffer.
97#[derive(Debug, PartialEq, Eq, Clone)]
98pub struct Assembler {
99    contigs: [Contig; ASSEMBLER_MAX_SEGMENT_COUNT],
100}
101
102impl fmt::Display for Assembler {
103    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
104        write!(f, "[ ")?;
105        for contig in self.contigs.iter() {
106            if !contig.has_data() {
107                break;
108            }
109            write!(f, "{contig} ")?;
110        }
111        write!(f, "]")?;
112        Ok(())
113    }
114}
115
116#[cfg(feature = "defmt")]
117impl defmt::Format for Assembler {
118    fn format(&self, fmt: defmt::Formatter) {
119        defmt::write!(fmt, "[ ");
120        for contig in self.contigs.iter() {
121            if !contig.has_data() {
122                break;
123            }
124            defmt::write!(fmt, "{} ", contig);
125        }
126        defmt::write!(fmt, "]");
127    }
128}
129
130// Invariant on Assembler::contigs:
131// - There's an index `i` where all contigs before have data, and all contigs after don't (are unused).
132// - All contigs with data must have hole_size != 0, except the first.
133
134impl Assembler {
135    /// Create a new buffer assembler.
136    pub const fn new() -> Assembler {
137        const EMPTY: Contig = Contig::empty();
138        Assembler {
139            contigs: [EMPTY; ASSEMBLER_MAX_SEGMENT_COUNT],
140        }
141    }
142
143    pub fn clear(&mut self) {
144        self.contigs.fill(Contig::empty());
145    }
146
147    fn front(&self) -> Contig {
148        self.contigs[0]
149    }
150
151    /// Return length of the front contiguous range without removing it from the assembler
152    pub fn peek_front(&self) -> usize {
153        let front = self.front();
154        if front.has_hole() { 0 } else { front.data_size }
155    }
156
157    fn back(&self) -> Contig {
158        self.contigs[self.contigs.len() - 1]
159    }
160
161    /// Return whether the assembler contains no data.
162    pub fn is_empty(&self) -> bool {
163        !self.front().has_data()
164    }
165
166    /// Remove a contig at the given index.
167    fn remove_contig_at(&mut self, at: usize) {
168        debug_assert!(self.contigs[at].has_data());
169
170        for i in at..self.contigs.len() - 1 {
171            if !self.contigs[i].has_data() {
172                return;
173            }
174            self.contigs[i] = self.contigs[i + 1];
175        }
176
177        // Removing the last one.
178        self.contigs[self.contigs.len() - 1] = Contig::empty();
179    }
180
181    /// Add a contig at the given index, and return a pointer to it.
182    fn add_contig_at(&mut self, at: usize) -> Result<&mut Contig, TooManyHolesError> {
183        if self.back().has_data() {
184            return Err(TooManyHolesError);
185        }
186
187        for i in (at + 1..self.contigs.len()).rev() {
188            self.contigs[i] = self.contigs[i - 1];
189        }
190
191        self.contigs[at] = Contig::empty();
192        Ok(&mut self.contigs[at])
193    }
194
195    /// Add a new contiguous range to the assembler,
196    /// or return `Err(TooManyHolesError)` if too many discontinuities are already recorded.
197    pub fn add(&mut self, mut offset: usize, size: usize) -> Result<(), TooManyHolesError> {
198        if size == 0 {
199            return Ok(());
200        }
201
202        let mut i = 0;
203
204        // Find index of the contig containing the start of the range.
205        loop {
206            if i == self.contigs.len() {
207                // The new range is after all the previous ranges, but there/s no space to add it.
208                return Err(TooManyHolesError);
209            }
210            let contig = &mut self.contigs[i];
211            if !contig.has_data() {
212                // The new range is after all the previous ranges. Add it.
213                *contig = Contig::hole_and_data(offset, size);
214                return Ok(());
215            }
216            if offset <= contig.total_size() {
217                break;
218            }
219            offset -= contig.total_size();
220            i += 1;
221        }
222
223        let contig = &mut self.contigs[i];
224        if offset < contig.hole_size {
225            // Range starts within the hole.
226
227            if offset + size < contig.hole_size {
228                // Range also ends within the hole.
229                let new_contig = self.add_contig_at(i)?;
230                new_contig.hole_size = offset;
231                new_contig.data_size = size;
232
233                // Previous contigs[index] got moved to contigs[index+1]
234                self.contigs[i + 1].shrink_hole_by(offset + size);
235                return Ok(());
236            }
237
238            // The range being added covers both a part of the hole and a part of the data
239            // in this contig, shrink the hole in this contig.
240            contig.shrink_hole_to(offset);
241        }
242
243        // coalesce contigs to the right.
244        let mut j = i + 1;
245        while j < self.contigs.len()
246            && self.contigs[j].has_data()
247            && offset + size >= self.contigs[i].total_size() + self.contigs[j].hole_size
248        {
249            self.contigs[i].data_size += self.contigs[j].total_size();
250            j += 1;
251        }
252        let shift = j - i - 1;
253        if shift != 0 {
254            for x in i + 1..self.contigs.len() {
255                if !self.contigs[x].has_data() {
256                    break;
257                }
258
259                self.contigs[x] = self
260                    .contigs
261                    .get(x + shift)
262                    .copied()
263                    .unwrap_or_else(Contig::empty);
264            }
265        }
266
267        if offset + size > self.contigs[i].total_size() {
268            // The added range still extends beyond the current contig. Increase data size.
269            let left = offset + size - self.contigs[i].total_size();
270            self.contigs[i].data_size += left;
271
272            // Decrease hole size of the next, if any.
273            if i + 1 < self.contigs.len() && self.contigs[i + 1].has_data() {
274                self.contigs[i + 1].hole_size -= left;
275            }
276        }
277
278        Ok(())
279    }
280
281    /// Remove a contiguous range from the front of the assembler.
282    /// If no such range, return 0.
283    pub fn remove_front(&mut self) -> usize {
284        let front = self.front();
285        if front.has_hole() || !front.has_data() {
286            0
287        } else {
288            self.remove_contig_at(0);
289            debug_assert!(front.data_size > 0);
290            front.data_size
291        }
292    }
293
294    /// Add a segment, then remove_front.
295    ///
296    /// This is equivalent to calling `add` then `remove_front` individually,
297    /// except it's guaranteed to not fail when offset = 0.
298    /// This is required for TCP: we must never drop the next expected segment, or
299    /// the protocol might get stuck.
300    pub fn add_then_remove_front(
301        &mut self,
302        offset: usize,
303        size: usize,
304    ) -> Result<usize, TooManyHolesError> {
305        // This is the only case where a segment at offset=0 would cause the
306        // total amount of contigs to rise (and therefore can potentially cause
307        // a TooManyHolesError). Handle it in a way that is guaranteed to succeed.
308        if offset == 0 && size < self.contigs[0].hole_size {
309            self.contigs[0].hole_size -= size;
310            return Ok(size);
311        }
312
313        self.add(offset, size)?;
314        Ok(self.remove_front())
315    }
316
317    /// Iterate over all of the contiguous data ranges.
318    ///
319    /// Returns `(offset, size)` tuples for each contiguous data range, where
320    /// offset is relative to the start of the assembler.
321    ///
322    ///    Data        Hole        Data
323    /// |--- 100 ---|--- 200 ---|--- 100 ---|
324    ///
325    /// Would return the ranges: ``(0, 100), (300, 400)``
326    pub fn iter_data(&self) -> impl Iterator<Item = (usize, usize)> + '_ {
327        let mut offset = 0;
328        self.contigs.iter().filter_map(move |contig| {
329            offset += contig.hole_size;
330            let left = offset;
331            offset += contig.data_size;
332            let right = offset;
333            if left < right {
334                Some((left, right))
335            } else {
336                None
337            }
338        })
339    }
340}
341
342#[cfg(test)]
343mod test {
344    use super::*;
345    use std::vec::Vec;
346
347    impl From<Vec<(usize, usize)>> for Assembler {
348        fn from(vec: Vec<(usize, usize)>) -> Assembler {
349            const EMPTY: Contig = Contig::empty();
350
351            let mut contigs = [EMPTY; ASSEMBLER_MAX_SEGMENT_COUNT];
352            for (i, &(hole_size, data_size)) in vec.iter().enumerate() {
353                contigs[i] = Contig {
354                    hole_size,
355                    data_size,
356                };
357            }
358            Assembler { contigs }
359        }
360    }
361
362    macro_rules! contigs {
363        [$( $x:expr ),*] => ({
364            Assembler::from(vec![$( $x ),*])
365        })
366    }
367
368    #[test]
369    fn test_new() {
370        let assr = Assembler::new();
371        assert_eq!(assr, contigs![]);
372    }
373
374    #[test]
375    fn test_empty_add_full() {
376        let mut assr = Assembler::new();
377        assert_eq!(assr.add(0, 16), Ok(()));
378        assert_eq!(assr, contigs![(0, 16)]);
379    }
380
381    #[test]
382    fn test_empty_add_front() {
383        let mut assr = Assembler::new();
384        assert_eq!(assr.add(0, 4), Ok(()));
385        assert_eq!(assr, contigs![(0, 4)]);
386    }
387
388    #[test]
389    fn test_empty_add_back() {
390        let mut assr = Assembler::new();
391        assert_eq!(assr.add(12, 4), Ok(()));
392        assert_eq!(assr, contigs![(12, 4)]);
393    }
394
395    #[test]
396    fn test_empty_add_mid() {
397        let mut assr = Assembler::new();
398        assert_eq!(assr.add(4, 8), Ok(()));
399        assert_eq!(assr, contigs![(4, 8)]);
400    }
401
402    #[test]
403    fn test_partial_add_front() {
404        let mut assr = contigs![(4, 8)];
405        assert_eq!(assr.add(0, 4), Ok(()));
406        assert_eq!(assr, contigs![(0, 12)]);
407    }
408
409    #[test]
410    fn test_partial_add_back() {
411        let mut assr = contigs![(4, 8)];
412        assert_eq!(assr.add(12, 4), Ok(()));
413        assert_eq!(assr, contigs![(4, 12)]);
414    }
415
416    #[test]
417    fn test_partial_add_front_overlap() {
418        let mut assr = contigs![(4, 8)];
419        assert_eq!(assr.add(0, 8), Ok(()));
420        assert_eq!(assr, contigs![(0, 12)]);
421    }
422
423    #[test]
424    fn test_partial_add_front_overlap_split() {
425        let mut assr = contigs![(4, 8)];
426        assert_eq!(assr.add(2, 6), Ok(()));
427        assert_eq!(assr, contigs![(2, 10)]);
428    }
429
430    #[test]
431    fn test_partial_add_back_overlap() {
432        let mut assr = contigs![(4, 8)];
433        assert_eq!(assr.add(8, 8), Ok(()));
434        assert_eq!(assr, contigs![(4, 12)]);
435    }
436
437    #[test]
438    fn test_partial_add_back_overlap_split() {
439        let mut assr = contigs![(4, 8)];
440        assert_eq!(assr.add(10, 4), Ok(()));
441        assert_eq!(assr, contigs![(4, 10)]);
442    }
443
444    #[test]
445    fn test_partial_add_both_overlap() {
446        let mut assr = contigs![(4, 8)];
447        assert_eq!(assr.add(0, 16), Ok(()));
448        assert_eq!(assr, contigs![(0, 16)]);
449    }
450
451    #[test]
452    fn test_partial_add_both_overlap_split() {
453        let mut assr = contigs![(4, 8)];
454        assert_eq!(assr.add(2, 12), Ok(()));
455        assert_eq!(assr, contigs![(2, 12)]);
456    }
457
458    #[test]
459    fn test_rejected_add_keeps_state() {
460        let mut assr = Assembler::new();
461        for c in 1..=ASSEMBLER_MAX_SEGMENT_COUNT {
462            assert_eq!(assr.add(c * 10, 3), Ok(()));
463        }
464        // Maximum of allowed holes is reached
465        let assr_before = assr.clone();
466        assert_eq!(assr.add(1, 3), Err(TooManyHolesError));
467        assert_eq!(assr_before, assr);
468    }
469
470    #[test]
471    fn test_empty_remove_front() {
472        let mut assr = contigs![];
473        assert_eq!(assr.remove_front(), 0);
474    }
475
476    #[test]
477    fn test_trailing_hole_remove_front() {
478        let mut assr = contigs![(0, 4)];
479        assert_eq!(assr.remove_front(), 4);
480        assert_eq!(assr, contigs![]);
481    }
482
483    #[test]
484    fn test_trailing_data_remove_front() {
485        let mut assr = contigs![(0, 4), (4, 4)];
486        assert_eq!(assr.remove_front(), 4);
487        assert_eq!(assr, contigs![(4, 4)]);
488    }
489
490    #[test]
491    fn test_boundary_case_remove_front() {
492        let mut vec = vec![(1, 1); ASSEMBLER_MAX_SEGMENT_COUNT];
493        vec[0] = (0, 2);
494        let mut assr = Assembler::from(vec);
495        assert_eq!(assr.remove_front(), 2);
496        let mut vec = vec![(1, 1); ASSEMBLER_MAX_SEGMENT_COUNT];
497        vec[ASSEMBLER_MAX_SEGMENT_COUNT - 1] = (0, 0);
498        let exp_assr = Assembler::from(vec);
499        assert_eq!(assr, exp_assr);
500    }
501
502    #[test]
503    fn test_shrink_next_hole() {
504        let mut assr = Assembler::new();
505        assert_eq!(assr.add(100, 10), Ok(()));
506        assert_eq!(assr.add(50, 10), Ok(()));
507        assert_eq!(assr.add(40, 30), Ok(()));
508        assert_eq!(assr, contigs![(40, 30), (30, 10)]);
509    }
510
511    #[test]
512    fn test_join_two() {
513        let mut assr = Assembler::new();
514        assert_eq!(assr.add(10, 10), Ok(()));
515        assert_eq!(assr.add(50, 10), Ok(()));
516        assert_eq!(assr.add(15, 40), Ok(()));
517        assert_eq!(assr, contigs![(10, 50)]);
518    }
519
520    #[test]
521    fn test_join_two_reversed() {
522        let mut assr = Assembler::new();
523        assert_eq!(assr.add(50, 10), Ok(()));
524        assert_eq!(assr.add(10, 10), Ok(()));
525        assert_eq!(assr.add(15, 40), Ok(()));
526        assert_eq!(assr, contigs![(10, 50)]);
527    }
528
529    #[test]
530    fn test_join_two_overlong() {
531        let mut assr = Assembler::new();
532        assert_eq!(assr.add(50, 10), Ok(()));
533        assert_eq!(assr.add(10, 10), Ok(()));
534        assert_eq!(assr.add(15, 60), Ok(()));
535        assert_eq!(assr, contigs![(10, 65)]);
536    }
537
538    #[test]
539    fn test_iter_empty() {
540        let assr = Assembler::new();
541        let segments: Vec<_> = assr.iter_data().collect();
542        assert_eq!(segments, vec![]);
543    }
544
545    #[test]
546    fn test_iter_full() {
547        let mut assr = Assembler::new();
548        assert_eq!(assr.add(0, 16), Ok(()));
549        let segments: Vec<_> = assr.iter_data().collect();
550        assert_eq!(segments, vec![(0, 16)]);
551    }
552
553    #[test]
554    fn test_iter_one_front() {
555        let mut assr = Assembler::new();
556        assert_eq!(assr.add(0, 4), Ok(()));
557        let segments: Vec<_> = assr.iter_data().collect();
558        assert_eq!(segments, vec![(0, 4)]);
559    }
560
561    #[test]
562    fn test_iter_one_back() {
563        let mut assr = Assembler::new();
564        assert_eq!(assr.add(12, 4), Ok(()));
565        let segments: Vec<_> = assr.iter_data().collect();
566        assert_eq!(segments, vec![(12, 16)]);
567    }
568
569    #[test]
570    fn test_iter_one_mid() {
571        let mut assr = Assembler::new();
572        assert_eq!(assr.add(4, 8), Ok(()));
573        let segments: Vec<_> = assr.iter_data().collect();
574        assert_eq!(segments, vec![(4, 12)]);
575    }
576
577    #[test]
578    fn test_iter_one_trailing_gap() {
579        let assr = contigs![(4, 8)];
580        let segments: Vec<_> = assr.iter_data().collect();
581        assert_eq!(segments, vec![(4, 12)]);
582    }
583
584    #[test]
585    fn test_iter_two_split() {
586        let assr = contigs![(2, 6), (4, 1)];
587        let segments: Vec<_> = assr.iter_data().collect();
588        assert_eq!(segments, vec![(2, 8), (12, 13)]);
589    }
590
591    #[test]
592    fn test_iter_three_split() {
593        let assr = contigs![(2, 6), (2, 1), (2, 2)];
594        let segments: Vec<_> = assr.iter_data().collect();
595        assert_eq!(segments, vec![(2, 8), (10, 11), (13, 15)]);
596    }
597
598    #[test]
599    fn test_issue_694() {
600        let mut assr = Assembler::new();
601        assert_eq!(assr.add(0, 1), Ok(()));
602        assert_eq!(assr.add(2, 1), Ok(()));
603        assert_eq!(assr.add(1, 1), Ok(()));
604    }
605
606    #[test]
607    fn test_add_then_remove_front() {
608        let mut assr = Assembler::new();
609        assert_eq!(assr.add(50, 10), Ok(()));
610        assert_eq!(assr.add_then_remove_front(10, 10), Ok(0));
611        assert_eq!(assr, contigs![(10, 10), (30, 10)]);
612    }
613
614    #[test]
615    fn test_add_then_remove_front_at_front() {
616        let mut assr = Assembler::new();
617        assert_eq!(assr.add(50, 10), Ok(()));
618        assert_eq!(assr.add_then_remove_front(0, 10), Ok(10));
619        assert_eq!(assr, contigs![(40, 10)]);
620    }
621
622    #[test]
623    fn test_add_then_remove_front_at_front_touch() {
624        let mut assr = Assembler::new();
625        assert_eq!(assr.add(50, 10), Ok(()));
626        assert_eq!(assr.add_then_remove_front(0, 50), Ok(60));
627        assert_eq!(assr, contigs![]);
628    }
629
630    #[test]
631    fn test_add_then_remove_front_at_front_full() {
632        let mut assr = Assembler::new();
633        for c in 1..=ASSEMBLER_MAX_SEGMENT_COUNT {
634            assert_eq!(assr.add(c * 10, 3), Ok(()));
635        }
636        // Maximum of allowed holes is reached
637        let assr_before = assr.clone();
638        assert_eq!(assr.add_then_remove_front(1, 3), Err(TooManyHolesError));
639        assert_eq!(assr_before, assr);
640    }
641
642    #[test]
643    fn test_add_then_remove_front_at_front_full_offset_0() {
644        let mut assr = Assembler::new();
645        for c in 1..=ASSEMBLER_MAX_SEGMENT_COUNT {
646            assert_eq!(assr.add(c * 10, 3), Ok(()));
647        }
648        assert_eq!(assr.add_then_remove_front(0, 3), Ok(3));
649    }
650
651    // Test against an obviously-correct but inefficient bitmap impl.
652    #[test]
653    fn test_random() {
654        use rand::Rng;
655
656        const MAX_INDEX: usize = 256;
657
658        for max_size in [2, 5, 10, 100] {
659            for _ in 0..300 {
660                //println!("===");
661                let mut assr = Assembler::new();
662                let mut map = [false; MAX_INDEX];
663
664                for _ in 0..60 {
665                    let offset = rand::thread_rng().gen_range(0..MAX_INDEX - max_size - 1);
666                    let size = rand::thread_rng().gen_range(1..=max_size);
667
668                    //println!("add {}..{} {}", offset, offset + size, size);
669                    // Real impl
670                    let res = assr.add(offset, size);
671
672                    // Bitmap impl
673                    let mut map2 = map;
674                    map2[offset..][..size].fill(true);
675
676                    let mut contigs = vec![];
677                    let mut hole: usize = 0;
678                    let mut data: usize = 0;
679                    for b in map2 {
680                        if b {
681                            data += 1;
682                        } else {
683                            if data != 0 {
684                                contigs.push((hole, data));
685                                hole = 0;
686                                data = 0;
687                            }
688                            hole += 1;
689                        }
690                    }
691
692                    // Compare.
693                    let wanted_res = if contigs.len() > ASSEMBLER_MAX_SEGMENT_COUNT {
694                        Err(TooManyHolesError)
695                    } else {
696                        Ok(())
697                    };
698                    assert_eq!(res, wanted_res);
699                    if res.is_ok() {
700                        map = map2;
701                        assert_eq!(assr, Assembler::from(contigs));
702                    }
703                }
704            }
705        }
706    }
707}
smoltcp/storage/assembler.rs

smoltcp/storage/
assembler.rs
