quick_protobuf/
writer.rs

1//! A module to manage protobuf serialization
2
3use crate::errors::{Error, Result};
4use crate::message::MessageWrite;
5
6use byteorder::{ByteOrder, LittleEndian as LE};
7
8#[cfg(feature = "std")]
9use byteorder::WriteBytesExt;
10
11/// A struct to write protobuf messages
12///
13/// # Examples
14///
15/// ```rust
16/// // an automatically generated module which is in a separate file in general
17/// mod foo_bar {
18///     # use quick_protobuf::{MessageWrite, Writer, WriterBackend, Result};
19///     # use std::borrow::Cow;
20///     pub struct Foo<'a> { pub name: Option<Cow<'a, str>>, }
21///     pub struct Bar { pub id: Option<u32> }
22///     pub struct FooBar<'a> { pub foos: Vec<Foo<'a>>, pub bars: Vec<Bar>, }
23///     impl<'a> MessageWrite for FooBar<'a> {
24///         // implements
25///         // fn get_size(&self) -> usize { ... }
26///         // fn write_message<W: WriterBackend>(&self, r: &mut Writer<W>) -> Result<()> { ... }
27///         # fn get_size(&self) -> usize { 0 }
28///         # fn write_message<W: WriterBackend>(&self, _: &mut Writer<W>) -> Result<()> { Ok(()) }
29///     }
30/// }
31///
32/// // FooBar is a message generated from a proto file
33/// // in parcicular it contains a `write_message` function
34/// use foo_bar::{FooBar, Foo, Bar};
35/// use std::borrow::Cow;
36/// use quick_protobuf::Writer;
37///
38/// fn main() {
39///     // let mut r = File::create("...").expect("Cannot create file");
40///     // for the sake of example, we'll use a simpler struct which impl `Write`
41///     let mut r = Vec::new();
42///     let mut writer = Writer::new(&mut r);
43///
44///     // manually generates a FooBar for the example
45///     let foobar = FooBar {
46///         foos: vec![Foo { name: Some(Cow::Borrowed("test!")) }, Foo { name: None }],
47///         bars: vec![Bar { id: Some(43) }, Bar { id: None }],
48///     };
49///
50///     // now using the generated module
51///     writer.write_message(&foobar).expect("Cannot write FooBar");
52/// }
53/// ```
54pub struct Writer<W: WriterBackend> {
55    inner: W,
56}
57
58impl<W: WriterBackend> Writer<W> {
59    /// Creates a new `ProtobufWriter`
60    pub fn new(w: W) -> Writer<W> {
61        Writer { inner: w }
62    }
63
64    /// Writes a byte which is NOT internally coded as a `varint`
65    pub fn write_u8(&mut self, byte: u8) -> Result<()> {
66        self.inner.pb_write_u8(byte)
67    }
68
69    /// Writes a `varint` (compacted `u64`)
70    pub fn write_varint(&mut self, mut v: u64) -> Result<()> {
71        while v > 0x7F {
72            self.inner.pb_write_u8(((v as u8) & 0x7F) | 0x80)?;
73            v >>= 7;
74        }
75        self.inner.pb_write_u8(v as u8)
76    }
77
78    /// Writes a tag, which represents both the field number and the wire type
79    #[cfg_attr(std, inline(always))]
80    pub fn write_tag(&mut self, tag: u32) -> Result<()> {
81        self.write_varint(tag as u64)
82    }
83
84    /// Writes a `int32` which is internally coded as a `varint`
85    #[cfg_attr(std, inline(always))]
86    pub fn write_int32(&mut self, v: i32) -> Result<()> {
87        self.write_varint(v as u64)
88    }
89
90    /// Writes a `int64` which is internally coded as a `varint`
91    #[cfg_attr(std, inline(always))]
92    pub fn write_int64(&mut self, v: i64) -> Result<()> {
93        self.write_varint(v as u64)
94    }
95
96    /// Writes a `uint32` which is internally coded as a `varint`
97    #[cfg_attr(std, inline(always))]
98    pub fn write_uint32(&mut self, v: u32) -> Result<()> {
99        self.write_varint(v as u64)
100    }
101
102    /// Writes a `uint64` which is internally coded as a `varint`
103    #[cfg_attr(std, inline(always))]
104    pub fn write_uint64(&mut self, v: u64) -> Result<()> {
105        self.write_varint(v)
106    }
107
108    /// Writes a `sint32` which is internally coded as a `varint`
109    #[cfg_attr(std, inline(always))]
110    pub fn write_sint32(&mut self, v: i32) -> Result<()> {
111        self.write_varint(((v << 1) ^ (v >> 31)) as u64)
112    }
113
114    /// Writes a `sint64` which is internally coded as a `varint`
115    #[cfg_attr(std, inline(always))]
116    pub fn write_sint64(&mut self, v: i64) -> Result<()> {
117        self.write_varint(((v << 1) ^ (v >> 63)) as u64)
118    }
119
120    /// Writes a `fixed64` which is little endian coded `u64`
121    #[cfg_attr(std, inline(always))]
122    pub fn write_fixed64(&mut self, v: u64) -> Result<()> {
123        self.inner.pb_write_u64(v)
124    }
125
126    /// Writes a `fixed32` which is little endian coded `u32`
127    #[cfg_attr(std, inline(always))]
128    pub fn write_fixed32(&mut self, v: u32) -> Result<()> {
129        self.inner.pb_write_u32(v)
130    }
131
132    /// Writes a `sfixed64` which is little endian coded `i64`
133    #[cfg_attr(std, inline(always))]
134    pub fn write_sfixed64(&mut self, v: i64) -> Result<()> {
135        self.inner.pb_write_i64(v)
136    }
137
138    /// Writes a `sfixed32` which is little endian coded `i32`
139    #[cfg_attr(std, inline(always))]
140    pub fn write_sfixed32(&mut self, v: i32) -> Result<()> {
141        self.inner.pb_write_i32(v)
142    }
143
144    /// Writes a `float`
145    #[cfg_attr(std, inline(always))]
146    pub fn write_float(&mut self, v: f32) -> Result<()> {
147        self.inner.pb_write_f32(v)
148    }
149
150    /// Writes a `double`
151    #[cfg_attr(std, inline(always))]
152    pub fn write_double(&mut self, v: f64) -> Result<()> {
153        self.inner.pb_write_f64(v)
154    }
155
156    /// Writes a `bool` 1 = true, 0 = false
157    #[cfg_attr(std, inline(always))]
158    pub fn write_bool(&mut self, v: bool) -> Result<()> {
159        self.inner.pb_write_u8(if v { 1 } else { 0 })
160    }
161
162    /// Writes an `enum` converting it to a `i32` first
163    #[cfg_attr(std, inline(always))]
164    pub fn write_enum(&mut self, v: i32) -> Result<()> {
165        self.write_int32(v)
166    }
167
168    /// Writes `bytes`: length first then the chunk of data
169    #[cfg_attr(std, inline(always))]
170    pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<()> {
171        self.write_varint(bytes.len() as u64)?;
172        self.inner.pb_write_all(bytes)
173    }
174
175    /// Writes `string`: length first then the chunk of data
176    #[cfg_attr(std, inline(always))]
177    pub fn write_string(&mut self, s: &str) -> Result<()> {
178        self.write_bytes(s.as_bytes())
179    }
180
181    /// Writes packed repeated field: length first then the chunk of data
182    pub fn write_packed<M, F, S>(&mut self, v: &[M], mut write: F, size: &S) -> Result<()>
183    where
184        F: FnMut(&mut Self, &M) -> Result<()>,
185        S: Fn(&M) -> usize,
186    {
187        if v.is_empty() {
188            return Ok(());
189        }
190        let len: usize = v.iter().map(|m| size(m)).sum();
191        self.write_varint(len as u64)?;
192        for m in v {
193            write(self, m)?;
194        }
195        Ok(())
196    }
197
198    /// Writes packed repeated field when we know the size of items
199    ///
200    /// `item_size` is internally used to compute the total length
201    /// As the length is fixed (and the same as rust internal representation, we can directly dump
202    /// all data at once
203    #[cfg_attr(std, inline)]
204    pub fn write_packed_fixed<M>(&mut self, v: &[M]) -> Result<()> {
205        let len = v.len() * ::core::mem::size_of::<M>();
206        let bytes = unsafe { ::core::slice::from_raw_parts(v.as_ptr() as *const u8, len) };
207        self.write_bytes(bytes)
208    }
209
210    /// Writes a message which implements `MessageWrite`
211    #[cfg_attr(std, inline)]
212    pub fn write_message<M: MessageWrite>(&mut self, m: &M) -> Result<()> {
213        let len = m.get_size();
214        self.write_varint(len as u64)?;
215        m.write_message(self)
216    }
217
218    /// Writes another item prefixed with tag
219    #[cfg_attr(std, inline)]
220    pub fn write_with_tag<F>(&mut self, tag: u32, mut write: F) -> Result<()>
221    where
222        F: FnMut(&mut Self) -> Result<()>,
223    {
224        self.write_tag(tag)?;
225        write(self)
226    }
227
228    /// Writes tag then repeated field
229    ///
230    /// If array is empty, then do nothing (do not even write the tag)
231    pub fn write_packed_with_tag<M, F, S>(
232        &mut self,
233        tag: u32,
234        v: &[M],
235        mut write: F,
236        size: &S,
237    ) -> Result<()>
238    where
239        F: FnMut(&mut Self, &M) -> Result<()>,
240        S: Fn(&M) -> usize,
241    {
242        if v.is_empty() {
243            return Ok(());
244        }
245
246        self.write_tag(tag)?;
247        let len: usize = v.iter().map(|m| size(m)).sum();
248        self.write_varint(len as u64)?;
249        for m in v {
250            write(self, m)?;
251        }
252        Ok(())
253    }
254
255    /// Writes tag then repeated field
256    ///
257    /// If array is empty, then do nothing (do not even write the tag)
258    pub fn write_packed_fixed_with_tag<M>(&mut self, tag: u32, v: &[M]) -> Result<()> {
259        if v.is_empty() {
260            return Ok(());
261        }
262
263        self.write_tag(tag)?;
264        let len = ::core::mem::size_of::<M>() * v.len();
265        let bytes = unsafe { ::core::slice::from_raw_parts(v.as_ptr() as *const u8, len) };
266        self.write_bytes(bytes)
267    }
268
269    /// Writes tag then repeated field with fixed length item size
270    ///
271    /// If array is empty, then do nothing (do not even write the tag)
272    pub fn write_packed_fixed_size_with_tag<M>(
273        &mut self,
274        tag: u32,
275        v: &[M],
276        item_size: usize,
277    ) -> Result<()> {
278        if v.is_empty() {
279            return Ok(());
280        }
281        self.write_tag(tag)?;
282        let len = v.len() * item_size;
283        let bytes =
284            unsafe { ::core::slice::from_raw_parts(v as *const [M] as *const M as *const u8, len) };
285        self.write_bytes(bytes)
286    }
287
288    /// Write entire map
289    pub fn write_map<FK, FV>(
290        &mut self,
291        size: usize,
292        tag_key: u32,
293        mut write_key: FK,
294        tag_val: u32,
295        mut write_val: FV,
296    ) -> Result<()>
297    where
298        FK: FnMut(&mut Self) -> Result<()>,
299        FV: FnMut(&mut Self) -> Result<()>,
300    {
301        self.write_varint(size as u64)?;
302        self.write_tag(tag_key)?;
303        write_key(self)?;
304        self.write_tag(tag_val)?;
305        write_val(self)
306    }
307}
308
309/// Serialize a `MessageWrite` into a `Vec`
310#[cfg(feature = "std")]
311pub fn serialize_into_vec<M: MessageWrite>(message: &M) -> Result<Vec<u8>> {
312    let len = message.get_size();
313    let mut v = Vec::with_capacity(len + crate::sizeofs::sizeof_len(len));
314    {
315        let mut writer = Writer::new(&mut v);
316        writer.write_message(message)?;
317    }
318    Ok(v)
319}
320
321/// Serialize a `MessageWrite` into a byte slice
322pub fn serialize_into_slice<M: MessageWrite>(message: &M, out: &mut [u8]) -> Result<()> {
323    let len = message.get_size();
324    if out.len() < len {
325        return Err(Error::OutputBufferTooSmall);
326    }
327    {
328        let mut writer = Writer::new(BytesWriter::new(out));
329        writer.write_message(message)?;
330    }
331
332    Ok(())
333}
334
335/// Writer backend abstraction
336pub trait WriterBackend {
337    /// Write a u8
338    fn pb_write_u8(&mut self, x: u8) -> Result<()>;
339
340    /// Write a u32
341    fn pb_write_u32(&mut self, x: u32) -> Result<()>;
342
343    /// Write a i32
344    fn pb_write_i32(&mut self, x: i32) -> Result<()>;
345
346    /// Write a f32
347    fn pb_write_f32(&mut self, x: f32) -> Result<()>;
348
349    /// Write a u64
350    fn pb_write_u64(&mut self, x: u64) -> Result<()>;
351
352    /// Write a i64
353    fn pb_write_i64(&mut self, x: i64) -> Result<()>;
354
355    /// Write a f64
356    fn pb_write_f64(&mut self, x: f64) -> Result<()>;
357
358    /// Write all bytes in buf
359    fn pb_write_all(&mut self, buf: &[u8]) -> Result<()>;
360}
361
362/// A writer backend for byte buffers
363pub struct BytesWriter<'a> {
364    buf: &'a mut [u8],
365    cursor: usize,
366}
367
368impl<'a> BytesWriter<'a> {
369    /// Create a new BytesWriter to write into `buf`
370    pub fn new(buf: &'a mut [u8]) -> BytesWriter<'a> {
371        BytesWriter { buf, cursor: 0 }
372    }
373}
374
375impl<'a> WriterBackend for BytesWriter<'a> {
376    #[cfg_attr(std, inline(always))]
377    fn pb_write_u8(&mut self, x: u8) -> Result<()> {
378        if self.buf.len() - self.cursor < 1 {
379            Err(Error::UnexpectedEndOfBuffer)
380        } else {
381            self.buf[self.cursor] = x;
382            self.cursor += 1;
383            Ok(())
384        }
385    }
386
387    #[cfg_attr(std, inline(always))]
388    fn pb_write_u32(&mut self, x: u32) -> Result<()> {
389        if self.buf.len() - self.cursor < 4 {
390            Err(Error::UnexpectedEndOfBuffer)
391        } else {
392            LE::write_u32(&mut self.buf[self.cursor..], x);
393            self.cursor += 4;
394            Ok(())
395        }
396    }
397
398    #[cfg_attr(std, inline(always))]
399    fn pb_write_i32(&mut self, x: i32) -> Result<()> {
400        if self.buf.len() - self.cursor < 4 {
401            Err(Error::UnexpectedEndOfBuffer)
402        } else {
403            LE::write_i32(&mut self.buf[self.cursor..], x);
404            self.cursor += 4;
405            Ok(())
406        }
407    }
408
409    #[cfg_attr(std, inline(always))]
410    fn pb_write_f32(&mut self, x: f32) -> Result<()> {
411        if self.buf.len() - self.cursor < 4 {
412            Err(Error::UnexpectedEndOfBuffer)
413        } else {
414            LE::write_f32(&mut self.buf[self.cursor..], x);
415            self.cursor += 4;
416            Ok(())
417        }
418    }
419
420    #[cfg_attr(std, inline(always))]
421    fn pb_write_u64(&mut self, x: u64) -> Result<()> {
422        if self.buf.len() - self.cursor < 8 {
423            Err(Error::UnexpectedEndOfBuffer)
424        } else {
425            LE::write_u64(&mut self.buf[self.cursor..], x);
426            self.cursor += 8;
427            Ok(())
428        }
429    }
430
431    #[cfg_attr(std, inline(always))]
432    fn pb_write_i64(&mut self, x: i64) -> Result<()> {
433        if self.buf.len() - self.cursor < 8 {
434            Err(Error::UnexpectedEndOfBuffer)
435        } else {
436            LE::write_i64(&mut self.buf[self.cursor..], x);
437            self.cursor += 8;
438            Ok(())
439        }
440    }
441
442    #[cfg_attr(std, inline(always))]
443    fn pb_write_f64(&mut self, x: f64) -> Result<()> {
444        if self.buf.len() - self.cursor < 8 {
445            Err(Error::UnexpectedEndOfBuffer)
446        } else {
447            LE::write_f64(&mut self.buf[self.cursor..], x);
448            self.cursor += 8;
449            Ok(())
450        }
451    }
452
453    #[cfg_attr(std, inline(always))]
454    fn pb_write_all(&mut self, buf: &[u8]) -> Result<()> {
455        if self.buf.len() - self.cursor < buf.len() {
456            Err(Error::UnexpectedEndOfBuffer)
457        } else {
458            self.buf[self.cursor..(self.cursor + buf.len())].copy_from_slice(buf);
459            self.cursor += buf.len();
460            Ok(())
461        }
462    }
463}
464
465#[cfg(feature = "std")]
466impl<W: std::io::Write> WriterBackend for W {
467    #[inline(always)]
468    fn pb_write_u8(&mut self, x: u8) -> Result<()> {
469        self.write_u8(x).map_err(|e| e.into())
470    }
471
472    #[inline(always)]
473    fn pb_write_u32(&mut self, x: u32) -> Result<()> {
474        self.write_u32::<LE>(x).map_err(|e| e.into())
475    }
476
477    #[inline(always)]
478    fn pb_write_i32(&mut self, x: i32) -> Result<()> {
479        self.write_i32::<LE>(x).map_err(|e| e.into())
480    }
481
482    #[inline(always)]
483    fn pb_write_f32(&mut self, x: f32) -> Result<()> {
484        self.write_f32::<LE>(x).map_err(|e| e.into())
485    }
486
487    #[inline(always)]
488    fn pb_write_u64(&mut self, x: u64) -> Result<()> {
489        self.write_u64::<LE>(x).map_err(|e| e.into())
490    }
491
492    #[inline(always)]
493    fn pb_write_i64(&mut self, x: i64) -> Result<()> {
494        self.write_i64::<LE>(x).map_err(|e| e.into())
495    }
496
497    #[inline(always)]
498    fn pb_write_f64(&mut self, x: f64) -> Result<()> {
499        self.write_f64::<LE>(x).map_err(|e| e.into())
500    }
501
502    #[inline(always)]
503    fn pb_write_all(&mut self, buf: &[u8]) -> Result<()> {
504        self.write_all(buf).map_err(|e| e.into())
505    }
506}