protobuf_codec/
message.rs

1//! Encoders, decoders and traits for messages.
2use crate::field::{FieldDecode, FieldEncode, UnknownFieldDecoder};
3use crate::value::{ValueDecode, ValueEncode};
4use crate::wire::{LengthDelimitedDecoder, LengthDelimitedEncoder, TagDecoder, WireType};
5use bytecodec::combinator::{Map, MapErr, MapFrom, PreEncode, TryMap, TryMapFrom};
6use bytecodec::{ByteCount, Decode, Encode, Eos, Error, ErrorKind, Result, SizedEncode};
7use std;
8
9/// This trait allows for decoding messages.
10pub trait MessageDecode: Decode {}
11impl<M, T, F> MessageDecode for Map<M, T, F>
12where
13    M: MessageDecode,
14    F: Fn(M::Item) -> T,
15{
16}
17impl<M, T, E, F> MessageDecode for TryMap<M, T, E, F>
18where
19    M: MessageDecode,
20    F: Fn(M::Item) -> std::result::Result<T, E>,
21    Error: From<E>,
22{
23}
24impl<M, E, F> MessageDecode for MapErr<M, E, F>
25where
26    M: MessageDecode,
27    F: Fn(Error) -> E,
28    Error: From<E>,
29{
30}
31
32/// This trait allows for encoding messages.
33pub trait MessageEncode: Encode {}
34impl<M: MessageEncode> MessageEncode for PreEncode<M> {}
35impl<M, T, F> MessageEncode for MapFrom<M, T, F>
36where
37    M: MessageEncode,
38    F: Fn(T) -> M::Item,
39{
40}
41impl<M, T, E, F> MessageEncode for TryMapFrom<M, T, E, F>
42where
43    M: MessageEncode,
44    F: Fn(T) -> std::result::Result<M::Item, E>,
45    Error: From<E>,
46{
47}
48impl<M, E, F> MessageEncode for MapErr<M, E, F>
49where
50    M: MessageEncode,
51    F: Fn(Error) -> E,
52    Error: From<E>,
53{
54}
55
56/// Decoder for messages.
57#[derive(Debug, Default)]
58pub struct MessageDecoder<F> {
59    tag: TagDecoder,
60    field: F,
61    unknown_field: UnknownFieldDecoder,
62    started: bool,
63    eos: bool, // end-of-stream
64    target: DecodeTarget,
65}
66impl<F: FieldDecode> MessageDecoder<F> {
67    /// Makes a new `MessageDecoder` instance.
68    pub fn new(field_decoder: F) -> Self {
69        MessageDecoder {
70            tag: TagDecoder::default(),
71            field: field_decoder,
72            unknown_field: UnknownFieldDecoder::default(),
73            started: false,
74            eos: false,
75            target: DecodeTarget::None,
76        }
77    }
78}
79impl<F: FieldDecode> Decode for MessageDecoder<F> {
80    type Item = F::Item;
81
82    fn decode(&mut self, buf: &[u8], eos: Eos) -> Result<usize> {
83        self.started = true;
84        if self.eos {
85            return Ok(0);
86        }
87
88        let mut offset = 0;
89        while offset < buf.len() {
90            match self.target {
91                DecodeTarget::None | DecodeTarget::Tag => {
92                    let size = track!(self.tag.decode(&buf[offset..], eos))?;
93                    offset += size;
94                    if size != 0 {
95                        self.target = DecodeTarget::Tag;
96                    }
97                    if self.tag.is_idle() {
98                        let tag = track!(self.tag.finish_decoding())?;
99                        let started = track!(self.field.start_decoding(tag))?;
100                        if started {
101                            self.target = DecodeTarget::KnownField;
102                        } else {
103                            self.target = DecodeTarget::UnknownField;
104                            track!(self.unknown_field.start_decoding(tag))?;
105                        }
106                    }
107                }
108                DecodeTarget::KnownField => {
109                    bytecodec_try_decode!(self.field, offset, buf, eos);
110                    self.target = DecodeTarget::None;
111                }
112                DecodeTarget::UnknownField => {
113                    bytecodec_try_decode!(self.unknown_field, offset, buf, eos);
114                    track!(self.unknown_field.finish_decoding())?;
115                    self.target = DecodeTarget::None;
116                }
117            }
118        }
119        self.eos = eos.is_reached();
120        Ok(offset)
121    }
122
123    fn finish_decoding(&mut self) -> Result<Self::Item> {
124        track_assert!(!self.started | self.eos, ErrorKind::IncompleteDecoding; self.target, self.started, self.eos);
125        track_assert_eq!(
126            self.target,
127            DecodeTarget::None,
128            ErrorKind::IncompleteDecoding
129        );
130        let v = track!(self.field.finish_decoding())?;
131        self.started = false;
132        self.eos = false;
133        Ok(v)
134    }
135
136    fn requiring_bytes(&self) -> ByteCount {
137        if self.eos {
138            ByteCount::Finite(0)
139        } else {
140            match self.target {
141                DecodeTarget::None => ByteCount::Unknown,
142                DecodeTarget::Tag => self.tag.requiring_bytes(),
143                DecodeTarget::KnownField => self.field.requiring_bytes(),
144                DecodeTarget::UnknownField => self.unknown_field.requiring_bytes(),
145            }
146        }
147    }
148
149    fn is_idle(&self) -> bool {
150        self.eos
151    }
152}
153impl<F: FieldDecode> MessageDecode for MessageDecoder<F> {}
154
155#[derive(Debug, PartialEq, Eq)]
156enum DecodeTarget {
157    None,
158    Tag,
159    KnownField,
160    UnknownField,
161}
162impl Default for DecodeTarget {
163    fn default() -> Self {
164        DecodeTarget::None
165    }
166}
167
168/// Decoder for embedded messages.
169#[derive(Debug, Default)]
170pub(crate) struct EmbeddedMessageDecoder<M>(LengthDelimitedDecoder<M>);
171impl<M: MessageDecode> EmbeddedMessageDecoder<M> {
172    /// Makes a new `EmbeddedMessageDecoder` instance.
173    pub(crate) fn new(message_decoder: M) -> Self {
174        EmbeddedMessageDecoder(LengthDelimitedDecoder::new(message_decoder))
175    }
176}
177impl<M: MessageDecode> Decode for EmbeddedMessageDecoder<M> {
178    type Item = M::Item;
179
180    fn decode(&mut self, buf: &[u8], eos: Eos) -> Result<usize> {
181        track!(self.0.decode(buf, eos))
182    }
183
184    fn finish_decoding(&mut self) -> Result<Self::Item> {
185        track!(self.0.finish_decoding())
186    }
187
188    fn requiring_bytes(&self) -> ByteCount {
189        self.0.requiring_bytes()
190    }
191
192    fn is_idle(&self) -> bool {
193        self.0.is_idle()
194    }
195}
196impl<M: MessageDecode> ValueDecode for EmbeddedMessageDecoder<M> {
197    fn wire_type(&self) -> WireType {
198        WireType::LengthDelimited
199    }
200}
201
202/// Encoder for messages.
203#[derive(Debug, Default)]
204pub struct MessageEncoder<F> {
205    field: F,
206}
207impl<F: FieldEncode> MessageEncoder<F> {
208    /// Makes a new `MessageEncoder` instance.
209    pub fn new(field_encoder: F) -> Self {
210        MessageEncoder {
211            field: field_encoder,
212        }
213    }
214}
215impl<F: FieldEncode> Encode for MessageEncoder<F> {
216    type Item = F::Item;
217
218    fn encode(&mut self, buf: &mut [u8], eos: Eos) -> Result<usize> {
219        track!(self.field.encode(buf, eos))
220    }
221
222    fn start_encoding(&mut self, item: Self::Item) -> Result<()> {
223        track!(self.field.start_encoding(item))
224    }
225
226    fn is_idle(&self) -> bool {
227        self.field.is_idle()
228    }
229
230    fn requiring_bytes(&self) -> ByteCount {
231        self.field.requiring_bytes()
232    }
233}
234impl<F: FieldEncode + SizedEncode> SizedEncode for MessageEncoder<F> {
235    fn exact_requiring_bytes(&self) -> u64 {
236        self.field.exact_requiring_bytes()
237    }
238}
239impl<F: FieldEncode> MessageEncode for MessageEncoder<F> {}
240
241/// Encoder for embedded messages.
242#[derive(Debug, Default)]
243pub(crate) struct EmbeddedMessageEncoder<M> {
244    message: LengthDelimitedEncoder<M>,
245}
246impl<M: MessageEncode + SizedEncode> EmbeddedMessageEncoder<M> {
247    /// Makes a new `EmbeddedMessageEncoder` instance.
248    pub(crate) fn new(message_encoder: M) -> Self {
249        EmbeddedMessageEncoder {
250            message: LengthDelimitedEncoder::new(message_encoder),
251        }
252    }
253}
254impl<M: MessageEncode + SizedEncode> Encode for EmbeddedMessageEncoder<M> {
255    type Item = M::Item;
256
257    fn encode(&mut self, buf: &mut [u8], eos: Eos) -> Result<usize> {
258        track!(self.message.encode(buf, eos))
259    }
260
261    fn start_encoding(&mut self, item: Self::Item) -> Result<()> {
262        track!(self.message.start_encoding(item))
263    }
264
265    fn is_idle(&self) -> bool {
266        self.message.is_idle()
267    }
268
269    fn requiring_bytes(&self) -> ByteCount {
270        self.message.requiring_bytes()
271    }
272}
273impl<M: MessageEncode + SizedEncode> SizedEncode for EmbeddedMessageEncoder<M> {
274    fn exact_requiring_bytes(&self) -> u64 {
275        self.message.exact_requiring_bytes()
276    }
277}
278impl<M: MessageEncode + SizedEncode> ValueEncode for EmbeddedMessageEncoder<M> {
279    fn wire_type(&self) -> WireType {
280        WireType::LengthDelimited
281    }
282}