Skip to content

Commit

Permalink
feat: add borsh::object_length helper (#236)
Browse files Browse the repository at this point in the history
* feat: add `borsh::object_length` helper

* chore: add bench to illustrate diff

* chore: check overflow on `usize` addition
  • Loading branch information
dj8yfo authored Sep 28, 2023
1 parent d2c63ac commit bb5248e
Show file tree
Hide file tree
Showing 6 changed files with 145 additions and 2 deletions.
4 changes: 4 additions & 0 deletions benchmarks/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
55 changes: 55 additions & 0 deletions benchmarks/benches/object_length.rs
Original file line number Diff line number Diff line change
@@ -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<T>(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<Vec<u8>> = 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::<ValidatorStake>("ser_account", 3, c);
}
criterion_group!(ser_length, ser_length_validator_stake,);
criterion_main!(ser_length);
2 changes: 1 addition & 1 deletion borsh/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
5 changes: 5 additions & 0 deletions borsh/src/nostd_io.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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",
}
}
}
Expand Down
34 changes: 33 additions & 1 deletion borsh/src/ser/helpers.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -21,3 +21,35 @@ where
{
value.serialize(&mut writer)
}

/// Serializes an object without allocation to compute and return its length
pub fn object_length<T>(value: &T) -> Result<usize>
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<usize> {
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)
}
47 changes: 47 additions & 0 deletions borsh/tests/test_simple_structs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> = BTreeMap::new();
map.insert("test".into(), "test".into());
let mut set: BTreeSet<u64> = 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);
}

0 comments on commit bb5248e

Please sign in to comment.