1 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2 
3 use crate::buf::GrpcSlice;
4 use crate::call::MessageReader;
5 use crate::error::Result;
6 
7 pub type DeserializeFn<T> = fn(MessageReader) -> Result<T>;
8 pub type SerializeFn<T> = fn(&T, &mut GrpcSlice);
9 
10 /// Defines how to serialize and deserialize between the specialized type and byte slice.
11 pub struct Marshaller<T> {
12     // Use function pointer here to simplify the signature.
13     // Compiler will probably inline the function so performance
14     // impact can be omitted.
15     //
16     // Using trait will require a trait object or generic, which will
17     // either have performance impact or make signature complicated.
18     //
19     // const function is not stable yet (rust-lang/rust#24111), hence
20     // make all fields public.
21     /// The serialize function.
22     pub ser: SerializeFn<T>,
23 
24     /// The deserialize function.
25     pub de: DeserializeFn<T>,
26 }
27 
28 #[cfg(feature = "protobuf-codec")]
29 pub mod pb_codec {
30     use protobuf::{CodedInputStream, CodedOutputStream, Message};
31 
32     use super::MessageReader;
33     use crate::buf::GrpcSlice;
34     use crate::error::Result;
35 
36     #[inline]
ser<T: Message>(t: &T, buf: &mut GrpcSlice)37     pub fn ser<T: Message>(t: &T, buf: &mut GrpcSlice) {
38         let cap = t.compute_size();
39         unsafe {
40             let bytes = buf.realloc(cap as usize);
41             let raw_bytes = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
42             let mut s = CodedOutputStream::bytes(raw_bytes);
43             t.write_to_with_cached_sizes(&mut s).unwrap();
44         }
45     }
46 
47     #[inline]
de<T: Message>(mut reader: MessageReader) -> Result<T>48     pub fn de<T: Message>(mut reader: MessageReader) -> Result<T> {
49         let mut s = CodedInputStream::from_buffered_reader(&mut reader);
50         let mut m = T::new();
51         m.merge_from(&mut s)?;
52         Ok(m)
53     }
54 }
55 
56 #[cfg(feature = "prost-codec")]
57 pub mod pr_codec {
58     use prost::Message;
59 
60     use super::MessageReader;
61     use crate::buf::GrpcSlice;
62     use crate::error::Result;
63 
64     #[inline]
ser<M: Message>(msg: &M, buf: &mut GrpcSlice)65     pub fn ser<M: Message>(msg: &M, buf: &mut GrpcSlice) {
66         let size = msg.encoded_len();
67         unsafe {
68             let bytes = buf.realloc(size);
69             let mut b = &mut *(bytes as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
70             msg.encode(&mut b)
71                 .expect("Writing message to buffer failed");
72             debug_assert!(b.is_empty());
73         }
74     }
75 
76     #[inline]
de<M: Message + Default>(mut reader: MessageReader) -> Result<M>77     pub fn de<M: Message + Default>(mut reader: MessageReader) -> Result<M> {
78         use bytes::buf::Buf;
79         reader.advance(0);
80         M::decode(reader).map_err(Into::into)
81     }
82 }
83