RUST-2239 Fix performance of bson 3.0

src/bson.rs

@@ -29,7 +29,14 @@ use std::{
 };
 
 pub use crate::document::Document;
-use crate::{oid, raw::CString, spec::ElementType, Binary, Decimal128};
+use crate::{
+    oid,
+    raw::{doc_writer::DocWriter, CString},
+    spec::ElementType,
+    Binary,
+    Decimal128,
+    RawBsonRef,
+};
 
 /// Possible BSON value types.
 #[derive(Clone, Default, PartialEq)]
@@ -815,6 +822,57 @@ impl Bson {
             Bson::Decimal128(_) => Unexpected::Other("decimal128"),
         }
     }
+
+    pub(crate) fn append_to(&self, buf: &mut Vec<u8>) -> crate::error::Result<()> {
+        match self {
+            Self::Int32(val) => RawBsonRef::Int32(*val).append_to(buf),
+            Self::Int64(val) => RawBsonRef::Int64(*val).append_to(buf),
+            Self::Double(val) => RawBsonRef::Double(*val).append_to(buf),
+            Self::Binary(bin) => RawBsonRef::Binary(crate::RawBinaryRef {
+                subtype: bin.subtype,
+                bytes: &bin.bytes,
+            })
+            .append_to(buf),
+            Self::String(s) => RawBsonRef::String(s).append_to(buf),
+            Self::Array(arr) => {
+                let mut writer = DocWriter::open(buf);
+                for (ix, v) in arr.iter().enumerate() {
+                    writer.append_key(
+                        v.element_type(),
+                        &crate::raw::CString::from_string_unchecked(ix.to_string()),
+                    );
+                    v.append_to(writer.buffer())?;
+                }
+            }
+            Self::Document(doc) => doc.append_to(buf)?,
+            Self::Boolean(b) => RawBsonRef::Boolean(*b).append_to(buf),
+            Self::RegularExpression(re) => RawBsonRef::RegularExpression(crate::RawRegexRef {
+                pattern: &re.pattern,
+                options: &re.options,
+            })
+            .append_to(buf),
+            Self::JavaScriptCode(js) => RawBsonRef::JavaScriptCode(js).append_to(buf),
+            Self::JavaScriptCodeWithScope(cws) => {
+                let start = buf.len();
+                buf.extend(0i32.to_le_bytes()); // placeholder
+                RawBsonRef::String(&cws.code).append_to(buf);
+                cws.scope.append_to(buf)?;
+                let len: i32 = (buf.len() - start) as i32;
+                buf[start..start + 4].copy_from_slice(&len.to_le_bytes());
+            }
+            Self::Timestamp(ts) => RawBsonRef::Timestamp(*ts).append_to(buf),
+            Self::ObjectId(oid) => RawBsonRef::ObjectId(*oid).append_to(buf),
+            Self::DateTime(dt) => RawBsonRef::DateTime(*dt).append_to(buf),
+            Self::Symbol(s) => RawBsonRef::Symbol(s).append_to(buf),
+            Self::Decimal128(d) => RawBsonRef::Decimal128(*d).append_to(buf),
+            Self::DbPointer(dbp) => {
+                RawBsonRef::String(&dbp.namespace).append_to(buf);
+                RawBsonRef::ObjectId(dbp.id).append_to(buf);
+            }
+            Self::Null | Self::Undefined | Self::MinKey | Self::MaxKey => {}
+        }
+        Ok(())
+    }
 }
 
 /// Value helpers

src/document.rs

@@ -732,6 +732,15 @@ impl Document {
         let raw = crate::raw::RawDocumentBuf::from_reader(reader)?;
         raw.try_into()
     }
+
+    pub(crate) fn append_to(&self, buf: &mut Vec<u8>) -> crate::error::Result<()> {
+        let mut writer = crate::raw::doc_writer::DocWriter::open(buf);
+        for (k, v) in self {
+            writer.append_key(v.element_type(), k.as_str().try_into()?);
+            v.append_to(writer.buffer())?;
+        }
+        Ok(())
+    }
 }
 
 /// A view into a single entry in a document, which may either be vacant or occupied.

src/macros.rs

