Skip to main content

conspire/io/vtk/read/
mod.rs

1use super::{invalid, unsupported};
2use crate::io::deflate::zlib_decode;
3use std::io::Result;
4
5#[derive(Clone, Copy)]
6pub struct DataArray<'a> {
7    pub name: Option<&'a str>,
8    pub data_type: &'a str,
9    pub format: &'a str,
10    pub tuples: usize,
11    pub text: &'a str,
12    pub information: &'a str,
13}
14
15#[derive(Clone, Copy)]
16pub struct Encoding {
17    pub header_bytes: usize,
18    pub compressed: bool,
19}
20
21pub fn encoding(header: &str) -> Result<Encoding> {
22    let compressor = attribute(header, "compressor");
23    if matches!(compressor, Some(other) if other != "vtkZLibDataCompressor") {
24        return Err(unsupported(
25            "only the vtkZLibDataCompressor compressor is supported",
26        ));
27    }
28    Ok(Encoding {
29        header_bytes: match attribute(header, "header_type") {
30            Some("UInt32") | None => 4,
31            Some("UInt64") => 8,
32            Some(other) => return Err(invalid(format!("unsupported header_type {other}"))),
33        },
34        compressed: compressor.is_some(),
35    })
36}
37
38pub fn tag<'a>(text: &'a str, open: &str) -> Result<&'a str> {
39    let start = text
40        .find(open)
41        .ok_or_else(|| invalid(format!("missing {open}")))?;
42    let end = text[start..]
43        .find('>')
44        .ok_or_else(|| invalid(format!("unterminated {open}")))?;
45    Ok(&text[start..start + end])
46}
47
48pub fn region<'a>(text: &'a str, name: &str) -> Result<&'a str> {
49    let open = format!("<{name}>");
50    let close = format!("</{name}>");
51    let start = text
52        .find(&open)
53        .ok_or_else(|| invalid(format!("missing <{name}>")))?;
54    let end = text
55        .find(&close)
56        .ok_or_else(|| invalid(format!("missing </{name}>")))?;
57    Ok(&text[start..end])
58}
59
60pub fn attribute<'a>(tag: &'a str, name: &str) -> Option<&'a str> {
61    let key = format!("{name}=\"");
62    let start = tag.find(&key)? + key.len();
63    let end = tag[start..].find('"')? + start;
64    Some(&tag[start..end])
65}
66
67pub fn data_arrays<'a>(region: &'a str) -> Result<Vec<DataArray<'a>>> {
68    let mut rest = region;
69    let mut arrays = Vec::new();
70    while let Some(open) = rest.find("<DataArray") {
71        let attributes_end = rest[open..]
72            .find('>')
73            .ok_or_else(|| invalid("unterminated DataArray".into()))?
74            + open;
75        let attributes = &rest[open..attributes_end];
76        let close = rest[attributes_end..]
77            .find("</DataArray>")
78            .ok_or_else(|| invalid("unclosed DataArray".into()))?
79            + attributes_end;
80        let inner = &rest[attributes_end + 1..close];
81        arrays.push(DataArray {
82            name: attribute(attributes, "Name"),
83            data_type: attribute(attributes, "type")
84                .ok_or_else(|| invalid("DataArray without type".into()))?,
85            format: attribute(attributes, "format").unwrap_or("ascii"),
86            tuples: attribute(attributes, "NumberOfTuples")
87                .and_then(|t| t.parse().ok())
88                .unwrap_or(0),
89            text: without_information(inner),
90            information: inner,
91        });
92        rest = &rest[close..];
93    }
94    Ok(arrays)
95}
96
97fn without_information(inner: &str) -> &str {
98    match inner.rfind("</InformationKey>") {
99        Some(end) => {
100            let after = inner[end + "</InformationKey>".len()..].trim();
101            if after.is_empty() {
102                inner[..inner.find("<InformationKey").unwrap_or(0)].trim()
103            } else {
104                after
105            }
106        }
107        None => inner.trim(),
108    }
109}
110
111pub fn information<'a>(array: &DataArray<'a>, name: &str) -> Option<&'a str> {
112    let mut rest = array.information;
113    while let Some(open) = rest.find("<InformationKey") {
114        let attributes_end = rest[open..].find('>')? + open;
115        let close = rest[attributes_end..].find("</InformationKey>")? + attributes_end;
116        if attribute(&rest[open..attributes_end], "name") == Some(name) {
117            return Some(rest[attributes_end + 1..close].trim());
118        }
119        rest = &rest[close..];
120    }
121    None
122}
123
124pub fn data_array<'a>(region: &'a str, name: Option<&str>) -> Result<DataArray<'a>> {
125    find_data_array(&data_arrays(region)?, name)
126}
127
128pub fn find_data_array<'a>(arrays: &[DataArray<'a>], name: Option<&str>) -> Result<DataArray<'a>> {
129    arrays
130        .iter()
131        .find(|array| name.is_none() || array.name == name)
132        .copied()
133        .ok_or_else(|| invalid(format!("missing DataArray {}", name.unwrap_or(""))))
134}
135
136pub fn floats(array: &DataArray, encoding: &Encoding) -> Result<Vec<f64>> {
137    if array.format == "ascii" {
138        return array.text.split_whitespace().map(parse).collect();
139    }
140    let bytes = decode(array, encoding)?;
141    Ok(match array.data_type {
142        "Float64" => bytes.chunks(8).map(le_f64).collect(),
143        "Float32" => bytes.chunks(4).map(le_f32).collect(),
144        other => return Err(invalid(format!("unsupported point type {other}"))),
145    })
146}
147
148pub fn integers(array: &DataArray, encoding: &Encoding) -> Result<Vec<i64>> {
149    if array.format == "ascii" {
150        return array.text.split_whitespace().map(parse).collect();
151    }
152    let bytes = decode(array, encoding)?;
153    Ok(match array.data_type {
154        "Int64" | "UInt64" => bytes.chunks(8).map(le_i64).collect(),
155        "Int32" | "UInt32" => bytes.chunks(4).map(le_i32).collect(),
156        "Int8" | "UInt8" => bytes.iter().map(|&b| b as i64).collect(),
157        other => return Err(invalid(format!("unsupported cell-data type {other}"))),
158    })
159}
160
161pub fn bits(array: &DataArray, encoding: &Encoding) -> Result<Vec<u8>> {
162    if array.format == "ascii" {
163        return array.text.split_whitespace().map(parse).collect();
164    }
165    let bytes = decode(array, encoding)?;
166    Ok((0..array.tuples)
167        .map(|i| bytes[i / 8] >> (7 - i % 8) & 1)
168        .collect())
169}
170
171pub fn decode(array: &DataArray, encoding: &Encoding) -> Result<Vec<u8>> {
172    if array.format != "binary" {
173        return Err(unsupported(
174            "only ascii and inline binary DataArrays are supported",
175        ));
176    }
177    if encoding.compressed {
178        return decode_compressed_blocks(array.text, encoding.header_bytes);
179    }
180    let bytes = unbase64(array.text);
181    if bytes.len() < encoding.header_bytes {
182        return Err(invalid("binary DataArray shorter than its header".into()));
183    }
184    Ok(bytes[encoding.header_bytes..].to_vec())
185}
186
187fn read_uint(bytes: &[u8], offset: usize, size: usize) -> Result<usize> {
188    let slice = bytes
189        .get(offset..offset + size)
190        .ok_or_else(|| invalid("compressed DataArray header is truncated".into()))?;
191    Ok(match size {
192        4 => u32::from_le_bytes(slice.try_into().unwrap()) as usize,
193        8 => u64::from_le_bytes(slice.try_into().unwrap()) as usize,
194        _ => unreachable!("header integer size is always 4 or 8"),
195    })
196}
197
198fn decode_compressed_blocks(text: &str, header_bytes: usize) -> Result<Vec<u8>> {
199    let encoded: String = text
200        .chars()
201        .filter(|&character| {
202            character.is_ascii_alphanumeric()
203                || character == '+'
204                || character == '/'
205                || character == '='
206        })
207        .collect();
208    let encoding_of = |bytes: usize| bytes.div_ceil(3) * 4;
209    let counted = encoding_of(3 * header_bytes);
210    let counts = encoded
211        .get(..counted)
212        .ok_or_else(|| invalid("compressed DataArray header is truncated".into()))?;
213    let num_blocks = read_uint(&unbase64(counts), 0, header_bytes)?;
214    let split = encoding_of(3 * header_bytes + num_blocks * header_bytes);
215    let header = unbase64(
216        encoded
217            .get(..split)
218            .ok_or_else(|| invalid("compressed DataArray header is truncated".into()))?,
219    );
220    let bytes = unbase64(&encoded[split..]);
221    let sizes_start = 3 * header_bytes;
222    let mut compressed_sizes = Vec::with_capacity(num_blocks);
223    for block in 0..num_blocks {
224        compressed_sizes.push(read_uint(
225            &header,
226            sizes_start + block * header_bytes,
227            header_bytes,
228        )?);
229    }
230    let mut offset = 0;
231    let mut out = Vec::new();
232    for size in compressed_sizes {
233        let block = bytes
234            .get(offset..offset + size)
235            .ok_or_else(|| invalid("compressed DataArray block is truncated".into()))?;
236        out.extend(zlib_decode(block)?);
237        offset += size;
238    }
239    Ok(out)
240}
241
242fn le_f64(b: &[u8]) -> f64 {
243    f64::from_le_bytes(b.try_into().unwrap())
244}
245fn le_f32(b: &[u8]) -> f64 {
246    f32::from_le_bytes(b.try_into().unwrap()) as f64
247}
248fn le_i64(b: &[u8]) -> i64 {
249    i64::from_le_bytes(b.try_into().unwrap())
250}
251fn le_i32(b: &[u8]) -> i64 {
252    i32::from_le_bytes(b.try_into().unwrap()) as i64
253}
254
255pub fn parse<T: std::str::FromStr>(token: &str) -> Result<T> {
256    token
257        .parse()
258        .map_err(|_| invalid(format!("could not parse '{token}'")))
259}
260
261const INVALID: u8 = 0xFF;
262
263const DECODE_TABLE: [u8; 256] = {
264    let mut table = [INVALID; 256];
265    let mut i = 0;
266    while i < 26 {
267        table[b'A' as usize + i] = i as u8;
268        table[b'a' as usize + i] = 26 + i as u8;
269        i += 1;
270    }
271    i = 0;
272    while i < 10 {
273        table[b'0' as usize + i] = 52 + i as u8;
274        i += 1;
275    }
276    table[b'+' as usize] = 62;
277    table[b'/' as usize] = 63;
278    table
279};
280
281pub fn unbase64(text: &str) -> Vec<u8> {
282    let mut out = Vec::with_capacity(text.len() / 4 * 3);
283    let mut buffer = 0u32;
284    let mut bits = 0u32;
285    for &byte in text.as_bytes() {
286        let value = DECODE_TABLE[byte as usize];
287        if value == INVALID {
288            continue;
289        }
290        buffer = (buffer << 6) | value as u32;
291        bits += 6;
292        if bits >= 8 {
293            bits -= 8;
294            out.push((buffer >> bits) as u8);
295        }
296    }
297    out
298}