1use crate::declaration::{Dtype, ValueDecl};
18
19const MAGIC: &[u8; 3] = b"BSP";
20const HEADER_LEN: usize = 3 + 1 + 4;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Kind {
25 Input,
26 Output,
27}
28
29impl Kind {
30 fn byte(self) -> u8 {
31 match self {
32 Self::Input => b'I',
33 Self::Output => b'O',
34 }
35 }
36}
37
38#[derive(Debug, Clone, PartialEq)]
40pub enum Values {
41 F32(Vec<f32>),
42 U8(Vec<u8>),
43}
44
45impl Values {
46 pub fn len(&self) -> usize {
47 match self {
48 Self::F32(values) => values.len(),
49 Self::U8(values) => values.len(),
50 }
51 }
52
53 pub fn is_empty(&self) -> bool {
54 self.len() == 0
55 }
56
57 pub fn dtype(&self) -> Dtype {
58 match self {
59 Self::F32(_) => Dtype::F32,
60 Self::U8(_) => Dtype::U8,
61 }
62 }
63
64 pub fn get(&self, index: usize) -> Option<f32> {
66 match self {
67 Self::F32(values) => values.get(index).copied(),
68 Self::U8(values) => values.get(index).map(|value| f32::from(*value)),
69 }
70 }
71
72 pub fn as_f32(&self) -> Option<&[f32]> {
73 match self {
74 Self::F32(values) => Some(values),
75 Self::U8(_) => None,
76 }
77 }
78
79 pub fn as_u8(&self) -> Option<&[u8]> {
80 match self {
81 Self::U8(values) => Some(values),
82 Self::F32(_) => None,
83 }
84 }
85}
86
87#[derive(Debug, Clone, PartialEq)]
89pub enum PayloadError {
90 TooShort,
91 BadHeader,
92 WrongKind,
93 WrongSchema {
94 expected: u32,
95 found: u32,
96 },
97 WrongLength {
98 value: String,
99 expected: usize,
100 found: usize,
101 },
102 NotFinite {
103 value: String,
104 index: usize,
105 },
106 OutOfBounds {
107 value: String,
108 index: usize,
109 found: f32,
110 min: f32,
111 max: f32,
112 },
113 TrailingBytes,
114 Mismatch(String),
116}
117
118impl std::fmt::Display for PayloadError {
119 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120 match self {
121 Self::TooShort => write!(f, "payload ends early"),
122 Self::BadHeader => write!(f, "payload header isn't BotStadia's"),
123 Self::WrongKind => write!(f, "payload is the wrong kind (input vs output)"),
124 Self::WrongSchema { expected, found } => {
125 write!(f, "payload schema {found}, expected {expected}")
126 }
127 Self::WrongLength {
128 value,
129 expected,
130 found,
131 } => write!(f, "`{value}` is {found} bytes, expected {expected}"),
132 Self::NotFinite { value, index } => {
133 write!(f, "`{value}`[{index}] isn't a finite number")
134 }
135 Self::OutOfBounds {
136 value,
137 index,
138 found,
139 min,
140 max,
141 } => {
142 write!(f, "`{value}`[{index}] = {found} is outside {min}..={max}")
143 }
144 Self::TrailingBytes => write!(f, "payload has bytes after its last value"),
145 Self::Mismatch(text) => write!(f, "{text}"),
146 }
147 }
148}
149
150impl std::error::Error for PayloadError {}
151
152pub fn encoded_len(decls: &[ValueDecl]) -> usize {
154 HEADER_LEN + decls.iter().map(|decl| 4 + decl.bytes()).sum::<usize>()
155}
156
157pub fn encode(
159 kind: Kind,
160 schema: u32,
161 decls: &[ValueDecl],
162 values: &[Values],
163) -> Result<Vec<u8>, PayloadError> {
164 if decls.len() != values.len() {
165 return Err(PayloadError::Mismatch(format!(
166 "{} values for {} declared",
167 values.len(),
168 decls.len()
169 )));
170 }
171 let mut bytes = Vec::with_capacity(encoded_len(decls));
172 bytes.extend_from_slice(MAGIC);
173 bytes.push(kind.byte());
174 bytes.extend_from_slice(&schema.to_le_bytes());
175 for (decl, value) in decls.iter().zip(values) {
176 if decl.dtype != value.dtype() || decl.elements() != value.len() {
177 return Err(PayloadError::Mismatch(format!(
178 "`{}` wants {} × {}, got {} × {}",
179 decl.name,
180 decl.elements(),
181 decl.dtype.name(),
182 value.len(),
183 value.dtype().name()
184 )));
185 }
186 let len = u32::try_from(decl.bytes())
187 .map_err(|_| PayloadError::Mismatch(format!("`{}` is too large", decl.name)))?;
188 bytes.extend_from_slice(&len.to_le_bytes());
189 match value {
190 Values::F32(elements) => elements
191 .iter()
192 .for_each(|element| bytes.extend_from_slice(&element.to_le_bytes())),
193 Values::U8(elements) => bytes.extend_from_slice(elements),
194 }
195 }
196 Ok(bytes)
197}
198
199pub fn decode(
201 kind: Kind,
202 schema: u32,
203 decls: &[ValueDecl],
204 bytes: &[u8],
205) -> Result<Vec<Values>, PayloadError> {
206 if bytes.len() < HEADER_LEN {
207 return Err(PayloadError::TooShort);
208 }
209 if &bytes[..3] != MAGIC {
210 return Err(PayloadError::BadHeader);
211 }
212 if bytes[3] != kind.byte() {
213 return Err(PayloadError::WrongKind);
214 }
215 let found = u32::from_le_bytes(bytes[4..8].try_into().expect("four bytes"));
216 if found != schema {
217 return Err(PayloadError::WrongSchema {
218 expected: schema,
219 found,
220 });
221 }
222 let mut rest = &bytes[HEADER_LEN..];
223 let mut values = Vec::with_capacity(decls.len());
224 for decl in decls {
225 if rest.len() < 4 {
226 return Err(PayloadError::TooShort);
227 }
228 let len = u32::from_le_bytes(rest[..4].try_into().expect("four bytes")) as usize;
229 rest = &rest[4..];
230 if len != decl.bytes() {
231 return Err(PayloadError::WrongLength {
232 value: decl.name.clone(),
233 expected: decl.bytes(),
234 found: len,
235 });
236 }
237 if rest.len() < len {
238 return Err(PayloadError::TooShort);
239 }
240 let (body, tail) = rest.split_at(len);
241 rest = tail;
242 values.push(match decl.dtype {
243 Dtype::F32 => Values::F32(
244 body.as_chunks::<4>()
245 .0
246 .iter()
247 .map(|chunk| f32::from_le_bytes(*chunk))
248 .collect(),
249 ),
250 Dtype::U8 => Values::U8(body.to_vec()),
251 });
252 }
253 if !rest.is_empty() {
254 return Err(PayloadError::TrailingBytes);
255 }
256 Ok(values)
257}
258
259pub fn check(
262 kind: Kind,
263 schema: u32,
264 decls: &[ValueDecl],
265 bytes: &[u8],
266) -> Result<Vec<Values>, PayloadError> {
267 let values = decode(kind, schema, decls, bytes)?;
268 for (decl, value) in decls.iter().zip(&values) {
269 for index in 0..value.len() {
270 let element = value.get(index).expect("index in range");
271 if !element.is_finite() {
272 return Err(PayloadError::NotFinite {
273 value: decl.name.clone(),
274 index,
275 });
276 }
277 let (min, max) = decl
278 .bounds
279 .range(index)
280 .expect("bounds cover every element");
281 if element < min || element > max {
282 return Err(PayloadError::OutOfBounds {
283 value: decl.name.clone(),
284 index,
285 found: element,
286 min,
287 max,
288 });
289 }
290 }
291 }
292 Ok(values)
293}
294
295pub fn check_output(
297 schema: u32,
298 decls: &[ValueDecl],
299 bytes: &[u8],
300) -> Result<Vec<Values>, PayloadError> {
301 check(Kind::Output, schema, decls, bytes)
302}
303
304pub fn sample_uniform(decls: &[ValueDecl], mut next: impl FnMut() -> u64) -> Vec<Values> {
308 let mut unit = move || (next() >> 11) as f64 / (1u64 << 53) as f64;
309 decls
310 .iter()
311 .map(|decl| {
312 let range = |index: usize| {
313 decl.bounds
314 .range(index)
315 .expect("bounds cover every element")
316 };
317 match decl.dtype {
318 Dtype::F32 => Values::F32(
319 (0..decl.elements())
320 .map(|index| {
321 let (min, max) = range(index);
322 (f64::from(min) + unit() * (f64::from(max) - f64::from(min)))
323 .clamp(f64::from(min), f64::from(max))
324 as f32
325 })
326 .collect(),
327 ),
328 Dtype::U8 => Values::U8(
329 (0..decl.elements())
330 .map(|index| {
331 let (min, max) = range(index);
332 let span = f64::from(max) - f64::from(min) + 1.0;
333 (f64::from(min) + (unit() * span).floor()).min(f64::from(max)) as u8
334 })
335 .collect(),
336 ),
337 }
338 })
339 .collect()
340}
341
342pub fn extremes(decls: &[ValueDecl], mut pick: impl FnMut() -> bool) -> Vec<Values> {
345 decls
346 .iter()
347 .map(|decl| {
348 let mut bound = |index: usize| {
349 let (min, max) = decl
350 .bounds
351 .range(index)
352 .expect("bounds cover every element");
353 if pick() { max } else { min }
354 };
355 match decl.dtype {
356 Dtype::F32 => Values::F32((0..decl.elements()).map(bound).collect()),
357 Dtype::U8 => Values::U8(
358 (0..decl.elements())
359 .map(|index| bound(index) as u8)
360 .collect(),
361 ),
362 }
363 })
364 .collect()
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370 use crate::declaration::tests::sample;
371
372 #[test]
373 fn round_trips_and_sizes_match() {
374 let declaration = sample();
375 let values = vec![Values::F32(vec![0.5, -0.5])];
376 let bytes = encode(Kind::Output, 1, &declaration.outputs, &values).unwrap();
377 assert_eq!(bytes.len(), declaration.output_payload_bytes());
378 assert_eq!(
379 check_output(1, &declaration.outputs, &bytes).unwrap(),
380 values
381 );
382 }
383
384 #[test]
385 fn rejects_without_clamping() {
386 let declaration = sample();
387 let outputs = &declaration.outputs;
388 let out_of_bounds =
389 encode(Kind::Output, 1, outputs, &[Values::F32(vec![1.5, 0.0])]).unwrap();
390 assert!(matches!(
391 check_output(1, outputs, &out_of_bounds),
392 Err(PayloadError::OutOfBounds { index: 0, .. })
393 ));
394 let nan = encode(
395 Kind::Output,
396 1,
397 outputs,
398 &[Values::F32(vec![0.0, f32::NAN])],
399 )
400 .unwrap();
401 assert!(matches!(
402 check_output(1, outputs, &nan),
403 Err(PayloadError::NotFinite { index: 1, .. })
404 ));
405 let good = encode(Kind::Output, 1, outputs, &[Values::F32(vec![0.0, 0.0])]).unwrap();
406 assert!(matches!(
407 check_output(2, outputs, &good),
408 Err(PayloadError::WrongSchema { .. })
409 ));
410 assert!(matches!(
411 check(Kind::Input, 1, outputs, &good),
412 Err(PayloadError::WrongKind)
413 ));
414 let mut long = good.clone();
415 long.push(0);
416 assert_eq!(
417 check_output(1, outputs, &long),
418 Err(PayloadError::TrailingBytes)
419 );
420 assert_eq!(
421 check_output(1, outputs, &good[..good.len() - 1]),
422 Err(PayloadError::TooShort)
423 );
424 }
425}