1use crate::SigmfError;
43use crate::datatype::{ByteType, ComponentType, Datatype, Endianness, SampleDomain};
44use num_complex::Complex;
45use std::io::{Read, Write};
46
47#[derive(Debug, PartialEq)]
53pub enum SampleValue {
54 ComplexUnsignedBytes(Complex<u8>),
56 ComplexSignedBytes(Complex<i8>),
58 ComplexF32(Complex<f32>),
60 ComplexF64(Complex<f64>),
62 ComplexI16(Complex<i16>),
64 ComplexI32(Complex<i32>),
66 ComplexU16(Complex<u16>),
68 ComplexU32(Complex<u32>),
70 RealUnsignedBytes(u8),
72 RealSignedBytes(i8),
74 RealF32(f32),
76 RealF64(f64),
78 RealI16(i16),
80 RealI32(i32),
82 RealU16(u16),
84 RealU32(u32),
86}
87
88#[allow(unknown_lints)]
100#[allow(clippy::manual_is_multiple_of)]
101pub fn decode(
102 data: &[u8],
103 datatype: &Datatype,
104 channel_count: usize,
105) -> Result<Vec<SampleValue>, SigmfError> {
106 if channel_count == 0 {
107 return Err(SigmfError::InvalidMetadata(
108 "core:num_channels must be greater than zero".to_string(),
109 ));
110 }
111
112 let components_per_frame = channel_count
113 .checked_mul(if datatype.is_complex() { 2 } else { 1 })
114 .ok_or_else(|| SigmfError::InvalidMetadata("frame size overflow".to_string()))?;
115 let bytes_per_frame = components_per_frame
116 .checked_mul(datatype.component_bytes())
117 .ok_or_else(|| SigmfError::InvalidMetadata("frame size overflow".to_string()))?;
118 if data.len() % bytes_per_frame != 0 {
119 return Err(SigmfError::InvalidMetadata(format!(
120 "dataset ends with an incomplete frame at byte {}; frame size is {bytes_per_frame}",
121 data.len() - data.len() % bytes_per_frame
122 )));
123 }
124
125 let components = decode_components(data, datatype)?;
126 if datatype.get_domain() == SampleDomain::Real {
127 Ok(components)
128 } else {
129 pair_complex_components(components, datatype)
130 }
131}
132
133#[allow(unknown_lints)]
146#[allow(clippy::manual_is_multiple_of)]
147pub fn encode(
148 samples: &[SampleValue],
149 datatype: &Datatype,
150 channel_count: usize,
151) -> Result<Vec<u8>, SigmfError> {
152 if channel_count == 0 {
153 return Err(SigmfError::InvalidMetadata(
154 "core:num_channels must be greater than zero".to_string(),
155 ));
156 }
157 if samples.len() % channel_count != 0 {
158 return Err(SigmfError::InvalidMetadata(format!(
159 "{} channel samples do not form complete {channel_count}-channel frames",
160 samples.len()
161 )));
162 }
163 let capacity = samples
164 .len()
165 .checked_mul(datatype.bytes_per_channel_sample())
166 .ok_or_else(|| SigmfError::ResourceLimit("encoded sample block is too large".into()))?;
167
168 macro_rules! encode_scalars {
169 ($ty:ty, $real_variant:ident, $complex_variant:ident, $to_bytes:ident) => {{
170 let mut encoded = Vec::with_capacity(capacity);
171 for sample in samples {
172 match (datatype.get_domain(), sample) {
173 (SampleDomain::Real, SampleValue::$real_variant(value)) => {
174 encoded.extend_from_slice(&value.$to_bytes());
175 }
176 (SampleDomain::Complex, SampleValue::$complex_variant(value)) => {
177 let real: $ty = value.re;
178 let imaginary: $ty = value.im;
179 encoded.extend_from_slice(&real.$to_bytes());
180 encoded.extend_from_slice(&imaginary.$to_bytes());
181 }
182 _ => {
183 return Err(SigmfError::InvalidDataType(format!(
184 "sample value does not match datatype {}",
185 datatype
186 )));
187 }
188 }
189 }
190 encoded
191 }};
192 }
193
194 let encoded = match (
195 datatype.get_byte(),
196 datatype.get_dtype(),
197 datatype.get_endian(),
198 ) {
199 (Some(ByteType::U8), None, None) => {
200 encode_scalars!(u8, RealUnsignedBytes, ComplexUnsignedBytes, to_ne_bytes)
201 }
202 (Some(ByteType::I8), None, None) => {
203 encode_scalars!(i8, RealSignedBytes, ComplexSignedBytes, to_ne_bytes)
204 }
205 (None, Some(ComponentType::F32), Some(Endianness::LE)) => {
206 encode_scalars!(f32, RealF32, ComplexF32, to_le_bytes)
207 }
208 (None, Some(ComponentType::F32), Some(Endianness::BE)) => {
209 encode_scalars!(f32, RealF32, ComplexF32, to_be_bytes)
210 }
211 (None, Some(ComponentType::F64), Some(Endianness::LE)) => {
212 encode_scalars!(f64, RealF64, ComplexF64, to_le_bytes)
213 }
214 (None, Some(ComponentType::F64), Some(Endianness::BE)) => {
215 encode_scalars!(f64, RealF64, ComplexF64, to_be_bytes)
216 }
217 (None, Some(ComponentType::I16), Some(Endianness::LE)) => {
218 encode_scalars!(i16, RealI16, ComplexI16, to_le_bytes)
219 }
220 (None, Some(ComponentType::I16), Some(Endianness::BE)) => {
221 encode_scalars!(i16, RealI16, ComplexI16, to_be_bytes)
222 }
223 (None, Some(ComponentType::I32), Some(Endianness::LE)) => {
224 encode_scalars!(i32, RealI32, ComplexI32, to_le_bytes)
225 }
226 (None, Some(ComponentType::I32), Some(Endianness::BE)) => {
227 encode_scalars!(i32, RealI32, ComplexI32, to_be_bytes)
228 }
229 (None, Some(ComponentType::U16), Some(Endianness::LE)) => {
230 encode_scalars!(u16, RealU16, ComplexU16, to_le_bytes)
231 }
232 (None, Some(ComponentType::U16), Some(Endianness::BE)) => {
233 encode_scalars!(u16, RealU16, ComplexU16, to_be_bytes)
234 }
235 (None, Some(ComponentType::U32), Some(Endianness::LE)) => {
236 encode_scalars!(u32, RealU32, ComplexU32, to_le_bytes)
237 }
238 (None, Some(ComponentType::U32), Some(Endianness::BE)) => {
239 encode_scalars!(u32, RealU32, ComplexU32, to_be_bytes)
240 }
241 _ => {
242 return Err(SigmfError::InvalidDataType(
243 "unsupported or invalid datatype representation".to_string(),
244 ));
245 }
246 };
247 Ok(encoded)
248}
249
250pub struct SampleWriter<W> {
283 writer: W,
284 datatype: Datatype,
285 channel_count: usize,
286}
287
288pub struct SampleReader<R> {
292 reader: R,
293 datatype: Datatype,
294 channel_count: usize,
295 frame_size: usize,
296 bytes_read: u64,
297}
298
299impl<R: Read> SampleReader<R> {
300 pub fn new(reader: R, datatype: Datatype, channel_count: usize) -> Result<Self, SigmfError> {
308 let frame_size = bytes_per_frame(&datatype, channel_count)?;
309 Ok(Self {
310 reader,
311 datatype,
312 channel_count,
313 frame_size,
314 bytes_read: 0,
315 })
316 }
317
318 #[allow(unknown_lints)]
319 #[allow(clippy::manual_is_multiple_of)]
320 pub fn read_frames(
336 &mut self,
337 maximum_frames: usize,
338 ) -> Result<Option<Vec<SampleValue>>, SigmfError> {
339 if maximum_frames == 0 {
340 return Err(SigmfError::InvalidRange(
341 "sample reader block size must be greater than zero".into(),
342 ));
343 }
344 let maximum_bytes = self
345 .frame_size
346 .checked_mul(maximum_frames)
347 .ok_or_else(|| SigmfError::ResourceLimit("sample block is too large".into()))?;
348 let mut bytes = vec![0; maximum_bytes];
349 let mut filled = 0usize;
350 while filled < maximum_bytes {
351 match self.reader.read(&mut bytes[filled..]) {
352 Ok(0) => break,
353 Ok(count) => {
354 filled += count;
355 self.bytes_read = self
356 .bytes_read
357 .checked_add(count as u64)
358 .ok_or_else(|| SigmfError::InvalidRange("byte offset overflow".into()))?;
359 }
360 Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
361 Err(error) => return Err(SigmfError::Io(error)),
362 }
363 }
364 if filled == 0 {
365 return Ok(None);
366 }
367 if filled % self.frame_size != 0 {
368 let incomplete_start = self.bytes_read - (filled % self.frame_size) as u64;
369 return Err(SigmfError::InvalidMetadata(format!(
370 "dataset ends with an incomplete frame at byte {incomplete_start}; frame size is {}",
371 self.frame_size
372 )));
373 }
374 bytes.truncate(filled);
375 decode(&bytes, &self.datatype, self.channel_count).map(Some)
376 }
377}
378
379impl<W: Write> SampleWriter<W> {
380 pub fn new(writer: W, datatype: Datatype, channel_count: usize) -> Result<Self, SigmfError> {
387 if channel_count == 0 {
388 return Err(SigmfError::InvalidMetadata(
389 "core:num_channels must be greater than zero".to_string(),
390 ));
391 }
392 Ok(Self {
393 writer,
394 datatype,
395 channel_count,
396 })
397 }
398
399 pub fn write_frames(&mut self, samples: &[SampleValue]) -> Result<(), SigmfError> {
406 let encoded = encode(samples, &self.datatype, self.channel_count)?;
407 self.writer.write_all(&encoded)?;
408 Ok(())
409 }
410
411 pub fn finish(mut self) -> Result<W, SigmfError> {
417 self.writer.flush()?;
418 Ok(self.writer)
419 }
420}
421
422fn decode_components(data: &[u8], datatype: &Datatype) -> Result<Vec<SampleValue>, SigmfError> {
423 macro_rules! scalar_values {
424 ($ty:ty, $width:expr, $decode:expr, $real:ident, $complex:ident) => {{
425 if data.len() % $width != 0 {
426 return Err(incomplete_component(data.len(), $width));
427 }
428 data.chunks_exact($width)
429 .map(|chunk| {
430 let value: $ty = ($decode)(chunk);
431 if datatype.get_domain() == SampleDomain::Real {
432 SampleValue::$real(value)
433 } else {
434 SampleValue::$complex(Complex::new(value, <$ty>::default()))
435 }
436 })
437 .collect()
438 }};
439 }
440
441 let values = match (
442 datatype.get_byte(),
443 datatype.get_dtype(),
444 datatype.get_endian(),
445 ) {
446 (Some(ByteType::U8), None, None) => data
447 .iter()
448 .map(|&value| {
449 if datatype.get_domain() == SampleDomain::Real {
450 SampleValue::RealUnsignedBytes(value)
451 } else {
452 SampleValue::ComplexUnsignedBytes(Complex::new(value, 0))
453 }
454 })
455 .collect(),
456 (Some(ByteType::I8), None, None) => data
457 .iter()
458 .map(|&value| {
459 let value = i8::from_ne_bytes([value]);
460 if datatype.get_domain() == SampleDomain::Real {
461 SampleValue::RealSignedBytes(value)
462 } else {
463 SampleValue::ComplexSignedBytes(Complex::new(value, 0))
464 }
465 })
466 .collect(),
467 (None, Some(ComponentType::F32), Some(Endianness::LE)) => {
468 scalar_values!(
469 f32,
470 4,
471 |b: &[u8]| f32::from_le_bytes([b[0], b[1], b[2], b[3]]),
472 RealF32,
473 ComplexF32
474 )
475 }
476 (None, Some(ComponentType::F32), Some(Endianness::BE)) => {
477 scalar_values!(
478 f32,
479 4,
480 |b: &[u8]| f32::from_be_bytes([b[0], b[1], b[2], b[3]]),
481 RealF32,
482 ComplexF32
483 )
484 }
485 (None, Some(ComponentType::F64), Some(Endianness::LE)) => {
486 scalar_values!(
487 f64,
488 8,
489 |b: &[u8]| f64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]),
490 RealF64,
491 ComplexF64
492 )
493 }
494 (None, Some(ComponentType::F64), Some(Endianness::BE)) => {
495 scalar_values!(
496 f64,
497 8,
498 |b: &[u8]| f64::from_be_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]),
499 RealF64,
500 ComplexF64
501 )
502 }
503 (None, Some(ComponentType::I16), Some(Endianness::LE)) => {
504 scalar_values!(
505 i16,
506 2,
507 |b: &[u8]| i16::from_le_bytes([b[0], b[1]]),
508 RealI16,
509 ComplexI16
510 )
511 }
512 (None, Some(ComponentType::I16), Some(Endianness::BE)) => {
513 scalar_values!(
514 i16,
515 2,
516 |b: &[u8]| i16::from_be_bytes([b[0], b[1]]),
517 RealI16,
518 ComplexI16
519 )
520 }
521 (None, Some(ComponentType::I32), Some(Endianness::LE)) => {
522 scalar_values!(
523 i32,
524 4,
525 |b: &[u8]| i32::from_le_bytes([b[0], b[1], b[2], b[3]]),
526 RealI32,
527 ComplexI32
528 )
529 }
530 (None, Some(ComponentType::I32), Some(Endianness::BE)) => {
531 scalar_values!(
532 i32,
533 4,
534 |b: &[u8]| i32::from_be_bytes([b[0], b[1], b[2], b[3]]),
535 RealI32,
536 ComplexI32
537 )
538 }
539 (None, Some(ComponentType::U16), Some(Endianness::LE)) => {
540 scalar_values!(
541 u16,
542 2,
543 |b: &[u8]| u16::from_le_bytes([b[0], b[1]]),
544 RealU16,
545 ComplexU16
546 )
547 }
548 (None, Some(ComponentType::U16), Some(Endianness::BE)) => {
549 scalar_values!(
550 u16,
551 2,
552 |b: &[u8]| u16::from_be_bytes([b[0], b[1]]),
553 RealU16,
554 ComplexU16
555 )
556 }
557 (None, Some(ComponentType::U32), Some(Endianness::LE)) => {
558 scalar_values!(
559 u32,
560 4,
561 |b: &[u8]| u32::from_le_bytes([b[0], b[1], b[2], b[3]]),
562 RealU32,
563 ComplexU32
564 )
565 }
566 (None, Some(ComponentType::U32), Some(Endianness::BE)) => {
567 scalar_values!(
568 u32,
569 4,
570 |b: &[u8]| u32::from_be_bytes([b[0], b[1], b[2], b[3]]),
571 RealU32,
572 ComplexU32
573 )
574 }
575 _ => {
576 return Err(SigmfError::InvalidDataType(
577 "unsupported or invalid datatype representation".to_string(),
578 ));
579 }
580 };
581 Ok(values)
582}
583
584#[allow(unknown_lints)]
585#[allow(clippy::manual_is_multiple_of)]
586fn pair_complex_components(
587 components: Vec<SampleValue>,
588 datatype: &Datatype,
589) -> Result<Vec<SampleValue>, SigmfError> {
590 if components.len() % 2 != 0 {
591 return Err(incomplete_component(
592 components.len() * datatype.component_bytes(),
593 datatype.component_bytes() * 2,
594 ));
595 }
596 let mut values = Vec::with_capacity(components.len() / 2);
597 let mut components = components.into_iter();
598 while let (Some(real), Some(imaginary)) = (components.next(), components.next()) {
599 let value = match (real, imaginary) {
600 (
601 SampleValue::ComplexUnsignedBytes(real),
602 SampleValue::ComplexUnsignedBytes(imaginary),
603 ) => SampleValue::ComplexUnsignedBytes(Complex::new(real.re, imaginary.re)),
604 (SampleValue::ComplexSignedBytes(real), SampleValue::ComplexSignedBytes(imaginary)) => {
605 SampleValue::ComplexSignedBytes(Complex::new(real.re, imaginary.re))
606 }
607 (SampleValue::ComplexF32(real), SampleValue::ComplexF32(imaginary)) => {
608 SampleValue::ComplexF32(Complex::new(real.re, imaginary.re))
609 }
610 (SampleValue::ComplexF64(real), SampleValue::ComplexF64(imaginary)) => {
611 SampleValue::ComplexF64(Complex::new(real.re, imaginary.re))
612 }
613 (SampleValue::ComplexI16(real), SampleValue::ComplexI16(imaginary)) => {
614 SampleValue::ComplexI16(Complex::new(real.re, imaginary.re))
615 }
616 (SampleValue::ComplexI32(real), SampleValue::ComplexI32(imaginary)) => {
617 SampleValue::ComplexI32(Complex::new(real.re, imaginary.re))
618 }
619 (SampleValue::ComplexU16(real), SampleValue::ComplexU16(imaginary)) => {
620 SampleValue::ComplexU16(Complex::new(real.re, imaginary.re))
621 }
622 (SampleValue::ComplexU32(real), SampleValue::ComplexU32(imaginary)) => {
623 SampleValue::ComplexU32(Complex::new(real.re, imaginary.re))
624 }
625 _ => {
626 return Err(SigmfError::InvalidDataType(
627 "inconsistent complex component types".to_string(),
628 ));
629 }
630 };
631 values.push(value);
632 }
633 Ok(values)
634}
635
636fn incomplete_component(actual: usize, expected: usize) -> SigmfError {
637 SigmfError::InvalidMetadata(format!(
638 "dataset ends with an incomplete sample at byte {actual}; component/frame size is {expected}"
639 ))
640}
641
642pub fn bytes_per_frame(datatype: &Datatype, channel_count: usize) -> Result<usize, SigmfError> {
665 if channel_count == 0 {
666 return Err(SigmfError::InvalidMetadata(
667 "core:num_channels must be greater than zero".to_string(),
668 ));
669 }
670 datatype
671 .bytes_per_channel_sample()
672 .checked_mul(channel_count)
673 .ok_or_else(|| SigmfError::InvalidMetadata("frame size overflow".to_string()))
674}
675
676#[allow(unknown_lints)]
685#[allow(clippy::manual_is_multiple_of)]
686pub fn frame_count(
687 dataset_bytes: u64,
688 datatype: &Datatype,
689 channel_count: usize,
690) -> Result<u64, SigmfError> {
691 let frame_size = bytes_per_frame(datatype, channel_count)? as u64;
692 if dataset_bytes % frame_size != 0 {
693 return Err(SigmfError::InvalidMetadata(format!(
694 "dataset size {dataset_bytes} is not aligned to {frame_size}-byte frames"
695 )));
696 }
697 Ok(dataset_bytes / frame_size)
698}
699
700#[cfg(test)]
701mod tests {
702 use super::*;
703 use std::io;
704 use std::str::FromStr;
705
706 #[derive(Default)]
707 struct ShortWriter(Vec<u8>);
708
709 impl Write for ShortWriter {
710 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
711 let Some(byte) = bytes.first() else {
712 return Ok(0);
713 };
714 self.0.push(*byte);
715 Ok(1)
716 }
717
718 fn flush(&mut self) -> io::Result<()> {
719 Ok(())
720 }
721 }
722
723 struct OneByteReader(io::Cursor<Vec<u8>>);
724
725 impl Read for OneByteReader {
726 fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
727 if buffer.is_empty() {
728 return Ok(0);
729 }
730 self.0.read(&mut buffer[..1])
731 }
732 }
733
734 #[test]
735 fn sample_reader_handles_one_byte_reads_and_rejects_truncated_frames() {
736 let datatype = Datatype::from_str("ru8").unwrap();
737 let mut reader =
738 SampleReader::new(OneByteReader(io::Cursor::new(vec![1, 2, 3])), datatype, 1).unwrap();
739 assert_eq!(
740 reader.read_frames(2).unwrap().unwrap(),
741 [
742 SampleValue::RealUnsignedBytes(1),
743 SampleValue::RealUnsignedBytes(2)
744 ]
745 );
746 assert_eq!(
747 reader.read_frames(2).unwrap().unwrap(),
748 [SampleValue::RealUnsignedBytes(3)]
749 );
750 assert!(reader.read_frames(2).unwrap().is_none());
751
752 let complex = Datatype::from_str("ci16_be").unwrap();
753 let mut reader =
754 SampleReader::new(OneByteReader(io::Cursor::new(vec![1, 2, 3])), complex, 1).unwrap();
755 assert!(matches!(
756 reader.read_frames(1),
757 Err(SigmfError::InvalidMetadata(message)) if message.contains("incomplete frame")
758 ));
759 }
760
761 #[test]
762 fn real_byte_samples_preserve_every_input_byte() {
763 let datatype = Datatype::from_str("ru8").unwrap();
764 assert_eq!(
765 decode(&[1, 2, 255], &datatype, 1).unwrap(),
766 vec![
767 SampleValue::RealUnsignedBytes(1),
768 SampleValue::RealUnsignedBytes(2),
769 SampleValue::RealUnsignedBytes(255)
770 ]
771 );
772 }
773
774 #[test]
775 fn complex_values_follow_declared_big_endian_and_iq_order() {
776 let datatype = Datatype::from_str("ci16_be").unwrap();
777 assert_eq!(
778 decode(&[0x01, 0x02, 0xfe, 0xdc], &datatype, 1).unwrap(),
779 vec![SampleValue::ComplexI16(Complex::new(
780 0x0102,
781 0xfedc_u16 as i16
782 ))]
783 );
784 }
785
786 #[test]
787 fn complex_and_multichannel_incomplete_frames_are_errors() {
788 let complex = Datatype::from_str("ci8").unwrap();
789 assert!(decode(&[1], &complex, 1).is_err());
790
791 let real = Datatype::from_str("ri16_le").unwrap();
792 assert!(decode(&[1, 2, 3, 4, 5, 6], &real, 2).is_err());
793 }
794
795 #[test]
796 fn frame_count_uses_channels_and_samples_not_components() {
797 let datatype = Datatype::from_str("ci16_le").unwrap();
798 assert_eq!(frame_count(16, &datatype, 2).unwrap(), 2);
799 assert!(frame_count(15, &datatype, 2).is_err());
800 }
801
802 #[test]
803 fn codecs_round_trip_every_core_datatype() {
804 let components = [
805 "i8", "u8", "f32_le", "f32_be", "f64_le", "f64_be", "i16_le", "i16_be", "i32_le",
806 "i32_be", "u16_le", "u16_be", "u32_le", "u32_be",
807 ];
808 for domain in ["r", "c"] {
809 for component in components {
810 let datatype = Datatype::from_str(&format!("{domain}{component}")).unwrap();
811 let channels = 2;
812 let frame_bytes = bytes_per_frame(&datatype, channels).unwrap();
813 let source: Vec<_> = (0..frame_bytes * 3)
814 .map(|index| (index.wrapping_mul(73).wrapping_add(19)) as u8)
815 .collect();
816 let decoded = decode(&source, &datatype, channels).unwrap();
817 assert_eq!(
818 encode(&decoded, &datatype, channels).unwrap(),
819 source,
820 "codec round trip failed for {datatype}"
821 );
822 }
823 }
824 }
825
826 #[test]
827 fn encoder_rejects_mismatched_values_and_incomplete_channel_frames() {
828 let real = Datatype::from_str("ri16_le").unwrap();
829 assert!(encode(&[SampleValue::RealI32(1)], &real, 1).is_err());
830 assert!(encode(&[SampleValue::RealI16(1)], &real, 2).is_err());
831 }
832
833 #[test]
834 fn sample_writer_handles_short_writes_and_flushes_on_finish() {
835 let datatype = Datatype::from_str("ri16_be").unwrap();
836 let mut writer = SampleWriter::new(ShortWriter::default(), datatype, 1).unwrap();
837 writer
838 .write_frames(&[SampleValue::RealI16(0x1234), SampleValue::RealI16(-2)])
839 .unwrap();
840 assert_eq!(writer.finish().unwrap().0, [0x12, 0x34, 0xff, 0xfe]);
841 }
842}