Skip to main content

idx_lib/
lib.rs

1pub use ndarray::ArrayD;
2use ndarray::IxDyn;
3use std::io::{Read, Seek};
4use std::convert::From;
5
6#[derive(Debug, Default, Clone)]
7pub enum IdxData {
8    #[default]
9    None,
10    UnsignedByte(u8),
11    SignedByte(i8),
12    Short(i16),
13    Int(i32),
14    Float(f32),
15    Double(f64),
16}
17
18impl IdxData {
19    pub fn cast_as<T: From<u8> + From<i8> + From<i16> + From<i32> + From<f32> + From<f64>>(&self) -> Option<T> {
20        match self {
21            Self::None => None,
22            Self::UnsignedByte(val) => Some(<u8  as Into<T>>::into(*val)),
23            Self::SignedByte(val)   => Some(<i8  as Into<T>>::into(*val)),
24            Self::Short(val)        => Some(<i16 as Into<T>>::into(*val)),
25            Self::Int(val)          => Some(<i32 as Into<T>>::into(*val)),
26            Self::Float(val)        => Some(<f32 as Into<T>>::into(*val)),
27            Self::Double(val)       => Some(<f64 as Into<T>>::into(*val)),
28        }
29    }
30}
31
32#[derive(Debug)]
33enum IdxType {
34    UnsignedByte,
35    SignedByte,
36    Short,
37    Int,
38    Float,
39    Double,
40}
41
42#[derive(Debug, Clone)]
43struct IdxError;
44
45impl std::fmt::Display for IdxError {
46    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
47        write!(f, "Error")
48    }
49}
50
51impl std::error::Error for IdxError {}
52
53fn recurser(
54    idx_source: &mut (impl Read + Seek),
55    data: &mut ArrayD<IdxData>,
56    dimension_sizes: &[usize],
57    data_type: &IdxType,
58    past_idxes: &mut Vec<usize>,
59) -> Result<(), Box<dyn std::error::Error>> {
60    let current_idx = past_idxes.len();
61    // If we are on the last dimension
62    if current_idx == dimension_sizes.len() {
63        // Add the data
64        match data_type {
65            IdxType::UnsignedByte => {
66                let mut buf = [0u8; 1];
67                idx_source.read_exact(&mut buf)?;
68                data[&past_idxes[..]] = IdxData::UnsignedByte(u8::from_be_bytes(buf));
69            }
70            IdxType::SignedByte => {
71                let mut buf = [0u8; 1];
72                idx_source.read_exact(&mut buf)?;
73                data[&past_idxes[..]] = IdxData::SignedByte(i8::from_be_bytes(buf));
74            }
75            IdxType::Short => {
76                let mut buf = [0u8; 2];
77                idx_source.read_exact(&mut buf)?;
78                data[&past_idxes[..]] = IdxData::Short(i16::from_be_bytes(buf));
79            }
80            IdxType::Int => {
81                let mut buf = [0u8; 4];
82                idx_source.read_exact(&mut buf)?;
83                data[&past_idxes[..]] = IdxData::Int(i32::from_be_bytes(buf));
84            }
85            IdxType::Float => {
86                let mut buf = [0u8; 4];
87                idx_source.read_exact(&mut buf)?;
88                data[&past_idxes[..]] = IdxData::Float(f32::from_be_bytes(buf));
89            }
90            IdxType::Double => {
91                let mut buf = [0u8; 8];
92                idx_source.read_exact(&mut buf)?;
93                data[&past_idxes[..]] = IdxData::Double(f64::from_be_bytes(buf));
94            }
95        }
96        Ok(())
97    } else {
98        let my_idx = past_idxes.len();
99        past_idxes.push(0);
100        // Not in the final dimension
101        for i in 0..dimension_sizes[current_idx] {
102            past_idxes[my_idx] = i;
103            recurser(idx_source, data, dimension_sizes, data_type, past_idxes)?;
104        }
105        // Remember to remove our index
106        past_idxes.pop();
107        Ok(())
108    }
109}
110
111fn process_dimensions(
112    idx_source: &mut (impl Read + Seek),
113    data: &mut ArrayD<IdxData>,
114    dimension_sizes: &[usize],
115    data_type: &IdxType,
116) -> Result<(), Box<dyn std::error::Error>> {
117    recurser(idx_source, data, dimension_sizes, data_type, &mut vec![])
118}
119
120pub fn read_idx(
121    idx_source: &mut (impl Read + Seek),
122) -> Result<ArrayD<IdxData>, Box<dyn std::error::Error>> {
123    // First 2 bytes are always 0
124    idx_source.seek_relative(2)?;
125
126    // Data in idx is stored in big endian format
127    let mut data_type_buf = [0u8; 1];
128    let mut dimension_count_buf = [0u8; 1];
129
130    idx_source.read_exact(&mut data_type_buf)?;
131    idx_source.read_exact(&mut dimension_count_buf)?;
132
133    let data_type = match u8::from_be_bytes(data_type_buf) {
134        0x08u8 => IdxType::UnsignedByte,
135        0x09u8 => IdxType::SignedByte,
136        0x0Bu8 => IdxType::Short,
137        0x0Cu8 => IdxType::Int,
138        0x0Du8 => IdxType::Float,
139        0x0Eu8 => IdxType::Double,
140        _ => return Err(IdxError.into()),
141    };
142    let dimension_count = u8::from_be_bytes(dimension_count_buf);
143
144    let mut dimension_sizes = Vec::with_capacity(dimension_count as usize);
145
146    for _ in 0..dimension_count {
147        let mut dimension_size_buf = [0u8; 4];
148        idx_source.read_exact(&mut dimension_size_buf)?;
149        dimension_sizes.push(i32::from_be_bytes(dimension_size_buf) as usize);
150    }
151
152    //println!("Data Type: {:#?}\nDimension Count: {}\nDimension Sizes: {:#?}", data_type, dimension_count, dimension_sizes);
153
154    let mut data: ArrayD<IdxData> = ArrayD::default(IxDyn(&dimension_sizes));
155
156    process_dimensions(idx_source, &mut data, &dimension_sizes, &data_type)?;
157
158    Ok(data)
159}