Skip to main content

sigmf/utils/
data.rs

1//! Encoding and decoding of raw SigMF sample data.
2//!
3//! A SigMF dataset is a sequence of *frames*. A frame holds one sample for
4//! each of the `core:num_channels` channels, stored channel after channel;
5//! complex samples store their real (I) component before their imaginary (Q)
6//! component. The layout of each component is given by a
7//! [`Datatype`].
8//!
9//! Decoded samples are represented as a flat list of [`SampleValue`]s in the
10//! same order as on disk: frame 0 channel 0, frame 0 channel 1, …, frame 1
11//! channel 0, and so on. A list of samples therefore always holds a multiple
12//! of the channel count.
13//!
14//! * [`encode`] and [`decode`] convert whole in-memory blocks.
15//! * [`SampleWriter`] and [`SampleReader`] stream frames to a [`Write`] or
16//!   from a [`Read`].
17//! * [`bytes_per_frame`] and [`frame_count`] compute dataset sizes.
18//!
19//! ```rust
20//! use num_complex::Complex;
21//! use sigmf::datatype::Datatype;
22//! use sigmf::utils::data::{decode, encode, SampleValue};
23//! use std::str::FromStr;
24//!
25//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
26//! let datatype = Datatype::from_str("ci16_le")?;
27//! // Two channels, two frames.
28//! let samples = vec![
29//!     SampleValue::ComplexI16(Complex::new(1, -1)),
30//!     SampleValue::ComplexI16(Complex::new(2, -2)),
31//!     SampleValue::ComplexI16(Complex::new(3, -3)),
32//!     SampleValue::ComplexI16(Complex::new(4, -4)),
33//! ];
34//! let bytes = encode(&samples, &datatype, 2)?;
35//! assert_eq!(bytes.len(), 4 * 2 * 2);
36//! assert_eq!(&bytes[..4], &[1, 0, 0xff, 0xff]);
37//! assert_eq!(decode(&bytes, &datatype, 2)?, samples);
38//! # Ok(())
39//! # }
40//! ```
41
42use crate::SigmfError;
43use crate::datatype::{ByteType, ComponentType, Datatype, Endianness, SampleDomain};
44use num_complex::Complex;
45use std::io::{Read, Write};
46
47/// One sample on one channel, typed by its datatype.
48///
49/// The variant must match the domain and component type of the
50/// [`Datatype`] it is encoded with; endianness is not part of the value. For
51/// example, both `cf32_le` and `cf32_be` use [`SampleValue::ComplexF32`].
52#[derive(Debug, PartialEq)]
53pub enum SampleValue {
54    /// A `cu8` sample.
55    ComplexUnsignedBytes(Complex<u8>),
56    /// A `ci8` sample.
57    ComplexSignedBytes(Complex<i8>),
58    /// A `cf32_le` or `cf32_be` sample.
59    ComplexF32(Complex<f32>),
60    /// A `cf64_le` or `cf64_be` sample.
61    ComplexF64(Complex<f64>),
62    /// A `ci16_le` or `ci16_be` sample.
63    ComplexI16(Complex<i16>),
64    /// A `ci32_le` or `ci32_be` sample.
65    ComplexI32(Complex<i32>),
66    /// A `cu16_le` or `cu16_be` sample.
67    ComplexU16(Complex<u16>),
68    /// A `cu32_le` or `cu32_be` sample.
69    ComplexU32(Complex<u32>),
70    /// An `ru8` sample.
71    RealUnsignedBytes(u8),
72    /// An `ri8` sample.
73    RealSignedBytes(i8),
74    /// An `rf32_le` or `rf32_be` sample.
75    RealF32(f32),
76    /// An `rf64_le` or `rf64_be` sample.
77    RealF64(f64),
78    /// An `ri16_le` or `ri16_be` sample.
79    RealI16(i16),
80    /// An `ri32_le` or `ri32_be` sample.
81    RealI32(i32),
82    /// An `ru16_le` or `ru16_be` sample.
83    RealU16(u16),
84    /// An `ru32_le` or `ru32_be` sample.
85    RealU32(u32),
86}
87
88/// Decodes a block of complete frames into samples.
89///
90/// `data` must contain a whole number of frames for `datatype` and
91/// `channel_count` channels. The result holds one [`SampleValue`] per channel
92/// per frame, in on-disk order. See the [module documentation](self) for an
93/// example.
94///
95/// # Errors
96///
97/// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero, the
98/// frame size overflows, or `data` ends with an incomplete frame.
99#[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/// Encodes complete frames of samples into raw bytes.
134///
135/// `samples` is in on-disk order (see the [module documentation](self)) and
136/// its length must be a multiple of `channel_count`. Every sample must use
137/// the [`SampleValue`] variant matching `datatype`.
138///
139/// # Errors
140///
141/// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero or the
142/// samples do not form complete frames, [`SigmfError::InvalidDataType`] if a
143/// sample does not match `datatype`, or [`SigmfError::ResourceLimit`] if the
144/// encoded size overflows.
145#[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
250/// Streams frames of samples into a [`Write`] destination.
251///
252/// Each call to [`write_frames`](Self::write_frames) encodes its samples with
253/// [`encode`] and writes them immediately; wrap the destination in a
254/// [`std::io::BufWriter`] when writing many small blocks to a file.
255///
256/// # Examples
257///
258/// ```rust
259/// use sigmf::datatype::Datatype;
260/// use sigmf::utils::data::{SampleReader, SampleValue, SampleWriter};
261/// use std::io::Cursor;
262/// use std::str::FromStr;
263///
264/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
265/// let datatype = Datatype::from_str("rf32_be")?;
266/// let mut writer = SampleWriter::new(Vec::new(), datatype.clone(), 1)?;
267/// writer.write_frames(&[SampleValue::RealF32(0.5), SampleValue::RealF32(-1.0)])?;
268/// writer.write_frames(&[SampleValue::RealF32(2.0)])?;
269/// let bytes = writer.finish()?;
270/// assert_eq!(bytes.len(), 12);
271///
272/// let mut reader = SampleReader::new(Cursor::new(bytes), datatype, 1)?;
273/// assert_eq!(
274///     reader.read_frames(2)?,
275///     Some(vec![SampleValue::RealF32(0.5), SampleValue::RealF32(-1.0)])
276/// );
277/// assert_eq!(reader.read_frames(2)?, Some(vec![SampleValue::RealF32(2.0)]));
278/// assert_eq!(reader.read_frames(2)?, None);
279/// # Ok(())
280/// # }
281/// ```
282pub struct SampleWriter<W> {
283    writer: W,
284    datatype: Datatype,
285    channel_count: usize,
286}
287
288/// Streams frames of samples out of a [`Read`] source.
289///
290/// See [`SampleWriter`] for an example.
291pub 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    /// Creates a reader that decodes `channel_count`-channel frames of
301    /// `datatype` from `reader`.
302    ///
303    /// # Errors
304    ///
305    /// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero or
306    /// the frame size overflows.
307    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    /// Reads and decodes up to `maximum_frames` frames.
321    ///
322    /// Short and interrupted reads are retried until `maximum_frames` frames
323    /// are available or the source reaches end of input. Returns
324    /// `Ok(Some(samples))` with between one and `maximum_frames` frames (one
325    /// [`SampleValue`] per channel per frame), or `Ok(None)` once the source is
326    /// exhausted. A buffer for `maximum_frames` frames is allocated on each
327    /// call.
328    ///
329    /// # Errors
330    ///
331    /// Returns [`SigmfError::InvalidRange`] if `maximum_frames` is zero,
332    /// [`SigmfError::ResourceLimit`] if the block size overflows,
333    /// [`SigmfError::InvalidMetadata`] if the source ends with an incomplete
334    /// frame, or [`SigmfError::Io`] if reading fails.
335    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    /// Creates a writer that encodes `channel_count`-channel frames of
381    /// `datatype` into `writer`.
382    ///
383    /// # Errors
384    ///
385    /// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero.
386    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    /// Encodes and writes one block of complete frames.
400    ///
401    /// # Errors
402    ///
403    /// Returns any error from [`encode`] (in which case nothing is written),
404    /// or [`SigmfError::Io`] if writing fails.
405    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    /// Flushes the destination and returns it.
412    ///
413    /// # Errors
414    ///
415    /// Returns [`SigmfError::Io`] if flushing fails.
416    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
642/// Returns the size in bytes of one frame of `channel_count` channels.
643///
644/// # Errors
645///
646/// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero or the
647/// size overflows.
648///
649/// # Examples
650///
651/// ```rust
652/// use sigmf::datatype::Datatype;
653/// use sigmf::utils::data::{bytes_per_frame, frame_count};
654/// use std::str::FromStr;
655///
656/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
657/// let datatype = Datatype::from_str("cf32_le")?;
658/// assert_eq!(bytes_per_frame(&datatype, 4)?, 32);
659/// assert_eq!(frame_count(320, &datatype, 4)?, 10);
660/// assert!(frame_count(321, &datatype, 4).is_err());
661/// # Ok(())
662/// # }
663/// ```
664pub 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/// Returns the number of frames in `dataset_bytes` bytes of sample data.
677///
678/// `dataset_bytes` must exclude any capture header and trailing bytes.
679///
680/// # Errors
681///
682/// Returns [`SigmfError::InvalidMetadata`] if `channel_count` is zero, the
683/// frame size overflows, or `dataset_bytes` is not a whole number of frames.
684#[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}