Feature: Merkle Tree implementation as Orion pcs building block (#143)

* init commit for merkle tree implementation

* benchmarked and optimized, shit runs fast

* allocate rather than concat

* use move for intermediate results that can be dropped later

* shorten bench time

* additional benchmarks

Author: tonyfloatersu

Cargo.lock

@@ -13,6 +13,7 @@ members = [
     "pcs",
     "sumcheck",
     "transcript",
+    "tree"
 ]
 resolver = "2"
 

Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "tree"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arith = { path = "../arith" }
+
+# babybear = { path = "../arith/babybear" }
+# poseidon = { path = "../poseidon" }
+
+ark-std.workspace = true
+sha2.workspace = true
+rayon.workspace = true
+
+[dev-dependencies]
+criterion.workspace = true
+gf2 = { path = "../arith/gf2" }
+gf2_128 = { path = "../arith/gf2_128" }
+mersenne31 = { path = "../arith/mersenne31" }
+tynm.workspace = true
+
+# p3-baby-bear.workspace = true
+# p3-field.workspace = true
+# p3-monty-31.workspace = true
+
+[[bench]]
+name = "tree"
+harness = false

tree/Cargo.toml

@@ -0,0 +1,120 @@
+use std::time::Duration;
+
+use arith::{Field, FieldSerde, SimdField};
+use ark_std::{rand::RngCore, test_rng};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use gf2::{GF2x128, GF2x64, GF2x8, GF2};
+use mersenne31::{M31x16, M31};
+use tree::{Leaf, Tree, LEAF_BYTES};
+use tynm::type_name;
+
+const FINAL_MT_LEAVES_LOG2: usize = 15;
+
+fn tree_building_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("SHA512-256 merkle tree");
+
+    let mut rng = test_rng();
+    let mut data_buffer = [0u8; LEAF_BYTES];
+    let leaves: Vec<_> = (0..(1 << FINAL_MT_LEAVES_LOG2))
+        .map(|_| {
+            Leaf::new({
+                rng.fill_bytes(&mut data_buffer);
+                data_buffer.clone()
+            })
+        })
+        .collect();
+
+    for i in 10..=FINAL_MT_LEAVES_LOG2 {
+        group
+            .bench_function(BenchmarkId::new(format!("2^{i} leaves"), i), |b| {
+                let leaves_benchmark = leaves[..(1 << i)].to_vec();
+
+                b.iter(|| {
+                    Tree::new_with_leaves(leaves_benchmark.clone());
+                })
+            })
+            .sample_size(10)
+            .measurement_time(Duration::from_secs(5));
+    }
+}
+
+fn compact_field_elem_tree_building_benchmark_generic<F, PackF>(c: &mut Criterion)
+where
+    F: Field + FieldSerde,
+    PackF: SimdField<Scalar = F>,
+{
+    let mut group = c.benchmark_group(format!(
+        "SHA512-256 merkle tree with field element {} packed by SIMD field element {}",
+        type_name::<F>(),
+        type_name::<PackF>()
+    ));
+    let num_of_elems_in_leaf = LEAF_BYTES * 8 / F::FIELD_SIZE;
+
+    let mut rng = test_rng();
+    let field_elems: Vec<_> = (0..(1 << FINAL_MT_LEAVES_LOG2) * num_of_elems_in_leaf)
+        .map(|_| F::random_unsafe(&mut rng))
+        .collect();
+
+    for i in 10..=FINAL_MT_LEAVES_LOG2 {
+        group
+            .bench_function(BenchmarkId::new(format!("2^{i} leaves"), i), |b| {
+                let field_elems_benchmark = field_elems[..(1 << i) * num_of_elems_in_leaf].to_vec();
+
+                b.iter(|| {
+                    Tree::compact_new_with_field_elems::<F, PackF>(&field_elems_benchmark);
+                })
+            })
+            .sample_size(10)
+            .measurement_time(Duration::from_secs(5));
+    }
+}
+
+fn compact_field_elem_tree_building_benchmark(c: &mut Criterion) {
+    compact_field_elem_tree_building_benchmark_generic::<GF2, GF2x8>(c);
+    compact_field_elem_tree_building_benchmark_generic::<GF2, GF2x64>(c);
+    compact_field_elem_tree_building_benchmark_generic::<GF2, GF2x128>(c);
+    compact_field_elem_tree_building_benchmark_generic::<M31, M31x16>(c)
+}
+
+fn compact_packed_field_elem_tree_building_benchmark_generic<F, PackF>(c: &mut Criterion)
+where
+    F: Field + FieldSerde,
+    PackF: SimdField<Scalar = F>,
+{
+    let mut group = c.benchmark_group(format!(
+        "SHA512-256 merkle tree with SIMD field element {}",
+        type_name::<PackF>()
+    ));
+    let num_of_elems_in_leaf = LEAF_BYTES / PackF::SIZE;
+
+    let mut rng = test_rng();
+    let field_elems: Vec<_> = (0..(1 << FINAL_MT_LEAVES_LOG2) * num_of_elems_in_leaf)
+        .map(|_| PackF::random_unsafe(&mut rng))
+        .collect();
+
+    for i in 10..=FINAL_MT_LEAVES_LOG2 {
+        group
+            .bench_function(BenchmarkId::new(format!("2^{i} leaves"), i), |b| {
+                let field_elems_benchmark = field_elems[..(1 << i) * num_of_elems_in_leaf].to_vec();
+
+                b.iter(|| {
+                    Tree::compact_new_with_packed_field_elems::<F, PackF>(&field_elems_benchmark);
+                })
+            })
+            .sample_size(10)
+            .measurement_time(Duration::from_secs(5));
+    }
+}
+
+fn compact_packed_field_elem_tree_building_benchmark(c: &mut Criterion) {
+    compact_packed_field_elem_tree_building_benchmark_generic::<GF2, GF2x8>(c);
+    compact_packed_field_elem_tree_building_benchmark_generic::<GF2, GF2x64>(c);
+}
+
+criterion_group!(
+    bench,
+    tree_building_benchmark,
+    compact_field_elem_tree_building_benchmark,
+    compact_packed_field_elem_tree_building_benchmark
+);
+criterion_main!(bench);

