From 8541b22f06e96b177877b9c291dbd9e1c9a24449 Mon Sep 17 00:00:00 2001
From: Jordan Doyle <jordan@doyle.la>
Date: Sat, 24 Jul 2021 15:36:23 +0100
Subject: [PATCH] Support nested documents, remove take_mut dependency

---
 Cargo.toml |  1 -
 src/lib.rs | 38 ++++++++++----------------------------
 src/ser.rs | 49 +++++++++++++++++++++++++++++++++++--------------
 3 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index f43bd78..2e2e303 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,7 +12,6 @@ license = "0BSD"
 [dependencies]
 serde = "1"
 bytes = "1"
-take_mut = "0.2"
 
 [dev-dependencies]
 serde = { version = "1", features = ["derive"] }
diff --git a/src/lib.rs b/src/lib.rs
index f42503f..78d00c1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,36 +3,11 @@ pub mod ser;
 
 pub use error::Error;
 
-use bytes::{BufMut, BytesMut};
+use bytes::BytesMut;
 use serde::Serialize;
 
 pub fn to_string<T: Serialize>(val: &T, output: &mut BytesMut) -> Result<(), Error> {
-    const SIZE_OF_SIZE: usize = std::mem::size_of::<i32>();
-
-    // essentially reserves a i32 we can prepend back onto the BytesMut later
-    // at the cost of an atomic increment
-    output.put_i32(0);
-    let mut size = output.split_to(SIZE_OF_SIZE);
-
-    val.serialize(ser::Serializer { key: None, output })?;
-
-    // writes the total length of the output to the i32 we split off before
-    for (i, byte) in ((output.len() + SIZE_OF_SIZE) as i32)
-        .to_le_bytes()
-        .iter()
-        .enumerate()
-    {
-        size[i] = *byte;
-    }
-
-    // this is safe because `unsplit` can't panic
-    take_mut::take(output, move |output| {
-        // O(1) prepend since `size` originally came from `output`.
-        size.unsplit(output);
-        size
-    });
-
-    Ok(())
+    val.serialize(ser::Serializer { key: None, output })
 }
 
 #[cfg(test)]
@@ -49,16 +24,23 @@ mod test {
             #[serde(with = "serde_bytes")]
             beans: &'a [u8],
             bro: &'a str,
+            b: B<'a>,
+        }
+
+        #[derive(Serialize)]
+        pub struct B<'a> {
+            s: &'a str,
         }
 
         let test = &A {
             cool: 999,
             beans: "so there was this one time at bandcamp".as_bytes(),
             bro: "the craziest thing happened",
+            b: B { s: "dddd" },
         };
 
         let mut ours = BytesMut::new();
-        to_string(&test, &mut ours);
+        to_string(&test, &mut ours).unwrap();
 
         let mut theirs = BytesMut::new().writer();
         bson::ser::to_document(&test)
diff --git a/src/ser.rs b/src/ser.rs
index 8fdb2e0..1276d7a 100644
--- a/src/ser.rs
+++ b/src/ser.rs
@@ -1,11 +1,11 @@
 use crate::Error;
-use bytes::BufMut;
+use bytes::{BufMut, BytesMut};
 use serde::Serialize;
 use std::convert::TryFrom;
 
-pub struct Serializer<'a, B: BufMut> {
+pub struct Serializer<'a> {
     pub key: Option<&'static str>,
-    pub output: &'a mut B,
+    pub output: &'a mut BytesMut,
 }
 
 macro_rules! write_key_or_error {
@@ -20,7 +20,7 @@ macro_rules! write_key_or_error {
     };
 }
 
-impl<'a, B: BufMut> serde::Serializer for Serializer<'a, B> {
+impl<'a> serde::Serializer for Serializer<'a> {
     type Ok = ();
     type Error = Error;
 
@@ -29,7 +29,7 @@ impl<'a, B: BufMut> serde::Serializer for Serializer<'a, B> {
     type SerializeTupleStruct = serde::ser::Impossible<Self::Ok, Self::Error>;
     type SerializeTupleVariant = serde::ser::Impossible<Self::Ok, Self::Error>;
     type SerializeMap = serde::ser::Impossible<Self::Ok, Self::Error>;
-    type SerializeStruct = StructSerializer<'a, B>;
+    type SerializeStruct = StructSerializer<'a>;
     type SerializeStructVariant = serde::ser::Impossible<Self::Ok, Self::Error>;
 
     fn serialize_bool(self, v: bool) -> Result<Self::Ok, Self::Error> {
@@ -208,11 +208,19 @@ impl<'a, B: BufMut> serde::Serializer for Serializer<'a, B> {
         _len: usize,
     ) -> Result<Self::SerializeStruct, Self::Error> {
         if self.key.is_some() {
-            todo!("nested struct: {:?}", self.key);
+            write_key_or_error!(0x03, self.key, self.output);
         }
 
+        // splits the output for the doc to be written to, this is appended back onto to the
+        // output when `StructSerializer::close` is called.
+        let mut doc_output = self.output.split_off(self.output.len());
+
+        // reserves a i32 we can write the document size to later
+        doc_output.put_i32(0);
+
         Ok(StructSerializer {
-            output: self.output,
+            original_output: self.output,
+            doc_output,
         })
     }
 
@@ -227,13 +235,14 @@ impl<'a, B: BufMut> serde::Serializer for Serializer<'a, B> {
     }
 }
 
-pub struct StructSerializer<'a, B: BufMut> {
-    output: &'a mut B,
+pub struct StructSerializer<'a> {
+    original_output: &'a mut BytesMut,
+    doc_output: BytesMut,
 }
 
-impl<'a, B: BufMut> serde::ser::SerializeStruct for StructSerializer<'a, B> {
+impl<'a> serde::ser::SerializeStruct for StructSerializer<'a> {
     type Ok = ();
-    type Error = <Serializer<'a, B> as serde::Serializer>::Error;
+    type Error = <Serializer<'a> as serde::Serializer>::Error;
 
     fn serialize_field<T: ?Sized>(
         &mut self,
@@ -245,12 +254,24 @@ impl<'a, B: BufMut> serde::ser::SerializeStruct for StructSerializer<'a, B> {
     {
         value.serialize(Serializer {
             key: Some(key),
-            output: &mut self.output,
+            output: &mut self.doc_output,
         })
     }
 
-    fn end(self) -> Result<Self::Ok, Self::Error> {
-        self.output.put_u8(0x00); // doc terminator
+    fn end(mut self) -> Result<Self::Ok, Self::Error> {
+        self.doc_output.put_u8(0x00); // doc terminator
+
+        // writes the total length of the output to the i32 we split off before
+        for (i, byte) in (self.doc_output.len() as i32)
+            .to_le_bytes()
+            .iter()
+            .enumerate()
+        {
+            self.doc_output[i] = *byte;
+        }
+
+        self.original_output.unsplit(self.doc_output);
+
         Ok(())
     }
 }
--
libgit2 1.7.2