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

Author: dj8yfo

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"]

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<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);

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;
 

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",
         }
     }
 }

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<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)
+}

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<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);
+}