reuse buffer for map serialization

cemoktra · Jul 10, 2024 · d05e804 · d05e804
1 parent 55e7762
commit d05e804
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 5 deletions.
diff --git a/core/src/wire/map.rs b/core/src/wire/map.rs
@@ -18,6 +18,21 @@ where
     WireType::LengthEncoded(map_buffer.freeze())
 }
 
+// #[inline(always)]
+// pub fn into_wire_pre_alloc<K, V>(key: K, value: V, buffer: &mut bytes::BytesMut) -> WireType
+// where
+//     K: IntoWire,
+//     V: IntoWire,
+// {
+//     let wire_type = key.into_wire();
+//     wire_type.serialize(1, buffer);
+
+//     let wire_type = value.into_wire();
+//     wire_type.serialize(2, buffer);
+
+//     WireType::LengthEncoded(buffer.clone().freeze())
+// }
+
 /// read a key-value pair from a [WireTypeView]
 #[inline(always)]
 pub fn from_wire<K, V>(wire_type: WireTypeView) -> Result<(K, V), Error>

diff --git a/core/src/wire/wire_type.rs b/core/src/wire/wire_type.rs
@@ -55,8 +55,6 @@ impl WireType {
 
                 writer.put_slice(&len_varint[0..len_size]);
                 writer.put_slice(data);
-
-                // );
             }
             WireType::SGroup => {
                 let tag = tag | 0b11;
@@ -103,3 +101,72 @@ impl WireType {
         }
     }
 }
+
+impl<'a> WireTypeView<'a> {
+    /// serialize a [WireType] to anything that implements [std::io::Write]
+    #[inline(always)]
+    pub fn serialize(&self, field_number: u32, writer: &mut impl bytes::BufMut) {
+        // let x = std::time::Instant::now();
+        let mut tag_varint = [0u8; 10];
+        let tag = field_number << 3;
+
+        match self {
+            WireTypeView::VarInt(data) => {
+                let tag_size = tag.encode_var(&mut tag_varint);
+                writer.put_slice(&tag_varint[0..tag_size]);
+                writer.put_slice(data);
+            }
+            WireTypeView::FixedI64(data) => {
+                let tag = tag | 0b1;
+                let tag_size = tag.encode_var(&mut tag_varint);
+                writer.put_slice(&tag_varint[0..tag_size]);
+                writer.put_slice(data);
+            }
+            WireTypeView::LengthEncoded(data) => {
+                let mut len_varint = [0u8; 10];
+
+                let tag = tag | 0b10;
+                let tag_size = tag.encode_var(&mut tag_varint);
+
+                writer.put_slice(&tag_varint[0..tag_size]);
+
+                let len: u32 = data.len().try_into().expect("this is good");
+                let len_size = len.encode_var(&mut len_varint);
+
+                writer.put_slice(&len_varint[0..len_size]);
+                writer.put_slice(data);
+            }
+            WireTypeView::SGroup => {
+                let tag = tag | 0b11;
+                let tag_size = tag.encode_var(&mut tag_varint);
+                writer.put_slice(&tag_varint[0..tag_size]);
+            }
+            WireTypeView::EGroup => {
+                let tag = tag | 0b100;
+                let tag_size = tag.encode_var(&mut tag_varint);
+                writer.put_slice(&tag_varint[0..tag_size]);
+            }
+
+            WireTypeView::FixedI32(data) => {
+                let tag = tag | 0b101;
+                let tag_size = tag.encode_var(&mut tag_varint);
+                writer.put_slice(&tag_varint[0..tag_size]);
+                writer.put_slice(data);
+            }
+        }
+    }
+
+    pub fn size_hint(&self, tag: u32) -> usize {
+        match self {
+            WireTypeView::VarInt(data) => tag.required_space() + data.len(),
+            WireTypeView::FixedI64(_) => tag.required_space() + 8,
+            WireTypeView::SGroup => tag.required_space(),
+            WireTypeView::EGroup => tag.required_space(),
+            WireTypeView::LengthEncoded(data) => {
+                let data_len = data.len();
+                tag.required_space() + data_len.required_space() + data_len
+            }
+            WireTypeView::FixedI32(_) => tag.required_space() + 4,
+        }
+    }
+}
diff --git a/derive/src/codegen.rs b/derive/src/codegen.rs
@@ -119,8 +119,22 @@ fn expand_message_message(
                 }
                 Kind::Map => {
                     serialize_impl.extend(quote_spanned! { span=>
+                        let mut buffer = Vec::new();
                         for (key, value) in self.#field_ident {
-                            let wire_type = #root::gin_tonic_core::map_into_wire(key, value);
+                            let size = key.size_hint(1) + value.size_hint(2);
+                            if size > buffer.capacity() {
+                                buffer.resize(size, 0);
+                            }
+
+                            let mut buffer_ref = buffer.as_mut_slice();
+
+                            let wire_type = key.into_wire();
+                            wire_type.serialize(1, &mut buffer_ref);
+
+                            let wire_type = value.into_wire();
+                            wire_type.serialize(2, &mut buffer_ref);
+
+                            let wire_type = #root::gin_tonic_core::WireTypeView::LengthEncoded(&buffer);
                             wire_type.serialize(#tag, writer);
                         }
                     });
@@ -233,8 +247,22 @@ fn expand_message_message(
                 Kind::Map => {
                     serialize_impl.extend(quote_spanned! { span=>
                         if let Some(value) = self.#field_ident {
-                            for (key, value) in value {
-                                let wire_type = #root::gin_tonic_core::map_into_wire(key, value);
+                            let mut buffer = Vec::new();
+                            for (key, value) in self.#field_ident {
+                                let size = key.size_hint(1) + value.size_hint(2);
+                                if size > buffer.capacity() {
+                                    buffer.resize(size, 0);
+                                }
+
+                                let mut buffer_ref = buffer.as_mut_slice();
+
+                                let wire_type = key.into_wire();
+                                wire_type.serialize(1, &mut buffer_ref);
+
+                                let wire_type = value.into_wire();
+                                wire_type.serialize(2, &mut buffer_ref);
+
+                                let wire_type = #root::gin_tonic_core::WireTypeView::LengthEncoded(&buffer);
                                 wire_type.serialize(#tag, writer);
                             }
                         }