Skip to main content

asdf_core/core/
elements.rs

1//! Decoding an ndarray's bytes into typed elements, and inlining them back
2//! into the tree.
3//!
4//! Inlining is what the ASDF Standard's reference corpus asks a reader to do
5//! before comparing a `.asdf` file against its expected `.yaml`: turn every
6//! block-backed array into the nested sequences the expected output carries.
7
8use asdf_yaml::{CollectionStyle, Document, Node, NodeData, NodeId, ScalarStyle};
9
10use crate::core::datatype::{ByteOrder, Datatype, ScalarType};
11use crate::core::ndarray::{Ndarray, Source};
12use crate::error::{Result, err};
13
14/// One decoded array element.
15#[derive(Clone, PartialEq, Debug)]
16pub enum Element {
17    /// A signed integer.
18    Int(i64),
19    /// An unsigned integer.
20    Uint(u64),
21    /// A float, including the half- and single-precision types widened to
22    /// `f64`.
23    Float(f64),
24    /// A boolean, from `bool8`.
25    Bool(bool),
26    /// Fixed-length text, with trailing NULs trimmed.
27    Text(String),
28    /// A complex number.
29    Complex(f64, f64),
30    /// One record of a compound array.
31    Record(Vec<Element>),
32}
33
34/// Read a big- or little-endian integer of `n` bytes.
35fn read_uint(bytes: &[u8], order: ByteOrder) -> u64 {
36    let mut acc = 0u64;
37    if order == ByteOrder::Big {
38        for b in bytes {
39            acc = (acc << 8) | u64::from(*b);
40        }
41    } else {
42        for b in bytes.iter().rev() {
43            acc = (acc << 8) | u64::from(*b);
44        }
45    }
46    acc
47}
48
49/// Sign-extend an `n`-byte two's-complement value.
50fn sign_extend(value: u64, bytes: usize) -> i64 {
51    let bits = bytes * 8;
52    if bits >= 64 {
53        return value as i64;
54    }
55    let shift = 64 - bits;
56    ((value << shift) as i64) >> shift
57}
58
59/// The order to actually use for a field: its own if set, else the array's.
60fn effective_order(field: ByteOrder, array: ByteOrder) -> ByteOrder {
61    match field {
62        ByteOrder::Big | ByteOrder::Little => field,
63        // The schema's default when nothing says otherwise.
64        _ => match array {
65            ByteOrder::Big | ByteOrder::Little => array,
66            _ => ByteOrder::Little,
67        },
68    }
69}
70
71/// Decode a single element of `datatype` from the front of `bytes`.
72fn decode_one(datatype: &Datatype, bytes: &[u8], array_order: ByteOrder) -> Result<Element> {
73    if datatype.is_structured() {
74        let mut fields = Vec::with_capacity(datatype.fields.len());
75        let mut offset = 0usize;
76        for field in &datatype.fields {
77            let width = field.datatype.item_size() as usize;
78            let slice = bytes.get(offset..offset + width).ok_or_else(|| {
79                err!(UnexpectedEof, "compound element truncated at field offset {offset}")
80            })?;
81            fields.push(decode_one(&field.datatype, slice, array_order)?);
82            offset += width;
83        }
84        return Ok(Element::Record(fields));
85    }
86
87    let order = effective_order(datatype.byteorder, array_order);
88    let width = datatype.item_size() as usize;
89    let raw = bytes.get(..width).ok_or_else(|| {
90        err!(UnexpectedEof, "element needs {width} bytes, {} available", bytes.len())
91    })?;
92
93    Ok(match datatype.scalar {
94        ScalarType::Bool8 => Element::Bool(raw[0] != 0),
95
96        ScalarType::Uint8 | ScalarType::Uint16 | ScalarType::Uint32 | ScalarType::Uint64 => {
97            Element::Uint(read_uint(raw, order))
98        }
99
100        ScalarType::Int8 | ScalarType::Int16 | ScalarType::Int32 | ScalarType::Int64 => {
101            Element::Int(sign_extend(read_uint(raw, order), width))
102        }
103
104        ScalarType::Float16 => {
105            let bits = read_uint(raw, order) as u16;
106            Element::Float(f64::from(half::f16::from_bits(bits)))
107        }
108        ScalarType::Float32 => {
109            let bits = read_uint(raw, order) as u32;
110            Element::Float(f64::from(f32::from_bits(bits)))
111        }
112        ScalarType::Float64 => Element::Float(f64::from_bits(read_uint(raw, order))),
113
114        ScalarType::Complex64 => {
115            let re = f32::from_bits(read_uint(&raw[..4], order) as u32);
116            let im = f32::from_bits(read_uint(&raw[4..], order) as u32);
117            Element::Complex(f64::from(re), f64::from(im))
118        }
119        ScalarType::Complex128 => {
120            let re = f64::from_bits(read_uint(&raw[..8], order));
121            let im = f64::from_bits(read_uint(&raw[8..], order));
122            Element::Complex(re, im)
123        }
124
125        ScalarType::Ascii => {
126            // Fixed-length text is NUL-padded to its declared width.
127            let end = raw.iter().position(|b| *b == 0).unwrap_or(raw.len());
128            Element::Text(String::from_utf8_lossy(&raw[..end]).into_owned())
129        }
130        ScalarType::Ucs4 => {
131            let mut out = String::new();
132            for chunk in raw.chunks_exact(4) {
133                let cp = read_uint(chunk, order) as u32;
134                if cp == 0 {
135                    break;
136                }
137                out.push(char::from_u32(cp).unwrap_or(char::REPLACEMENT_CHARACTER));
138            }
139            Element::Text(out)
140        }
141
142        ScalarType::Unknown | ScalarType::Structured => {
143            return Err(err!(
144                InvalidArgument,
145                "cannot decode a {} element",
146                datatype.scalar.name()
147            ));
148        }
149    })
150}
151
152/// Decode every element of an array, in C order.
153///
154/// `shape` must be fully resolved. Strides, when present, are honoured, so a
155/// tile view or a Fortran-order array reads correctly.
156pub fn decode_all(nd: &Ndarray, shape: &[u64], bytes: &[u8]) -> Result<Vec<Element>> {
157    let item = nd.datatype.item_size();
158    if item == 0 {
159        return Err(err!(InvalidArgument, "cannot decode elements of zero width"));
160    }
161
162    let count: u64 = shape.iter().product();
163    let count = usize::try_from(count)
164        .map_err(|_| err!(OverLimit, "array has too many elements for this platform"))?;
165
166    let strides = match &nd.strides {
167        Some(s) if s.len() == shape.len() => s.clone(),
168        Some(s) => {
169            return Err(err!(
170                InvalidArgument,
171                "strides have {} entries but the shape has {}",
172                s.len(),
173                shape.len()
174            ));
175        }
176        None => Ndarray::c_strides(shape, item),
177    };
178
179    let base = usize::try_from(nd.offset)
180        .map_err(|_| err!(InvalidArgument, "ndarray offset overflows this platform"))?;
181
182    let mut out = Vec::with_capacity(count);
183    let mut index = vec![0u64; shape.len()];
184
185    for _ in 0..count {
186        // Byte position of this element, from the per-dimension strides.
187        let mut pos = base as i64;
188        for (dim, idx) in index.iter().enumerate() {
189            pos += strides[dim] * (*idx as i64);
190        }
191        let pos = usize::try_from(pos)
192            .map_err(|_| err!(InvalidArgument, "strides address a negative offset"))?;
193
194        let slice = bytes.get(pos..).ok_or_else(|| {
195            err!(UnexpectedEof, "element at byte {pos} is past the end of the block")
196        })?;
197        out.push(decode_one(&nd.datatype, slice, nd.byteorder)?);
198
199        // Odometer step, last dimension fastest.
200        for dim in (0..shape.len()).rev() {
201            index[dim] += 1;
202            if index[dim] < shape[dim] {
203                break;
204            }
205            index[dim] = 0;
206        }
207    }
208    Ok(out)
209}
210
211/// Decode an array whose data is inline in the tree rather than in a block.
212///
213/// The elements are already text in the document, so this is the mirror of
214/// [`inline_ndarray`]: walk the nested sequences under `data` in row-major
215/// order and read each scalar as the array's datatype.
216///
217/// The descent is driven by `shape`, not by how deeply the sequences happen
218/// to nest. That distinction matters for a compound array, whose innermost
219/// "sequence" is one record's fields rather than another dimension.
220///
221/// A `mask`ed value written as `null` decodes to the datatype's zero, since
222/// [`Element`] carries no missing-value marker; `mask` itself is available on
223/// the [`Ndarray`] for a caller that needs to know which those were.
224pub fn decode_inline(doc: &Document, array: &Ndarray, shape: &[u64]) -> Result<Vec<Element>> {
225    let Source::Inline(root) = array.source else {
226        return Err(err!(InvalidArgument, "this array's data is not inline"));
227    };
228
229    let expected: u64 = shape.iter().copied().product();
230    let mut out = Vec::with_capacity(usize::try_from(expected).unwrap_or(0));
231    collect_inline(doc, root, &array.datatype, shape, &mut out)?;
232
233    if out.len() as u64 != expected {
234        return Err(err!(
235            InvalidArgument,
236            "inline data holds {} elements but the shape calls for {expected}",
237            out.len()
238        ));
239    }
240    Ok(out)
241}
242
243/// Walk `shape.len()` levels of sequences, reading each leaf as `datatype`.
244fn collect_inline(
245    doc: &Document,
246    node: NodeId,
247    datatype: &Datatype,
248    shape: &[u64],
249    out: &mut Vec<Element>,
250) -> Result<()> {
251    let resolved = doc.resolve(node);
252
253    let Some((dim, rest)) = shape.split_first() else {
254        // Past the last dimension: this is one element.
255        out.push(leaf_element(doc, resolved, datatype)?);
256        return Ok(());
257    };
258
259    let items = doc.sequence_items(resolved).map(<[_]>::to_vec).ok_or_else(|| {
260        err!(InvalidArgument, "inline array data is not nested {} deep", shape.len())
261    })?;
262    if items.len() as u64 != *dim {
263        return Err(err!(
264            InvalidArgument,
265            "inline dimension holds {} entries but the shape calls for {dim}",
266            items.len()
267        ));
268    }
269    for item in items {
270        collect_inline(doc, item, datatype, rest, out)?;
271    }
272    Ok(())
273}
274
275/// Read one element: a scalar, or a record for a compound datatype.
276fn leaf_element(doc: &Document, node: NodeId, datatype: &Datatype) -> Result<Element> {
277    if !datatype.fields.is_empty() {
278        let items = doc.sequence_items(node).map(<[_]>::to_vec).ok_or_else(|| {
279            err!(InvalidArgument, "a compound element must be a sequence of its fields")
280        })?;
281        if items.len() != datatype.fields.len() {
282            return Err(err!(
283                InvalidArgument,
284                "a compound element holds {} values but the datatype has {} fields",
285                items.len(),
286                datatype.fields.len()
287            ));
288        }
289        let mut record = Vec::with_capacity(items.len());
290        for (item, field) in items.iter().zip(datatype.fields.iter()) {
291            record.push(leaf_element(doc, doc.resolve(*item), &field.datatype)?);
292        }
293        return Ok(Element::Record(record));
294    }
295
296    let text = doc
297        .resolved(node)
298        .as_str()
299        .ok_or_else(|| err!(InvalidArgument, "inline array data holds a non-scalar leaf"))?;
300    scalar_element(text, datatype.scalar)
301}
302
303/// Read one inline scalar as `scalar`.
304fn scalar_element(text: &str, scalar: ScalarType) -> Result<Element> {
305    use ScalarType as S;
306
307    // A masked element is written `null`; there is no missing marker in
308    // `Element`, so it reads as the type's zero.
309    if matches!(text, "null" | "~" | "") {
310        return Ok(match scalar {
311            S::Float16 | S::Float32 | S::Float64 => Element::Float(0.0),
312            S::Complex64 | S::Complex128 => Element::Complex(0.0, 0.0),
313            S::Bool8 => Element::Bool(false),
314            S::Ascii | S::Ucs4 => Element::Text(String::new()),
315            S::Uint8 | S::Uint16 | S::Uint32 | S::Uint64 => Element::Uint(0),
316            _ => Element::Int(0),
317        });
318    }
319
320    let bad = |what: &str| err!(InvalidArgument, "inline {what} value {text:?} does not parse");
321    match scalar {
322        S::Uint8 | S::Uint16 | S::Uint32 | S::Uint64 => {
323            Ok(Element::Uint(text.parse::<u64>().map_err(|_| bad("unsigned"))?))
324        }
325        S::Int8 | S::Int16 | S::Int32 | S::Int64 => {
326            Ok(Element::Int(text.parse::<i64>().map_err(|_| bad("integer"))?))
327        }
328        S::Float16 | S::Float32 | S::Float64 => Ok(Element::Float(parse_inline_float(text)?)),
329        S::Complex64 | S::Complex128 => {
330            let (re, im) = parse_inline_complex(text)?;
331            Ok(Element::Complex(re, im))
332        }
333        S::Bool8 => Ok(Element::Bool(matches!(text, "true" | "True" | "1"))),
334        S::Ascii | S::Ucs4 => Ok(Element::Text(text.to_string())),
335        S::Unknown | S::Structured => {
336            Err(err!(InvalidArgument, "inline data needs a known scalar datatype"))
337        }
338    }
339}
340
341/// Parse a float, accepting YAML's non-finite spellings and Python's.
342fn parse_inline_float(text: &str) -> Result<f64> {
343    match text {
344        ".nan" | ".NaN" | ".NAN" | "nan" => return Ok(f64::NAN),
345        ".inf" | ".Inf" | ".INF" | "inf" => return Ok(f64::INFINITY),
346        "-.inf" | "-.Inf" | "-.INF" | "-inf" => return Ok(f64::NEG_INFINITY),
347        _ => {}
348    }
349    text.parse::<f64>()
350        .map_err(|_| err!(InvalidArgument, "inline float value {text:?} does not parse"))
351}
352
353/// Parse a complex number in the spellings the `core/complex` schema allows.
354///
355/// Both `(1+2j)` and `1+2j` are valid, as is a pure imaginary `3j` or a pure
356/// real `1`. The imaginary unit may be `i` or `j`, either case.
357fn parse_inline_complex(text: &str) -> Result<(f64, f64)> {
358    let body = text.trim();
359    let body = body.strip_prefix('(').map_or(body, |rest| rest.strip_suffix(')').unwrap_or(rest));
360
361    let imaginary_unit = |c: char| matches!(c, 'i' | 'I' | 'j' | 'J');
362    let Some(unit) = body.char_indices().rev().find(|(_, c)| imaginary_unit(*c)) else {
363        // No imaginary part at all: a plain real number.
364        return Ok((parse_inline_float(body)?, 0.0));
365    };
366    // The unit must be last; anything after it is not a complex number.
367    if unit.0 + unit.1.len_utf8() != body.len() {
368        return Err(err!(InvalidArgument, "inline complex value {text:?} does not parse"));
369    }
370    let without_unit = &body[..unit.0];
371
372    // Split off the imaginary part at the sign that separates the two, which
373    // is the last `+`/`-` not part of an exponent and not the leading sign.
374    let split = without_unit
375        .char_indices()
376        .rev()
377        .find(|(index, c)| {
378            (*c == '+' || *c == '-')
379                && *index > 0
380                && !matches!(without_unit.as_bytes()[index - 1], b'e' | b'E')
381        })
382        .map(|(index, _)| index);
383
384    match split {
385        None => Ok((0.0, parse_inline_float(without_unit)?)),
386        Some(index) => {
387            let (real, imaginary) = without_unit.split_at(index);
388            // The imaginary part keeps its sign; a bare sign means one.
389            let imaginary = match imaginary {
390                "+" => "1",
391                "-" => "-1",
392                other => other,
393            };
394            Ok((parse_inline_float(real)?, parse_inline_float(imaginary)?))
395        }
396    }
397}
398
399/// The tag Python asdf puts on every inline complex value.
400const COMPLEX_TAG: &str = "tag:stsci.edu:asdf/core/complex-1.0.0";
401
402/// Format a float the way libasdf's emitter does.
403///
404/// `%.17g` for doubles, with YAML's own spellings for the non-finite values.
405pub fn format_float(value: f64) -> String {
406    if value.is_nan() {
407        return ".nan".to_string();
408    }
409    if value.is_infinite() {
410        return if value.is_sign_negative() { "-.inf".into() } else { ".inf".into() };
411    }
412    // Shortest representation that round-trips, which is what Rust's
413    // formatter gives and what a reader will parse back identically.
414    let mut s = format!("{value}");
415    if !s.contains('.') && !s.contains('e') && !s.contains("inf") && !s.contains("nan") {
416        s.push_str(".0");
417    }
418    s
419}
420
421/// Render one element as a tree node.
422fn element_to_node(doc: &mut Document, element: &Element) -> NodeId {
423    match element {
424        Element::Int(v) => doc.add_scalar(v.to_string()),
425        Element::Uint(v) => doc.add_scalar(v.to_string()),
426        Element::Bool(v) => doc.add_scalar(if *v { "true" } else { "false" }),
427        Element::Float(v) => doc.add_scalar(format_float(*v)),
428        // Text is quoted so it round-trips as a string rather than being
429        // re-resolved as a number.
430        Element::Text(s) => doc.add_scalar_styled(s.clone(), ScalarStyle::SingleQuoted),
431        Element::Complex(re, im) => {
432            // The core/complex schema leaves the spelling open; Python asdf
433            // writes CPython's `repr(complex)` and tags each value, so that
434            // is the canonical form to reproduce.
435            let node = Node::scalar(crate::core::pyrepr::repr_complex(*re, *im))
436                .with_tag(asdf_yaml::Tag::parse(COMPLEX_TAG));
437            doc.add(node)
438        }
439        Element::Record(fields) => {
440            let items: Vec<NodeId> = fields.iter().map(|f| element_to_node(doc, f)).collect();
441            doc.add_sequence(items)
442        }
443    }
444}
445
446/// Build nested sequences for `elements` laid out in `shape`.
447pub fn nest(doc: &mut Document, elements: &[Element], shape: &[u64]) -> NodeId {
448    fn build(
449        doc: &mut Document,
450        elements: &[Element],
451        shape: &[u64],
452        cursor: &mut usize,
453    ) -> NodeId {
454        match shape.split_first() {
455            None => {
456                let node = element_to_node(doc, &elements[*cursor]);
457                *cursor += 1;
458                node
459            }
460            Some((dim, rest)) => {
461                let mut items = Vec::with_capacity(*dim as usize);
462                for _ in 0..*dim {
463                    items.push(build(doc, elements, rest, cursor));
464                }
465                let id = doc.add_sequence(items);
466                // Inline data reads better in flow style, which is how both
467                // other implementations write it.
468                if let NodeData::Sequence { style, .. } = &mut doc.node_mut(id).data {
469                    *style = CollectionStyle::Flow;
470                }
471                id
472            }
473        }
474    }
475
476    let mut cursor = 0;
477    build(doc, elements, shape, &mut cursor)
478}
479
480/// Replace an ndarray's `source` with the inline `data` it stands for.
481///
482/// This is the transformation the reference corpus prescribes. The array's
483/// `byteorder`, `offset` and `strides` describe how bytes sit in a block and
484/// become meaningless once the data is inline, so they are removed too.
485pub fn inline_ndarray(
486    doc: &mut Document,
487    id: NodeId,
488    elements: &[Element],
489    shape: &[u64],
490) -> Result<()> {
491    let data = nest(doc, elements, shape);
492    let target = doc.resolve(id);
493
494    if !doc.node(target).is_mapping() {
495        // The bare-sequence shorthand is already inline.
496        return Ok(());
497    }
498
499    doc.mapping_remove(target, "source");
500    for key in ["byteorder", "offset", "strides"] {
501        doc.mapping_remove(target, key);
502    }
503    // A compound datatype's fields carry their own byteorder, which is just
504    // as meaningless once the data is inline.
505    if let Some(dt) = doc.mapping_get(target, "datatype")
506        && let Some(fields) = doc.sequence_items(dt).map(<[_]>::to_vec)
507    {
508        for field in fields {
509            let field = doc.resolve(field);
510            if doc.node(field).is_mapping() {
511                doc.mapping_remove(field, "byteorder");
512            }
513        }
514    }
515    doc.mapping_set(target, "data", data);
516
517    // Record the shape the data actually has, replacing any '*'.
518    let dims: Vec<NodeId> = shape.iter().map(|d| doc.add_scalar(d.to_string())).collect();
519    let shape_node = doc.add_sequence(dims);
520    if let NodeData::Sequence { style, .. } = &mut doc.node_mut(shape_node).data {
521        *style = CollectionStyle::Flow;
522    }
523    doc.mapping_set(target, "shape", shape_node);
524    Ok(())
525}
526
527/// Build a node holding an element, for tests and callers wanting one value.
528pub fn element_node(doc: &mut Document, element: &Element) -> NodeId {
529    element_to_node(doc, element)
530}
531
532/// A node with a tag applied.
533pub fn tagged(doc: &mut Document, node: Node, tag: asdf_yaml::Tag) -> NodeId {
534    doc.add(node.with_tag(tag))
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540    use asdf_yaml::parse_document;
541
542    fn ndarray(yaml: &str) -> Ndarray {
543        let doc = parse_document(yaml).unwrap();
544        let root = doc.root().unwrap();
545        Ndarray::parse(&doc, doc.mapping_get(root, "a").unwrap()).unwrap()
546    }
547
548    #[test]
549    fn inline_integers_decode_from_the_tree() {
550        let doc = parse_document(
551            "a:\n  data: [[1, 2, 3], [4, 5, 6]]\n  datatype: int32\n  shape: [2, 3]\n",
552        )
553        .unwrap();
554        let root = doc.root().unwrap();
555        let nd = Ndarray::parse(&doc, doc.mapping_get(root, "a").unwrap()).unwrap();
556        let shape = nd.resolved_shape(None).unwrap();
557        assert_eq!(shape, vec![2, 3]);
558
559        let els = decode_inline(&doc, &nd, &shape).unwrap();
560        assert_eq!(
561            els,
562            (1..=6).map(Element::Int).collect::<Vec<_>>(),
563            "row-major order, flattened"
564        );
565    }
566
567    #[test]
568    fn inline_floats_accept_yamls_non_finite_spellings() {
569        let doc = parse_document(
570            "a:\n  data: [1.5, .inf, -.inf, .nan]\n  datatype: float64\n  shape: [4]\n",
571        )
572        .unwrap();
573        let root = doc.root().unwrap();
574        let nd = Ndarray::parse(&doc, doc.mapping_get(root, "a").unwrap()).unwrap();
575        let els = decode_inline(&doc, &nd, &[4]).unwrap();
576
577        assert_eq!(els[0], Element::Float(1.5));
578        assert_eq!(els[1], Element::Float(f64::INFINITY));
579        assert_eq!(els[2], Element::Float(f64::NEG_INFINITY));
580        let Element::Float(nan) = els[3] else { panic!("{:?}", els[3]) };
581        assert!(nan.is_nan());
582    }
583
584    /// The schema allows a family of spellings; all of them must read.
585    #[test]
586    fn inline_complex_accepts_every_spelling_the_schema_allows() {
587        let cases = [
588            ("0j", (0.0, 0.0)),
589            ("(1+2j)", (1.0, 2.0)),
590            ("1+2j", (1.0, 2.0)),
591            ("(1-2j)", (1.0, -2.0)),
592            ("-1j", (0.0, -1.0)),
593            ("(-0+0j)", (-0.0, 0.0)),
594            ("3", (3.0, 0.0)),
595            ("2i", (0.0, 2.0)),
596            ("(1.5e-3+2.5e+4j)", (1.5e-3, 2.5e4)),
597            // A bare sign before the unit means one.
598            ("(1+j)", (1.0, 1.0)),
599            ("(1-j)", (1.0, -1.0)),
600        ];
601        for (text, (re, im)) in cases {
602            let got = parse_inline_complex(text).unwrap_or_else(|e| panic!("{text}: {e}"));
603            assert_eq!(got.0, re, "real part of {text}");
604            assert_eq!(got.1, im, "imaginary part of {text}");
605        }
606
607        // The non-finites, which the corpus does carry.
608        let (re, im) = parse_inline_complex("(nan-infj)").unwrap();
609        assert!(re.is_nan());
610        assert_eq!(im, f64::NEG_INFINITY);
611    }
612
613    /// Everything we write must read back, which is the property that
614    /// matters for a round trip through inline form.
615    #[test]
616    fn complex_spellings_round_trip_through_the_parser() {
617        let values = [
618            (0.0, 0.0),
619            (-0.0, 0.0),
620            (1.0, 2.0),
621            (1.0, -2.0),
622            (0.0, -1.0),
623            (1.5e-3, 2.5e4),
624            (f64::MAX, f64::MIN_POSITIVE),
625        ];
626        for (re, im) in values {
627            let text = crate::core::pyrepr::repr_complex(re, im);
628            let (back_re, back_im) = parse_inline_complex(&text).unwrap();
629            assert_eq!(back_re.to_bits(), re.to_bits(), "{text}");
630            assert_eq!(back_im.to_bits(), im.to_bits(), "{text}");
631        }
632    }
633
634    #[test]
635    fn inline_compound_records_stay_grouped() {
636        let doc = parse_document(
637            "a:\n  data: [[1, 2.5], [3, 4.5]]\n  shape: [2]\n  datatype:\n  \
638             - {name: n, datatype: int32}\n  - {name: x, datatype: float64}\n",
639        )
640        .unwrap();
641        let root = doc.root().unwrap();
642        let nd = Ndarray::parse(&doc, doc.mapping_get(root, "a").unwrap()).unwrap();
643        let els = decode_inline(&doc, &nd, &[2]).unwrap();
644        assert_eq!(
645            els,
646            vec![
647                Element::Record(vec![Element::Int(1), Element::Float(2.5)]),
648                Element::Record(vec![Element::Int(3), Element::Float(4.5)]),
649            ]
650        );
651    }
652
653    #[test]
654    fn inline_data_must_match_the_declared_shape() {
655        let doc =
656            parse_document("a:\n  data: [1, 2, 3]\n  datatype: int32\n  shape: [4]\n").unwrap();
657        let root = doc.root().unwrap();
658        let nd = Ndarray::parse(&doc, doc.mapping_get(root, "a").unwrap()).unwrap();
659        let err = decode_inline(&doc, &nd, &[4]).unwrap_err();
660        assert!(err.message().contains("shape calls for 4"), "{err}");
661    }
662
663    /// A block-backed array decoded and re-read through inline form must
664    /// come back the same, which is the property `tree_inlined` relies on.
665    #[test]
666    fn a_block_array_survives_a_trip_through_inline_form() {
667        let nd =
668            ndarray("a:\n  source: 0\n  shape: [5]\n  datatype: float64\n  byteorder: little\n");
669        let values = [1.5f64, -2.25, 0.0, f64::MAX, -0.125];
670        let bytes: Vec<u8> = values.iter().flat_map(|v| v.to_le_bytes()).collect();
671        let original = decode_all(&nd, &[5], &bytes).unwrap();
672
673        let mut doc = parse_document(
674            "a:\n  source: 0\n  shape: [5]\n  datatype: float64\n  byteorder: little\n",
675        )
676        .unwrap();
677        let root = doc.root().unwrap();
678        let node = doc.mapping_get(root, "a").unwrap();
679        inline_ndarray(&mut doc, node, &original, &[5]).unwrap();
680
681        let inlined = Ndarray::parse(&doc, node).unwrap();
682        let read_back = decode_inline(&doc, &inlined, &[5]).unwrap();
683        assert_eq!(read_back, original);
684    }
685
686    #[test]
687    fn decodes_little_endian_integers() {
688        let nd = ndarray("a:\n  source: 0\n  shape: [4]\n  datatype: int32\n  byteorder: little\n");
689        let mut bytes = Vec::new();
690        for v in [1i32, -1, 256, i32::MIN] {
691            bytes.extend_from_slice(&v.to_le_bytes());
692        }
693        let els = decode_all(&nd, &[4], &bytes).unwrap();
694        assert_eq!(
695            els,
696            vec![
697                Element::Int(1),
698                Element::Int(-1),
699                Element::Int(256),
700                Element::Int(i64::from(i32::MIN)),
701            ]
702        );
703    }
704
705    #[test]
706    fn decodes_big_endian_integers() {
707        let nd = ndarray("a:\n  source: 0\n  shape: [3]\n  datatype: int16\n  byteorder: big\n");
708        let mut bytes = Vec::new();
709        for v in [1i16, -2, 1000] {
710            bytes.extend_from_slice(&v.to_be_bytes());
711        }
712        let els = decode_all(&nd, &[3], &bytes).unwrap();
713        assert_eq!(els, vec![Element::Int(1), Element::Int(-2), Element::Int(1000)]);
714    }
715
716    #[test]
717    fn byte_order_actually_changes_the_value() {
718        let bytes = [0x01u8, 0x00];
719        let le =
720            ndarray("a:\n  source: 0\n  shape: [1]\n  datatype: uint16\n  byteorder: little\n");
721        let be = ndarray("a:\n  source: 0\n  shape: [1]\n  datatype: uint16\n  byteorder: big\n");
722        assert_eq!(decode_all(&le, &[1], &bytes).unwrap(), vec![Element::Uint(1)]);
723        assert_eq!(decode_all(&be, &[1], &bytes).unwrap(), vec![Element::Uint(256)]);
724    }
725
726    #[test]
727    fn decodes_floats_of_every_width() {
728        let nd =
729            ndarray("a:\n  source: 0\n  shape: [2]\n  datatype: float64\n  byteorder: little\n");
730        let mut bytes = Vec::new();
731        bytes.extend_from_slice(&1.5f64.to_le_bytes());
732        bytes.extend_from_slice(&(-0.25f64).to_le_bytes());
733        assert_eq!(
734            decode_all(&nd, &[2], &bytes).unwrap(),
735            vec![Element::Float(1.5), Element::Float(-0.25)]
736        );
737
738        let nd =
739            ndarray("a:\n  source: 0\n  shape: [1]\n  datatype: float32\n  byteorder: little\n");
740        assert_eq!(
741            decode_all(&nd, &[1], &2.5f32.to_le_bytes()).unwrap(),
742            vec![Element::Float(2.5)]
743        );
744
745        let nd =
746            ndarray("a:\n  source: 0\n  shape: [1]\n  datatype: float16\n  byteorder: little\n");
747        let h = half::f16::from_f32(0.5);
748        assert_eq!(
749            decode_all(&nd, &[1], &h.to_bits().to_le_bytes()).unwrap(),
750            vec![Element::Float(0.5)]
751        );
752    }
753
754    #[test]
755    fn decodes_bools_and_text() {
756        let nd = ndarray("a:\n  source: 0\n  shape: [2]\n  datatype: bool8\n  byteorder: little\n");
757        assert_eq!(
758            decode_all(&nd, &[2], &[0u8, 1]).unwrap(),
759            vec![Element::Bool(false), Element::Bool(true)]
760        );
761
762        // Fixed-length ASCII is NUL-padded and trimmed on the way out.
763        let nd = ndarray(
764            "a:\n  source: 0\n  shape: [2]\n  datatype: ['ascii', 4]\n  byteorder: little\n",
765        );
766        let bytes = b"M31\0Cas\0";
767        assert_eq!(
768            decode_all(&nd, &[2], bytes).unwrap(),
769            vec![Element::Text("M31".into()), Element::Text("Cas".into())]
770        );
771    }
772
773    #[test]
774    fn decodes_ucs4_text() {
775        let nd = ndarray(
776            "a:\n  source: 0\n  shape: [1]\n  datatype: ['ucs4', 3]\n  byteorder: little\n",
777        );
778        let mut bytes = Vec::new();
779        for cp in ['a' as u32, 0x00E9 /* é */, 0] {
780            bytes.extend_from_slice(&cp.to_le_bytes());
781        }
782        assert_eq!(decode_all(&nd, &[1], &bytes).unwrap(), vec![Element::Text("aé".into())]);
783    }
784
785    #[test]
786    fn honours_offset() {
787        let nd = ndarray(
788            "a:\n  source: 0\n  shape: [2]\n  datatype: uint8\n  byteorder: little\n  offset: 3\n",
789        );
790        let bytes = [9u8, 9, 9, 1, 2];
791        assert_eq!(
792            decode_all(&nd, &[2], &bytes).unwrap(),
793            vec![Element::Uint(1), Element::Uint(2)]
794        );
795    }
796
797    #[test]
798    fn honours_strides_for_a_fortran_order_array() {
799        // A 2x3 array stored column-major: strides are [1, 2] elements.
800        let nd = ndarray(
801            "a:\n  source: 0\n  shape: [2, 3]\n  datatype: uint8\n  byteorder: little\n  \
802             strides: [1, 2]\n",
803        );
804        // Column-major layout of [[1,2,3],[4,5,6]].
805        let bytes = [1u8, 4, 2, 5, 3, 6];
806        let els = decode_all(&nd, &[2, 3], &bytes).unwrap();
807        let values: Vec<u64> = els
808            .iter()
809            .map(|e| match e {
810                Element::Uint(v) => *v,
811                _ => unreachable!(),
812            })
813            .collect();
814        // Read back in C order it must be the logical array.
815        assert_eq!(values, vec![1, 2, 3, 4, 5, 6]);
816    }
817
818    #[test]
819    fn honours_strides_for_a_tile_view() {
820        // A 2x2 tile of a 4x4 uint8 image, starting at row 1 column 1.
821        let nd = ndarray(
822            "a:\n  source: 0\n  shape: [2, 2]\n  datatype: uint8\n  byteorder: little\n  \
823             strides: [4, 1]\n  offset: 5\n",
824        );
825        let bytes: Vec<u8> = (0..16).collect();
826        let els = decode_all(&nd, &[2, 2], &bytes).unwrap();
827        let values: Vec<u64> = els
828            .iter()
829            .map(|e| match e {
830                Element::Uint(v) => *v,
831                _ => unreachable!(),
832            })
833            .collect();
834        assert_eq!(values, vec![5, 6, 9, 10]);
835    }
836
837    #[test]
838    fn decodes_compound_records() {
839        let nd = ndarray(
840            "a:\n  source: 0\n  shape: [2]\n  byteorder: little\n  \
841             datatype:\n    - name: id\n      datatype: uint16\n    \
842             - name: value\n      datatype: float32\n",
843        );
844        let mut bytes = Vec::new();
845        for (id, value) in [(1u16, 1.5f32), (2, -2.5)] {
846            bytes.extend_from_slice(&id.to_le_bytes());
847            bytes.extend_from_slice(&value.to_le_bytes());
848        }
849        let els = decode_all(&nd, &[2], &bytes).unwrap();
850        assert_eq!(
851            els,
852            vec![
853                Element::Record(vec![Element::Uint(1), Element::Float(1.5)]),
854                Element::Record(vec![Element::Uint(2), Element::Float(-2.5)]),
855            ]
856        );
857    }
858
859    #[test]
860    fn truncated_data_is_an_error_not_a_panic() {
861        let nd = ndarray("a:\n  source: 0\n  shape: [4]\n  datatype: int64\n  byteorder: little\n");
862        assert!(decode_all(&nd, &[4], &[0u8; 8]).is_err());
863    }
864
865    #[test]
866    fn nesting_reproduces_the_shape() {
867        let mut doc = Document::new();
868        let els: Vec<Element> = (0..6).map(Element::Uint).collect();
869        let node = nest(&mut doc, &els, &[2, 3]);
870        doc.set_root(node);
871
872        assert_eq!(doc.container_len(node), Some(2));
873        let first = doc.sequence_get(node, 0).unwrap();
874        assert_eq!(doc.container_len(first), Some(3));
875        assert_eq!(doc.resolved(doc.sequence_get(first, 2).unwrap()).as_str(), Some("2"));
876    }
877
878    #[test]
879    fn float_formatting_uses_yaml_spellings() {
880        assert_eq!(format_float(f64::NAN), ".nan");
881        assert_eq!(format_float(f64::INFINITY), ".inf");
882        assert_eq!(format_float(f64::NEG_INFINITY), "-.inf");
883        // A whole float keeps a decimal point so it does not read as an int.
884        assert_eq!(format_float(1.0), "1.0");
885        assert_eq!(format_float(1.5), "1.5");
886    }
887
888    #[test]
889    fn inlining_replaces_source_with_data() {
890        let mut doc = parse_document(
891            "a:\n  source: 0\n  shape: [4]\n  datatype: uint8\n  byteorder: little\n  offset: 0\n",
892        )
893        .unwrap();
894        let root = doc.root().unwrap();
895        let nd_id = doc.mapping_get(root, "a").unwrap();
896
897        let els: Vec<Element> = (0..4).map(Element::Uint).collect();
898        inline_ndarray(&mut doc, nd_id, &els, &[4]).unwrap();
899
900        assert!(doc.mapping_get(nd_id, "source").is_none(), "source must be removed");
901        assert!(doc.mapping_get(nd_id, "byteorder").is_none(), "byteorder is meaningless inline");
902        assert!(doc.mapping_get(nd_id, "offset").is_none(), "offset is meaningless inline");
903
904        let data = doc.mapping_get(nd_id, "data").expect("data must be added");
905        assert_eq!(doc.container_len(data), Some(4));
906        // The datatype survives; it still describes the values.
907        assert!(doc.mapping_get(nd_id, "datatype").is_some());
908    }
909}