Skip to main content

leash_harness/
stream_processing.rs

1use std::{
2    collections::HashMap,
3    hash::Hash,
4    time::{Duration, Instant},
5};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct LatestValue<T> {
9    value: Option<T>,
10    dropped: u64,
11}
12
13impl<T> Default for LatestValue<T> {
14    fn default() -> Self {
15        Self {
16            value: None,
17            dropped: 0,
18        }
19    }
20}
21
22impl<T> LatestValue<T> {
23    pub fn push(&mut self, value: T) {
24        if self.value.replace(value).is_some() {
25            self.dropped += 1;
26        }
27    }
28
29    pub fn take(&mut self) -> Option<T> {
30        self.value.take()
31    }
32
33    pub fn latest(&self) -> Option<&T> {
34        self.value.as_ref()
35    }
36
37    pub fn dropped(&self) -> u64 {
38        self.dropped
39    }
40}
41
42#[derive(Debug, Clone)]
43pub struct RateLimiter<K> {
44    min_interval: Duration,
45    last_emit: HashMap<K, Instant>,
46}
47
48impl<K> RateLimiter<K>
49where
50    K: Eq + Hash + Clone,
51{
52    pub fn new(min_interval: Duration) -> Self {
53        Self {
54            min_interval,
55            last_emit: HashMap::new(),
56        }
57    }
58
59    pub fn allow(&mut self, key: K, now: Instant) -> bool {
60        let Some(last) = self.last_emit.get(&key).copied() else {
61            self.last_emit.insert(key, now);
62            return true;
63        };
64        if now.duration_since(last) < self.min_interval {
65            return false;
66        }
67        self.last_emit.insert(key, now);
68        true
69    }
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
73#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
74pub enum FrameQuality {
75    Drop,
76    Low,
77    Medium,
78    High,
79}
80
81#[derive(Debug, Clone, PartialEq, Eq)]
82#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
83pub struct QualityDecision {
84    pub accept: bool,
85    pub quality: FrameQuality,
86    pub reason: Option<String>,
87}
88
89impl QualityDecision {
90    pub fn accept(quality: FrameQuality) -> Self {
91        Self {
92            accept: quality != FrameQuality::Drop,
93            quality,
94            reason: None,
95        }
96    }
97
98    pub fn reject(reason: impl Into<String>) -> Self {
99        Self {
100            accept: false,
101            quality: FrameQuality::Drop,
102            reason: Some(reason.into()),
103        }
104    }
105}
106
107pub trait QualityFilter<T> {
108    fn evaluate(&self, frame: &T) -> QualityDecision;
109}
110
111impl<T, F> QualityFilter<T> for F
112where
113    F: Fn(&T) -> QualityDecision,
114{
115    fn evaluate(&self, frame: &T) -> QualityDecision {
116        self(frame)
117    }
118}
119
120pub fn select_best_frame<T>(
121    frames: impl IntoIterator<Item = T>,
122    filter: impl QualityFilter<T>,
123) -> Option<(T, QualityDecision)> {
124    frames
125        .into_iter()
126        .filter_map(|frame| {
127            let decision = filter.evaluate(&frame);
128            decision.accept.then_some((frame, decision))
129        })
130        .max_by_key(|(_, decision)| decision.quality)
131}
132
133pub trait Timestamped {
134    fn ts_ms(&self) -> u128;
135}
136
137impl<T> Timestamped for &T
138where
139    T: Timestamped,
140{
141    fn ts_ms(&self) -> u128 {
142        (*self).ts_ms()
143    }
144}
145
146#[derive(Debug, Clone, PartialEq, Eq)]
147pub struct TimestampPair<A, B> {
148    pub left: A,
149    pub right: B,
150    pub delta_ms: u128,
151}
152
153pub fn pair_by_timestamp<A, B>(
154    left: impl IntoIterator<Item = A>,
155    right: impl IntoIterator<Item = B>,
156    tolerance_ms: u128,
157) -> Vec<TimestampPair<A, B>>
158where
159    A: Timestamped,
160    B: Timestamped,
161{
162    let mut right = right.into_iter().collect::<Vec<_>>();
163    let mut pairs = Vec::new();
164
165    for left_item in left {
166        let Some((index, delta_ms)) = right
167            .iter()
168            .enumerate()
169            .map(|(index, right_item)| (index, abs_delta(left_item.ts_ms(), right_item.ts_ms())))
170            .filter(|(_, delta_ms)| *delta_ms <= tolerance_ms)
171            .min_by_key(|(_, delta_ms)| *delta_ms)
172        else {
173            continue;
174        };
175        let right_item = right.remove(index);
176        pairs.push(TimestampPair {
177            left: left_item,
178            right: right_item,
179            delta_ms,
180        });
181    }
182
183    pairs
184}
185
186fn abs_delta(left: u128, right: u128) -> u128 {
187    left.abs_diff(right)
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    #[derive(Debug, Clone, PartialEq, Eq)]
195    struct Sample {
196        ts_ms: u128,
197        score: u8,
198        corrupt: bool,
199    }
200
201    impl Timestamped for Sample {
202        fn ts_ms(&self) -> u128 {
203            self.ts_ms
204        }
205    }
206
207    #[test]
208    fn latest_value_keeps_only_newest_frame() {
209        let mut latest = LatestValue::default();
210
211        latest.push("old");
212        latest.push("newer");
213        latest.push("newest");
214
215        assert_eq!(latest.dropped(), 2);
216        assert_eq!(latest.take(), Some("newest"));
217        assert_eq!(latest.take(), None);
218    }
219
220    #[test]
221    fn rate_limiter_blocks_until_interval_elapses() {
222        let start = Instant::now();
223        let mut limiter = RateLimiter::new(Duration::from_millis(100));
224
225        assert!(limiter.allow("camera", start));
226        assert!(!limiter.allow("camera", start + Duration::from_millis(99)));
227        assert!(limiter.allow("camera", start + Duration::from_millis(100)));
228        assert!(limiter.allow("lidar", start + Duration::from_millis(1)));
229    }
230
231    #[test]
232    fn quality_filter_selects_best_accepted_frame() {
233        let frames = vec![
234            Sample {
235                ts_ms: 1,
236                score: 80,
237                corrupt: true,
238            },
239            Sample {
240                ts_ms: 2,
241                score: 30,
242                corrupt: false,
243            },
244            Sample {
245                ts_ms: 3,
246                score: 90,
247                corrupt: false,
248            },
249        ];
250
251        let (frame, decision) = select_best_frame(frames, |frame: &Sample| {
252            if frame.corrupt {
253                return QualityDecision::reject("corrupt");
254            }
255            QualityDecision::accept(if frame.score >= 75 {
256                FrameQuality::High
257            } else {
258                FrameQuality::Low
259            })
260        })
261        .unwrap();
262
263        assert_eq!(frame.ts_ms, 3);
264        assert_eq!(decision.quality, FrameQuality::High);
265    }
266
267    #[test]
268    fn timestamp_pairing_honors_tolerance_and_consumes_matches() {
269        let left = vec![
270            Sample {
271                ts_ms: 100,
272                score: 1,
273                corrupt: false,
274            },
275            Sample {
276                ts_ms: 150,
277                score: 2,
278                corrupt: false,
279            },
280            Sample {
281                ts_ms: 300,
282                score: 3,
283                corrupt: false,
284            },
285        ];
286        let right = vec![
287            Sample {
288                ts_ms: 110,
289                score: 10,
290                corrupt: false,
291            },
292            Sample {
293                ts_ms: 140,
294                score: 20,
295                corrupt: false,
296            },
297            Sample {
298                ts_ms: 500,
299                score: 30,
300                corrupt: false,
301            },
302        ];
303
304        let pairs = pair_by_timestamp(left, right, 15);
305
306        assert_eq!(pairs.len(), 2);
307        assert_eq!(pairs[0].left.ts_ms, 100);
308        assert_eq!(pairs[0].right.ts_ms, 110);
309        assert_eq!(pairs[0].delta_ms, 10);
310        assert_eq!(pairs[1].left.ts_ms, 150);
311        assert_eq!(pairs[1].right.ts_ms, 140);
312        assert_eq!(pairs[1].delta_ms, 10);
313    }
314}