data_rw/
data.rs

1use crate::serde::error::DataError;
2use anyhow::{ensure, Result};
3use auto_impl::auto_impl;
4use serde::Serialize;
5use std::mem::size_of;
6use std::ops::{Deref, DerefMut};
7
8#[auto_impl(&, &mut,Box)]
9pub trait WriteNumberFixed {
10    fn write(&self, data: &mut Data);
11    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()>;
12}
13#[auto_impl(&, &mut,Box)]
14pub trait WriteNumberVar {
15    fn write(&self, data: &mut Data);
16}
17
18#[derive(Debug, Default)]
19pub struct Data {
20    buff: Vec<u8>,
21    pub mode: u8,
22}
23
24macro_rules! impl_number_fixed {
25    ($type:ty) => {
26        impl WriteNumberFixed for $type {
27            #[cfg(not(feature = "big_endian"))]
28            #[inline]
29            fn write(&self, data: &mut Data) {
30                unsafe {
31                    let size = size_of::<$type>();
32                    let len = data.check_reserve(size);
33                    std::ptr::copy_nonoverlapping(
34                        self.to_le_bytes().as_ptr(),
35                        data.as_mut_ptr().add(len),
36                        size,
37                    );
38                    data.buff.set_len(len.wrapping_add(size));
39                }
40            }
41            #[cfg(not(feature = "big_endian"))]
42            #[inline]
43            fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
44                unsafe {
45                    let size = size_of::<$type>();
46                    ensure!(
47                        idx.wrapping_add(size) <= data.len(),
48                        "idx too max {}>{}",
49                        idx.wrapping_add(size),
50                        data.len()
51                    );
52                    std::ptr::copy_nonoverlapping(
53                        self.to_le_bytes().as_ptr(),
54                        data.as_mut_ptr().add(idx),
55                        size,
56                    );
57                    Ok(())
58                }
59            }
60
61            #[cfg(feature = "big_endian")]
62            #[inline]
63            fn write(&self, data: &mut Data) {
64                unsafe {
65                    let size = size_of::<$type>();
66                    let len = data.check_reserve(size);
67                    std::ptr::copy_nonoverlapping(
68                        self.to_be_bytes().as_ptr(),
69                        data.as_mut_ptr().add(len),
70                        size,
71                    );
72                    data.buff.set_len(len.wrapping_add(size));
73                }
74            }
75            #[cfg(feature = "big_endian")]
76            #[inline]
77            fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
78                unsafe {
79                    let size = size_of::<$type>();
80                    ensure!(
81                        idx.wrapping_add(size) <= data.len(),
82                        "idx too max {}>{}",
83                        idx.wrapping_add(size),
84                        data.len()
85                    );
86                    std::ptr::copy_nonoverlapping(
87                        self.to_be_bytes().as_ptr(),
88                        data.as_mut_ptr().add(idx),
89                        size,
90                    );
91                    Ok(())
92                }
93            }
94        }
95    };
96}
97
98impl_number_fixed!(u8);
99impl_number_fixed!(i8);
100impl_number_fixed!(i16);
101impl_number_fixed!(u16);
102impl_number_fixed!(i32);
103impl_number_fixed!(u32);
104impl_number_fixed!(i64);
105impl_number_fixed!(u64);
106impl_number_fixed!(i128);
107impl_number_fixed!(u128);
108impl_number_fixed!(f32);
109impl_number_fixed!(f64);
110
111impl WriteNumberFixed for bool {
112    #[inline]
113    fn write(&self, data: &mut Data) {
114        let v = if *self { 1u8 } else { 0u8 };
115        data.write_fixed(v);
116    }
117
118    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
119        let v = if *self { 1u8 } else { 0u8 };
120        data.write_fixed_at(idx, v)
121    }
122}
123
124impl WriteNumberFixed for String {
125    #[inline]
126    fn write(&self, data: &mut Data) {
127        let buff = self.as_bytes();
128        data.write_fixed(buff.len() as u32);
129        data.write_buf(buff);
130    }
131
132    #[inline]
133    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
134        let buff = self.as_bytes();
135        data.write_fixed_at(idx, buff.len() as u32)?;
136        data.write_buf_at(idx + 4, buff)?;
137        Ok(())
138    }
139}
140
141impl WriteNumberFixed for &str {
142    #[inline]
143    fn write(&self, data: &mut Data) {
144        let buff = self.as_bytes();
145        data.write_fixed(buff.len() as u32);
146        data.write_buf(buff);
147    }
148
149    #[inline]
150    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
151        let buff = self.as_bytes();
152        data.write_fixed_at(idx, buff.len() as u32)?;
153        data.write_buf_at(idx + 4, buff)?;
154        Ok(())
155    }
156}
157
158impl<'a> WriteNumberFixed for &'a [u8] {
159    #[inline]
160    fn write(&self, data: &mut Data) {
161        let len = self.len() as u32;
162        data.write_fixed(len);
163        data.write_buf(self);
164    }
165
166    #[inline]
167    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
168        let len = self.len() as u32;
169        data.write_fixed_at(idx, len)?;
170        data.write_buf_at(idx + 4, self)?;
171        Ok(())
172    }
173}
174
175impl WriteNumberFixed for Vec<u8> {
176    fn write(&self, data: &mut Data) {
177        data.write_fixed(self.len() as u32);
178        data.write_buf(self);
179    }
180
181    fn write_at(&self, idx: usize, data: &mut Data) -> Result<()> {
182        data.write_fixed_at(idx, self.len() as u32)?;
183        data.write_buf_at(idx + 4, self)?;
184        Ok(())
185    }
186}
187
188impl WriteNumberVar for u16 {
189    #[inline]
190    fn write(&self, data: &mut Data) {
191        let mut value = *self;
192        let size = compute_raw_varint64_size(value as u64);
193        let current_len = data.check_reserve(size);
194        unsafe {
195            let mut len: usize = 1;
196            let mut ptr = data.as_mut_ptr().add(current_len);
197            while value >= 1 << 7 {
198                ptr.write((value & 0x7f | 0x80) as u8);
199                ptr = ptr.offset(1);
200                len += 1;
201                value >>= 7;
202            }
203            ptr.write(value as u8);
204            data.set_len(current_len + len);
205        }
206    }
207}
208impl WriteNumberVar for i16 {
209    #[inline]
210    fn write(&self, data: &mut Data) {
211        WriteNumberVar::write(&zig_zag_encode_u16(self), data);
212    }
213}
214impl WriteNumberVar for u32 {
215    #[inline]
216    fn write(&self, data: &mut Data) {
217        let mut value = *self;
218        let size = compute_raw_varint32_size(value);
219        let current_len = data.check_reserve(size);
220        unsafe {
221            let mut len: usize = 1;
222            let mut ptr = data.as_mut_ptr().add(current_len);
223            while value >= 1 << 7 {
224                ptr.write((value & 0x7f | 0x80) as u8);
225                ptr = ptr.offset(1);
226                len += 1;
227                value >>= 7;
228            }
229            ptr.write(value as u8);
230            data.set_len(current_len + len);
231        }
232    }
233}
234impl WriteNumberVar for i32 {
235    #[inline]
236    fn write(&self, data: &mut Data) {
237        WriteNumberVar::write(&zig_zag_encode_u32(self), data);
238    }
239}
240impl WriteNumberVar for u64 {
241    #[inline]
242    fn write(&self, data: &mut Data) {
243        let mut value = *self;
244        let size = compute_raw_varint64_size(value);
245        let current_len = data.check_reserve(size);
246        unsafe {
247            let mut len: usize = 1;
248            let mut ptr = data.as_mut_ptr().add(current_len);
249            while value >= 1 << 7 {
250                ptr.write((value & 0x7f | 0x80) as u8);
251                ptr = ptr.offset(1);
252                len += 1;
253                value >>= 7;
254            }
255            ptr.write(value as u8);
256            data.set_len(current_len + len);
257        }
258    }
259}
260impl WriteNumberVar for i64 {
261    #[inline]
262    fn write(&self, data: &mut Data) {
263        WriteNumberVar::write(&zig_zag_encode_u64(self), data);
264    }
265}
266
267#[inline(always)]
268fn zig_zag_encode_u16(v: &i16) -> u16 {
269    ((v << 1) ^ (v >> 15)) as u16
270}
271#[inline(always)]
272fn zig_zag_encode_u32(v: &i32) -> u32 {
273    ((v << 1) ^ (v >> 31)) as u32
274}
275#[inline(always)]
276fn zig_zag_encode_u64(v: &i64) -> u64 {
277    ((v << 1) ^ (v >> 63)) as u64
278}
279
280impl WriteNumberVar for String {
281    #[inline]
282    fn write(&self, data: &mut Data) {
283        let buff = self.as_bytes();
284        data.write_var_integer(buff.len() as u64);
285        data.write_buf(buff);
286    }
287}
288impl WriteNumberVar for &str {
289    #[inline]
290    fn write(&self, data: &mut Data) {
291        let buff = self.as_bytes();
292        data.write_var_integer(buff.len() as u64);
293        data.write_buf(buff);
294    }
295}
296
297impl WriteNumberVar for &[u8] {
298    #[inline]
299    fn write(&self, data: &mut Data) {
300        let len = self.len() as u64;
301        data.write_var_integer(len);
302        data.write_buf(self);
303    }
304}
305
306impl WriteNumberVar for Vec<u8> {
307    #[inline]
308    fn write(&self, data: &mut Data) {
309        let len = self.len() as u64;
310        data.write_var_integer(len);
311        data.write_buf(self);
312    }
313}
314
315impl Data {
316    #[inline]
317    pub fn new() -> Self {
318        Data {
319            buff: Vec::with_capacity(4096),
320            mode: 0,
321        }
322    }
323
324    #[inline]
325    pub fn with_len(len: usize, default: u8) -> Self {
326        let mut data = Data::new();
327        data.resize(len, default);
328        data
329    }
330
331    #[inline]
332    pub fn with_capacity(cap: usize) -> Data {
333        Data {
334            buff: Vec::with_capacity(cap),
335            mode: 0,
336        }
337    }
338
339    #[inline]
340    pub fn serialize<T: Serialize>(value: T) -> Result<Self, DataError> {
341        let mut data = Data::new();
342        data.serde_serialize(value)?;
343        Ok(data)
344    }
345
346    #[inline]
347    pub fn into_inner(self) -> Vec<u8> {
348        self.buff
349    }
350
351    #[inline]
352    pub fn write_buf(&mut self, buff: &[u8]) {
353        unsafe {
354            let size = buff.len();
355            let len = self.check_reserve(size);
356            std::ptr::copy_nonoverlapping(buff.as_ptr(), self.as_mut_ptr().add(len), size);
357            self.set_len(len.wrapping_add(size));
358        }
359    }
360
361    #[inline]
362    pub fn write_buf_at(&mut self, idx: usize, buff: &[u8]) -> Result<()> {
363        let size = buff.len();
364        ensure!(
365            idx.wrapping_add(size) <= self.len(),
366            "idx too max {}>{}",
367            idx.wrapping_add(size),
368            self.len()
369        );
370        unsafe {
371            std::ptr::copy_nonoverlapping(buff.as_ptr(), self.as_mut_ptr().add(idx), size);
372        }
373        Ok(())
374    }
375
376    #[inline]
377    pub fn write_fixed(&mut self, v: impl WriteNumberFixed) {
378        v.write(self)
379    }
380
381    #[inline]
382    pub fn write_fixed_at(&mut self, idx: usize, v: impl WriteNumberFixed) -> Result<()> {
383        v.write_at(idx, self)
384    }
385
386    #[inline]
387    pub fn write_var_integer(&mut self, v: impl WriteNumberVar) {
388        v.write(self);
389    }
390
391    #[inline]
392    pub fn check_reserve(&mut self, size: usize) -> usize {
393        let len = self.len();
394        let cap = self.capacity();
395        if size.wrapping_add(len) > cap {
396            self.reserve(size.wrapping_add(len));
397        }
398        len
399    }
400}
401
402impl From<Data> for Vec<u8> {
403    #[inline]
404    fn from(val: Data) -> Self {
405        val.buff
406    }
407}
408
409impl Deref for Data {
410    type Target = Vec<u8>;
411
412    #[inline]
413    fn deref(&self) -> &Self::Target {
414        &self.buff
415    }
416}
417
418impl DerefMut for Data {
419    #[inline]
420    fn deref_mut(&mut self) -> &mut Self::Target {
421        &mut self.buff
422    }
423}
424
425impl AsRef<[u8]> for Data {
426    #[inline]
427    fn as_ref(&self) -> &[u8] {
428        &self.buff
429    }
430}
431
432/// Given `u64` value compute varint encoded length.
433#[inline(always)]
434pub fn compute_raw_varint64_size(value: u64) -> usize {
435    if (value & (0xffffffffffffffffu64 << 7)) == 0 {
436        return 1;
437    }
438    if (value & (0xffffffffffffffffu64 << 14)) == 0 {
439        return 2;
440    }
441    if (value & (0xffffffffffffffffu64 << 21)) == 0 {
442        return 3;
443    }
444    if (value & (0xffffffffffffffffu64 << 28)) == 0 {
445        return 4;
446    }
447    if (value & (0xffffffffffffffffu64 << 35)) == 0 {
448        return 5;
449    }
450    if (value & (0xffffffffffffffffu64 << 42)) == 0 {
451        return 6;
452    }
453    if (value & (0xffffffffffffffffu64 << 49)) == 0 {
454        return 7;
455    }
456    if (value & (0xffffffffffffffffu64 << 56)) == 0 {
457        return 8;
458    }
459    if (value & (0xffffffffffffffffu64 << 63)) == 0 {
460        return 9;
461    }
462    10
463}
464
465/// Given `u32` value compute varint encoded length.
466#[inline(always)]
467pub fn compute_raw_varint32_size(value: u32) -> usize {
468    compute_raw_varint64_size(value as u64)
469}
470
471impl From<Vec<u8>> for Data {
472    fn from(buff: Vec<u8>) -> Self {
473        Data { buff, mode: 0 }
474    }
475}
data_rw/data.rs

data_rw/
data.rs
