Add MerkleSet struct to Rust

crates/chia-consensus/Cargo.toml

@@ -32,6 +32,7 @@ num-traits = "0.2.15"
 rstest = "0.16.0"
 text-diff = "0.4.0"
 criterion = "0.5.1"
+rand = { version = "0.8.5", features = [ "small_rng" ] }
 
 [lib]
 bench = false
@@ -43,3 +44,7 @@ harness = false
 [[bench]]
 name = "tree-hash"
 harness = false
+
+[[bench]]
+name = "merkle-set"
+harness = false

crates/chia-consensus/benches/merkle-set.rs

@@ -0,0 +1,60 @@
+use chia_consensus::merkle_tree::MerkleSet;
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use rand::rngs::SmallRng;
+use rand::{Rng, SeedableRng};
+use std::time::Instant;
+
+fn run(c: &mut Criterion) {
+    let mut group = c.benchmark_group("merkle-set");
+
+    let mut rng = SmallRng::seed_from_u64(1337);
+
+    let mut leafs = Vec::<[u8; 32]>::with_capacity(1000);
+    for _ in 0..1000 {
+        let mut item = [0_u8; 32];
+        rng.fill(&mut item);
+        leafs.push(item);
+    }
+
+    group.bench_function("from_leafs", |b| {
+        b.iter(|| {
+            let start = Instant::now();
+            let _ = black_box(MerkleSet::from_leafs(&mut leafs));
+            start.elapsed()
+        })
+    });
+
+    let tree = MerkleSet::from_leafs(&mut leafs);
+
+    group.bench_function("generate_proof", |b| {
+        b.iter(|| {
+            let start = Instant::now();
+            for item in &leafs {
+                let _ = black_box(tree.generate_proof(&item));
+            }
+            start.elapsed()
+        })
+    });
+
+    let mut proofs = Vec::<Vec<u8>>::with_capacity(leafs.len());
+    for item in &leafs {
+        proofs.push(
+            tree.generate_proof(&item)
+                .expect("failed to generate proof")
+                .expect("item not found"),
+        );
+    }
+
+    group.bench_function("deserialize_proof", |b| {
+        b.iter(|| {
+            let start = Instant::now();
+            for p in &proofs {
+                let _ = black_box(MerkleSet::from_proof(&p));
+            }
+            start.elapsed()
+        })
+    });
+}
+
+criterion_group!(merkle_set, run);
+criterion_main!(merkle_set);

crates/chia-consensus/fuzz/Cargo.toml

@@ -98,3 +98,17 @@ path = "fuzz_targets/parse-spend-id.rs"
 test = false
 doc = false
 bench = false
+
+[[bin]]
+name = "deserialize-proof"
+path = "fuzz_targets/deserialize-proof.rs"
+test = false
+doc = false
+bench = false
+
+[[bin]]
+name = "merkle-set"
+path = "fuzz_targets/merkle-set.rs"
+test = false
+doc = false
+bench = false

crates/chia-consensus/fuzz/fuzz_targets/deserialize-proof.rs

@@ -0,0 +1,8 @@
+#![no_main]
+use libfuzzer_sys::fuzz_target;
+
+use chia_consensus::merkle_tree::MerkleSet;
+
+fuzz_target!(|data: &[u8]| {
+    let _r = MerkleSet::from_proof(data);
+});

crates/chia-consensus/fuzz/fuzz_targets/merkle-set.rs

@@ -0,0 +1,28 @@
+#![no_main]
+use chia_consensus::merkle_tree::MerkleSet;
+use libfuzzer_sys::fuzz_target;
+
+fuzz_target!(|data: &[u8]| {
+    let mut input = data;
+    let num_leafs = input.len() / 32;
+    let mut leafs = Vec::<[u8; 32]>::with_capacity(num_leafs);
+    for _ in 0..num_leafs {
+        leafs.push(input[..32].try_into().unwrap());
+        input = &input[32..];
+    }
+
+    let tree = MerkleSet::from_leafs(&mut leafs);
+
+    for item in &leafs {
+        let proof = tree
+            .generate_proof(item)
+            .expect("failed to generate proof")
+            .expect("item not found");
+        let rebuilt = MerkleSet::from_proof(&proof).expect("failed to parse proof");
+        assert!(rebuilt
+            .generate_proof(item)
+            .expect("failed to validate proof")
+            .is_some());
+        assert_eq!(rebuilt.get_root(), tree.get_root());
+    }
+});

crates/chia-consensus/src/lib.rs

@@ -7,3 +7,4 @@ pub mod fast_forward;
 pub mod gen;
 pub mod generator_rom;
 pub mod merkle_set;
+pub mod merkle_tree;

crates/chia-consensus/src/merkle_set.rs

@@ -7,7 +7,7 @@ fn get_bit(val: &[u8; 32], bit: u8) -> u8 {
 
 #[repr(u8)]
 #[derive(PartialEq, Eq, Copy, Clone)]
-enum NodeType {
+pub(crate) enum NodeType {
     Empty,
     Term,
     Mid,
@@ -27,7 +27,12 @@ fn encode_type(t: NodeType) -> u8 {
     }
 }
 
-fn hash(ltype: NodeType, rtype: NodeType, left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
+pub(crate) fn hash(
+    ltype: NodeType,
+    rtype: NodeType,
+    left: &[u8; 32],
+    right: &[u8; 32],
+) -> [u8; 32] {
     let mut hasher = Sha256::new();
     hasher.update(hex!(
         "000000000000000000000000000000000000000000000000000000000000"
@@ -38,7 +43,8 @@ fn hash(ltype: NodeType, rtype: NodeType, left: &[u8; 32], right: &[u8; 32]) ->
     hasher.finalize().into()
 }
 
-const BLANK: [u8; 32] = hex!("0000000000000000000000000000000000000000000000000000000000000000");
+pub(crate) const BLANK: [u8; 32] =
+    hex!("0000000000000000000000000000000000000000000000000000000000000000");
 
 // this function performs an in-place, recursive radix sort of the range.
 // as each level returns, values are hashed pair-wise and as a hash tree.
@@ -171,7 +177,7 @@ pub fn compute_merkle_set_root(leafs: &mut [[u8; 32]]) -> [u8; 32] {
 }
 
 #[cfg(test)]
-mod test {
+pub mod test {
     use super::*;
 
     fn h2(buf1: &[u8], buf2: &[u8]) -> [u8; 32] {
@@ -399,7 +405,7 @@ mod test {
         //                 c   b
     }
 
-    fn merkle_set_test_cases() -> Vec<([u8; 32], Vec<[u8; 32]>)> {
+    pub fn merkle_set_test_cases() -> Vec<([u8; 32], Vec<[u8; 32]>)> {
         let a = hex!("7000000000000000000000000000000000000000000000000000000000000000");
         let b = hex!("7100000000000000000000000000000000000000000000000000000000000000");
         let c = hex!("8000000000000000000000000000000000000000000000000000000000000000");