Skip to main content

lutra_bin/
shape.rs

1use crate::ArrayReader;
2use crate::ir::{self, TyKind};
3
4/// Structural shape of a value.
5#[derive(Debug, Clone, Copy)]
6pub struct Shape {
7    /// Number of array items; `Some` only when the top-level type is `Array`.
8    pub items: Option<usize>,
9    /// Number of tuple fields; `Some` when the top-level type is `Tuple` or
10    /// `Array` whose item type is a `Tuple`.
11    pub fields: Option<usize>,
12}
13
14/// Extract the [`Shape`] of `data` without decoding individual values.
15///
16/// Item count is read from the array header in O(1).
17/// Field count is read from the type definition without touching `data`.
18pub fn get_shape(data: &[u8], ty: &ir::Ty, ty_defs: &[ir::TyDef]) -> Shape {
19    let ty_mat = get_ty_mat(ty, ty_defs);
20    match &ty_mat.kind {
21        TyKind::Array(item_ty) => {
22            let (_, items) = ArrayReader::<&[u8]>::read_head(data);
23            Shape {
24                items: Some(items),
25                fields: get_shape_tuple(item_ty, ty_defs),
26            }
27        }
28        TyKind::Tuple(_) => Shape {
29            items: None,
30            fields: get_shape_tuple(ty_mat, ty_defs),
31        },
32        _ => Shape {
33            items: None,
34            fields: None,
35        },
36    }
37}
38
39pub fn get_shape_tuple(ty: &ir::Ty, ty_defs: &[ir::TyDef]) -> Option<usize> {
40    let ty_mat = get_ty_mat(ty, ty_defs);
41    match &ty_mat.kind {
42        TyKind::Tuple(fields) => Some(
43            fields
44                .iter()
45                .map(|f| get_shape_tuple(&f.ty, ty_defs).unwrap_or(1))
46                .sum(),
47        ),
48        _ => None,
49    }
50}
51
52fn get_ty_mat<'a>(mut ty: &'a ir::Ty, ty_defs: &'a [ir::TyDef]) -> &'a ir::Ty {
53    while let TyKind::Ident(path) = &ty.kind {
54        if ir::TyStd::try_new(path).is_some() {
55            return ty;
56        }
57        let Some(def) = ty_defs.iter().find(|d| &d.name == path) else {
58            return ty;
59        };
60        ty = &def.ty;
61    }
62    ty
63}