1extern crate alloc;
10
11use alloc::vec::Vec;
12use serde::{Deserialize, Serialize};
13
14use crate::types::{Class, TimingSample};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub enum SampleClass {
19 Fixed,
21 Random,
23}
24
25#[derive(Debug, Clone, Default, Serialize, Deserialize)]
37pub struct AcquisitionStream {
38 samples: Vec<(SampleClass, f64)>,
40}
41
42impl AcquisitionStream {
43 pub fn new() -> Self {
45 Self {
46 samples: Vec::new(),
47 }
48 }
49
50 pub fn with_capacity(capacity: usize) -> Self {
52 Self {
53 samples: Vec::with_capacity(capacity),
54 }
55 }
56
57 #[inline]
59 pub fn push(&mut self, class: SampleClass, timing: f64) {
60 self.samples.push((class, timing));
61 }
62
63 pub fn push_batch_interleaved(&mut self, fixed: &[f64], random: &[f64]) {
68 debug_assert_eq!(
69 fixed.len(),
70 random.len(),
71 "Fixed and random batches must have same length"
72 );
73
74 self.samples.reserve(fixed.len() + random.len());
75 for (f, r) in fixed.iter().zip(random.iter()) {
76 self.samples.push((SampleClass::Fixed, *f));
77 self.samples.push((SampleClass::Random, *r));
78 }
79 }
80
81 #[inline]
83 pub fn len(&self) -> usize {
84 self.samples.len()
85 }
86
87 #[inline]
89 pub fn is_empty(&self) -> bool {
90 self.samples.is_empty()
91 }
92
93 #[inline]
95 pub fn n_per_class(&self) -> usize {
96 self.samples.len() / 2
97 }
98
99 pub fn split_by_class(&self) -> (Vec<f64>, Vec<f64>) {
103 let mut fixed = Vec::with_capacity(self.samples.len() / 2);
104 let mut random = Vec::with_capacity(self.samples.len() / 2);
105
106 for &(class, timing) in &self.samples {
107 match class {
108 SampleClass::Fixed => fixed.push(timing),
109 SampleClass::Random => random.push(timing),
110 }
111 }
112
113 (fixed, random)
114 }
115
116 pub fn timings(&self) -> impl Iterator<Item = f64> + '_ {
120 self.samples.iter().map(|&(_, t)| t)
121 }
122
123 pub fn as_slice(&self) -> &[(SampleClass, f64)] {
125 &self.samples
126 }
127
128 pub fn as_mut_slice(&mut self) -> &mut [(SampleClass, f64)] {
130 &mut self.samples
131 }
132
133 pub fn iter(&self) -> impl Iterator<Item = &(SampleClass, f64)> {
135 self.samples.iter()
136 }
137
138 pub fn clear(&mut self) {
140 self.samples.clear();
141 }
142
143 pub fn truncate(&mut self, len: usize) {
145 self.samples.truncate(len);
146 }
147
148 #[inline]
150 pub fn get(&self, index: usize) -> Option<&(SampleClass, f64)> {
151 self.samples.get(index)
152 }
153
154 pub fn from_raw_interleaved(baseline: &[u64], sample: &[u64], ns_per_tick: f64) -> Self {
158 debug_assert_eq!(
159 baseline.len(),
160 sample.len(),
161 "Baseline and sample must have same length"
162 );
163
164 let mut stream = Self::with_capacity(baseline.len() + sample.len());
165 for (b, s) in baseline.iter().zip(sample.iter()) {
166 stream.push(SampleClass::Fixed, *b as f64 * ns_per_tick);
167 stream.push(SampleClass::Random, *s as f64 * ns_per_tick);
168 }
169 stream
170 }
171
172 pub fn to_timing_samples(&self) -> Vec<TimingSample> {
179 self.samples
180 .iter()
181 .map(|&(class, time_ns)| TimingSample {
182 time_ns,
183 class: match class {
184 SampleClass::Fixed => Class::Baseline,
185 SampleClass::Random => Class::Sample,
186 },
187 })
188 .collect()
189 }
190}
191
192#[cfg(test)]
193mod tests {
194 use super::*;
195
196 #[test]
197 fn test_empty_stream() {
198 let stream = AcquisitionStream::new();
199 assert!(stream.is_empty());
200 assert_eq!(stream.len(), 0);
201 assert_eq!(stream.n_per_class(), 0);
202 }
203
204 #[test]
205 fn test_push_and_split() {
206 let mut stream = AcquisitionStream::new();
207
208 stream.push(SampleClass::Fixed, 100.0);
210 stream.push(SampleClass::Random, 105.0);
211 stream.push(SampleClass::Fixed, 101.0);
212 stream.push(SampleClass::Random, 106.0);
213
214 assert_eq!(stream.len(), 4);
215 assert_eq!(stream.n_per_class(), 2);
216
217 let (fixed, random) = stream.split_by_class();
218 assert_eq!(fixed, vec![100.0, 101.0]);
219 assert_eq!(random, vec![105.0, 106.0]);
220 }
221
222 #[test]
223 fn test_push_batch_interleaved() {
224 let mut stream = AcquisitionStream::new();
225
226 let fixed = vec![100.0, 101.0, 102.0];
227 let random = vec![200.0, 201.0, 202.0];
228 stream.push_batch_interleaved(&fixed, &random);
229
230 assert_eq!(stream.len(), 6);
231
232 assert_eq!(stream.samples[0], (SampleClass::Fixed, 100.0));
234 assert_eq!(stream.samples[1], (SampleClass::Random, 200.0));
235 assert_eq!(stream.samples[2], (SampleClass::Fixed, 101.0));
236 assert_eq!(stream.samples[3], (SampleClass::Random, 201.0));
237 }
238
239 #[test]
240 fn test_timings_iterator() {
241 let mut stream = AcquisitionStream::new();
242 stream.push(SampleClass::Fixed, 1.0);
243 stream.push(SampleClass::Random, 2.0);
244 stream.push(SampleClass::Fixed, 3.0);
245
246 let timings: Vec<f64> = stream.timings().collect();
247 assert_eq!(timings, vec![1.0, 2.0, 3.0]);
248 }
249
250 #[test]
251 fn test_from_raw_interleaved() {
252 let baseline = vec![100u64, 110, 120];
253 let sample = vec![200u64, 210, 220];
254 let ns_per_tick = 2.0;
255
256 let stream = AcquisitionStream::from_raw_interleaved(&baseline, &sample, ns_per_tick);
257
258 assert_eq!(stream.len(), 6);
259 assert_eq!(stream.n_per_class(), 3);
260
261 let (fixed, random) = stream.split_by_class();
262 assert_eq!(fixed, vec![200.0, 220.0, 240.0]); assert_eq!(random, vec![400.0, 420.0, 440.0]); }
265
266 #[test]
267 fn test_clear_and_truncate() {
268 let mut stream = AcquisitionStream::new();
269 stream.push(SampleClass::Fixed, 1.0);
270 stream.push(SampleClass::Random, 2.0);
271 stream.push(SampleClass::Fixed, 3.0);
272 stream.push(SampleClass::Random, 4.0);
273
274 stream.truncate(2);
275 assert_eq!(stream.len(), 2);
276
277 stream.clear();
278 assert!(stream.is_empty());
279 }
280}