From 175d64363231f9f844650729318eee3610f29140 Mon Sep 17 00:00:00 2001 From: nuskey8 Date: Sat, 29 Aug 2026 20:41:49 +0900 Subject: [PATCH 1/4] optimize: add `ToMessagePack::size()` --- zerompk/src/impl.rs | 164 +++++++++++++++++++++++++++++++++++--- zerompk/src/lib.rs | 35 +++++++- zerompk/src/write.rs | 28 +++++-- zerompk_derive/src/lib.rs | 97 ++++++++++++++++++++-- 4 files changed, 294 insertions(+), 30 deletions(-) diff --git a/zerompk/src/impl.rs b/zerompk/src/impl.rs index 1422679..d964148 100644 --- a/zerompk/src/impl.rs +++ b/zerompk/src/impl.rs @@ -9,7 +9,7 @@ use core::hash::Hash; // ------------------------------------------------------------------------------- macro_rules! impl_scalar { - ($ty:ty, $write_fn:ident, $write_slice_fn:ident, $read_fn:ident) => { + ($ty:ty, $write_fn:ident, $write_slice_fn:ident, $read_fn:ident, $size:expr) => { impl<'a> FromMessagePack<'a> for $ty { #[inline(always)] fn read>(reader: &mut R) -> crate::Result @@ -30,21 +30,126 @@ macro_rules! impl_scalar { fn write_slice(values: &[Self], writer: &mut W) -> crate::Result<()> { writer.$write_slice_fn(values) } + + #[inline(always)] + unsafe fn size(&self) -> Option { + Some(($size)(*self)) + } } }; } -impl_scalar!(bool, write_boolean, write_boolean_slice, read_boolean); -impl_scalar!(i8, write_i8, write_i8_slice, read_i8); -impl_scalar!(i16, write_i16, write_i16_slice, read_i16); -impl_scalar!(i32, write_i32, write_i32_slice, read_i32); -impl_scalar!(i64, write_i64, write_i64_slice, read_i64); -impl_scalar!(u8, write_u8, write_u8_slice, read_u8); -impl_scalar!(u16, write_u16, write_u16_slice, read_u16); -impl_scalar!(u32, write_u32, write_u32_slice, read_u32); -impl_scalar!(u64, write_u64, write_u64_slice, read_u64); -impl_scalar!(f32, write_f32, write_f32_slice, read_f32); -impl_scalar!(f64, write_f64, write_f64_slice, read_f64); +impl_scalar!( + bool, + write_boolean, + write_boolean_slice, + read_boolean, + |_| 1 +); +impl_scalar!( + i8, + write_i8, + write_i8_slice, + read_i8, + |v: i8| if (-32..=127).contains(&v) { 1 } else { 2 } +); +impl_scalar!( + i16, + write_i16, + write_i16_slice, + read_i16, + |v: i16| if (-32..=127).contains(&v) { + 1 + } else if (-128..=127).contains(&v) { + 2 + } else { + 3 + } +); +impl_scalar!( + i32, + write_i32, + write_i32_slice, + read_i32, + |v: i32| if (-32..=127).contains(&v) { + 1 + } else if (-128..=127).contains(&v) { + 2 + } else if (-32768..=32767).contains(&v) { + 3 + } else { + 5 + } +); +impl_scalar!( + i64, + write_i64, + write_i64_slice, + read_i64, + |v: i64| if (-32..=127).contains(&v) { + 1 + } else if (-128..=127).contains(&v) { + 2 + } else if (-32768..=32767).contains(&v) { + 3 + } else if (-2147483648..=2147483647).contains(&v) { + 5 + } else { + 9 + } +); +impl_scalar!(u8, write_u8, write_u8_slice, read_u8, |v: u8| if v <= 127 { + 1 +} else { + 2 +}); +impl_scalar!( + u16, + write_u16, + write_u16_slice, + read_u16, + |v: u16| if v <= 127 { + 1 + } else if v <= 255 { + 2 + } else { + 3 + } +); +impl_scalar!( + u32, + write_u32, + write_u32_slice, + read_u32, + |v: u32| if v <= 127 { + 1 + } else if v <= 255 { + 2 + } else if v <= 65535 { + 3 + } else { + 5 + } +); +impl_scalar!( + u64, + write_u64, + write_u64_slice, + read_u64, + |v: u64| if v <= 127 { + 1 + } else if v <= 255 { + 2 + } else if v <= 65535 { + 3 + } else if v <= 4294967295 { + 5 + } else { + 9 + } +); +impl_scalar!(f32, write_f32, write_f32_slice, read_f32, |_| 5); +impl_scalar!(f64, write_f64, write_f64_slice, read_f64, |_| 9); impl<'a> FromMessagePack<'a> for usize { #[inline(always)] @@ -187,6 +292,23 @@ impl ToMessagePack for [T] { writer.write_array_len(self.len())?; T::write_slice(self, writer) } + + #[inline] + unsafe fn size(&self) -> Option { + let mut size: usize = if self.len() < 16 { + 1 + } else if self.len() <= u16::MAX as usize { + 3 + } else { + 5 + }; + for value in self { + // SAFETY: callers of this method uphold each element's size contract. + let value_size = unsafe { value.size()? }; + size = size.checked_add(value_size)?; + } + Some(size) + } } impl<'a, T: FromMessagePack<'a>, const N: usize> FromMessagePack<'a> for [T; N] { @@ -234,6 +356,12 @@ impl ToMessagePack for [T; N] { writer.write_array_len(N)?; T::write_slice(self, writer) } + + #[inline] + unsafe fn size(&self) -> Option { + // SAFETY: arrays use the same representation as slices. + unsafe { self.as_slice().size() } + } } impl<'a> FromMessagePack<'a> for alloc::string::String { @@ -305,6 +433,12 @@ impl ToMessagePack for &T { fn write(&self, writer: &mut W) -> crate::Result<()> { T::write(self, writer) } + + #[inline] + unsafe fn size(&self) -> Option { + // SAFETY: forwarded from the caller. + unsafe { T::size(self) } + } } impl ToMessagePack for &mut T { @@ -395,6 +529,12 @@ impl ToMessagePack for alloc::vec::Vec { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_slice().write(writer) } + + #[inline] + unsafe fn size(&self) -> Option { + // SAFETY: vectors use the same representation as slices. + unsafe { self.as_slice().size() } + } } impl<'a, T: FromMessagePack<'a>> FromMessagePack<'a> for alloc::collections::VecDeque { diff --git a/zerompk/src/lib.rs b/zerompk/src/lib.rs index b2bcae1..0bfcf5f 100644 --- a/zerompk/src/lib.rs +++ b/zerompk/src/lib.rs @@ -51,6 +51,17 @@ pub trait ToMessagePack { /// Writes the MessagePack representation of this value into the provided writer. fn write(&self, writer: &mut W) -> Result<()>; + /// Returns the exact number of bytes written by the next call to [`Self::write`], + /// or `None` when the size cannot be determined cheaply. + /// + /// # Safety + /// When an override returns `Some`, that size and the following `write` call + /// must describe exactly the same output. + #[inline] + unsafe fn size(&self) -> Option { + None + } + /// Writes the MessagePack representation of a slice of values into the provided writer. #[inline(always)] fn write_slice(values: &[Self], writer: &mut W) -> Result<()> @@ -113,7 +124,11 @@ pub fn from_msgpack<'a, T: FromMessagePack<'a>>(data: &'a [u8]) -> Result { /// } /// ``` pub fn to_msgpack_vec(value: &T) -> Result> { - let mut writer = write::VecWriter::new(); + // SAFETY: a returned size is required to match the immediately following write. + let mut writer = match unsafe { value.size() } { + Some(size) => write::VecWriter::with_capacity(size), + None => write::VecWriter::new(), + }; value.write(&mut writer)?; Ok(writer.into_vec()) } @@ -144,9 +159,21 @@ pub fn to_msgpack_vec(value: &T) -> Result> { /// } /// ``` pub fn to_msgpack(value: &T, buf: &mut [u8]) -> Result { - let mut writer = write::SliceWriter::new(buf); - value.write(&mut writer)?; - Ok(writer.position()) + // SAFETY: `size` guarantees the byte count of the immediately following write. + if let Some(size) = unsafe { value.size() } { + if size > buf.len() { + return Err(Error::BufferTooSmall); + } + // SAFETY: the exact output range was validated above. + let mut writer = unsafe { write::SliceWriter::new_unchecked(&mut buf[..size]) }; + value.write(&mut writer)?; + debug_assert_eq!(writer.position(), size); + Ok(writer.position()) + } else { + let mut writer = write::SliceWriter::new(buf); + value.write(&mut writer)?; + Ok(writer.position()) + } } /// Serializes a value of type `T` into the I/O stream. diff --git a/zerompk/src/write.rs b/zerompk/src/write.rs index 847bc41..398f931 100644 --- a/zerompk/src/write.rs +++ b/zerompk/src/write.rs @@ -77,7 +77,7 @@ macro_rules! impl_unsigned_slice { let Some(max_len) = values.len().checked_mul($max_len) else { return Err(Error::BufferTooSmall); }; - if max_len > self.buffer.len() - self.pos { + if CHECKED && max_len > self.buffer.len() - self.pos { for &value in values { self.$scalar(value)?; } @@ -104,7 +104,7 @@ macro_rules! impl_signed_slice { let Some(max_len) = values.len().checked_mul($max_len) else { return Err(Error::BufferTooSmall); }; - if max_len > self.buffer.len() - self.pos { + if CHECKED && max_len > self.buffer.len() - self.pos { for &value in values { self.$scalar(value)?; } @@ -339,19 +339,27 @@ pub trait Write { fn write_ext(&mut self, type_id: i8, data: &[u8]) -> Result<()>; } -pub struct SliceWriter<'a> { +pub struct SliceWriter<'a, const CHECKED: bool = true> { buffer: &'a mut [u8], pos: usize, } -impl<'a> SliceWriter<'a> { +impl<'a> SliceWriter<'a, true> { pub fn new(buffer: &'a mut [u8]) -> Self { SliceWriter { buffer, pos: 0 } } +} + +impl<'a> SliceWriter<'a, false> { + pub(crate) unsafe fn new_unchecked(buffer: &'a mut [u8]) -> Self { + SliceWriter { buffer, pos: 0 } + } +} +impl<'a, const CHECKED: bool> SliceWriter<'a, CHECKED> { #[inline(always)] fn take_array(&mut self) -> Result<&mut [u8; N]> { - if N > self.buffer.len() - self.pos { + if CHECKED && N > self.buffer.len() - self.pos { cold_path(); return Err(Error::BufferTooSmall); } @@ -362,7 +370,7 @@ impl<'a> SliceWriter<'a> { #[inline(always)] fn take_slice(&mut self, len: usize) -> Result<&mut [u8]> { - if len > self.buffer.len() - self.pos { + if CHECKED && len > self.buffer.len() - self.pos { cold_path(); return Err(Error::BufferTooSmall); } @@ -378,7 +386,7 @@ impl<'a> SliceWriter<'a> { } } -impl<'a> Write for SliceWriter<'a> { +impl<'a, const CHECKED: bool> Write for SliceWriter<'a, CHECKED> { #[inline(always)] fn write_static_string(&mut self, _value: &'static str, encoded: &'static [u8]) -> Result<()> { self.take_slice(encoded.len())?.copy_from_slice(encoded); @@ -543,6 +551,12 @@ impl VecWriter { VecWriter { buffer: Vec::new() } } + pub fn with_capacity(capacity: usize) -> Self { + VecWriter { + buffer: Vec::with_capacity(capacity), + } + } + pub fn into_vec(self) -> Vec { self.buffer } diff --git a/zerompk_derive/src/lib.rs b/zerompk_derive/src/lib.rs index 2974dd8..c030470 100644 --- a/zerompk_derive/src/lib.rs +++ b/zerompk_derive/src/lib.rs @@ -977,7 +977,7 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result quote! { @@ -986,6 +986,8 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result(&self, writer: &mut W) -> ::core::result::Result<(), ::zerompk::Error> { #write } + + #size } }, DeriveKind::From => { @@ -1026,6 +1028,7 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result Result { @@ -1149,7 +1152,36 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - Ok(ImplBody { write, read }) + let size = if is_dense_sequential + && field_configs.iter().all(|cfg| cfg.as_bytes.is_none()) + && tys + .iter() + .zip(field_configs.iter()) + .all(|(ty, cfg)| !(is_bin_type(ty) && should_use_bin(ty, Some(cfg)))) + { + let header_size = if array_len < 16 { + 1usize + } else if array_len <= u16::MAX as usize { + 3 + } else { + 5 + }; + quote! { + #[inline] + unsafe fn size(&self) -> ::core::option::Option { + let mut __size = #header_size; + #( + let __field_size = unsafe { ::zerompk::ToMessagePack::size(&self.#names)? }; + __size = __size.checked_add(__field_size)?; + )* + ::core::option::Option::Some(__size) + } + } + } else { + quote! {} + }; + + Ok(ImplBody { write, read, size }) } Fields::Unnamed(fields) => { let count = fields.unnamed.len(); @@ -1180,7 +1212,11 @@ fn expand_array_struct(data: &DataStruct) -> Result { Ok(Self(__f0)) }; - return Ok(ImplBody { write, read }); + return Ok(ImplBody { + write, + read, + size: quote! {}, + }); } let idx: Vec<_> = (0..count).map(syn::Index::from).collect(); @@ -1276,7 +1312,36 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - Ok(ImplBody { write, read }) + let size = if is_dense_sequential + && field_configs.iter().all(|cfg| cfg.as_bytes.is_none()) + && tys + .iter() + .zip(field_configs.iter()) + .all(|(ty, cfg)| !(is_bin_type(ty) && should_use_bin(ty, Some(cfg)))) + { + let header_size = if array_len < 16 { + 1usize + } else if array_len <= u16::MAX as usize { + 3 + } else { + 5 + }; + quote! { + #[inline] + unsafe fn size(&self) -> ::core::option::Option { + let mut __size = #header_size; + #( + let __field_size = unsafe { ::zerompk::ToMessagePack::size(&self.#idx)? }; + __size = __size.checked_add(__field_size)?; + )* + ::core::option::Option::Some(__size) + } + } + } else { + quote! {} + }; + + Ok(ImplBody { write, read, size }) } Fields::Unit => Ok(ImplBody { write: quote! { @@ -1287,6 +1352,12 @@ fn expand_array_struct(data: &DataStruct) -> Result { reader.read_nil()?; Ok(Self) }, + size: quote! { + #[inline] + unsafe fn size(&self) -> ::core::option::Option { + ::core::option::Option::Some(1) + } + }, }), } } @@ -1486,7 +1557,11 @@ fn expand_map_struct(data: &DataStruct, allow_unknown_fields: bool) -> Result Result { @@ -1550,7 +1625,11 @@ fn expand_c_enum(data: &DataEnum, repr: CEnumRepr) -> Result { } }; - Ok(ImplBody { write, read }) + Ok(ImplBody { + write, + read, + size: quote! {}, + }) } fn read_tag_dispatch( @@ -1702,7 +1781,11 @@ fn expand_enum(data: &DataEnum, repr: Repr) -> Result { } }; - Ok(ImplBody { write, read }) + Ok(ImplBody { + write, + read, + size: quote! {}, + }) } fn build_enum_variant_payload( From 9e2623f3c95a6fc400be76ad375e674420fe654b Mon Sep 17 00:00:00 2001 From: nuskey8 Date: Sat, 29 Aug 2026 21:05:58 +0900 Subject: [PATCH 2/4] fix: make serialization size proofs safe to use --- zerompk/src/impl.rs | 32 +++++++++++++++----------------- zerompk/src/lib.rs | 33 ++++++++++++++++++++++++++------- zerompk_derive/src/lib.rs | 23 +++++++++++++---------- 3 files changed, 54 insertions(+), 34 deletions(-) diff --git a/zerompk/src/impl.rs b/zerompk/src/impl.rs index d964148..3e93a7d 100644 --- a/zerompk/src/impl.rs +++ b/zerompk/src/impl.rs @@ -32,8 +32,9 @@ macro_rules! impl_scalar { } #[inline(always)] - unsafe fn size(&self) -> Option { - Some(($size)(*self)) + fn size(&self) -> Option { + // SAFETY: primitive encodings are completely determined by their value. + Some(unsafe { crate::ExactSize::new_unchecked(($size)(*self)) }) } } }; @@ -294,20 +295,20 @@ impl ToMessagePack for [T] { } #[inline] - unsafe fn size(&self) -> Option { - let mut size: usize = if self.len() < 16 { + fn size(&self) -> Option { + let header_size: usize = if self.len() < 16 { 1 } else if self.len() <= u16::MAX as usize { 3 } else { 5 }; + let mut size = header_size; for value in self { - // SAFETY: callers of this method uphold each element's size contract. - let value_size = unsafe { value.size()? }; - size = size.checked_add(value_size)?; + size = size.checked_add(value.size()?.get())?; } - Some(size) + // SAFETY: the array header and every element have exact-size proofs. + Some(unsafe { crate::ExactSize::new_unchecked(size) }) } } @@ -358,9 +359,8 @@ impl ToMessagePack for [T; N] { } #[inline] - unsafe fn size(&self) -> Option { - // SAFETY: arrays use the same representation as slices. - unsafe { self.as_slice().size() } + fn size(&self) -> Option { + self.as_slice().size() } } @@ -435,9 +435,8 @@ impl ToMessagePack for &T { } #[inline] - unsafe fn size(&self) -> Option { - // SAFETY: forwarded from the caller. - unsafe { T::size(self) } + fn size(&self) -> Option { + T::size(self) } } @@ -531,9 +530,8 @@ impl ToMessagePack for alloc::vec::Vec { } #[inline] - unsafe fn size(&self) -> Option { - // SAFETY: vectors use the same representation as slices. - unsafe { self.as_slice().size() } + fn size(&self) -> Option { + self.as_slice().size() } } diff --git a/zerompk/src/lib.rs b/zerompk/src/lib.rs index 0bfcf5f..34feee6 100644 --- a/zerompk/src/lib.rs +++ b/zerompk/src/lib.rs @@ -46,6 +46,27 @@ pub trait FromMessagePackOwned: for<'a> FromMessagePack<'a> {} impl FromMessagePackOwned for T where T: for<'a> FromMessagePack<'a> {} +/// Proof that the next serialization of a value has an exact encoded size. +pub struct ExactSize(usize); + +impl ExactSize { + /// Creates an exact-size proof. + /// + /// # Safety + /// The next serialization of the value this proof describes must write exactly + /// `size` bytes. Serialization-relevant state must not change in between. + #[doc(hidden)] + #[inline(always)] + pub const unsafe fn new_unchecked(size: usize) -> Self { + Self(size) + } + + #[inline(always)] + pub const fn get(&self) -> usize { + self.0 + } +} + /// A data structure that can be serialized into MessagePack format. pub trait ToMessagePack { /// Writes the MessagePack representation of this value into the provided writer. @@ -54,11 +75,8 @@ pub trait ToMessagePack { /// Returns the exact number of bytes written by the next call to [`Self::write`], /// or `None` when the size cannot be determined cheaply. /// - /// # Safety - /// When an override returns `Some`, that size and the following `write` call - /// must describe exactly the same output. #[inline] - unsafe fn size(&self) -> Option { + fn size(&self) -> Option { None } @@ -125,8 +143,8 @@ pub fn from_msgpack<'a, T: FromMessagePack<'a>>(data: &'a [u8]) -> Result { /// ``` pub fn to_msgpack_vec(value: &T) -> Result> { // SAFETY: a returned size is required to match the immediately following write. - let mut writer = match unsafe { value.size() } { - Some(size) => write::VecWriter::with_capacity(size), + let mut writer = match value.size() { + Some(size) => write::VecWriter::with_capacity(size.get()), None => write::VecWriter::new(), }; value.write(&mut writer)?; @@ -160,7 +178,8 @@ pub fn to_msgpack_vec(value: &T) -> Result> { /// ``` pub fn to_msgpack(value: &T, buf: &mut [u8]) -> Result { // SAFETY: `size` guarantees the byte count of the immediately following write. - if let Some(size) = unsafe { value.size() } { + if let Some(size) = value.size() { + let size = size.get(); if size > buf.len() { return Err(Error::BufferTooSmall); } diff --git a/zerompk_derive/src/lib.rs b/zerompk_derive/src/lib.rs index c030470..4a2f8c7 100644 --- a/zerompk_derive/src/lib.rs +++ b/zerompk_derive/src/lib.rs @@ -1168,13 +1168,14 @@ fn expand_array_struct(data: &DataStruct) -> Result { }; quote! { #[inline] - unsafe fn size(&self) -> ::core::option::Option { + fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { let mut __size = #header_size; #( - let __field_size = unsafe { ::zerompk::ToMessagePack::size(&self.#names)? }; - __size = __size.checked_add(__field_size)?; + let __field_size = ::zerompk::ToMessagePack::size(&self.#names)?; + __size = __size.checked_add(__field_size.get())?; )* - ::core::option::Option::Some(__size) + // SAFETY: the header and every serialized field have exact-size proofs. + ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(__size) }) } } } else { @@ -1328,13 +1329,14 @@ fn expand_array_struct(data: &DataStruct) -> Result { }; quote! { #[inline] - unsafe fn size(&self) -> ::core::option::Option { + fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { let mut __size = #header_size; #( - let __field_size = unsafe { ::zerompk::ToMessagePack::size(&self.#idx)? }; - __size = __size.checked_add(__field_size)?; + let __field_size = ::zerompk::ToMessagePack::size(&self.#idx)?; + __size = __size.checked_add(__field_size.get())?; )* - ::core::option::Option::Some(__size) + // SAFETY: the header and every serialized field have exact-size proofs. + ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(__size) }) } } } else { @@ -1354,8 +1356,9 @@ fn expand_array_struct(data: &DataStruct) -> Result { }, size: quote! { #[inline] - unsafe fn size(&self) -> ::core::option::Option { - ::core::option::Option::Some(1) + fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { + // SAFETY: unit structs always encode as one nil byte. + ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(1) }) } }, }), From dffbba2aaa30d1477a5dd65373174a8d88eda006 Mon Sep 17 00:00:00 2001 From: nuskey8 Date: Sat, 29 Aug 2026 22:16:47 +0900 Subject: [PATCH 3/4] feat: replace exact size with trusted size hints --- zerompk/src/impl.rs | 322 ++++++++++++++++++++++++++++++++----- zerompk/src/lib.rs | 93 +++++++---- zerompk/tests/serialize.rs | 19 +++ zerompk_derive/src/lib.rs | 192 +++++++++++++++++++--- 4 files changed, 526 insertions(+), 100 deletions(-) diff --git a/zerompk/src/impl.rs b/zerompk/src/impl.rs index 3e93a7d..dc210ca 100644 --- a/zerompk/src/impl.rs +++ b/zerompk/src/impl.rs @@ -4,12 +4,66 @@ use alloc::string::ToString; #[cfg(feature = "std")] use core::hash::Hash; +#[inline(always)] +fn array_header_size(len: usize) -> usize { + if len < 16 { + 1 + } else if len <= u16::MAX as usize { + 3 + } else { + 5 + } +} + +#[inline] +fn sequence_size_hint(len: usize) -> Option { + let element = T::max_size()?.upper_bound(); + let size = element + .checked_mul(len)? + .checked_add(array_header_size(len))?; + // SAFETY: the header is exact and every element is bounded by `element`. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) +} + +#[inline] +fn map_size_hint(len: usize) -> Option { + let pair = K::max_size()? + .upper_bound() + .checked_add(V::max_size()?.upper_bound())?; + let header = if len < 16 { + 1 + } else if len <= u16::MAX as usize { + 3 + } else { + 5 + }; + let size = pair.checked_mul(len)?.checked_add(header)?; + // SAFETY: the header is exact and every key/value pair is bounded by `pair`. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) +} + +#[inline] +fn string_size_hint(len: usize) -> Option { + let header = if len < 32 { + 1 + } else if len <= u8::MAX as usize { + 2 + } else if len <= u16::MAX as usize { + 3 + } else { + 5 + }; + let size = len.checked_add(header)?; + // SAFETY: MessagePack string headers depend only on the byte length. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) +} + // ------------------------------------------------------------------------------- // primitive types // ------------------------------------------------------------------------------- macro_rules! impl_scalar { - ($ty:ty, $write_fn:ident, $write_slice_fn:ident, $read_fn:ident, $size:expr) => { + ($ty:ty, $write_fn:ident, $write_slice_fn:ident, $read_fn:ident, $size:expr, $max:expr) => { impl<'a> FromMessagePack<'a> for $ty { #[inline(always)] fn read>(reader: &mut R) -> crate::Result @@ -32,9 +86,15 @@ macro_rules! impl_scalar { } #[inline(always)] - fn size(&self) -> Option { + fn size_hint(&self) -> Option { // SAFETY: primitive encodings are completely determined by their value. - Some(unsafe { crate::ExactSize::new_unchecked(($size)(*self)) }) + Some(unsafe { crate::TrustedSizeHint::new_unchecked(($size)(*self)) }) + } + + #[inline(always)] + fn max_size() -> Option { + // SAFETY: this is the largest encoding emitted for this primitive type. + Some(unsafe { crate::TrustedSizeHint::new_unchecked($max) }) } } }; @@ -45,14 +105,16 @@ impl_scalar!( write_boolean, write_boolean_slice, read_boolean, - |_| 1 + |_| 1, + 1 ); impl_scalar!( i8, write_i8, write_i8_slice, read_i8, - |v: i8| if (-32..=127).contains(&v) { 1 } else { 2 } + |v: i8| if (-32..=127).contains(&v) { 1 } else { 2 }, + 2 ); impl_scalar!( i16, @@ -65,7 +127,8 @@ impl_scalar!( 2 } else { 3 - } + }, + 3 ); impl_scalar!( i32, @@ -80,7 +143,8 @@ impl_scalar!( 3 } else { 5 - } + }, + 5 ); impl_scalar!( i64, @@ -97,13 +161,17 @@ impl_scalar!( 5 } else { 9 - } + }, + 9 ); -impl_scalar!(u8, write_u8, write_u8_slice, read_u8, |v: u8| if v <= 127 { - 1 -} else { +impl_scalar!( + u8, + write_u8, + write_u8_slice, + read_u8, + |v: u8| if v <= 127 { 1 } else { 2 }, 2 -}); +); impl_scalar!( u16, write_u16, @@ -115,7 +183,8 @@ impl_scalar!( 2 } else { 3 - } + }, + 3 ); impl_scalar!( u32, @@ -130,7 +199,8 @@ impl_scalar!( 3 } else { 5 - } + }, + 5 ); impl_scalar!( u64, @@ -147,10 +217,11 @@ impl_scalar!( 5 } else { 9 - } + }, + 9 ); -impl_scalar!(f32, write_f32, write_f32_slice, read_f32, |_| 5); -impl_scalar!(f64, write_f64, write_f64_slice, read_f64, |_| 9); +impl_scalar!(f32, write_f32, write_f32_slice, read_f32, |_| 5, 5); +impl_scalar!(f64, write_f64, write_f64_slice, read_f64, |_| 9, 9); impl<'a> FromMessagePack<'a> for usize { #[inline(always)] @@ -175,6 +246,22 @@ impl ToMessagePack for usize { writer.write_u64(*self as u64) } } + + fn size_hint(&self) -> Option { + if usize::BITS <= 32 { + (*self as u32).size_hint() + } else { + (*self as u64).size_hint() + } + } + + fn max_size() -> Option { + if usize::BITS <= 32 { + u32::max_size() + } else { + u64::max_size() + } + } } impl<'a> FromMessagePack<'a> for isize { @@ -200,6 +287,22 @@ impl ToMessagePack for isize { writer.write_i64(*self as i64) } } + + fn size_hint(&self) -> Option { + if isize::BITS <= 32 { + (*self as i32).size_hint() + } else { + (*self as i64).size_hint() + } + } + + fn max_size() -> Option { + if isize::BITS <= 32 { + i32::max_size() + } else { + i64::max_size() + } + } } impl<'a> FromMessagePack<'a> for char { @@ -221,6 +324,14 @@ impl ToMessagePack for char { fn write(&self, writer: &mut W) -> crate::Result<()> { writer.write_u32(*self as u32) } + + fn size_hint(&self) -> Option { + (*self as u32).size_hint() + } + + fn max_size() -> Option { + u32::max_size() + } } // ------------------------------------------------------------------------------- @@ -242,6 +353,16 @@ impl ToMessagePack for core::marker::PhantomData { fn write(&self, _: &mut W) -> crate::Result<()> { Ok(()) } + + fn size_hint(&self) -> Option { + // SAFETY: PhantomData writes no bytes. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(0) }) + } + + fn max_size() -> Option { + // SAFETY: PhantomData writes no bytes. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(0) }) + } } // ------------------------------------------------------------------------------- @@ -269,6 +390,10 @@ impl ToMessagePack for &str { fn write(&self, writer: &mut W) -> crate::Result<()> { writer.write_string(self) } + + fn size_hint(&self) -> Option { + string_size_hint(self.len()) + } } impl<'de, 'a> FromMessagePack<'de> for &'a [u8] @@ -294,21 +419,8 @@ impl ToMessagePack for [T] { T::write_slice(self, writer) } - #[inline] - fn size(&self) -> Option { - let header_size: usize = if self.len() < 16 { - 1 - } else if self.len() <= u16::MAX as usize { - 3 - } else { - 5 - }; - let mut size = header_size; - for value in self { - size = size.checked_add(value.size()?.get())?; - } - // SAFETY: the array header and every element have exact-size proofs. - Some(unsafe { crate::ExactSize::new_unchecked(size) }) + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) } } @@ -358,9 +470,12 @@ impl ToMessagePack for [T; N] { T::write_slice(self, writer) } - #[inline] - fn size(&self) -> Option { - self.as_slice().size() + fn size_hint(&self) -> Option { + sequence_size_hint::(N) + } + + fn max_size() -> Option { + sequence_size_hint::(N) } } @@ -382,6 +497,10 @@ impl ToMessagePack for alloc::string::String { fn write(&self, writer: &mut W) -> crate::Result<()> { writer.write_string(self) } + + fn size_hint(&self) -> Option { + string_size_hint(self.len()) + } } impl<'a> FromMessagePack<'a> for alloc::borrow::Cow<'a, str> { @@ -399,6 +518,10 @@ impl ToMessagePack for alloc::borrow::Cow<'_, str> { fn write(&self, writer: &mut W) -> crate::Result<()> { writer.write_string(self) } + + fn size_hint(&self) -> Option { + string_size_hint(self.len()) + } } impl<'de, 'a, T> FromMessagePack<'de> for alloc::borrow::Cow<'a, [T]> @@ -422,6 +545,10 @@ impl ToMessagePack for alloc::borrow::Cow<'_, [T]> { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_ref().write(writer) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } // ------------------------------------------------------------------------------- @@ -435,8 +562,8 @@ impl ToMessagePack for &T { } #[inline] - fn size(&self) -> Option { - T::size(self) + fn size_hint(&self) -> Option { + T::size_hint(self) } } @@ -445,6 +572,10 @@ impl ToMessagePack for &mut T { fn write(&self, writer: &mut W) -> crate::Result<()> { T::write(self, writer) } + + fn size_hint(&self) -> Option { + T::size_hint(self) + } } // ------------------------------------------------------------------------------- @@ -469,6 +600,21 @@ impl ToMessagePack for Option { None => writer.write_nil(), } } + + #[inline] + fn size_hint(&self) -> Option { + match self { + Some(value) => value.size_hint(), + // SAFETY: None always encodes as one nil byte. + None => Some(unsafe { crate::TrustedSizeHint::new_unchecked(1) }), + } + } + + fn max_size() -> Option { + let upper = T::max_size()?.upper_bound().max(1); + // SAFETY: Option is either a one-byte nil or a T value. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(upper) }) + } } impl<'a, T: FromMessagePack<'a>, E: FromMessagePack<'a>> FromMessagePack<'a> @@ -505,6 +651,27 @@ impl ToMessagePack for core::result::Result< } } } + + #[inline] + fn size_hint(&self) -> Option { + let payload = match self { + Ok(value) => value.size_hint()?, + Err(error) => error.size_hint()?, + }; + // fixarray header and boolean tag are one byte each. + let size = payload.upper_bound().checked_add(2)?; + // SAFETY: the wrapper and payload sizes are exact. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) + } + + fn max_size() -> Option { + let payload = T::max_size()? + .upper_bound() + .max(E::max_size()?.upper_bound()); + let upper = payload.checked_add(2)?; + // SAFETY: the wrapper is two bytes and the payload is bounded above. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(upper) }) + } } // ------------------------------------------------------------------------------- @@ -529,9 +696,8 @@ impl ToMessagePack for alloc::vec::Vec { self.as_slice().write(writer) } - #[inline] - fn size(&self) -> Option { - self.as_slice().size() + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) } } @@ -558,6 +724,10 @@ impl ToMessagePack for alloc::collections::VecDeque { T::write_slice(front, writer)?; T::write_slice(back, writer) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } impl<'a, T: FromMessagePack<'a>> FromMessagePack<'a> for alloc::collections::LinkedList { @@ -582,6 +752,10 @@ impl ToMessagePack for alloc::collections::LinkedList { } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } impl<'a, T: Ord + FromMessagePack<'a>> FromMessagePack<'a> for alloc::collections::BTreeSet { @@ -606,6 +780,10 @@ impl ToMessagePack for alloc::collections::BTreeSet { } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } impl<'a, K: Ord + FromMessagePack<'a>, V: FromMessagePack<'a>> FromMessagePack<'a> @@ -635,6 +813,10 @@ impl ToMessagePack for alloc::collections::B } Ok(()) } + + fn size_hint(&self) -> Option { + map_size_hint::(self.len()) + } } impl<'a, T: FromMessagePack<'a> + Ord> FromMessagePack<'a> for alloc::collections::BinaryHeap { @@ -662,6 +844,10 @@ impl ToMessagePack for alloc::collections::BinaryHeap } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } #[cfg(feature = "std")] @@ -691,6 +877,10 @@ impl ToMessagePack for std::collections::HashSet { } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } #[cfg(feature = "std")] @@ -725,6 +915,10 @@ impl ToMessagePack for std::collections::Has } Ok(()) } + + fn size_hint(&self) -> Option { + map_size_hint::(self.len()) + } } // ------------------------------------------------------------------------------- @@ -746,6 +940,14 @@ impl ToMessagePack for alloc::boxed::Box { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_ref().write(writer) } + + fn size_hint(&self) -> Option { + self.as_ref().size_hint() + } + + fn max_size() -> Option { + T::max_size() + } } #[cfg(feature = "std")] @@ -763,6 +965,14 @@ impl ToMessagePack for std::sync::Arc { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_ref().write(writer) } + + fn size_hint(&self) -> Option { + self.as_ref().size_hint() + } + + fn max_size() -> Option { + T::max_size() + } } impl<'a, T: FromMessagePack<'a>> FromMessagePack<'a> for alloc::rc::Rc { @@ -778,6 +988,14 @@ impl ToMessagePack for alloc::rc::Rc { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_ref().write(writer) } + + fn size_hint(&self) -> Option { + self.as_ref().size_hint() + } + + fn max_size() -> Option { + T::max_size() + } } // ------------------------------------------------------------------------------- @@ -797,6 +1015,16 @@ impl ToMessagePack for () { fn write(&self, writer: &mut W) -> crate::Result<()> { writer.write_nil() } + + fn size_hint(&self) -> Option { + // SAFETY: unit always encodes as one nil byte. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(1) }) + } + + fn max_size() -> Option { + // SAFETY: unit always encodes as one nil byte. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(1) }) + } } macro_rules! impl_tuple_message_packable { @@ -817,6 +1045,20 @@ macro_rules! impl_tuple_message_packable { $(self.$idx.write(writer)?;)+ Ok(()) } + + fn size_hint(&self) -> Option { + let mut size = array_header_size($len); + $(size = size.checked_add(self.$idx.size_hint()?.upper_bound())?;)+ + // SAFETY: the tuple header and every element have exact-size proofs. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) + } + + fn max_size() -> Option { + let mut size = array_header_size($len); + $(size = size.checked_add($t::max_size()?.upper_bound())?;)+ + // SAFETY: the tuple header is exact and every field is bounded above. + Some(unsafe { crate::TrustedSizeHint::new_unchecked(size) }) + } } }; } diff --git a/zerompk/src/lib.rs b/zerompk/src/lib.rs index 34feee6..2da1728 100644 --- a/zerompk/src/lib.rs +++ b/zerompk/src/lib.rs @@ -46,23 +46,23 @@ pub trait FromMessagePackOwned: for<'a> FromMessagePack<'a> {} impl FromMessagePackOwned for T where T: for<'a> FromMessagePack<'a> {} -/// Proof that the next serialization of a value has an exact encoded size. -pub struct ExactSize(usize); +/// A trusted upper bound for the next serialization of a value. +pub struct TrustedSizeHint(usize); -impl ExactSize { - /// Creates an exact-size proof. +impl TrustedSizeHint { + /// Creates a trusted encoded-size upper bound. /// /// # Safety - /// The next serialization of the value this proof describes must write exactly - /// `size` bytes. Serialization-relevant state must not change in between. + /// The next serialization of the value this hint describes must write at most + /// `upper_bound` bytes. Serialization-relevant state must not change in between. #[doc(hidden)] #[inline(always)] - pub const unsafe fn new_unchecked(size: usize) -> Self { - Self(size) + pub const unsafe fn new_unchecked(upper_bound: usize) -> Self { + Self(upper_bound) } #[inline(always)] - pub const fn get(&self) -> usize { + pub const fn upper_bound(&self) -> usize { self.0 } } @@ -72,14 +72,6 @@ pub trait ToMessagePack { /// Writes the MessagePack representation of this value into the provided writer. fn write(&self, writer: &mut W) -> Result<()>; - /// Returns the exact number of bytes written by the next call to [`Self::write`], - /// or `None` when the size cannot be determined cheaply. - /// - #[inline] - fn size(&self) -> Option { - None - } - /// Writes the MessagePack representation of a slice of values into the provided writer. #[inline(always)] fn write_slice(values: &[Self], writer: &mut W) -> Result<()> @@ -91,6 +83,38 @@ pub trait ToMessagePack { } Ok(()) } + + /// Returns a trusted upper bound for the number of bytes written by the next + /// call to [`Self::write`], or `None` when it cannot be determined cheaply. + /// + /// Implementations must run in O(1) with respect to the serialized value's + /// runtime-sized contents. Return `None` if determining the size requires + /// traversing a slice, collection, string contents, or another variable-size value. + /// + /// ## Safety + /// + /// The returned upper bound must never be smaller than the serialized representation. + #[inline] + fn size_hint(&self) -> Option { + None + } + + /// Returns a trusted upper bound valid for every value of this type. + /// + /// This is used to compute O(1) hints for runtime-sized homogeneous containers. + /// The returned upper bound must never be smaller than any serialized value + /// of this type. + /// + /// ## Safety + /// + /// The returned upper bound must never be smaller than any serialized value of this type. + #[inline] + fn max_size() -> Option + where + Self: Sized, + { + None + } } /// Deserializes a value of type `T` from a MessagePack-encoded byte slice. @@ -142,9 +166,8 @@ pub fn from_msgpack<'a, T: FromMessagePack<'a>>(data: &'a [u8]) -> Result { /// } /// ``` pub fn to_msgpack_vec(value: &T) -> Result> { - // SAFETY: a returned size is required to match the immediately following write. - let mut writer = match value.size() { - Some(size) => write::VecWriter::with_capacity(size.get()), + let mut writer = match value.size_hint() { + Some(hint) => write::VecWriter::with_capacity(hint.upper_bound()), None => write::VecWriter::new(), }; value.write(&mut writer)?; @@ -177,22 +200,24 @@ pub fn to_msgpack_vec(value: &T) -> Result> { /// } /// ``` pub fn to_msgpack(value: &T, buf: &mut [u8]) -> Result { - // SAFETY: `size` guarantees the byte count of the immediately following write. - if let Some(size) = value.size() { - let size = size.get(); - if size > buf.len() { - return Err(Error::BufferTooSmall); - } - // SAFETY: the exact output range was validated above. - let mut writer = unsafe { write::SliceWriter::new_unchecked(&mut buf[..size]) }; - value.write(&mut writer)?; - debug_assert_eq!(writer.position(), size); - Ok(writer.position()) - } else { - let mut writer = write::SliceWriter::new(buf); + if let Some(hint) = value.size_hint() + && hint.upper_bound() <= buf.len() + { + // SAFETY: the trusted hint guarantees that the write fits in `buf`. + let mut writer = unsafe { write::SliceWriter::new_unchecked(buf) }; value.write(&mut writer)?; - Ok(writer.position()) + debug_assert!(writer.position() <= hint.upper_bound()); + return Ok(writer.position()); } + to_msgpack_checked(value, buf) +} + +#[cold] +#[inline(never)] +fn to_msgpack_checked(value: &T, buf: &mut [u8]) -> Result { + let mut writer = write::SliceWriter::new(buf); + value.write(&mut writer)?; + Ok(writer.position()) } /// Serializes a value of type `T` into the I/O stream. diff --git a/zerompk/tests/serialize.rs b/zerompk/tests/serialize.rs index c53d7ad..ab08b2e 100644 --- a/zerompk/tests/serialize.rs +++ b/zerompk/tests/serialize.rs @@ -336,6 +336,25 @@ fn test_to_msgpack_with_small_buffer() { assert!(matches!(err, zerompk::Error::BufferTooSmall)); } +#[test] +fn test_trusted_size_hint_bounds_runtime_sized_container() { + let values = vec![0u64, 127, 128, u64::MAX]; + let hint = zerompk::ToMessagePack::size_hint(&values).unwrap(); + let max_element = ::max_size().unwrap(); + assert_eq!( + hint.upper_bound(), + 1 + values.len() * max_element.upper_bound() + ); + + let mut buf = vec![0u8; hint.upper_bound()]; + let written = zerompk::to_msgpack(&values, &mut buf).unwrap(); + assert!(written <= hint.upper_bound()); + assert_eq!( + zerompk::from_msgpack::>(&buf[..written]).unwrap(), + values + ); +} + #[test] #[cfg(feature = "std")] fn test_write_read_msgpack_std_io() { diff --git a/zerompk_derive/src/lib.rs b/zerompk_derive/src/lib.rs index 4a2f8c7..507b977 100644 --- a/zerompk_derive/src/lib.rs +++ b/zerompk_derive/src/lib.rs @@ -778,6 +778,38 @@ fn build_write_expr( } } +fn build_size_expr( + value: proc_macro2::TokenStream, + ty: &Type, + cfg: Option<&FieldConfig>, +) -> proc_macro2::TokenStream { + if is_bin_type(ty) && should_use_bin(ty, cfg) { + quote! {{ + let __len = (#value).len(); + let __header = if __len <= u8::MAX as usize { + 2usize + } else if __len <= u16::MAX as usize { + 3usize + } else { + 5usize + }; + let __size = __len.checked_add(__header)?; + // SAFETY: MessagePack binary headers depend only on payload length. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) + }} + } else { + quote! { ::zerompk::ToMessagePack::size_hint(&(#value)) } + } +} + +fn build_max_size_expr(ty: &Type, cfg: Option<&FieldConfig>) -> proc_macro2::TokenStream { + if is_bin_type(ty) && should_use_bin(ty, cfg) { + quote! { ::core::option::Option::<::zerompk::TrustedSizeHint>::None } + } else { + quote! { <#ty as ::zerompk::ToMessagePack>::max_size() } + } +} + fn build_named_array_slots( fields: &syn::FieldsNamed, configs: &[FieldConfig], @@ -1152,13 +1184,18 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - let size = if is_dense_sequential - && field_configs.iter().all(|cfg| cfg.as_bytes.is_none()) - && tys + let size = if is_dense_sequential { + let field_sizes: Vec<_> = names + .iter() + .zip(tys.iter()) + .zip(field_configs.iter()) + .map(|((name, ty), cfg)| build_size_expr(quote! { self.#name }, ty, Some(cfg))) + .collect(); + let field_max_sizes: Vec<_> = tys .iter() .zip(field_configs.iter()) - .all(|(ty, cfg)| !(is_bin_type(ty) && should_use_bin(ty, Some(cfg)))) - { + .map(|(ty, cfg)| build_max_size_expr(ty, Some(cfg))) + .collect(); let header_size = if array_len < 16 { 1usize } else if array_len <= u16::MAX as usize { @@ -1168,14 +1205,24 @@ fn expand_array_struct(data: &DataStruct) -> Result { }; quote! { #[inline] - fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { let mut __size = #header_size; #( - let __field_size = ::zerompk::ToMessagePack::size(&self.#names)?; - __size = __size.checked_add(__field_size.get())?; + let __field_size = (#field_sizes)?; + __size = __size.checked_add(__field_size.upper_bound())?; )* // SAFETY: the header and every serialized field have exact-size proofs. - ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(__size) }) + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) + } + + #[inline] + fn max_size() -> ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + __size = __size.checked_add((#field_max_sizes)?.upper_bound())?; + )* + // SAFETY: the header is exact and every field is bounded above. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) } } } else { @@ -1202,6 +1249,8 @@ fn expand_array_struct(data: &DataStruct) -> Result { let cfg = &field_configs[0]; let write_expr = build_write_expr(quote! { self.0 }, &ty, Some(cfg)); let read_expr = build_read_expr(&ty, Some(cfg)); + let size_expr = build_size_expr(quote! { self.0 }, &ty, Some(cfg)); + let max_size_expr = build_max_size_expr(&ty, Some(cfg)); let write = quote! { #write_expr @@ -1216,7 +1265,17 @@ fn expand_array_struct(data: &DataStruct) -> Result { return Ok(ImplBody { write, read, - size: quote! {}, + size: quote! { + #[inline] + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { + #size_expr + } + + #[inline] + fn max_size() -> ::core::option::Option<::zerompk::TrustedSizeHint> { + #max_size_expr + } + }, }); } @@ -1313,13 +1372,20 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - let size = if is_dense_sequential - && field_configs.iter().all(|cfg| cfg.as_bytes.is_none()) - && tys + let size = if is_dense_sequential { + let field_sizes: Vec<_> = idx + .iter() + .zip(tys.iter()) + .zip(field_configs.iter()) + .map(|((index, ty), cfg)| { + build_size_expr(quote! { self.#index }, ty, Some(cfg)) + }) + .collect(); + let field_max_sizes: Vec<_> = tys .iter() .zip(field_configs.iter()) - .all(|(ty, cfg)| !(is_bin_type(ty) && should_use_bin(ty, Some(cfg)))) - { + .map(|(ty, cfg)| build_max_size_expr(ty, Some(cfg))) + .collect(); let header_size = if array_len < 16 { 1usize } else if array_len <= u16::MAX as usize { @@ -1329,14 +1395,24 @@ fn expand_array_struct(data: &DataStruct) -> Result { }; quote! { #[inline] - fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { let mut __size = #header_size; #( - let __field_size = ::zerompk::ToMessagePack::size(&self.#idx)?; - __size = __size.checked_add(__field_size.get())?; + let __field_size = (#field_sizes)?; + __size = __size.checked_add(__field_size.upper_bound())?; )* // SAFETY: the header and every serialized field have exact-size proofs. - ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(__size) }) + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) + } + + #[inline] + fn max_size() -> ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + __size = __size.checked_add((#field_max_sizes)?.upper_bound())?; + )* + // SAFETY: the header is exact and every field is bounded above. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) } } } else { @@ -1356,9 +1432,14 @@ fn expand_array_struct(data: &DataStruct) -> Result { }, size: quote! { #[inline] - fn size(&self) -> ::core::option::Option<::zerompk::ExactSize> { + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { + // SAFETY: unit structs always encode as one nil byte. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(1) }) + } + #[inline] + fn max_size() -> ::core::option::Option<::zerompk::TrustedSizeHint> { // SAFETY: unit structs always encode as one nil byte. - ::core::option::Option::Some(unsafe { ::zerompk::ExactSize::new_unchecked(1) }) + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(1) }) } }, }), @@ -1389,6 +1470,21 @@ fn expand_map_struct(data: &DataStruct, allow_unknown_fields: bool) -> Result>()?; let (field_indices, key_lits) = parse_named_map_keys(fields, &field_configs)?; let encoded_key_lits: Vec<_> = key_lits.iter().map(encode_static_string).collect(); + let encoded_key_sizes: Vec<_> = key_lits + .iter() + .map(|key| { + let len = key.value().len(); + len + if len < 32 { + 1 + } else if len <= u8::MAX as usize { + 2 + } else if len <= u16::MAX as usize { + 3 + } else { + 5 + } + }) + .collect(); let count = field_indices.len(); let names: Vec<_> = field_indices .iter() @@ -1439,6 +1535,23 @@ fn expand_map_struct(data: &DataStruct, allow_unknown_fields: bool) -> Result = names + .iter() + .zip(tys.iter()) + .enumerate() + .map(|(idx, (name, ty))| { + let cfg = &field_configs[field_indices[idx]]; + build_size_expr(quote! { self.#name }, ty, Some(cfg)) + }) + .collect(); + let value_max_sizes: Vec<_> = tys + .iter() + .enumerate() + .map(|(idx, ty)| { + let cfg = &field_configs[field_indices[idx]]; + build_max_size_expr(ty, Some(cfg)) + }) + .collect(); let write = quote! { writer.write_map_len(#count)?; @@ -1560,11 +1673,38 @@ fn expand_map_struct(data: &DataStruct, allow_unknown_fields: bool) -> Result ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + __size = __size.checked_add(#encoded_key_sizes)?; + __size = __size.checked_add((#value_sizes)?.upper_bound())?; + )* + // SAFETY: the map header, static keys, and field values have exact sizes. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) + } + + #[inline] + fn max_size() -> ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + __size = __size.checked_add(#encoded_key_sizes)?; + __size = __size.checked_add((#value_max_sizes)?.upper_bound())?; + )* + // SAFETY: the map header and keys are exact and every value is bounded above. + ::core::option::Option::Some(unsafe { ::zerompk::TrustedSizeHint::new_unchecked(__size) }) + } + }; + + Ok(ImplBody { write, read, size }) } fn expand_c_enum(data: &DataEnum, repr: CEnumRepr) -> Result { From c8e264b2d60939122ee51ae7606d9568e2210253 Mon Sep 17 00:00:00 2001 From: nuskey8 Date: Sat, 29 Aug 2026 22:50:24 +0900 Subject: [PATCH 4/4] fix: limit trusted size hint preallocation --- fuzz/Cargo.lock | 4 ++-- fuzz/fuzz_targets/roundtrip.rs | 7 +++++++ zerompk/src/lib.rs | 8 +++++++- zerompk/src/write.rs | 10 ++++++---- zerompk/tests/serialize.rs | 18 ++++++++++++++++++ 5 files changed, 40 insertions(+), 7 deletions(-) diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index 3410a74..8841eb4 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -157,14 +157,14 @@ checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" [[package]] name = "zerompk" -version = "0.4.2" +version = "0.7.1" dependencies = [ "zerompk_derive", ] [[package]] name = "zerompk_derive" -version = "0.4.2" +version = "0.7.1" dependencies = [ "proc-macro2", "quote", diff --git a/fuzz/fuzz_targets/roundtrip.rs b/fuzz/fuzz_targets/roundtrip.rs index 542859f..03c34e8 100644 --- a/fuzz/fuzz_targets/roundtrip.rs +++ b/fuzz/fuzz_targets/roundtrip.rs @@ -24,6 +24,13 @@ struct FuzzNested { fuzz_target!(|value: FuzzValue| { if let Ok(buf) = zerompk::to_msgpack_vec(&value) { + if let Some(hint) = value.size_hint() { + assert!(buf.len() <= hint.upper_bound()); + let mut hinted = vec![0u8; hint.upper_bound()]; + let written = zerompk::to_msgpack(&value, &mut hinted) + .expect("trusted size hint must bound unchecked serialization"); + assert_eq!(&hinted[..written], buf); + } let decoded = zerompk::from_msgpack::(&buf) .expect("roundtrip decode must succeed for encoded data"); assert_eq!(decoded, value); diff --git a/zerompk/src/lib.rs b/zerompk/src/lib.rs index 2da1728..e65bb33 100644 --- a/zerompk/src/lib.rs +++ b/zerompk/src/lib.rs @@ -22,6 +22,7 @@ use alloc::vec::Vec; pub use error::{Error, Result}; pub use read::{Read, SliceReader, Tag}; pub use value::Value; + pub use write::Write; extern crate alloc; @@ -166,8 +167,13 @@ pub fn from_msgpack<'a, T: FromMessagePack<'a>>(data: &'a [u8]) -> Result { /// } /// ``` pub fn to_msgpack_vec(value: &T) -> Result> { + /// The maximum size hint preallocation to avoid excessive memory usage. + const MAX_SIZE_HINT_PREALLOC: usize = 16 * 1024 * 1024; + let mut writer = match value.size_hint() { - Some(hint) => write::VecWriter::with_capacity(hint.upper_bound()), + Some(hint) => { + write::VecWriter::with_capacity_hint(hint.upper_bound().min(MAX_SIZE_HINT_PREALLOC)) + } None => write::VecWriter::new(), }; value.write(&mut writer)?; diff --git a/zerompk/src/write.rs b/zerompk/src/write.rs index 398f931..494a15b 100644 --- a/zerompk/src/write.rs +++ b/zerompk/src/write.rs @@ -551,10 +551,12 @@ impl VecWriter { VecWriter { buffer: Vec::new() } } - pub fn with_capacity(capacity: usize) -> Self { - VecWriter { - buffer: Vec::with_capacity(capacity), - } + pub(crate) fn with_capacity_hint(capacity: usize) -> Self { + let mut buffer = Vec::new(); + // A hint is only an optimization. Allocation failure must not abort + // serialization before any output has actually required this memory. + let _ = buffer.try_reserve(capacity); + VecWriter { buffer } } pub fn into_vec(self) -> Vec { diff --git a/zerompk/tests/serialize.rs b/zerompk/tests/serialize.rs index ab08b2e..d4d1653 100644 --- a/zerompk/tests/serialize.rs +++ b/zerompk/tests/serialize.rs @@ -355,6 +355,24 @@ fn test_trusted_size_hint_bounds_runtime_sized_container() { ); } +#[test] +fn test_huge_trusted_size_hint_does_not_force_huge_preallocation() { + struct TinyValue; + + impl zerompk::ToMessagePack for TinyValue { + fn write(&self, writer: &mut W) -> zerompk::Result<()> { + writer.write_nil() + } + + fn size_hint(&self) -> Option { + // SAFETY: one byte is less than this deliberately loose upper bound. + Some(unsafe { zerompk::TrustedSizeHint::new_unchecked(usize::MAX) }) + } + } + + assert_eq!(zerompk::to_msgpack_vec(&TinyValue).unwrap(), [0xc0]); +} + #[test] #[cfg(feature = "std")] fn test_write_read_msgpack_std_io() {