Skip to main content

lutra_bin/
encode.rs

1use crate::layout;
2use crate::string;
3use crate::vec;
4
5use bytes::{BufMut, BytesMut};
6
7use crate::{Layout, Result};
8
9pub trait Encode {
10    type HeadPtr;
11
12    fn encode(&self) -> vec::Vec<u8> {
13        let mut buf = BytesMut::new();
14        let ptr = self.encode_head(&mut buf);
15        self.encode_body(ptr, &mut buf);
16        buf.to_vec()
17    }
18
19    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr;
20
21    fn encode_body(&self, head: Self::HeadPtr, buf: &mut BytesMut);
22}
23
24impl<T> Encode for &T
25where
26    T: Encode,
27{
28    type HeadPtr = T::HeadPtr;
29
30    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
31        T::encode_head(self, buf)
32    }
33
34    fn encode_body(&self, head: Self::HeadPtr, buf: &mut BytesMut) {
35        T::encode_body(self, head, buf)
36    }
37}
38
39impl Encode for bool {
40    type HeadPtr = ();
41
42    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
43        buf.put_u8(*self as u8);
44    }
45
46    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
47}
48
49impl Encode for i8 {
50    type HeadPtr = ();
51
52    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
53        buf.put_i8(*self)
54    }
55
56    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
57}
58impl Encode for i16 {
59    type HeadPtr = ();
60
61    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
62        buf.put_i16_le(*self)
63    }
64
65    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
66}
67impl Encode for i32 {
68    type HeadPtr = ();
69
70    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
71        buf.put_i32_le(*self);
72    }
73
74    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
75}
76impl Encode for i64 {
77    type HeadPtr = ();
78
79    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
80        buf.put_i64_le(*self);
81    }
82
83    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
84}
85impl Encode for u8 {
86    type HeadPtr = ();
87
88    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
89        buf.put_u8(*self);
90    }
91
92    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
93}
94impl Encode for u16 {
95    type HeadPtr = ();
96
97    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
98        buf.put_u16_le(*self);
99    }
100
101    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
102}
103impl Encode for u32 {
104    type HeadPtr = ();
105
106    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
107        buf.put_u32_le(*self);
108    }
109
110    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
111}
112impl Encode for u64 {
113    type HeadPtr = ();
114
115    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
116        buf.put_u64_le(*self);
117    }
118
119    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
120}
121impl Encode for f32 {
122    type HeadPtr = ();
123
124    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
125        buf.put_f32_le(*self);
126    }
127
128    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
129}
130impl Encode for f64 {
131    type HeadPtr = ();
132
133    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
134        buf.put_f64_le(*self);
135    }
136
137    fn encode_body(&self, _head: Self::HeadPtr, _buf: &mut BytesMut) {}
138}
139
140impl Encode for str {
141    type HeadPtr = ReversePointer;
142
143    fn encode_head(&self, buf: &mut BytesMut) -> ReversePointer {
144        let offset = ReversePointer::new(buf);
145        buf.put_u32_le(self.len() as u32);
146        offset
147    }
148
149    fn encode_body(&self, offset: Self::HeadPtr, buf: &mut BytesMut) {
150        offset.write_cur_len(buf);
151        buf.put_slice(self.as_bytes());
152    }
153}
154impl Encode for string::String {
155    type HeadPtr = ReversePointer;
156
157    fn encode_head(&self, buf: &mut BytesMut) -> ReversePointer {
158        (self as &str).encode_head(buf)
159    }
160
161    fn encode_body(&self, offset: Self::HeadPtr, buf: &mut BytesMut) {
162        (self as &str).encode_body(offset, buf)
163    }
164}
165impl<E: Encode> Encode for vec::Vec<E> {
166    type HeadPtr = ReversePointer;
167
168    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
169        self.as_slice().encode_head(buf)
170    }
171
172    fn encode_body(&self, head: Self::HeadPtr, buf: &mut BytesMut) {
173        self.as_slice().encode_body(head, buf)
174    }
175}
176impl<E: Encode> Encode for [E] {
177    type HeadPtr = ReversePointer;
178
179    fn encode_head(&self, buf: &mut BytesMut) -> ReversePointer {
180        let offset = ReversePointer::new(buf);
181        buf.put_u32_le(self.len() as u32);
182        offset
183    }
184
185    fn encode_body(&self, offset_ptr: Self::HeadPtr, buf: &mut BytesMut) {
186        offset_ptr.write_cur_len(buf);
187
188        let mut heads: vec::Vec<_> = vec::Vec::with_capacity(self.len());
189        for i in self {
190            heads.push(i.encode_head(buf));
191        }
192
193        for (i, h) in self.iter().zip(heads.into_iter()) {
194            i.encode_body(h, buf);
195        }
196    }
197}
198impl<E: Encode + Layout> Encode for Option<E> {
199    type HeadPtr = Option<Result<E::HeadPtr, ReversePointer>>;
200
201    #[allow(clippy::collapsible_else_if)]
202    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
203        let tag: u8 = if self.is_none() { 0 } else { 1 };
204        buf.put_u8(tag);
205
206        let inner_head_size = E::head_size();
207        let has_ptr = inner_head_size > 32;
208
209        if has_ptr {
210            if self.is_some() {
211                let offset = ReversePointer::new(buf);
212                Some(Err(offset))
213            } else {
214                buf.put_bytes(0, 4);
215                None
216            }
217        } else {
218            if let Some(inner) = self {
219                let inner_head_ptr = inner.encode_head(buf);
220                Some(Ok(inner_head_ptr))
221            } else {
222                buf.put_bytes(0, inner_head_size.div_ceil(8));
223                None
224            }
225        }
226    }
227
228    fn encode_body(&self, head: Self::HeadPtr, buf: &mut BytesMut) {
229        let Some(inner) = self else { return };
230
231        let inner_head_size = E::head_size();
232        let has_ptr = inner_head_size > 32;
233
234        let head = head.unwrap();
235
236        if has_ptr {
237            let Err(offset_ptr) = head else {
238                unreachable!()
239            };
240            offset_ptr.write_cur_len(buf);
241
242            let head_ptr = inner.encode_head(buf);
243            inner.encode_body(head_ptr, buf)
244        } else {
245            let Ok(head) = head else { unreachable!() };
246            inner.encode_body(head, buf)
247        }
248    }
249}
250
251impl Encode for () {
252    type HeadPtr = ();
253
254    fn encode_head(&self, _buf: &mut BytesMut) {}
255
256    fn encode_body(&self, _: (), _buf: &mut BytesMut) {}
257}
258
259/// Pointer to a location where an offset should be written to.
260#[derive(Debug, Clone, Copy)]
261pub struct ReversePointer {
262    /// Location offset within the buffer.
263    location: usize,
264}
265
266impl ReversePointer {
267    pub fn new(buf: &mut BytesMut) -> ReversePointer {
268        let r = ReversePointer {
269            location: buf.len(),
270        };
271        buf.put_bytes(0, 4);
272        r
273    }
274
275    pub fn new_at(location: usize) -> ReversePointer {
276        ReversePointer { location }
277    }
278
279    pub fn unwrap(&self) -> usize {
280        self.location
281    }
282
283    pub fn write_cur_len(self, w: &mut [u8]) {
284        self.write(w, w.len())
285    }
286
287    pub fn write(self, w: &mut [u8], absolute_ptr: usize) {
288        let relative = absolute_ptr - self.location;
289        w[self.location..(self.location + 4)].copy_from_slice(&(relative as u32).to_le_bytes());
290    }
291}
292
293pub fn encode_enum_head_tag(
294    format: &layout::EnumFormat,
295    tag: u64,
296    buf: &mut BytesMut,
297) -> Option<ReversePointer> {
298    // tag
299    let tag_slice = &tag.to_le_bytes()[0..format.tag_bytes as usize];
300    buf.put_slice(tag_slice);
301
302    if format.has_ptr {
303        // inner pointer
304        let variant = &format.variants[tag as usize];
305        if !variant.is_unit {
306            Some(ReversePointer::new(buf))
307        } else {
308            // this is unit variant, no need for a pointer
309            None
310        }
311    } else {
312        // no pointer, head of inner should follow directly
313        None
314    }
315}
316
317pub fn encode_enum_head_padding(format: &layout::EnumFormat, tag: u64, buf: &mut BytesMut) {
318    let variant = &format.variants[tag as usize];
319    if variant.padding_bytes > 0 {
320        buf.put_bytes(0, variant.padding_bytes as usize);
321    }
322}