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#[derive(Debug, Clone, Copy)]
326pub struct ReversePointer {
327 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 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 let variant = &format.variants[tag as usize];
370 if !variant.is_unit {
371 Some(ReversePointer::new(buf))
372 } else {
373 None
375 }
376 } else {
377 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}