Skip to main content

lutra_bin/
reader.rs

1use crate::borrow;
2
3use crate::ir;
4use crate::layout;
5
6pub trait ReaderExt {
7    fn read_n(&self, n: usize) -> &[u8];
8
9    fn read_const<const N: usize>(&self) -> [u8; N];
10
11    #[must_use]
12    fn skip(self, bytes: usize) -> Self;
13}
14
15impl ReaderExt for &[u8] {
16    fn read_n(&self, n: usize) -> &[u8] {
17        &self[..n]
18    }
19
20    fn read_const<const N: usize>(&self) -> [u8; N] {
21        self[0..N].try_into().unwrap()
22    }
23
24    fn skip(self, bytes: usize) -> Self {
25        &self[bytes..]
26    }
27}
28
29#[derive(Clone)]
30pub struct ArrayReader<B: bytes::Buf + Clone> {
31    item_head_bytes: usize,
32    buf_next: B,
33    remaining: usize,
34}
35
36impl<B> ArrayReader<B>
37where
38    B: bytes::Buf + Clone,
39{
40    pub fn new_for_ty(buf: B, ty: &ir::Ty) -> Self {
41        let ir::TyKind::Array(items_ty) = &ty.kind else {
42            panic!()
43        };
44        let item_head_size = items_ty.layout.as_ref().unwrap().head_size as usize;
45
46        Self::new(buf, item_head_size.div_ceil(8))
47    }
48
49    pub fn new(buf: B, item_head_bytes: usize) -> Self {
50        let (offset, len) = Self::read_head(buf.chunk());
51
52        let mut body = buf.clone();
53        body.advance(offset);
54
55        ArrayReader {
56            item_head_bytes,
57            buf_next: body,
58            remaining: len,
59        }
60    }
61
62    pub fn remaining(&self) -> usize {
63        self.remaining
64    }
65
66    pub fn read_head(buffer: &[u8]) -> (usize, usize) {
67        let offset = u32::from_le_bytes(buffer.read_const::<4>());
68
69        let len = u32::from_le_bytes(buffer.skip(4).read_const::<4>());
70        (offset as usize, len as usize)
71    }
72
73    pub fn get(&self, index: usize) -> Option<B> {
74        if index >= self.remaining {
75            return None;
76        }
77
78        let mut data = self.buf_next.clone();
79        data.advance(self.item_head_bytes * index);
80        Some(data)
81    }
82}
83
84impl<B> Iterator for ArrayReader<B>
85where
86    B: bytes::Buf + Clone,
87{
88    type Item = B;
89
90    fn next(&mut self) -> Option<Self::Item> {
91        if self.remaining > 0 {
92            let this = self.buf_next.clone();
93            self.buf_next.advance(self.item_head_bytes);
94            self.remaining -= 1;
95            Some(this)
96        } else {
97            None
98        }
99    }
100
101    fn nth(&mut self, n: usize) -> Option<Self::Item> {
102        let n = if n > self.remaining {
103            self.remaining
104        } else {
105            n
106        };
107        self.buf_next.advance(self.item_head_bytes * n);
108        self.remaining -= n;
109        self.next()
110    }
111
112    fn size_hint(&self) -> (usize, Option<usize>) {
113        (self.remaining, Some(self.remaining))
114    }
115
116    fn count(self) -> usize
117    where
118        Self: Sized,
119    {
120        self.remaining
121    }
122}
123
124pub struct TupleReader<'t, B> {
125    buf: B,
126    field_offsets: borrow::Cow<'t, [u32]>,
127}
128
129impl<'t, B> TupleReader<'t, B>
130where
131    B: bytes::Buf + Clone,
132{
133    pub fn new(buf: B, field_offsets: borrow::Cow<'t, [u32]>) -> Self {
134        TupleReader { buf, field_offsets }
135    }
136
137    pub fn new_for_ty(buf: B, ty: &ir::Ty) -> Self {
138        let field_offsets = layout::tuple_field_offsets(ty);
139        Self::new(buf, field_offsets.into())
140    }
141
142    pub fn get_field(&self, index: usize) -> B {
143        let mut r = self.buf.clone();
144        r.advance(self.field_offsets[index] as usize);
145        r
146    }
147}