Skip to main content

leodos_protocols/transport/cfdp/class2/machine/
tracker.rs

1//! A helper for tracking received file segments and identifying gaps.
2
3use heapless::Vec;
4
5use crate::transport::cfdp::CfdpError;
6
7/// The maximum number of distinct missing-data gaps we can track.
8/// This limits the size of a NAK PDU.
9const MAX_GAPS: usize = 4;
10
11/// Sentinel value representing an unused slot in the received ranges array.
12const EMPTY_SLOT: (u64, u64) = (u64::MAX, u64::MAX);
13
14/// A data structure to track received segments of a file.
15///
16/// This implementation uses a sorted, merged list of received ranges to efficiently
17/// track progress and identify missing data.
18#[derive(Debug, Clone)]
19pub struct SegmentTracker {
20    /// A sorted list of disjoint ranges representing received data.
21    /// Invariant: For any two ranges (s1, e1) and (s2, e2) in the list, e1 < s2.
22    received_ranges: [(u64, u64); MAX_GAPS],
23    /// The total expected size of the file in bytes.
24    file_size: u64,
25}
26
27impl SegmentTracker {
28    /// Creates a new tracker for a file of a given size.
29    pub fn new(file_size: u64) -> Self {
30        Self {
31            received_ranges: [(u64::MAX, u64::MAX); MAX_GAPS],
32            file_size,
33        }
34    }
35
36    /// Helper to get an iterator over the valid ranges.
37    fn valid_ranges(&self) -> impl Iterator<Item = &(u64, u64)> {
38        self.received_ranges
39            .iter()
40            .take_while(|&&r| r != EMPTY_SLOT)
41    }
42
43    /// Records that a segment of the file has been received.
44    /// This method will merge overlapping or adjacent ranges to maintain a minimal
45    /// list of disjoint received segments.
46    pub fn add_segment(&mut self, offset: u64, len: u64) -> Result<(), CfdpError> {
47        if len == 0 {
48            return Ok(());
49        }
50
51        let mut new_start = offset;
52        let mut new_end = offset + len;
53
54        let mut next_ranges: Vec<(u64, u64), MAX_GAPS, u8> = Vec::new();
55
56        for &(start, end) in self.valid_ranges() {
57            if end < new_start {
58                // Case 1: Disjoint and before. Keep it.
59                next_ranges.push((start, end)).ok();
60            } else if start > new_end {
61                // Case 2: Disjoint and after. Keep it.
62                next_ranges.push((start, end)).ok();
63            } else {
64                // Case 3: Overlap. Merge by expanding the new range's bounds.
65                new_start = new_start.min(start);
66                new_end = new_end.max(end);
67            }
68        }
69
70        if next_ranges.len() == MAX_GAPS {
71            return Ok(());
72        }
73        next_ranges.push((new_start, new_end)).unwrap();
74
75        // Sort the temporary Vec to ensure the final array is sorted.
76        next_ranges.sort_unstable_by_key(|(start, _)| *start);
77
78        // Now, copy the result from the temporary Vec back into our array.
79        let mut final_ranges = [EMPTY_SLOT; MAX_GAPS];
80        for (i, range) in next_ranges.iter().enumerate() {
81            final_ranges[i] = *range;
82        }
83        self.received_ranges = final_ranges;
84        Ok(())
85    }
86
87    /// Returns a list of ranges (offset, length) that are still missing.
88    pub fn get_missing_ranges(&self) -> impl Iterator<Item = (u64, u64)> + '_ {
89        let mut last_offset = 0;
90        self.received_ranges
91            .iter()
92            .map(move |(start, end)| {
93                let missing_start = last_offset;
94                last_offset = *end;
95                if *start > missing_start {
96                    Some((missing_start, *start))
97                } else {
98                    None
99                }
100            })
101            .filter_map(|x| x)
102            .chain(if last_offset < self.file_size {
103                Some((last_offset, self.file_size))
104            } else {
105                None
106            })
107    }
108
109    /// Checks if the entire file has been received.
110    pub fn is_complete(&self) -> bool {
111        // The file is complete if there is exactly one range,
112        // and that range is [0, file_size).
113        if self.received_ranges.len() == 1 {
114            if let Some((start, end)) = self.received_ranges.first() {
115                return *start == 0 && *end >= self.file_size;
116            }
117        }
118        false
119    }
120}