Skip to main content

acktor_ipc/codec/
prost_codec.rs

1use bytes::{Bytes, BytesMut};
2
3use acktor_ipc_proto::utils::{
4    Tuple, VecBool, VecDouble, VecFloat, VecInt32, VecInt64, VecUint32, VecUint64,
5};
6
7use super::errors::{DecodeError, EncodeError};
8use super::protobuf_helper::LENGTH_DELIMITED_TAGS;
9use super::{Decode, DecodeContext, Encode, EncodeContext};
10
11macro_rules! impl_encode_decode_for {
12    ($type:ty) => {
13        impl Encode for $type {
14            #[inline]
15            fn encoded_len(&self) -> usize {
16                prost::Message::encoded_len(self)
17            }
18
19            #[inline]
20            fn encode(
21                &self,
22                buf: &mut BytesMut,
23                _ctx: Option<&EncodeContext>,
24            ) -> Result<(), EncodeError> {
25                prost::Message::encode(self, buf).map_err(Into::into)
26            }
27        }
28
29        impl Decode for $type {
30            #[inline]
31            fn decode(buf: Bytes, _ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
32                prost::Message::decode(buf).map_err(Into::into)
33            }
34        }
35    };
36    ($type:ty, $wire_type:ty) => {
37        impl Encode for $type {
38            #[inline]
39            fn encoded_len(&self) -> usize {
40                prost::Message::encoded_len(&(*self as $wire_type))
41            }
42
43            #[inline]
44            fn encode(
45                &self,
46                buf: &mut BytesMut,
47                _ctx: Option<&EncodeContext>,
48            ) -> Result<(), EncodeError> {
49                prost::Message::encode(&(*self as $wire_type), buf).map_err(Into::into)
50            }
51        }
52
53        impl Decode for $type {
54            #[inline]
55            fn decode(buf: Bytes, _ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
56                let value = <$wire_type as prost::Message>::decode(buf)?;
57                <$type>::try_from(value).map_err(|e| e.to_string().into())
58            }
59        }
60    };
61}
62
63impl_encode_decode_for!(());
64impl_encode_decode_for!(bool);
65impl_encode_decode_for!(i8, i32);
66impl_encode_decode_for!(i16, i32);
67impl_encode_decode_for!(i32);
68impl_encode_decode_for!(i64);
69impl_encode_decode_for!(u8, u32);
70impl_encode_decode_for!(u16, u32);
71impl_encode_decode_for!(u32);
72impl_encode_decode_for!(u64);
73impl_encode_decode_for!(f32);
74impl_encode_decode_for!(f64);
75impl_encode_decode_for!(isize, i64);
76impl_encode_decode_for!(usize, u64);
77impl_encode_decode_for!(String);
78impl_encode_decode_for!(Vec<u8>);
79
80macro_rules! impl_encode_decode_for_vec {
81    (varint, $type:ty, $msg:ident) => {
82        impl Encode for Vec<$type> {
83            #[inline]
84            fn encoded_len(&self) -> usize {
85                if self.is_empty() {
86                    return 0;
87                }
88                let data_len: usize = self
89                    .iter()
90                    .map(|&v| prost::encoding::encoded_len_varint(v as u64))
91                    .sum();
92                1 + prost::length_delimiter_len(data_len) + data_len
93            }
94
95            #[inline]
96            fn encode(
97                &self,
98                buf: &mut BytesMut,
99                _ctx: Option<&EncodeContext>,
100            ) -> Result<(), EncodeError> {
101                if self.is_empty() {
102                    return Ok(());
103                }
104
105                let data_len: usize = self
106                    .iter()
107                    .map(|&v| prost::encoding::encoded_len_varint(v as u64))
108                    .sum();
109                buf.reserve(1 + prost::length_delimiter_len(data_len) + data_len);
110                buf.extend_from_slice(&[LENGTH_DELIMITED_TAGS[1]]);
111                prost::encoding::encode_varint(data_len as u64, buf);
112                for value in self.iter() {
113                    prost::encoding::encode_varint(*value as u64, buf);
114                }
115
116                Ok(())
117            }
118        }
119
120        impl Decode for Vec<$type> {
121            #[inline]
122            fn decode(buf: Bytes, _ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
123                let message = <$msg as prost::Message>::decode(buf)?;
124
125                Ok(message.values)
126            }
127        }
128    };
129    (varint, $type:ty, $wire_type:ty, $msg:ident) => {
130        impl Encode for Vec<$type> {
131            #[inline]
132            fn encoded_len(&self) -> usize {
133                if self.is_empty() {
134                    return 0;
135                }
136                let data_len: usize = self
137                    .iter()
138                    .map(|&v| prost::encoding::encoded_len_varint(v as $wire_type as u64))
139                    .sum();
140                1 + prost::length_delimiter_len(data_len) + data_len
141            }
142
143            #[inline]
144            fn encode(
145                &self,
146                buf: &mut BytesMut,
147                _ctx: Option<&EncodeContext>,
148            ) -> Result<(), EncodeError> {
149                if self.is_empty() {
150                    return Ok(());
151                }
152
153                let data_len: usize = self
154                    .iter()
155                    .map(|&v| prost::encoding::encoded_len_varint(v as $wire_type as u64))
156                    .sum();
157                buf.reserve(1 + prost::length_delimiter_len(data_len) + data_len);
158                buf.extend_from_slice(&[LENGTH_DELIMITED_TAGS[1]]);
159                prost::encoding::encode_varint(data_len as u64, buf);
160                for value in self.iter() {
161                    prost::encoding::encode_varint(*value as $wire_type as u64, buf);
162                }
163
164                Ok(())
165            }
166        }
167
168        impl Decode for Vec<$type> {
169            #[inline]
170            fn decode(buf: Bytes, _ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
171                let message = <$msg as prost::Message>::decode(buf)?;
172                message
173                    .values
174                    .into_iter()
175                    .map(|v| <$type>::try_from(v).map_err(|e| e.to_string().into()))
176                    .collect()
177            }
178        }
179    };
180    (fixed, $type:ty, $msg:ident) => {
181        impl Encode for Vec<$type> {
182            #[inline]
183            fn encoded_len(&self) -> usize {
184                if self.is_empty() {
185                    return 0;
186                }
187                let data_len = self.len() * std::mem::size_of::<$type>();
188                1 + prost::length_delimiter_len(data_len) + data_len
189            }
190
191            #[inline]
192            fn encode(
193                &self,
194                buf: &mut BytesMut,
195                _ctx: Option<&EncodeContext>,
196            ) -> Result<(), EncodeError> {
197                if self.is_empty() {
198                    return Ok(());
199                }
200
201                let data_len = self.len() * std::mem::size_of::<$type>();
202                buf.reserve(1 + prost::length_delimiter_len(data_len) + data_len);
203                buf.extend_from_slice(&[LENGTH_DELIMITED_TAGS[1]]);
204                prost::encoding::encode_varint(data_len as u64, buf);
205                for value in self.iter() {
206                    buf.extend_from_slice(&value.to_le_bytes());
207                }
208
209                Ok(())
210            }
211        }
212
213        impl Decode for Vec<$type> {
214            #[inline]
215            fn decode(buf: Bytes, _ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
216                let message = <$msg as prost::Message>::decode(buf)?;
217
218                Ok(message.values)
219            }
220        }
221    };
222}
223
224impl_encode_decode_for_vec!(varint, bool, VecBool);
225impl_encode_decode_for_vec!(varint, i8, i32, VecInt32);
226impl_encode_decode_for_vec!(varint, i16, i32, VecInt32);
227impl_encode_decode_for_vec!(varint, i32, VecInt32);
228impl_encode_decode_for_vec!(varint, i64, VecInt64);
229// `Vec<u8>` is handled above by `impl_encode_decode_for!(Vec<u8>)` via prost's bytes field.
230impl_encode_decode_for_vec!(varint, u16, u32, VecUint32);
231impl_encode_decode_for_vec!(varint, u32, VecUint32);
232impl_encode_decode_for_vec!(varint, u64, VecUint64);
233impl_encode_decode_for_vec!(fixed, f32, VecFloat);
234impl_encode_decode_for_vec!(fixed, f64, VecDouble);
235impl_encode_decode_for_vec!(varint, isize, i64, VecInt64);
236impl_encode_decode_for_vec!(varint, usize, u64, VecUint64);
237
238macro_rules! impl_encode_decode_for_tuple {
239    ($($type:ident<$index:tt>($field:ident)),+) => {
240        impl<$($type,)+> Encode for ($($type,)+)
241        where
242            $($type: Encode,)+
243        {
244            #[inline]
245            fn encoded_len(&self) -> usize {
246                let mut len = 0;
247                // 1 byte tag + varint length + data
248                $({
249                    let element_len = self.$index.encoded_len();
250                    len += 1 + prost::length_delimiter_len(element_len) + element_len;
251                })+
252                len
253            }
254
255            #[inline]
256            fn encode(
257                &self,
258                buf: &mut BytesMut,
259                ctx: Option<&EncodeContext>,
260            ) -> Result<(), EncodeError> {
261                buf.reserve(self.encoded_len());
262                $({
263                    buf.extend_from_slice(&[LENGTH_DELIMITED_TAGS[$index + 1]]);
264                    prost::encoding::encode_varint(self.$index.encoded_len() as u64, buf);
265                    self.$index.encode(buf, ctx)?;
266                })+
267
268                Ok(())
269            }
270        }
271
272        impl<$($type,)+> Decode for ($($type,)+)
273        where
274            $($type: Decode,)+
275        {
276            #[inline]
277            fn decode(buf: Bytes, ctx: Option<&DecodeContext>) -> Result<Self, DecodeError> {
278                let message = <Tuple as prost::Message>::decode(buf)?;
279
280                Ok((
281                    $($type::decode(
282                        message.$field.ok_or(DecodeError::from(
283                            concat!("missing field ", stringify!($field), " in tuple message")
284                        ))?,
285                        ctx,
286                    )?,)+
287                ))
288            }
289        }
290    };
291}
292
293impl_encode_decode_for_tuple!(T0<0>(t0), T1<1>(t1));
294impl_encode_decode_for_tuple!(T0<0>(t0), T1<1>(t1), T2<2>(t2));
295impl_encode_decode_for_tuple!(T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3));
296impl_encode_decode_for_tuple!(T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4));
297impl_encode_decode_for_tuple!(
298    T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4), T5<5>(t5)
299);
300impl_encode_decode_for_tuple!(
301    T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4), T5<5>(t5), T6<6>(t6)
302);
303impl_encode_decode_for_tuple!(
304    T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4), T5<5>(t5), T6<6>(t6), T7<7>(t7)
305);
306impl_encode_decode_for_tuple!(
307    T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4), T5<5>(t5), T6<6>(t6), T7<7>(t7),
308    T8<8>(t8)
309);
310impl_encode_decode_for_tuple!(
311    T0<0>(t0), T1<1>(t1), T2<2>(t2), T3<3>(t3), T4<4>(t4), T5<5>(t5), T6<6>(t6), T7<7>(t7),
312    T8<8>(t8), T9<9>(t9)
313);