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(_) | ir::TyKind::Tuple(_) | ir::TyKind::Enum(_) => {
62                r.rem_items = 1;
63            }
64            ir::TyKind::Array(item) => {
65                let (offset, len) =
66                    crate::ArrayReader::<&[u8]>::read_head(TableCell { data, ty, ty_defs }.data);
67                r.inner.data = &TableCell { data, ty, ty_defs }.data[offset..];
68                r.rem_items = len;
69                r.array_item_size = item.layout.as_ref().unwrap().head_size.div_ceil(8) as usize;
70            }
71            ir::TyKind::Function(_) | ir::TyKind::Ident(_) => unreachable!(),
72        }
73        r
74    }
75
76    pub fn ty(&self) -> &'t ir::Ty {
77        self.inner.ty()
78    }
79
80    pub fn get_ty_mat(&self, ty: &'t ir::Ty) -> &'t ir::Ty {
81        let mut ty = ty;
82        while let ir::TyKind::Ident(path) = &ty.kind {
83            ty = self.types.get(path).unwrap();
84        }
85        ty
86    }
87
88    pub fn remaining(&self) -> usize {
89        self.rem_items
90    }
91
92    pub fn column_names(&self) -> Vec<String> {
93        self.column_names_of_ty(self.inner.ty)
94    }
95
96    fn column_names_of_ty(&self, ty: &ir::Ty) -> Vec<String> {
97        match &self.get_ty_mat(ty).kind {
98            // arrays are iterated over, columns come from inner type
99            ir::TyKind::Array(item) => self.column_names_of_ty(item),
100
101            // tuple fields become columns
102            ir::TyKind::Tuple(fields) => fields
103                .iter()
104                .enumerate()
105                .map(|(i, f)| {
106                    if let Some(name) = &f.name {
107                        name.clone()
108                    } else {
109                        i.to_string()
110                    }
111                })
112                .collect(),
113
114            // primitives and enums become a single column (we also infer name from ident)
115            ir::TyKind::Primitive(_) | ir::TyKind::Enum(_) => {
116                if let ir::TyKind::Ident(path) = &ty.kind {
117                    vec![path.0.last().unwrap().clone()]
118                } else {
119                    vec!["value".into()]
120                }
121            }
122            ir::TyKind::Ident(_) | ir::TyKind::Function(_) => unreachable!(),
123        }
124    }
125
126    pub fn advance_by_(&mut self, n: usize) -> Result<(), NonZero<usize>> {
127        let items = n.min(self.rem_items);
128        self.rem_items -= items;
129        self.inner.data = &self.inner.data[(items * self.array_item_size)..];
130        match NonZero::new(n - items) {
131            Some(k) => Err(k),
132            None => Ok(()),
133        }
134    }
135}
136
137impl<'d, 't> Iterator for TabularReader<'d, 't> {
138    type Item = vec::Vec<TableCell<'d, 't>>;
139
140    fn next(&mut self) -> Option<Self::Item> {
141        if self.rem_items == 0 {
142            return None;
143        }
144        let mut row = self.inner;
145        if let ir::TyKind::Array(item) = &row.ty.kind {
146            row.ty = item.as_ref();
147        }
148
149        // advance
150        self.rem_items -= 1;
151        if let ir::TyKind::Array(_) = &self.inner.ty.kind {
152            self.inner.data = &self.inner.data[self.array_item_size..];
153        }
154
155        // unpack row
156        let row_ty_mat = self.get_ty_mat(row.ty);
157        Some(match &row_ty_mat.kind {
158            ir::TyKind::Primitive(_) | ir::TyKind::Array(_) | ir::TyKind::Enum(_) => {
159                vec![row]
160            }
161            ir::TyKind::Tuple(fields) => {
162                let mut cells = Vec::with_capacity(fields.len());
163                let reader = TupleReader::new_for_ty(row.data, row_ty_mat);
164                for (i, f) in fields.iter().enumerate() {
165                    cells.push(TableCell {
166                        data: reader.get_field(i),
167                        ty: &f.ty,
168                        ty_defs: row.ty_defs,
169                    })
170                }
171                cells
172            }
173            ir::TyKind::Function(_) | ir::TyKind::Ident(_) => unreachable!(),
174        })
175    }
176}