diff --git a/core/src/wire/map.rs b/core/src/wire/map.rs index 5bbd7c1..1ac661d 100644 --- 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(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(wire_type: WireTypeView) -> Result<(K, V), Error> diff --git a/core/src/wire/wire_type.rs b/core/src/wire/wire_type.rs index 2ea80e7..67519e0 100644 --- 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 index a5cb4e7..67285cf 100644 --- 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); } }