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) {
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<T: Encode + Layout, E: Encode + Layout> Encode for core::result::Result<T, E> {
252    type HeadPtr =
253        core::result::Result<core::result::Result<T::HeadPtr, E::HeadPtr>, ReversePointer>;
254
255    #[allow(clippy::collapsible_else_if)]
256    fn encode_head(&self, buf: &mut BytesMut) -> Self::HeadPtr {
257        let t_head_bits = T::head_size();
258        let e_head_bits = E::head_size();
259        let max_head_bits = t_head_bits.max(e_head_bits);
260        let has_ptr = max_head_bits > 32;
261        let inner_bytes = if has_ptr {
262            4
263        } else {
264            max_head_bits.div_ceil(8)
265        };
266
267        let tag: u8 = if self.is_ok() { 0 } else { 1 };
268        buf.put_u8(tag);
269
270        if has_ptr {
271            let ptr = ReversePointer::new(buf);
272            Err(ptr)
273        } else {
274            match self {
275                Ok(inner) => {
276                    let head_ptr = inner.encode_head(buf);
277                    buf.put_bytes(0, inner_bytes - t_head_bits.div_ceil(8));
278                    Ok(Ok(head_ptr))
279                }
280                Err(inner) => {
281                    let head_ptr = inner.encode_head(buf);
282                    buf.put_bytes(0, inner_bytes - e_head_bits.div_ceil(8));
283                    Ok(Err(head_ptr))
284                }
285            }
286        }
287    }
288
289    fn encode_body(&self, head: Self::HeadPtr, buf: &mut BytesMut) {
290        match head {
291            Err(ptr) => {
292                ptr.write_cur_len(buf);
293                match self {
294                    Ok(inner) => {
295                        let hp = inner.encode_head(buf);
296                        inner.encode_body(hp, buf);
297                    }
298                    Err(inner) => {
299                        let hp = inner.encode_head(buf);
300                        inner.encode_body(hp, buf);
301                    }
302                }
303            }
304            Ok(Ok(head)) => {
305                let Ok(inner) = self else { unreachable!() };
306                inner.encode_body(head, buf);
307            }
308            Ok(Err(head)) => {
309                let Err(inner) = self else { unreachable!() };
310                inner.encode_body(head, buf);
311            }
312        }
313    }
314}
315
316impl Encode for () {
317    type HeadPtr = ();
318
319    fn encode_head(&self, _buf: &mut BytesMut) {}
320
321    fn encode_body(&self, _: (), _buf: &mut BytesMut) {}
322}
323
324/// Pointer to a location where an offset should be written to.
325#[derive(Debug, Clone, Copy)]
326pub struct ReversePointer {
327    /// Location offset within the buffer.
328    location: usize,
329}
330
331impl ReversePointer {
332    pub fn new(buf: &mut BytesMut) -> ReversePointer {
333        let r = ReversePointer {
334            location: buf.len(),
335        };
336        buf.put_bytes(0, 4);
337        r
338    }
339
340    pub fn new_at(location: usize) -> ReversePointer {
341        ReversePointer { location }
342    }
343
344    pub fn unwrap(&self) -> usize {
345        self.location
346    }
347
348    pub fn write_cur_len(self, w: &mut [u8]) {
349        self.write(w, w.len())
350    }
351
352    pub fn write(self, w: &mut [u8], absolute_ptr: usize) {
353        let relative = absolute_ptr - self.location;
354        w[self.location..(self.location + 4)].copy_from_slice(&(relative as u32).to_le_bytes());
355    }
356}
357
358pub fn encode_enum_head_tag(
359    format: &layout::EnumFormat,
360    tag: u64,
361    buf: &mut BytesMut,
362) -> Option<ReversePointer> {
363    // tag
364    let tag_slice = &tag.to_le_bytes()[0..format.tag_bytes as usize];
365    buf.put_slice(tag_slice);
366
367    if format.has_ptr {
368        // inner pointer
369        let variant = &format.variants[tag as usize];
370        if !variant.is_unit {
371            Some(ReversePointer::new(buf))
372        } else {
373            // this is unit variant, no need for a pointer
374            None
375        }
376    } else {
377        // no pointer, head of inner should follow directly
378        None
379    }
380}
381
382pub fn encode_enum_head_padding(format: &layout::EnumFormat, tag: u64, buf: &mut BytesMut) {
383    let variant = &format.variants[tag as usize];
384    if variant.padding_bytes > 0 {
385        buf.put_bytes(0, variant.padding_bytes as usize);
386    }
387}