From bb5248e3bde451e992f5636dbc74ecdcfc59645b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?dj8yf0=CE=BCl?= <26653921+dj8yfo@users.noreply.github.com> Date: Thu, 28 Sep 2023 22:58:46 +0300 Subject: [PATCH] feat: add `borsh::object_length` helper (#236) * feat: add `borsh::object_length` helper * chore: add bench to illustrate diff * chore: check overflow on `usize` addition --- benchmarks/Cargo.toml | 4 +++ benchmarks/benches/object_length.rs | 55 +++++++++++++++++++++++++++++ borsh/src/lib.rs | 2 +- borsh/src/nostd_io.rs | 5 +++ borsh/src/ser/helpers.rs | 34 +++++++++++++++++- borsh/tests/test_simple_structs.rs | 47 ++++++++++++++++++++++++ 6 files changed, 145 insertions(+), 2 deletions(-) create mode 100644 benchmarks/benches/object_length.rs diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 4e4da3193..aad6202b9 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -36,5 +36,9 @@ harness = false name = "maps_sets_inner_de" harness = false +[[bench]] +name = "object_length" +harness = false + [features] default = ["borsh/std", "borsh/derive"] diff --git a/benchmarks/benches/object_length.rs b/benchmarks/benches/object_length.rs new file mode 100644 index 000000000..e69e62a78 --- /dev/null +++ b/benchmarks/benches/object_length.rs @@ -0,0 +1,55 @@ +use benchmarks::{Generate, ValidatorStake}; +use borsh::{to_vec, BorshSerialize}; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use rand::SeedableRng; + +fn ser_obj_length(group_name: &str, num_samples: usize, c: &mut Criterion) +where + for<'a> T: Generate + BorshSerialize + 'static, +{ + let mut rng = rand_xorshift::XorShiftRng::from_seed([0u8; 16]); + let mut group = c.benchmark_group(group_name); + + let objects: Vec<_> = (0..num_samples).map(|_| T::generate(&mut rng)).collect(); + let borsh_datas: Vec> = objects.iter().map(|t| to_vec(t).unwrap()).collect(); + let borsh_sizes: Vec<_> = borsh_datas.iter().map(|d| d.len()).collect(); + + for i in 0..objects.len() { + let size = borsh_sizes[i]; + let obj = &objects[i]; + assert_eq!( + borsh::to_vec(obj).unwrap().len(), + borsh::object_length(obj).unwrap() + ); + + let benchmark_param_display = format!("idx={}; size={}", i, size); + + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input( + BenchmarkId::new( + "borsh::to_vec(obj).unwrap().len()", + benchmark_param_display.clone(), + ), + obj, + |b, d| { + b.iter(|| borsh::to_vec(d).unwrap().len()); + }, + ); + group.bench_with_input( + BenchmarkId::new( + "borsh::object_length(obj).unwrap()", + benchmark_param_display.clone(), + ), + obj, + |b, d| { + b.iter(|| borsh::object_length(d).unwrap()); + }, + ); + } + group.finish(); +} +fn ser_length_validator_stake(c: &mut Criterion) { + ser_obj_length::("ser_account", 3, c); +} +criterion_group!(ser_length, ser_length_validator_stake,); +criterion_main!(ser_length); diff --git a/borsh/src/lib.rs b/borsh/src/lib.rs index 3aff5665e..87af12d7b 100644 --- a/borsh/src/lib.rs +++ b/borsh/src/lib.rs @@ -93,7 +93,7 @@ pub use schema::BorshSchema; pub use schema_helpers::{ max_serialized_size, schema_container_of, try_from_slice_with_schema, try_to_vec_with_schema, }; -pub use ser::helpers::{to_vec, to_writer}; +pub use ser::helpers::{object_length, to_vec, to_writer}; pub use ser::BorshSerialize; pub mod error; diff --git a/borsh/src/nostd_io.rs b/borsh/src/nostd_io.rs index 27fba509c..52061b817 100644 --- a/borsh/src/nostd_io.rs +++ b/borsh/src/nostd_io.rs @@ -151,6 +151,10 @@ pub enum ErrorKind { /// particular number of bytes but only a smaller number of bytes could be /// read. UnexpectedEof, + + /// An operation could not be completed, because it failed + /// to allocate enough memory. + OutOfMemory, } impl ErrorKind { @@ -174,6 +178,7 @@ impl ErrorKind { ErrorKind::Interrupted => "operation interrupted", ErrorKind::Other => "other os error", ErrorKind::UnexpectedEof => "unexpected end of file", + ErrorKind::OutOfMemory => "out of memory", } } } diff --git a/borsh/src/ser/helpers.rs b/borsh/src/ser/helpers.rs index 73761d72b..4f00d5ab4 100644 --- a/borsh/src/ser/helpers.rs +++ b/borsh/src/ser/helpers.rs @@ -1,6 +1,6 @@ use crate::BorshSerialize; use crate::__private::maybestd::vec::Vec; -use crate::io::{Result, Write}; +use crate::io::{ErrorKind, Result, Write}; pub(super) const DEFAULT_SERIALIZER_CAPACITY: usize = 1024; @@ -21,3 +21,35 @@ where { value.serialize(&mut writer) } + +/// Serializes an object without allocation to compute and return its length +pub fn object_length(value: &T) -> Result +where + T: BorshSerialize + ?Sized, +{ + // copy-paste of solution provided by @matklad + // in https://github.com/near/borsh-rs/issues/23#issuecomment-816633365 + struct LengthWriter { + len: usize, + } + impl Write for LengthWriter { + #[inline] + fn write(&mut self, buf: &[u8]) -> Result { + let res = self.len.checked_add(buf.len()); + self.len = match res { + Some(res) => res, + None => { + return Err(ErrorKind::OutOfMemory.into()); + } + }; + Ok(buf.len()) + } + #[inline] + fn flush(&mut self) -> Result<()> { + Ok(()) + } + } + let mut w = LengthWriter { len: 0 }; + value.serialize(&mut w)?; + Ok(w.len) +} diff --git a/borsh/tests/test_simple_structs.rs b/borsh/tests/test_simple_structs.rs index 5dad88eb1..647bf3fe3 100644 --- a/borsh/tests/test_simple_structs.rs +++ b/borsh/tests/test_simple_structs.rs @@ -176,3 +176,50 @@ fn test_ultimate_combined_all_features() { assert_eq!(decoded_f2.aa.len(), 2); assert!(decoded_f2.aa.iter().all(|f2_a| f2_a == &expected_a)); } + +#[test] +fn test_object_length() { + let mut map: BTreeMap = BTreeMap::new(); + map.insert("test".into(), "test".into()); + let mut set: BTreeSet = BTreeSet::new(); + set.insert(u64::MAX); + set.insert(100); + set.insert(103); + set.insert(109); + let cow_arr = [ + borrow::Cow::Borrowed("Hello1"), + borrow::Cow::Owned("Hello2".to_string()), + ]; + let a = A { + x: 1, + b: B { + x: 2, + y: 3, + c: C::C5(D { x: 1 }), + }, + y: 4.0, + z: "123".to_string(), + t: ("Hello".to_string(), 10), + btree_map_string: map.clone(), + btree_set_u64: set.clone(), + linked_list_string: vec!["a".to_string(), "b".to_string()].into_iter().collect(), + vec_deque_u64: vec![1, 2, 3].into_iter().collect(), + bytes: vec![5, 4, 3, 2, 1].into(), + bytes_mut: BytesMut::from(&[1, 2, 3, 4, 5][..]), + v: vec!["qwe".to_string(), "zxc".to_string()], + w: vec![0].into_boxed_slice(), + box_str: Box::from("asd"), + i: [4u8; 32], + u: Ok("Hello".to_string()), + lazy: Some(5), + c: borrow::Cow::Borrowed("Hello"), + cow_arr: borrow::Cow::Borrowed(&cow_arr), + range_u32: 12..71, + skipped: Some(6), + }; + let encoded_a_len = to_vec(&a).unwrap().len(); + + let len_helper_result = borsh::object_length(&a).unwrap(); + + assert_eq!(encoded_a_len, len_helper_result); +}