RUST-2258 Use borrowed buffer in raw serializer (#590)

Co-authored-by: Isabel Atkinson <isabel.atkinson@mongodb.com>

Author: LeWimbes

src/lib.rs

@@ -474,6 +474,7 @@ pub use self::{
     },
     ser::{
         serialize_to_bson,
+        serialize_to_buffer,
         serialize_to_document,
         serialize_to_raw_document_buf,
         serialize_to_vec,

src/ser.rs

@@ -89,7 +89,19 @@ pub fn serialize_to_vec<T>(value: &T) -> Result<Vec<u8>>
 where
     T: Serialize,
 {
-    let mut serializer = raw::Serializer::new();
+    let mut bytes = Vec::new();
+    serialize_to_buffer(value, &mut bytes)?;
+    Ok(bytes)
+}
+
+/// Serialize the given `T` as a BSON byte vector into the provided byte buffer.
+/// This allows reusing the same buffer for multiple serializations.
+#[inline]
+pub fn serialize_to_buffer<T>(value: &T, buffer: &mut Vec<u8>) -> Result<()>
+where
+    T: Serialize,
+{
+    let mut serializer = raw::Serializer::new(buffer);
     #[cfg(feature = "serde_path_to_error")]
     {
         serde_path_to_error::serialize(value, &mut serializer).map_err(Error::with_path)?;
@@ -98,7 +110,7 @@ where
     {
         value.serialize(&mut serializer)?;
     }
-    Ok(serializer.into_vec())
+    Ok(())
 }
 
 /// Serialize the given `T` as a [`RawDocumentBuf`].

src/ser/raw.rs

@@ -22,8 +22,11 @@ use crate::{
 use document_serializer::DocumentSerializer;
 
 /// Serializer used to convert a type `T` into raw BSON bytes.
-pub(crate) struct Serializer {
-    bytes: Vec<u8>,
+pub(crate) struct Serializer<'a> {
+    bytes: &'a mut Vec<u8>,
+
+    /// The index into `bytes` where the current serialization started.
+    start_index: usize,
 
     /// The index into `bytes` where the current element type will need to be stored.
     /// This needs to be set retroactively because in BSON, the element type comes before the key,
@@ -58,21 +61,18 @@ impl SerializerHint {
     }
 }
 
-impl Serializer {
-    pub(crate) fn new() -> Self {
+impl<'a> Serializer<'a> {
+    pub(crate) fn new(bytes: &'a mut Vec<u8>) -> Self {
+        let start_index = bytes.len();
         Self {
-            bytes: Vec::new(),
-            type_index: 0,
+            bytes,
+            start_index,
+            type_index: start_index,
             hint: SerializerHint::None,
             human_readable: false,
         }
     }
 
-    /// Convert this serializer into the vec of the serialized bytes.
-    pub(crate) fn into_vec(self) -> Vec<u8> {
-        self.bytes
-    }
-
     /// Reserve a spot for the element type to be set retroactively via `update_element_type`.
     #[inline]
     fn reserve_element_type(&mut self) {
@@ -83,7 +83,7 @@ impl Serializer {
     /// Retroactively set the element type of the most recently serialized element.
     #[inline]
     fn update_element_type(&mut self, t: ElementType) -> Result<()> {
-        if self.type_index == 0 {
+        if self.type_index == self.start_index {
             if matches!(t, ElementType::EmbeddedDocument) {
                 // don't need to set the element type for the top level document
                 return Ok(());
@@ -108,22 +108,22 @@ impl Serializer {
 
     fn serialize_raw(&mut self, v: RawBsonRef) -> Result<()> {
         self.update_element_type(v.element_type())?;
-        v.append_to(&mut self.bytes);
+        v.append_to(self.bytes);
         Ok(())
     }
 }
 
-impl<'a> serde::Serializer for &'a mut Serializer {
+impl<'a, 'b> serde::Serializer for &'a mut Serializer<'b> {
     type Ok = ();
     type Error = Error;
 
-    type SerializeSeq = DocumentSerializer<'a>;
-    type SerializeTuple = DocumentSerializer<'a>;
-    type SerializeTupleStruct = DocumentSerializer<'a>;
-    type SerializeTupleVariant = VariantSerializer<'a>;
-    type SerializeMap = DocumentSerializer<'a>;
-    type SerializeStruct = StructSerializer<'a>;
-    type SerializeStructVariant = VariantSerializer<'a>;
+    type SerializeSeq = DocumentSerializer<'a, 'b>;
+    type SerializeTuple = DocumentSerializer<'a, 'b>;
+    type SerializeTupleStruct = DocumentSerializer<'a, 'b>;
+    type SerializeTupleVariant = VariantSerializer<'a, 'b>;
+    type SerializeMap = DocumentSerializer<'a, 'b>;
+    type SerializeStruct = StructSerializer<'a, 'b>;
+    type SerializeStructVariant = VariantSerializer<'a, 'b>;
 
     fn is_human_readable(&self) -> bool {
         self.human_readable
@@ -385,15 +385,15 @@ impl<'a> serde::Serializer for &'a mut Serializer {
     }
 }
 
-pub(crate) enum StructSerializer<'a> {
+pub(crate) enum StructSerializer<'a, 'b> {
     /// Serialize a BSON value currently represented in serde as a struct (e.g. ObjectId)
-    Value(ValueSerializer<'a>),
+    Value(ValueSerializer<'a, 'b>),
 
     /// Serialize the struct as a document.
-    Document(DocumentSerializer<'a>),
+    Document(DocumentSerializer<'a, 'b>),
 }
 
-impl SerializeStruct for StructSerializer<'_> {
+impl SerializeStruct for StructSerializer<'_, '_> {
     type Ok = ();
     type Error = Error;
 
@@ -424,8 +424,8 @@ enum VariantInnerType {
 
 /// Serializer used for enum variants, including both tuple (e.g. Foo::Bar(1, 2, 3)) and
 /// struct (e.g. Foo::Bar { a: 1 }).
-pub(crate) struct VariantSerializer<'a> {
-    root_serializer: &'a mut Serializer,
+pub(crate) struct VariantSerializer<'a, 'b> {
+    root_serializer: &'a mut Serializer<'b>,
 
     /// Variants are serialized as documents of the form `{ <variant name>: <document or array> }`,
     /// and `doc_start` indicates the index at which the outer document begins.
@@ -438,22 +438,26 @@ pub(crate) struct VariantSerializer<'a> {
     num_elements_serialized: usize,
 }
 
-impl<'a> VariantSerializer<'a> {
-    fn start(rs: &'a mut Serializer, variant: &'static CStr, inner_type: VariantInnerType) -> Self {
+impl<'a, 'b> VariantSerializer<'a, 'b> {
+    fn start(
+        rs: &'a mut Serializer<'b>,
+        variant: &'static CStr,
+        inner_type: VariantInnerType,
+    ) -> Self {
         let doc_start = rs.bytes.len();
         // write placeholder length for document, will be updated at end
         static ZERO: RawBsonRef = RawBsonRef::Int32(0);
-        ZERO.append_to(&mut rs.bytes);
+        ZERO.append_to(rs.bytes);
 
         let inner = match inner_type {
             VariantInnerType::Struct => ElementType::EmbeddedDocument,
             VariantInnerType::Tuple => ElementType::Array,
         };
         rs.bytes.push(inner as u8);
-        variant.append_to(&mut rs.bytes);
+        variant.append_to(rs.bytes);
         let inner_start = rs.bytes.len();
         // write placeholder length for inner, will be updated at end
-        ZERO.append_to(&mut rs.bytes);
+        ZERO.append_to(rs.bytes);
 
         Self {
             root_serializer: rs,
@@ -469,7 +473,7 @@ impl<'a> VariantSerializer<'a> {
         T: Serialize + ?Sized,
     {
         self.root_serializer.reserve_element_type();
-        CStr::from_str(k)?.append_to(&mut self.root_serializer.bytes);
+        CStr::from_str(k)?.append_to(self.root_serializer.bytes);
         v.serialize(&mut *self.root_serializer)?;
 
         self.num_elements_serialized += 1;
@@ -492,7 +496,7 @@ impl<'a> VariantSerializer<'a> {
     }
 }
 
-impl serde::ser::SerializeTupleVariant for VariantSerializer<'_> {
+impl serde::ser::SerializeTupleVariant for VariantSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -511,7 +515,7 @@ impl serde::ser::SerializeTupleVariant for VariantSerializer<'_> {
     }
 }
 
-impl serde::ser::SerializeStructVariant for VariantSerializer<'_> {
+impl serde::ser::SerializeStructVariant for VariantSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;

src/ser/raw/document_serializer.rs

@@ -7,21 +7,21 @@ use crate::{
 
 use super::Serializer;
 
-pub(crate) struct DocumentSerializationResult<'a> {
-    pub(crate) root_serializer: &'a mut Serializer,
+pub(crate) struct DocumentSerializationResult<'a, 'b> {
+    pub(crate) root_serializer: &'a mut Serializer<'b>,
 }
 
 /// Serializer used to serialize document or array bodies.
-pub(crate) struct DocumentSerializer<'a> {
-    root_serializer: &'a mut Serializer,
+pub(crate) struct DocumentSerializer<'a, 'b> {
+    root_serializer: &'a mut Serializer<'b>,
     num_keys_serialized: usize,
     start: usize,
 }
 
-impl<'a> DocumentSerializer<'a> {
-    pub(crate) fn start(rs: &'a mut Serializer) -> Self {
+impl<'a, 'b> DocumentSerializer<'a, 'b> {
+    pub(crate) fn start(rs: &'a mut Serializer<'b>) -> Self {
         let start = rs.bytes.len();
-        RawBsonRef::Int32(0).append_to(&mut rs.bytes);
+        RawBsonRef::Int32(0).append_to(rs.bytes);
         Self {
             root_serializer: rs,
             num_keys_serialized: 0,
@@ -30,7 +30,7 @@ impl<'a> DocumentSerializer<'a> {
     }
 
     /// Serialize a document key using the provided closure.
-    fn serialize_doc_key_custom<F: FnOnce(&mut Serializer) -> Result<()>>(
+    fn serialize_doc_key_custom<F: FnOnce(&mut Serializer<'b>) -> Result<()>>(
         &mut self,
         f: F,
     ) -> Result<()> {
@@ -55,7 +55,7 @@ impl<'a> DocumentSerializer<'a> {
         Ok(())
     }
 
-    pub(crate) fn end_doc(self) -> crate::ser::Result<DocumentSerializationResult<'a>> {
+    pub(crate) fn end_doc(self) -> crate::ser::Result<DocumentSerializationResult<'a, 'b>> {
         self.root_serializer.bytes.push(0);
         let length = (self.root_serializer.bytes.len() - self.start) as i32;
         self.root_serializer.replace_i32(self.start, length);
@@ -65,7 +65,7 @@ impl<'a> DocumentSerializer<'a> {
     }
 }
 
-impl serde::ser::SerializeSeq for DocumentSerializer<'_> {
+impl serde::ser::SerializeSeq for DocumentSerializer<'_, '_> {
     type Ok = ();
     type Error = Error;
 
@@ -90,7 +90,7 @@ impl serde::ser::SerializeSeq for DocumentSerializer<'_> {
     }
 }
 
-impl serde::ser::SerializeMap for DocumentSerializer<'_> {
+impl serde::ser::SerializeMap for DocumentSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -116,7 +116,7 @@ impl serde::ser::SerializeMap for DocumentSerializer<'_> {
     }
 }
 
-impl serde::ser::SerializeStruct for DocumentSerializer<'_> {
+impl serde::ser::SerializeStruct for DocumentSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -136,7 +136,7 @@ impl serde::ser::SerializeStruct for DocumentSerializer<'_> {
     }
 }
 
-impl serde::ser::SerializeTuple for DocumentSerializer<'_> {
+impl serde::ser::SerializeTuple for DocumentSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -156,7 +156,7 @@ impl serde::ser::SerializeTuple for DocumentSerializer<'_> {
     }
 }
 
-impl serde::ser::SerializeTupleStruct for DocumentSerializer<'_> {
+impl serde::ser::SerializeTupleStruct for DocumentSerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -178,11 +178,11 @@ impl serde::ser::SerializeTupleStruct for DocumentSerializer<'_> {
 
 /// Serializer used specifically for serializing document keys.
 /// Only keys that serialize to strings will be accepted.
-struct KeySerializer<'a> {
-    root_serializer: &'a mut Serializer,
+struct KeySerializer<'a, 'b> {
+    root_serializer: &'a mut Serializer<'b>,
 }
 
-impl serde::Serializer for KeySerializer<'_> {
+impl serde::Serializer for KeySerializer<'_, '_> {
     type Ok = ();
 
     type Error = Error;
@@ -257,7 +257,7 @@ impl serde::Serializer for KeySerializer<'_> {
 
     #[inline]
     fn serialize_str(self, v: &str) -> Result<Self::Ok> {
-        crate::raw::CStr::from_str(v)?.append_to(&mut self.root_serializer.bytes);
+        crate::raw::CStr::from_str(v)?.append_to(self.root_serializer.bytes);
         Ok(())
     }