Skip to main content

lutra_bin/
tabular.rs

1#![cfg(feature = "std")]
2
3//! Low-level tabular data reader.
4
5use std::collections::HashMap;
6use std::num::NonZero;
7use std::rc::Rc;
8
9use crate::TupleReader;
10use crate::ir;
11
12use crate::string::ToString;
13use crate::vec;
14
15/// Utility for iterating over arbitrary data in tabular manner (as rows and columns).
16#[derive(Clone)]
17pub struct TabularReader<'d, 't> {
18    inner: TableCell<'d, 't>,
19
20    rem_items: usize,
21    array_item_size: usize,
22
23    types: Rc<HashMap<&'t ir::Path, &'t ir::Ty>>,
24}
25
26#[derive(Clone, Copy)]
27pub struct TableCell<'d, 't> {
28    data: &'d [u8],
29    ty: &'t ir::Ty,
30    ty_defs: &'t [ir::TyDef],
31}
32
33impl<'d, 't> TableCell<'d, 't> {
34    pub fn new(data: &'d [u8], ty: &'t ir::Ty, ty_defs: &'t [ir::TyDef]) -> Self {
35        Self { data, ty, ty_defs }
36    }
37
38    pub fn data(&self) -> &'d [u8] {
39        self.data
40    }
41
42    pub fn ty(&self) -> &'t ir::Ty {
43        self.ty
44    }
45
46    pub fn ty_defs(&self) -> &'t [ir::TyDef] {
47        self.ty_defs
48    }
49}
50
51impl<'d, 't> TabularReader<'d, 't> {
52    pub fn new(data: &'d [u8], ty: &'t ir::Ty, ty_defs: &'t [ir::TyDef]) -> Self {
53        let mut r = TabularReader {
54            inner: TableCell { data, ty, ty_defs },
55            rem_items: 0,
56            array_item_size: 0,
57            types: Rc::new(HashMap::from_iter(ty_defs.iter().map(|d| (&d.name, &d.ty)))),
58        };
59
60        match &r.get_ty_mat(ty).kind {
61            ir::TyKind::Primitive(_)
62            | ir::TyKind::Ident(_)
63            | ir::TyKind::Tuple(_)
64            | ir::TyKind::Enum(_) => {
65                r.rem_items = 1;
66            }
67            ir::TyKind::Array(item) => {
68                let (offset, len) =
69                    crate::ArrayReader::<&[u8]>::read_head(TableCell { data, ty, ty_defs }.data);
70                r.inner.data = &TableCell { data, ty, ty_defs }.data[offset..];
71                r.rem_items = len;
72                r.array_item_size = item.layout.as_ref().unwrap().head_size.div_ceil(8) as usize;
73            }
74            ir::TyKind::Function(_) => unreachable!(),
75        }
76        r
77    }
78
79    pub fn ty(&self) -> &'t ir::Ty {
80        self.inner.ty()
81    }
82
83    pub fn get_ty(&self, path: &'t ir::Path) -> &'t ir::Ty {
84        self.types.get(path).unwrap()
85    }
86
87    pub fn get_ty_mat(&self, ty: &'t ir::Ty) -> &'t ir::Ty {
88        let mut ty = ty;
89        while let ir::TyKind::Ident(path) = &ty.kind {
90            if ir::TyStd::try_new(path).is_some() {
91                return ty;
92            }
93            ty = self.get_ty(path);
94        }
95        ty
96    }
97
98    pub fn remaining(&self) -> usize {
99        self.rem_items
100    }
101
102    pub fn column_names(&self) -> Vec<String> {
103        self.column_names_of_ty(self.inner.ty)
104    }
105
106    fn column_names_of_ty(&self, ty: &ir::Ty) -> Vec<String> {
107        match &self.get_ty_mat(ty).kind {
108            // arrays are iterated over, columns come from inner type
109            ir::TyKind::Array(item) => self.column_names_of_ty(item),
110
111            // tuple fields become columns
112            ir::TyKind::Tuple(fields) => fields
113                .iter()
114                .enumerate()
115                .map(|(i, f)| {
116                    if let Some(name) = &f.name {
117                        name.clone()
118                    } else {
119                        i.to_string()
120                    }
121                })
122                .collect(),
123
124            // primitives and enums become a single column (we also infer name from ident)
125            ir::TyKind::Primitive(_) | ir::TyKind::Ident(_) | ir::TyKind::Enum(_) => {
126                if let ir::TyKind::Ident(path) = &ty.kind {
127                    vec![path.0.last().unwrap().clone()]
128                } else {
129                    vec!["value".into()]
130                }
131            }
132            ir::TyKind::Function(_) => unreachable!(),
133        }
134    }
135
136    pub fn advance_by_(&mut self, n: usize) -> Result<(), NonZero<usize>> {
137        let items = n.min(self.rem_items);
138        self.rem_items -= items;
139        self.inner.data = &self.inner.data[(items * self.array_item_size)..];
140        match NonZero::new(n - items) {
141            Some(k) => Err(k),
142            None => Ok(()),
143        }
144    }
145}
146
147impl<'d, 't> Iterator for TabularReader<'d, 't> {
148    type Item = vec::Vec<TableCell<'d, 't>>;
149
150    fn next(&mut self) -> Option<Self::Item> {
151        if self.rem_items == 0 {
152            return None;
153        }
154        let mut row = self.inner;
155        if let ir::TyKind::Array(item) = &row.ty.kind {
156            row.ty = item.as_ref();
157        }
158
159        // advance
160        self.rem_items -= 1;
161        if let ir::TyKind::Array(_) = &self.inner.ty.kind {
162            self.inner.data = &self.inner.data[self.array_item_size..];
163        }
164
165        // unpack row
166        let row_ty_mat = self.get_ty_mat(row.ty);
167        Some(match &row_ty_mat.kind {
168            ir::TyKind::Primitive(_)
169            | ir::TyKind::Ident(_)
170            | ir::TyKind::Array(_)
171            | ir::TyKind::Enum(_) => {
172                vec![row]
173            }
174            ir::TyKind::Tuple(fields) => {
175                let mut cells = Vec::with_capacity(fields.len());
176                let reader = TupleReader::new_for_ty(row.data, row_ty_mat);
177                for (i, f) in fields.iter().enumerate() {
178                    cells.push(TableCell {
179                        data: reader.get_field(i),
180                        ty: &f.ty,
181                        ty_defs: row.ty_defs,
182                    })
183                }
184                cells
185            }
186            ir::TyKind::Function(_) => unreachable!(),
187        })
188    }
189}