@@ -439,6 +439,7 @@ macro_rules! rawdoc {
 ///
 /// This macro generates a `SerializeAs`/`DeserializeAs` implementation for a given type,
 /// with optional struct-level attributes like `#[derive(...)]` or `/// doc comments`.
+#[allow(unused_macros)]
 macro_rules! serde_conv_doc {
     ($(#[$meta:meta])* $vis:vis $m:ident, $t:ty, $ser:expr, $de:expr) => {
         #[allow(non_camel_case_types)]
@@ -490,4 +491,5 @@ macro_rules! serde_conv_doc {
     };
 }
 
+#[allow(unused_imports)]
 pub(crate) use serde_conv_doc;

src/raw.rs

@@ -129,6 +129,7 @@ mod array_buf;
 mod bson;
 mod bson_ref;
 mod cstr;
+pub(crate) mod doc_writer;
 mod document;
 mod document_buf;
 mod iter;

src/raw/doc_writer.rs

@@ -0,0 +1,35 @@
+use crate::{raw::CStr, spec::ElementType};
+
+pub(crate) struct DocWriter<'a> {
+    data: &'a mut Vec<u8>,
+    start: usize,
+}
+
+impl<'a> DocWriter<'a> {
+    pub(crate) fn open(data: &'a mut Vec<u8>) -> Self {
+        let start = data.len();
+        data.extend(crate::raw::MIN_BSON_DOCUMENT_SIZE.to_le_bytes());
+        Self { data, start }
+    }
+
+    pub(crate) fn resume(data: &'a mut Vec<u8>, start: usize) -> Self {
+        Self { data, start }
+    }
+
+    pub(crate) fn append_key(&mut self, element_type: ElementType, name: &CStr) {
+        self.data.push(element_type as u8);
+        name.append_to(self.data);
+    }
+
+    pub(crate) fn buffer(&mut self) -> &mut Vec<u8> {
+        self.data
+    }
+}
+
+impl<'a> Drop for DocWriter<'a> {
+    fn drop(&mut self) {
+        self.data.push(0);
+        let new_len = ((self.data.len() - self.start) as i32).to_le_bytes();
+        self.data[self.start..self.start + 4].copy_from_slice(&new_len);
+    }
+}

src/raw/document_buf.rs

@@ -5,14 +5,12 @@ use std::{
 };
 
 use crate::{
-    raw::{CStr, MIN_BSON_DOCUMENT_SIZE},
+    raw::{doc_writer::DocWriter, CStr},
     Document,
 };
 
 use super::{bson::RawBson, iter::Iter, RawBsonRef, RawDocument, RawIter, Result};
 
-mod raw_writer;
-
 /// An owned BSON document (akin to [`std::path::PathBuf`]), backed by a buffer of raw BSON bytes.
 /// This can be created from a `Vec<u8>` or a [`crate::Document`].
 ///
@@ -58,8 +56,7 @@ impl RawDocumentBuf {
     /// Creates a new, empty [`RawDocumentBuf`].
     pub fn new() -> Self {
         let mut data = Vec::new();
-        data.extend(MIN_BSON_DOCUMENT_SIZE.to_le_bytes());
-        data.push(0);
+        DocWriter::open(&mut data);
         Self { data }
     }
 
@@ -200,8 +197,13 @@ impl RawDocumentBuf {
     /// # Ok::<(), Error>(())
     /// ```
     pub fn append(&mut self, key: impl AsRef<CStr>, value: impl BindRawBsonRef) {
+        self.data.pop();
         let key = key.as_ref();
-        value.bind(|value_ref| raw_writer::RawWriter::new(&mut self.data).append(key, value_ref));
+        value.bind(|value_ref| {
+            let mut writer = DocWriter::resume(&mut self.data, 0);
+            writer.append_key(value_ref.element_type(), key);
+            value_ref.append_to(writer.buffer());
+        });
     }
 }
 
@@ -266,27 +268,17 @@ impl TryFrom<&Document> for RawDocumentBuf {
     type Error = crate::error::Error;
 
     fn try_from(doc: &Document) -> std::result::Result<Self, Self::Error> {
-        let mut out = RawDocumentBuf::new();
-        for (k, v) in doc {
-            let k: &CStr = k.as_str().try_into()?;
-            let val: RawBson = v.clone().try_into()?;
-            out.append(k, val);
-        }
-        Ok(out)
+        let mut out = vec![];
+        doc.append_to(&mut out)?;
+        RawDocumentBuf::from_bytes(out)
     }
 }
 
 impl TryFrom<Document> for RawDocumentBuf {
     type Error = crate::error::Error;
 
     fn try_from(doc: Document) -> std::result::Result<Self, Self::Error> {
-        let mut out = RawDocumentBuf::new();
-        for (k, v) in doc {
-            let k: &CStr = k.as_str().try_into()?;
-            let val: RawBson = v.try_into()?;
-            out.append(k, val);
-        }
-        Ok(out)
+        RawDocumentBuf::try_from(&doc)
     }
 }