nuvec_io/
idx.rs

1//! The IDX format (MNIST and Fashion-MNIST): a big-endian header (two zero bytes, a type
2//! code, the number of dimensions, then one `u32` per dimension) followed by row-major
3//! big-endian data.
4
5use std::path::Path;
6
7use nuvec_core::{Array, Error, Result};
8
9fn invalid(msg: impl Into<String>) -> Error {
10    Error::InvalidArgument(format!("idx: {}", msg.into()))
11}
12
13/// Parses an IDX file's bytes. Element types map to `uint8`, `int8`, `int16`, `int32`,
14/// `float32` and `float64`.
15pub fn parse_idx(bytes: &[u8]) -> Result<Array> {
16    let header = bytes
17        .get(..4)
18        .ok_or_else(|| invalid("file shorter than its header"))?;
19    if header[..2] != [0, 0] {
20        return Err(invalid("bad magic number"));
21    }
22    let ndim = usize::from(header[3]);
23    let dims_end = 4 + 4 * ndim;
24    let shape: Vec<usize> = bytes
25        .get(4..dims_end)
26        .ok_or_else(|| invalid("truncated dimensions"))?
27        .as_chunks::<4>()
28        .0
29        .iter()
30        .map(|&d| u32::from_be_bytes(d) as usize)
31        .collect();
32    let count = shape
33        .iter()
34        .try_fold(1usize, |n, &d| n.checked_mul(d))
35        .ok_or_else(|| invalid("dimensions overflow"))?;
36    let data = &bytes[dims_end..];
37    macro_rules! decode {
38        ($t:ty) => {{
39            let nbytes = count
40                .checked_mul(std::mem::size_of::<$t>())
41                .ok_or_else(|| invalid("dimensions overflow"))?;
42            if data.len() != nbytes {
43                return Err(invalid(format!(
44                    "expected {nbytes} data bytes for shape {shape:?}, got {}",
45                    data.len()
46                )));
47            }
48            let size = std::mem::size_of::<$t>();
49            let values: Vec<$t> = data
50                .chunks_exact(size)
51                .map(|c| <$t>::from_be_bytes(c.try_into().expect("element bytes")))
52                .collect();
53            Array::from_slice(&values, &shape[..]).check()
54        }};
55    }
56    match header[2] {
57        0x08 => decode!(u8),
58        0x09 => decode!(i8),
59        0x0B => decode!(i16),
60        0x0C => decode!(i32),
61        0x0D => decode!(f32),
62        0x0E => decode!(f64),
63        code => Err(invalid(format!("unknown type code {code:#04x}"))),
64    }
65}
66
67/// Reads an (uncompressed) IDX file.
68pub fn read_idx(path: impl AsRef<Path>) -> Result<Array> {
69    let path = path.as_ref();
70    let bytes = std::fs::read(path).map_err(|e| Error::Io(format!("{}: {e}", path.display())))?;
71    parse_idx(&bytes)
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    fn file(code: u8, shape: &[u32], data: &[u8]) -> Vec<u8> {
79        let mut b = vec![0, 0, code, shape.len() as u8];
80        for d in shape {
81            b.extend(d.to_be_bytes());
82        }
83        b.extend(data);
84        b
85    }
86
87    #[test]
88    fn images_and_labels() {
89        let images = parse_idx(&file(
90            0x08,
91            &[2, 2, 3],
92            &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 255],
93        ))
94        .unwrap();
95        assert_eq!(images.shape(), [2, 2, 3]);
96        assert_eq!(images.to_vec::<u8>().unwrap()[11], 255);
97        let labels = parse_idx(&file(0x08, &[3], &[7, 2, 1])).unwrap();
98        assert_eq!(labels.to_vec::<u8>().unwrap(), [7, 2, 1]);
99    }
100
101    #[test]
102    fn big_endian_wider_types() {
103        let mut data = Vec::new();
104        data.extend((-2i32).to_be_bytes());
105        data.extend(300i32.to_be_bytes());
106        assert_eq!(
107            parse_idx(&file(0x0C, &[2], &data))
108                .unwrap()
109                .to_vec::<i32>()
110                .unwrap(),
111            [-2, 300]
112        );
113        let data: Vec<u8> = [1.5f32, -0.25]
114            .iter()
115            .flat_map(|v| v.to_be_bytes())
116            .collect();
117        assert_eq!(
118            parse_idx(&file(0x0D, &[2], &data))
119                .unwrap()
120                .to_vec::<f32>()
121                .unwrap(),
122            [1.5, -0.25]
123        );
124    }
125
126    #[test]
127    fn malformed_files_are_errors() {
128        assert!(parse_idx(&[0, 0]).is_err());
129        assert!(parse_idx(&[1, 0, 8, 1, 0, 0, 0, 1, 5]).is_err());
130        assert!(parse_idx(&file(0x08, &[3], &[1, 2])).is_err());
131        assert!(parse_idx(&file(0x42, &[1], &[1])).is_err());
132        assert!(matches!(
133            read_idx("/nonexistent/train-images-idx3-ubyte"),
134            Err(Error::Io(_))
135        ));
136    }
137}
nuvec_io/idx.rs

nuvec_io/
idx.rs
