make tree_hash() support CLVM with backrefs

benches/tree-hash.rs

@@ -1,15 +1,21 @@
 use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, node_to_bytes_backrefs};
 use clvmr::Allocator;
-use criterion::{criterion_group, criterion_main, Criterion, SamplingMode};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use std::fs::read_to_string;
 use std::time::Instant;
 
 fn run(c: &mut Criterion) {
     let mut group = c.benchmark_group("tree-hash");
-    group.sample_size(20);
-    group.sampling_mode(SamplingMode::Flat);
 
-    for name in &["block-4671894", "block-834752", "block-225758"] {
+    for name in &[
+        "block-4671894",
+        "block-834752",
+        "block-225758",
+        "block-c2a8df0d",
+        "block-1ee588dc",
+        "block-6fe59b24",
+        "block-b45268ac",
+    ] {
         let filename = format!("generator-tests/{name}.txt");
         println!("file: {filename}");
         let test_file = read_to_string(filename).expect("test file not found");
@@ -22,14 +28,24 @@ fn run(c: &mut Criterion) {
             node_to_bytes_backrefs(&a, input).expect("failed to compress generator")
         };
 
-        for (gen, name_suffix) in &[(generator, ""), (compressed_generator, "-compressed")] {
+        for (gen, name_suffix) in &[(&generator, ""), (&compressed_generator, "-compressed")] {
             let mut a = Allocator::new();
             let gen = node_from_bytes_backrefs(&mut a, &gen).expect("parse generator");
 
             group.bench_function(format!("tree-hash {name}{name_suffix}"), |b| {
                 b.iter(|| {
                     let start = Instant::now();
-                    clvm_utils::tree_hash(&a, gen);
+                    let _ = black_box(clvm_utils::tree_hash(&a, gen));
+                    start.elapsed()
+                })
+            });
+        }
+
+        for (gen, name_suffix) in &[(&generator, ""), (&compressed_generator, "-compressed")] {
+            group.bench_function(format!("tree-hash-from-stream {name}{name_suffix}"), |b| {
+                b.iter(|| {
+                    let start = Instant::now();
+                    let _ = black_box(clvm_utils::tree_hash_from_bytes(gen));
                     start.elapsed()
                 })
             });

chia-protocol/src/program.rs

@@ -320,8 +320,7 @@ impl Program {
     }
 
     fn get_tree_hash(&self) -> crate::Bytes32 {
-        let mut cursor = Cursor::new(self.0.as_ref());
-        clvmr::serde::tree_hash_from_stream(&mut cursor)
+        clvm_utils::tree_hash_from_bytes(self.0.as_ref())
             .unwrap()
             .into()
     }

clvm-utils/src/tree_hash.rs

@@ -1,5 +1,7 @@
 use clvmr::allocator::{Allocator, NodePtr, SExp};
+use clvmr::serde::node_from_bytes_backrefs;
 use clvmr::sha2::{Digest, Sha256};
+use std::io;
 
 enum TreeOp {
     SExp(NodePtr),
@@ -49,6 +51,12 @@ pub fn tree_hash(a: &Allocator, node: NodePtr) -> [u8; 32] {
     hashes[0]
 }
 
+pub fn tree_hash_from_bytes(buf: &[u8]) -> io::Result<[u8; 32]> {
+    let mut a = Allocator::new();
+    let node = node_from_bytes_backrefs(&mut a, buf)?;
+    Ok(tree_hash(&a, node))
+}
+
 #[test]
 fn test_tree_hash() {
     let mut a = Allocator::new();
@@ -110,3 +118,31 @@ fn test_tree_hash() {
         assert_eq!(tree_hash(&a, root2), sha256.finalize().as_slice());
     }
 }
+
+#[test]
+fn test_tree_hash_from_bytes() {
+    use clvmr::serde::{node_to_bytes, node_to_bytes_backrefs};
+
+    let mut a = Allocator::new();
+    let atom1 = a.new_atom(&[1, 2, 3]).unwrap();
+    let atom2 = a.new_atom(&[4, 5, 6]).unwrap();
+    let node1 = a.new_pair(atom1, atom2).unwrap();
+    let node2 = a.new_pair(atom2, atom1).unwrap();
+
+    let node1 = a.new_pair(node1, node1).unwrap();
+    let node2 = a.new_pair(node2, node2).unwrap();
+
+    let root = a.new_pair(node1, node2).unwrap();
+
+    let serialized_clvm = node_to_bytes(&a, root).expect("node_to_bytes");
+    let serialized_clvm_backrefs =
+        node_to_bytes_backrefs(&a, root).expect("node_to_bytes_backrefs");
+
+    let hash1 = tree_hash_from_bytes(&serialized_clvm).expect("tree_hash_from_bytes");
+    let hash2 = tree_hash_from_bytes(&serialized_clvm_backrefs).expect("tree_hash_from_bytes");
+    let hash3 = tree_hash(&a, root);
+
+    assert!(serialized_clvm.len() > serialized_clvm_backrefs.len());
+    assert_eq!(hash1, hash2);
+    assert_eq!(hash1, hash3);
+}

wheel/Cargo.toml

@@ -24,5 +24,6 @@ chia-bls = { version = "=0.5.1", path = "../chia-bls", features = ["py-bindings"
 chia-protocol = { version = "=0.5.1", path = "../chia-protocol", features = ["py-bindings"]  }
 chia-traits = { version = "=0.5.2", path = "../chia-traits", features = ["py-bindings"]  }
 clvm-traits = { version = "=0.5.2", path = "../clvm-traits", features = ["derive", "py-bindings"] }
+clvm-utils = { version = "=0.5.1", path = "../clvm-utils" }
 chia_py_streamable_macro = { version = "=0.5.1", path = "../chia_py_streamable_macro" }
 chia_streamable_macro = { version = "=0.3.0", path = "../chia_streamable_macro" }

wheel/src/api.rs

@@ -35,7 +35,7 @@ use chia_protocol::{
     SubEpochSegments, SubEpochSummary, SubSlotData, SubSlotProofs, TimestampedPeerInfo,
     TransactionAck, TransactionsInfo, UnfinishedBlock, VDFInfo, VDFProof, WeightProof,
 };
-use clvmr::serde::tree_hash_from_stream;
+use clvm_utils::tree_hash_from_bytes;
 use clvmr::{ENABLE_BLS_OPS_OUTSIDE_GUARD, ENABLE_FIXED_DIV, LIMIT_HEAP, NO_UNKNOWN_OPS};
 use pyo3::buffer::PyBuffer;
 use pyo3::exceptions::PyRuntimeError;
@@ -87,8 +87,7 @@ pub fn tree_hash(py: Python, blob: PyBuffer<u8>) -> PyResult<&PyBytes> {
     }
     let slice =
         unsafe { std::slice::from_raw_parts(blob.buf_ptr() as *const u8, blob.len_bytes()) };
-    let mut input = std::io::Cursor::<&[u8]>::new(slice);
-    Ok(PyBytes::new(py, &tree_hash_from_stream(&mut input)?))
+    Ok(PyBytes::new(py, &tree_hash_from_bytes(slice)?))
 }
 
 #[allow(clippy::too_many_arguments)]