tree/benches/tree.rs

@@ -0,0 +1,77 @@
+use std::fmt;
+use std::fmt::{Debug, Display};
+
+use arith::Field;
+use sha2::{Digest, Sha512_256};
+
+use crate::Node;
+
+/// Each leaf should have 64 bytes or 512 bits
+pub const LEAF_BYTES: usize = 64;
+
+/// Each leaf hash should have 32 bytes
+pub const LEAF_HASH_BYTES: usize = 32;
+
+#[inline(always)]
+pub const fn leaf_adic<F: Field>() -> usize {
+    LEAF_BYTES * 8 / F::FIELD_SIZE
+}
+
+/// Represents a leaf in the Merkle tree, containing 64 bytes of data stored in a BabyBearx16.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct Leaf {
+    pub data: [u8; LEAF_BYTES],
+}
+
+impl Default for Leaf {
+    fn default() -> Self {
+        Self {
+            data: [0u8; LEAF_BYTES],
+        }
+    }
+}
+
+impl Display for Leaf {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        // Display the first and last byte of the leaf data for brevity
+        write!(f, "leaf: 0x{:02x?}", self.data)
+    }
+}
+
+impl Leaf {
+    /// Creates a new Leaf with the given data.
+    pub fn new(data: [u8; LEAF_BYTES]) -> Self {
+        Self { data }
+    }
+
+    pub fn leaf_hash(&self) -> Node {
+        let mut hasher = Sha512_256::new();
+        hasher.update(self.data);
+        let res: [u8; LEAF_HASH_BYTES] = hasher.finalize().into();
+
+        Node { data: res }
+    }
+
+    // /// Computes the hash of the leaf using Poseidon hash function.
+    // ///
+    // /// # Arguments
+    // ///
+    // /// * `hash_param` - The Poseidon hash parameters
+    // ///
+    // /// # Returns
+    // ///
+    // /// A Node containing the hash of the leaf data.
+    // pub fn leaf_hash(&self, hash_param: &PoseidonBabyBearParams) -> Node {
+    //     // Use Poseidon hash for leaf nodes
+    //     // Note: This could be replaced with SHA2 if performance requires
+    //     let mut state = PoseidonBabyBearState { state: self.data };
+    //     hash_param.permute(&mut state);
+    //     Node {
+    //         data: unsafe {
+    //             transmute::<BabyBearx16, [u8; 64]>(state.state)[..32]
+    //                 .try_into()
+    //                 .unwrap()
+    //         },
+    //     }
+    // }
+}

tree/src/leaf.rs

@@ -0,0 +1,17 @@
+//! This module defines the core components of a Merkle tree implementation.
+//! It includes definitions for tree structures, nodes, leaves, and paths.
+
+mod tree;
+pub use tree::*;
+
+mod node;
+pub use node::*;
+
+mod leaf;
+pub use leaf::*;
+
+mod path;
+pub use path::*;
+
+#[cfg(test)]
+mod tests;

tree/src/lib.rs

@@ -0,0 +1,54 @@
+use std::fmt;
+use std::fmt::Display;
+
+use sha2::{Digest, Sha512_256};
+
+use crate::LEAF_HASH_BYTES;
+
+/// A node in the Merkle tree, representing 32 bytes of data.
+#[derive(Debug, Copy, Clone, PartialEq, Default)]
+pub struct Node {
+    pub(crate) data: [u8; LEAF_HASH_BYTES],
+}
+
+impl Display for Node {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        // Display the first and last byte of the node for brevity
+        write!(f, "node: 0x{:02x?}...{:02x?}", self.data[0], self.data[31])
+    }
+}
+
+impl Node {
+    /// Creates a new Node with the given data.
+    pub fn new(data: [u8; LEAF_HASH_BYTES]) -> Self {
+        Self { data }
+    }
+
+    /// Computes the hash of two child nodes to create a parent node.
+    ///
+    /// This function uses SHA-512 for hashing and takes the first 32 bytes of the result.
+    ///
+    /// # Arguments
+    ///
+    /// * `left` - The left child node
+    /// * `right` - The right child node
+    ///
+    /// # Returns
+    ///
+    /// A new Node containing the hash of the two input nodes.
+    #[inline]
+    pub fn node_hash(left: &Node, right: &Node) -> Node {
+        let mut hasher = Sha512_256::new();
+        hasher.update(left.data);
+        hasher.update(right.data);
+        let result = hasher.finalize();
+        Node {
+            data: result.into(),
+        }
+    }
+
+    /// Returns the data of the node as a slice of bytes.
+    pub fn as_bytes(&self) -> &[u8] {
+        &self.data
+    }
+}