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/impl.rs b/zerompk/src/impl.rs index 1422679..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) => { + ($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 @@ -30,21 +84,144 @@ macro_rules! impl_scalar { fn write_slice(values: &[Self], writer: &mut W) -> crate::Result<()> { writer.$write_slice_fn(values) } + + #[inline(always)] + fn size_hint(&self) -> Option { + // SAFETY: primitive encodings are completely determined by their value. + 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) }) + } } }; } -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, + 1 +); +impl_scalar!( + i8, + write_i8, + write_i8_slice, + read_i8, + |v: i8| if (-32..=127).contains(&v) { 1 } else { 2 }, + 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 + }, + 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 + }, + 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 + }, + 9 +); +impl_scalar!( + u8, + write_u8, + write_u8_slice, + read_u8, + |v: u8| if v <= 127 { 1 } else { 2 }, + 2 +); +impl_scalar!( + u16, + write_u16, + write_u16_slice, + read_u16, + |v: u16| if v <= 127 { + 1 + } else if v <= 255 { + 2 + } else { + 3 + }, + 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 + }, + 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 + }, + 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)] @@ -69,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 { @@ -94,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 { @@ -115,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() + } } // ------------------------------------------------------------------------------- @@ -136,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) }) + } } // ------------------------------------------------------------------------------- @@ -163,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] @@ -187,6 +418,10 @@ impl ToMessagePack for [T] { writer.write_array_len(self.len())?; T::write_slice(self, writer) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } impl<'a, T: FromMessagePack<'a>, const N: usize> FromMessagePack<'a> for [T; N] { @@ -234,6 +469,14 @@ impl ToMessagePack for [T; N] { writer.write_array_len(N)?; T::write_slice(self, writer) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(N) + } + + fn max_size() -> Option { + sequence_size_hint::(N) + } } impl<'a> FromMessagePack<'a> for alloc::string::String { @@ -254,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> { @@ -271,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]> @@ -294,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()) + } } // ------------------------------------------------------------------------------- @@ -305,6 +560,11 @@ impl ToMessagePack for &T { fn write(&self, writer: &mut W) -> crate::Result<()> { T::write(self, writer) } + + #[inline] + fn size_hint(&self) -> Option { + T::size_hint(self) + } } impl ToMessagePack for &mut T { @@ -312,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) + } } // ------------------------------------------------------------------------------- @@ -336,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> @@ -372,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) }) + } } // ------------------------------------------------------------------------------- @@ -395,6 +695,10 @@ impl ToMessagePack for alloc::vec::Vec { fn write(&self, writer: &mut W) -> crate::Result<()> { self.as_slice().write(writer) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } impl<'a, T: FromMessagePack<'a>> FromMessagePack<'a> for alloc::collections::VecDeque { @@ -420,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 { @@ -444,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 { @@ -468,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> @@ -497,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 { @@ -524,6 +844,10 @@ impl ToMessagePack for alloc::collections::BinaryHeap } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } #[cfg(feature = "std")] @@ -553,6 +877,10 @@ impl ToMessagePack for std::collections::HashSet { } Ok(()) } + + fn size_hint(&self) -> Option { + sequence_size_hint::(self.len()) + } } #[cfg(feature = "std")] @@ -587,6 +915,10 @@ impl ToMessagePack for std::collections::Has } Ok(()) } + + fn size_hint(&self) -> Option { + map_size_hint::(self.len()) + } } // ------------------------------------------------------------------------------- @@ -608,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")] @@ -625,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 { @@ -640,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() + } } // ------------------------------------------------------------------------------- @@ -659,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 { @@ -679,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 b2bcae1..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; @@ -46,6 +47,27 @@ pub trait FromMessagePackOwned: for<'a> FromMessagePack<'a> {} impl FromMessagePackOwned for T where T: for<'a> FromMessagePack<'a> {} +/// A trusted upper bound for the next serialization of a value. +pub struct TrustedSizeHint(usize); + +impl TrustedSizeHint { + /// Creates a trusted encoded-size upper bound. + /// + /// # Safety + /// 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(upper_bound: usize) -> Self { + Self(upper_bound) + } + + #[inline(always)] + pub const fn upper_bound(&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. @@ -62,6 +84,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. @@ -113,7 +167,15 @@ 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(); + /// 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(hint.upper_bound().min(MAX_SIZE_HINT_PREALLOC)) + } + None => write::VecWriter::new(), + }; value.write(&mut writer)?; Ok(writer.into_vec()) } @@ -144,6 +206,21 @@ pub fn to_msgpack_vec(value: &T) -> Result> { /// } /// ``` pub fn to_msgpack(value: &T, buf: &mut [u8]) -> Result { + 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)?; + 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()) diff --git a/zerompk/src/write.rs b/zerompk/src/write.rs index 847bc41..494a15b 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,14 @@ impl VecWriter { VecWriter { buffer: Vec::new() } } + 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 { self.buffer } diff --git a/zerompk/tests/serialize.rs b/zerompk/tests/serialize.rs index c53d7ad..d4d1653 100644 --- a/zerompk/tests/serialize.rs +++ b/zerompk/tests/serialize.rs @@ -336,6 +336,43 @@ 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] +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() { diff --git a/zerompk_derive/src/lib.rs b/zerompk_derive/src/lib.rs index 2974dd8..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], @@ -977,7 +1009,7 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result quote! { @@ -986,6 +1018,8 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result(&self, writer: &mut W) -> ::core::result::Result<(), ::zerompk::Error> { #write } + + #size } }, DeriveKind::From => { @@ -1026,6 +1060,7 @@ fn expand(input: DeriveInput, kind: DeriveKind) -> Result Result { @@ -1149,7 +1184,52 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - Ok(ImplBody { write, read }) + 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()) + .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 { + 3 + } else { + 5 + }; + quote! { + #[inline] + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + 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::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 { + quote! {} + }; + + Ok(ImplBody { write, read, size }) } Fields::Unnamed(fields) => { let count = fields.unnamed.len(); @@ -1169,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 @@ -1180,7 +1262,21 @@ fn expand_array_struct(data: &DataStruct) -> Result { Ok(Self(__f0)) }; - return Ok(ImplBody { write, read }); + return Ok(ImplBody { + write, + read, + 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 + } + }, + }); } let idx: Vec<_> = (0..count).map(syn::Index::from).collect(); @@ -1276,7 +1372,54 @@ fn expand_array_struct(data: &DataStruct) -> Result { } }; - Ok(ImplBody { write, read }) + 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()) + .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 { + 3 + } else { + 5 + }; + quote! { + #[inline] + fn size_hint(&self) -> ::core::option::Option<::zerompk::TrustedSizeHint> { + let mut __size = #header_size; + #( + 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::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 { + quote! {} + }; + + Ok(ImplBody { write, read, size }) } Fields::Unit => Ok(ImplBody { write: quote! { @@ -1287,6 +1430,18 @@ fn expand_array_struct(data: &DataStruct) -> Result { reader.read_nil()?; Ok(Self) }, + size: quote! { + #[inline] + 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::TrustedSizeHint::new_unchecked(1) }) + } + }, }), } } @@ -1315,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() @@ -1365,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)?; @@ -1486,7 +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 { @@ -1550,7 +1768,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 +1924,11 @@ fn expand_enum(data: &DataEnum, repr: Repr) -> Result { } }; - Ok(ImplBody { write, read }) + Ok(ImplBody { + write, + read, + size: quote! {}, + }) } fn build_enum_variant_payload(