Merge pull request #523 from Chia-Network/serialized-length

move serialized_length functions into their own file

src/serde/mod.rs

@@ -9,6 +9,7 @@ mod parse_atom;
 mod read_cache_lookup;
 mod ser;
 mod ser_br;
+mod serialized_length;
 mod tools;
 mod utils;
 pub mod write_atom;
@@ -23,6 +24,7 @@ pub use incremental::{Serializer, UndoState};
 pub use object_cache::{serialized_length, treehash, ObjectCache};
 pub use ser::{node_to_bytes, node_to_bytes_limit};
 pub use ser_br::{node_to_bytes_backrefs, node_to_bytes_backrefs_limit};
+pub use serialized_length::{serialized_length_atom, serialized_length_small_number};
 pub use tools::{
     serialized_length_from_bytes, serialized_length_from_bytes_trusted, tree_hash_from_stream,
 };

src/serde/object_cache.rs

@@ -10,6 +10,7 @@ use crate::allocator::{Allocator, NodePtr, SExp};
 use std::collections::HashMap;
 type CachedFunction<T> = fn(&mut ObjectCache<T>, &Allocator, NodePtr) -> Option<T>;
 use super::bytes32::{hash_blobs, Bytes32};
+use crate::serde::serialized_length_atom;
 
 pub struct ObjectCache<T> {
     cache: HashMap<NodePtr, T>,
@@ -130,20 +131,7 @@ pub fn serialized_length(
         },
         SExp::Atom => {
             let buf = allocator.atom(node);
-            let lb: u64 = buf.as_ref().len().try_into().unwrap_or(u64::MAX);
-            Some(if lb == 0 || (lb == 1 && buf.as_ref()[0] < 128) {
-                1
-            } else if lb < 0x40 {
-                1 + lb
-            } else if lb < 0x2000 {
-                2 + lb
-            } else if lb < 0x100000 {
-                3 + lb
-            } else if lb < 0x8000000 {
-                4 + lb
-            } else {
-                5 + lb
-            })
+            Some(serialized_length_atom(buf.as_ref()) as u64)
         }
     }
 }

src/serde/serialized_length.rs

@@ -0,0 +1,56 @@
+use crate::allocator::len_for_value;
+
+pub fn serialized_length_atom(buf: &[u8]) -> u32 {
+    let lb = buf.len() as u32;
+    if lb == 0 || (lb == 1 && buf[0] < 128) {
+        1
+    } else if lb < 0x40 {
+        1 + lb
+    } else if lb < 0x2000 {
+        2 + lb
+    } else if lb < 0x100000 {
+        3 + lb
+    } else if lb < 0x8000000 {
+        4 + lb
+    } else {
+        5 + lb
+    }
+}
+
+pub fn serialized_length_small_number(val: u32) -> u32 {
+    len_for_value(val) as u32 + 1
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case(&[], 1)]
+    #[case(&[1], 1)]
+    #[case(&[0x7f], 1)]
+    #[case(&[0x80], 2)]
+    #[case(&[0x81], 2)]
+    #[case(&[0x80, 0], 3)]
+    #[case(&[1; 0x3f], 0x40)]
+    #[case(&[1; 0x40], 0x42)]
+    fn test_serialized_length_atom(#[case] atom: &[u8], #[case] expect: u32) {
+        assert_eq!(serialized_length_atom(atom), expect);
+    }
+
+    #[rstest]
+    #[case(0, 1)]
+    #[case(1, 2)]
+    #[case(0x7f, 2)]
+    #[case(0x80, 3)]
+    #[case(0x7fff, 3)]
+    #[case(0x7fffff, 4)]
+    #[case(0x800000, 5)]
+    #[case(0x7fffffff, 5)]
+    #[case(0x80000000, 6)]
+    #[case(0xffffffff, 6)]
+    fn test_serialized_length_small_number(#[case] value: u32, #[case] expect: u32) {
+        assert_eq!(serialized_length_small_number(value), expect);
+    }
+}