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#[derive(Debug, Clone, Copy)]
261pub struct ReversePointer {
262 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 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 let variant = &format.variants[tag as usize];
305 if !variant.is_unit {
306 Some(ReversePointer::new(buf))
307 } else {
308 None
310 }
311 } else {
312 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}