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}