tars_stream/
tars_encoder.rs

1use bytes::{BufMut, Bytes, BytesMut};
2use errors::EncodeErr;
3use std::collections::BTreeMap;
4use std::convert::TryFrom;
5use std::mem;
6use tars_trait::{EnumToI32, StructToTars};
7use tars_type::TarsTypeMark::*;
8use tars_type::*;
9
10const MAX_HEADER_LEN: usize = 2;
11const MAX_SIZE_LEN: usize = 4;
12
13#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
14pub struct TarsEncoder {
15    buf: BytesMut,
16}
17
18impl TarsEncoder {
19    pub fn new() -> Self {
20        TarsEncoder {
21            buf: BytesMut::new(),
22        }
23    }
24
25    pub fn individual_encode<T>(ele: &T) -> Result<Bytes, EncodeErr>
26    where
27        T: EncodeTars,
28    {
29        let mut encoder = TarsEncoder::new();
30        ele._encode(&mut encoder, 0)?;
31        Ok(encoder.to_bytes())
32    }
33
34    // move out buf
35    pub fn to_bytes(self) -> Bytes {
36        self.buf.freeze()
37    }
38
39    pub fn to_bytes_mut(self) -> BytesMut {
40        self.buf
41    }
42
43    pub fn check_maybe_resize(&mut self, len: usize) {
44        if self.buf.remaining_mut() < len {
45            let new_len = self.buf.remaining_mut() + len + 1024;
46            self.buf.reserve(new_len)
47        }
48    }
49
50    fn put_head(&mut self, tag: u8, tars_type: TarsTypeMark) -> Result<(), EncodeErr> {
51        self.check_maybe_resize(MAX_HEADER_LEN);
52        if tag > u8::max_value() {
53            Err(EncodeErr::TooBigTagErr)
54        } else {
55            if tag < 15 {
56                let head = (tag << 4) | tars_type.value();
57                self.buf.put_u8(head);
58            } else {
59                let head: u16 = u16::from((0xF0u8) | tars_type.value()) << 8 | u16::from(tag);
60                self.buf.put_u16_be(head)
61            }
62            Ok(())
63        }
64    }
65}
66
67// all write_xxxx method will move value into TarsDecoder
68
69pub trait TarsEncoderNormalTrait {
70    fn write_int8(&mut self, tag: u8, ele: i8) -> Result<(), EncodeErr>;
71    fn write_boolean(&mut self, tag: u8, ele: bool) -> Result<(), EncodeErr>;
72
73    fn write_int16(&mut self, tag: u8, ele: i16) -> Result<(), EncodeErr>;
74    fn write_int32(&mut self, tag: u8, ele: i32) -> Result<(), EncodeErr>;
75    fn write_int64(&mut self, tag: u8, ele: i64) -> Result<(), EncodeErr>;
76
77    fn write_uint8(&mut self, tag: u8, ele: u8) -> Result<(), EncodeErr>;
78    fn write_uint16(&mut self, tag: u8, ele: u16) -> Result<(), EncodeErr>;
79    fn write_uint32(&mut self, tag: u8, ele: u32) -> Result<(), EncodeErr>;
80
81    fn write_float(&mut self, tag: u8, ele: f32) -> Result<(), EncodeErr>;
82    fn write_double(&mut self, tag: u8, ele: f64) -> Result<(), EncodeErr>;
83
84    fn write_string(&mut self, tag: u8, ele: &String) -> Result<(), EncodeErr>;
85
86    fn write_bytes(&mut self, tag: u8, ele: &Bytes) -> Result<(), EncodeErr>;
87
88    fn write_map<K, V>(&mut self, tag: u8, ele: &BTreeMap<K, V>) -> Result<(), EncodeErr>
89    where
90        K: EncodeTars + Ord,
91        V: EncodeTars;
92
93    fn write_enum<T>(&mut self, tag: u8, ele: &T) -> Result<(), EncodeErr>
94    where
95        T: EnumToI32;
96
97    fn write_struct<T>(&mut self, tag: u8, ele: &T) -> Result<(), EncodeErr>
98    where
99        T: StructToTars;
100}
101
102pub trait TarsEncodeListTrait<T>
103where
104    T: EncodeTars,
105{
106    fn write_list(&mut self, tag: u8, ele: &Vec<T>) -> Result<(), EncodeErr>;
107}
108
109impl TarsEncoderNormalTrait for TarsEncoder {
110    fn write_int8(&mut self, tag: u8, ele: i8) -> Result<(), EncodeErr> {
111        if ele == 0 {
112            self.put_head(tag, EnZero)
113        } else {
114            self.put_head(tag, EnInt8)?;
115            self.check_maybe_resize(mem::size_of::<i8>());
116            self.buf.put_i8(ele);
117            Ok(())
118        }
119    }
120
121    fn write_boolean(&mut self, tag: u8, ele: bool) -> Result<(), EncodeErr> {
122        self.write_int8(tag, ele as i8)
123    }
124
125    fn write_int16(&mut self, tag: u8, ele: i16) -> Result<(), EncodeErr> {
126        if ele >= i16::from(i8::min_value()) && ele <= i16::from(i8::max_value()) {
127            self.write_int8(tag, ele as i8)
128        } else {
129            self.put_head(tag, EnInt16)?;
130            self.check_maybe_resize(mem::size_of::<i16>());
131            self.buf.put_i16_be(ele);
132            Ok(())
133        }
134    }
135
136    fn write_int32(&mut self, tag: u8, ele: i32) -> Result<(), EncodeErr> {
137        if ele >= i32::from(i16::min_value()) && ele <= i32::from(i16::max_value()) {
138            self.write_int16(tag, ele as i16)
139        } else {
140            self.put_head(tag, EnInt32)?;
141            self.check_maybe_resize(mem::size_of::<i32>());
142            self.buf.put_i32_be(ele);
143            Ok(())
144        }
145    }
146
147    fn write_int64(&mut self, tag: u8, ele: i64) -> Result<(), EncodeErr> {
148        if ele >= i64::from(i32::min_value()) && ele <= i64::from(i32::max_value()) {
149            self.write_int32(tag, ele as i32)
150        } else {
151            self.put_head(tag, EnInt64)?;
152            self.check_maybe_resize(mem::size_of::<i64>());
153            self.buf.put_i64_be(ele);
154            Ok(())
155        }
156    }
157
158    fn write_uint8(&mut self, tag: u8, ele: u8) -> Result<(), EncodeErr> {
159        self.write_int16(tag, ele as i16)
160    }
161
162    fn write_uint16(&mut self, tag: u8, ele: u16) -> Result<(), EncodeErr> {
163        self.write_int32(tag, ele as i32)
164    }
165
166    fn write_uint32(&mut self, tag: u8, ele: u32) -> Result<(), EncodeErr> {
167        self.write_int64(tag, ele as i64)
168    }
169
170    fn write_float(&mut self, tag: u8, ele: f32) -> Result<(), EncodeErr> {
171        if ele == 0.0 {
172            self.put_head(tag, EnZero)?;
173        } else {
174            self.put_head(tag, EnFloat)?;
175            self.check_maybe_resize(mem::size_of::<f32>());
176            self.buf.put_f32_be(ele)
177        }
178        Ok(())
179    }
180    fn write_double(&mut self, tag: u8, ele: f64) -> Result<(), EncodeErr> {
181        if ele == 0.0 {
182            self.put_head(tag, EnZero)?;
183        } else {
184            self.put_head(tag, EnDouble)?;
185            self.check_maybe_resize(mem::size_of::<f64>());
186            self.buf.put_f64_be(ele)
187        }
188        Ok(())
189    }
190    fn write_string(&mut self, tag: u8, ele: &String) -> Result<(), EncodeErr> {
191        let len = ele.len();
192        self.check_maybe_resize(MAX_SIZE_LEN + len);
193
194        if len <= usize::from(u8::max_value()) {
195            // encode as string1
196            self.put_head(tag, EnString1)?;
197            match u8::try_from(len) {
198                Ok(l) => {
199                    self.buf.put_u8(l);
200                    self.buf.put(ele);
201                    Ok(())
202                }
203                Err(_) => Err(EncodeErr::ConvertU8Err),
204            }
205        } else if len <= u32::max_value() as usize {
206            // encode as string4
207            self.put_head(tag, EnString4)?;
208            self.buf.put_u32_be(len as u32);
209            self.buf.put(ele);
210            Ok(())
211        } else {
212            Err(EncodeErr::DataTooBigErr)
213        }
214    }
215
216    fn write_bytes(&mut self, tag: u8, ele: &Bytes) -> Result<(), EncodeErr> {
217        let len = ele.len();
218        if len > i32::max_value() as usize {
219            Err(EncodeErr::DataTooBigErr)
220        } else {
221            self.put_head(tag, EnSimplelist)?;
222            self.put_head(0, EnInt8)?;
223            self.write_int32(0, len as i32)?;
224            self.buf.extend_from_slice(ele);
225            Ok(())
226        }
227    }
228
229    fn write_map<K, V>(&mut self, tag: u8, ele: &BTreeMap<K, V>) -> Result<(), EncodeErr>
230    where
231        K: EncodeTars + Ord,
232        V: EncodeTars,
233    {
234        let len = ele.len();
235        if len > i32::max_value() as usize {
236            Err(EncodeErr::DataTooBigErr)
237        } else {
238            self.put_head(tag, EnMaps)?;
239            self.write_int32(0, len as i32)?;
240            for (key, value) in ele.iter() {
241                key._encode(self, 0)?;
242                value._encode(self, 1)?;
243            }
244            Ok(())
245        }
246    }
247
248    fn write_enum<T>(&mut self, tag: u8, ele: &T) -> Result<(), EncodeErr>
249    where
250        T: EnumToI32,
251    {
252        self.write_int32(tag, ele._to_i32())
253    }
254
255    fn write_struct<T>(&mut self, tag: u8, ele: &T) -> Result<(), EncodeErr>
256    where
257        T: StructToTars,
258    {
259        self.put_head(tag, EnStructBegin)?;
260        ele._encode_to(self)?;
261        self.put_head(0, EnStructEnd)
262    }
263}
264
265impl<T> TarsEncodeListTrait<T> for TarsEncoder
266where
267    T: EncodeTars,
268{
269    default fn write_list(&mut self, tag: u8, ele: &Vec<T>) -> Result<(), EncodeErr> {
270        let len = ele.len();
271        if len > i32::max_value() as usize {
272            Err(EncodeErr::DataTooBigErr)
273        } else {
274            self.put_head(tag, EnList)?;
275            self.write_int32(0, len as i32)?;
276            for ele in ele.into_iter() {
277                ele._encode(self, 0)?;
278            }
279            Ok(())
280        }
281    }
282}
283
284impl TarsEncodeListTrait<i8> for TarsEncoder {
285    fn write_list(&mut self, tag: u8, ele: &Vec<i8>) -> Result<(), EncodeErr> {
286        let len = ele.len();
287        if len > i32::max_value() as usize {
288            Err(EncodeErr::DataTooBigErr)
289        } else {
290            self.put_head(tag, EnSimplelist)?;
291            self.put_head(0, EnInt8)?;
292            self.write_int32(0, len as i32)?;
293            self.buf
294                .extend_from_slice(unsafe { mem::transmute(ele.as_slice()) });
295            Ok(())
296        }
297    }
298}
299
300impl TarsEncodeListTrait<bool> for TarsEncoder {
301    fn write_list(&mut self, tag: u8, ele: &Vec<bool>) -> Result<(), EncodeErr> {
302        let len = ele.len();
303        if len > i32::max_value() as usize {
304            Err(EncodeErr::DataTooBigErr)
305        } else {
306            self.put_head(tag, EnSimplelist)?;
307            self.put_head(0, EnInt8)?;
308            self.write_int32(0, len as i32)?;
309            self.buf
310                .extend_from_slice(unsafe { mem::transmute(ele.as_slice()) });
311            Ok(())
312        }
313    }
314}
315
316// EncodeTars Trait, 各类型将自身写入 TarsEncoder 中
317pub trait EncodeTars {
318    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr>;
319}
320
321impl EncodeTars for i8 {
322    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
323        encoder.write_int8(tag, *self)
324    }
325}
326
327impl EncodeTars for i16 {
328    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
329        encoder.write_int16(tag, *self)
330    }
331}
332
333impl EncodeTars for i32 {
334    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
335        encoder.write_int32(tag, *self)
336    }
337}
338
339impl EncodeTars for i64 {
340    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
341        encoder.write_int64(tag, *self)
342    }
343}
344
345impl EncodeTars for u8 {
346    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
347        encoder.write_uint8(tag, *self)
348    }
349}
350
351impl EncodeTars for u16 {
352    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
353        encoder.write_uint16(tag, *self)
354    }
355}
356
357impl EncodeTars for u32 {
358    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
359        encoder.write_uint32(tag, *self)
360    }
361}
362
363impl EncodeTars for f32 {
364    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
365        encoder.write_float(tag, *self)
366    }
367}
368
369impl EncodeTars for f64 {
370    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
371        encoder.write_double(tag, *self)
372    }
373}
374
375impl EncodeTars for bool {
376    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
377        encoder.write_boolean(tag, *self)
378    }
379}
380
381impl EncodeTars for String {
382    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
383        encoder.write_string(tag, self)
384    }
385}
386
387impl<K, V> EncodeTars for BTreeMap<K, V>
388where
389    K: EncodeTars + Ord,
390    V: EncodeTars,
391{
392    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
393        encoder.write_map(tag, self)
394    }
395}
396
397impl<T> EncodeTars for Vec<T>
398where
399    T: EncodeTars,
400{
401    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
402        encoder.write_list(tag, self)
403    }
404}
405
406impl EncodeTars for Bytes {
407    fn _encode(&self, encoder: &mut TarsEncoder, tag: u8) -> Result<(), EncodeErr> {
408        encoder.write_bytes(tag, self)
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415
416    #[test]
417    fn test_encode_i8() {
418        let mut encoder = TarsEncoder::new();
419        let i0: i8 = -127;
420        let i3: i8 = 0;
421        let i1: i8 = 127;
422        let i2: i8 = -1;
423
424        encoder.write_int8(0, i0).unwrap();
425        encoder.write_int8(14, i1).unwrap();
426        encoder.write_int8(255, i2).unwrap();
427        encoder.write_int8(3, i3).unwrap();
428
429        assert_eq!(
430            &encoder.to_bytes(),
431            &b"\x00\x81\xe0\x7f\xf0\xff\xff\x3c"[..]
432        );
433    }
434
435    #[test]
436    fn test_encode_i16() {
437        let mut encoder = TarsEncoder::new();
438        let i0: i16 = -32768;
439        let i1: i16 = -127;
440        let i2: i16 = 32767;
441
442        encoder.write_int16(0, i0).unwrap();
443        encoder.write_int16(15, i1).unwrap();
444        encoder.write_int16(19, i2).unwrap();
445
446        assert_eq!(
447            &encoder.to_bytes(),
448            &b"\x01\x80\x00\xf0\x0f\x81\xf1\x13\x7f\xff"[..]
449        );
450    }
451
452    #[test]
453    fn test_encode_i32() {
454        let mut encoder = TarsEncoder::new();
455
456        let i0: i32 = 90909;
457        let i1: i32 = 255;
458        let i2: i32 = -127;
459        let i3: i32 = -95234;
460
461        encoder.write_int32(0, i0).unwrap();
462        encoder.write_int32(15, i1).unwrap();
463        encoder.write_int32(14, i2).unwrap();
464        encoder.write_int32(14, i3).unwrap();
465
466        assert_eq!(
467            &encoder.to_bytes(),
468            &b"\x02\x00\x01\x63\x1d\xf1\x0f\x00\xff\xe0\x81\xe2\xff\xfe\x8b\xfe"[..]
469        );
470    }
471
472    #[test]
473    fn test_encode_i64() {
474        let mut encoder = TarsEncoder::new();
475
476        let i0: i64 = -1;
477        let i1: i64 = -129;
478        let i2: i64 = -32769;
479        let i3: i64 = -2147483649;
480
481        encoder.write_int64(0, i0).unwrap();
482        encoder.write_int64(0, i1).unwrap();
483        encoder.write_int64(0, i2).unwrap();
484        encoder.write_int64(0, i3).unwrap();
485
486        assert_eq!(
487            &encoder.to_bytes(),
488            &b"\x00\xff\x01\xff\x7f\x02\xff\xff\x7f\xff\x03\xff\xff\xff\xff\x7f\xff\xff\xff"[..]
489        );
490    }
491
492    #[test]
493    fn test_encode_u8() {
494        let mut encoder = TarsEncoder::new();
495        let u0: u8 = 127;
496        let u1: u8 = 255;
497        let u2: u8 = 0;
498
499        encoder.write_uint8(0, u0).unwrap();
500        encoder.write_uint8(14, u1).unwrap();
501        encoder.write_uint8(255, u2).unwrap();
502
503        assert_eq!(&encoder.to_bytes(), &b"\x00\x7f\xe1\x00\xff\xfc\xff"[..]);
504    }
505
506    #[test]
507    fn test_encode_u16() {
508        let mut encoder = TarsEncoder::new();
509
510        let i0: u16 = 32768;
511        let i1: u16 = 255;
512        let i2: u16 = 65535;
513
514        encoder.write_uint16(0, i0).unwrap();
515        encoder.write_uint16(15, i1).unwrap();
516
517        encoder.write_uint16(19, i2).unwrap();
518        assert_eq!(
519            &encoder.to_bytes(),
520            &b"\x02\x00\x00\x80\x00\xf1\x0f\x00\xff\xf2\x13\x00\x00\xff\xff"[..]
521        );
522    }
523
524    #[test]
525    fn test_encode_u32() {
526        let mut encoder = TarsEncoder::new();
527        let u0: u32 = 88888;
528        let u1: u32 = 254;
529        let u2: u32 = 256;
530
531        encoder.write_uint32(0, u0).unwrap();
532        encoder.write_uint32(14, u1).unwrap();
533        encoder.write_uint32(14, u2).unwrap();
534
535        assert_eq!(
536            &encoder.to_bytes(),
537            &b"\x02\x00\x01\x5b\x38\xe1\x00\xfe\xe1\x01\x00"[..]
538        );
539    }
540
541    #[test]
542    fn test_encode_f32() {
543        let mut encoder = TarsEncoder::new();
544        let f1: f32 = 0.1472;
545        encoder.write_float(0, f1).unwrap();
546        assert_eq!(&encoder.to_bytes(), &b"\x04\x3e\x16\xbb\x99"[..]);
547    }
548
549    #[test]
550    fn test_encode_f64() {
551        let mut encoder = TarsEncoder::new();
552        let f1: f64 = 0.14723333;
553        encoder.write_double(0, f1).unwrap();
554        assert_eq!(
555            &encoder.to_bytes(),
556            &b"\x05\x3f\xc2\xd8\x8a\xb0\x9d\x97\x2a"[..]
557        );
558    }
559
560    #[test]
561    fn test_encode_bool() {
562        let mut encoder = TarsEncoder::new();
563        encoder.write_boolean(0, false).unwrap();
564        encoder.write_boolean(1, true).unwrap();
565        assert_eq!(&encoder.to_bytes(), &b"\x0c\x10\x01"[..]);
566    }
567
568    #[test]
569    fn test_encode_string() {
570        let mut encoder = TarsEncoder::new();
571        let s: String = "hello wrold!".to_string();
572        let expect_buf = "\x06\x0c".to_string() + &s;
573        encoder.write_string(0, &s).unwrap();
574        assert_eq!(&encoder.to_bytes(), &expect_buf);
575
576        let mut encoder = TarsEncoder::new();
577        let mut s1: String = String::new();
578        for _ in 0..0xf7f7f {
579            s1.push('z');
580        }
581        let expect_buf = "\x07\x00\x0f\x7f\x7f".to_string() + &s1;
582        encoder.write_string(0, &s1).unwrap();
583        assert_eq!(&encoder.to_bytes(), &expect_buf);
584    }
585
586    #[test]
587    fn test_encode_vec() {
588        let mut v2: Vec<i8> = Vec::with_capacity(0xf7f7f);
589        for _ in 0..0xf7f7f {
590            v2.push(-127);
591        }
592        let mut encoder = TarsEncoder::new();
593        encoder.write_list(0, &v2).unwrap();
594        let mut header_v: Vec<u8> = Vec::from(&b"\x0d\x00\x02\x00\x0f\x7f\x7f"[..]);
595        header_v.extend_from_slice(unsafe { mem::transmute(v2.as_slice()) });
596        assert_eq!(&encoder.to_bytes(), &header_v);
597
598        let mut v3: Vec<bool> = Vec::with_capacity(0xf6f7f);
599        let mut b = false;
600        for _ in 0..0xf6f7f {
601            v3.push(b);
602            b = !b;
603        }
604
605        let mut encoder = TarsEncoder::new();
606        encoder.write_list(0, &v3).unwrap();
607        let mut header_v: Vec<u8> = Vec::from(&b"\x0d\x00\x02\x00\x0f\x6f\x7f"[..]);
608        header_v.extend_from_slice(unsafe { mem::transmute(v3.as_slice()) });
609        assert_eq!(&encoder.to_bytes(), &header_v);
610
611        let mut v4: Vec<String> = Vec::with_capacity(0xf6f7e);
612        let str4 = "hello".repeat(128);
613        let str1 = "hello".to_string();
614        let times = 0xf6f7e / 2;
615        for _ in 0..times {
616            v4.push(str4.clone());
617        }
618        for _ in 0..times {
619            v4.push(str1.clone());
620        }
621
622        let mut encoder = TarsEncoder::new();
623        encoder.write_list(10, &v4).unwrap();
624        let buf = encoder.to_bytes();
625        assert_eq!(&buf[0..2], &b"\xa9\x02"[..]);
626        let len_in_u8: [u8; 4] = [buf[2], buf[3], buf[4], buf[5]];
627        let len: i32 = i32::from_be(unsafe { mem::transmute(len_in_u8) });
628        assert_eq!(len, v4.len() as i32);
629    }
630
631    #[test]
632    fn test_encode_map() {
633        let mut map: BTreeMap<String, i32> = BTreeMap::new();
634        map.insert("hello".to_string(), 32);
635        map.insert("world".to_string(), 42);
636
637        let mut encoder = TarsEncoder::new();
638        encoder.write_map(0, &map).unwrap();
639        assert_eq!(
640            &encoder.to_bytes(),
641            &b"\x08\x00\x02\x06\x05hello\x10\x20\x06\x05world\x10\x2a"[..]
642        );
643    }
644
645    #[test]
646    fn test_encode_bytes() {
647        let b = Bytes::from(&b"hello world!"[..]);
648        let mut encoder = TarsEncoder::new();
649        encoder.write_bytes(9, &b).unwrap();
650        assert_eq!(&encoder.to_bytes(), &b"\x9d\x00\x00\x0chello world!"[..]);
651    }
652}
tars_stream/tars_encoder.rs

tars_stream/
tars_encoder.rs
