From 7364f4a60c8b18d5a6dfd2ea6f4aecf2c89412d4 Mon Sep 17 00:00:00 2001 From: zhenfei Date: Tue, 1 Oct 2024 18:30:46 -0400 Subject: [PATCH] docs --- tree/src/leaf.rs | 16 +++++++++++++++- tree/src/lib.rs | 3 +++ tree/src/node.rs | 17 +++++++++++++++-- tree/src/path.rs | 22 +++++++++++++++++----- tree/src/tests.rs | 24 +++++++++++++++++++++--- tree/src/tree.rs | 46 +++++++++++++++++++++++----------------------- 6 files changed, 94 insertions(+), 34 deletions(-) diff --git a/tree/src/leaf.rs b/tree/src/leaf.rs index 3cadb165..6c9a9005 100644 --- a/tree/src/leaf.rs +++ b/tree/src/leaf.rs @@ -7,7 +7,7 @@ use poseidon::{PoseidonBabyBearParams, PoseidonBabyBearState}; use crate::Node; -/// A leaf is a blob of 64 bytes of data, stored in a BabyBearx16 +/// Represents a leaf in the Merkle tree, containing 64 bytes of data stored in a BabyBearx16. #[derive(Debug, Copy, Clone, PartialEq, Default)] pub struct Leaf { pub(crate) data: BabyBearx16, @@ -15,17 +15,30 @@ pub struct Leaf { 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 let t = unsafe { transmute::(self.data) }; write!(f, "leaf: 0x{:02x?}...{:02x?}", t[0], t[63]) } } impl Leaf { + /// Creates a new Leaf with the given data. pub fn new(data: BabyBearx16) -> Self { Self { data } } + /// 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 { @@ -39,6 +52,7 @@ impl Leaf { } impl From for Leaf { + /// Implements the From trait to allow creation of a Leaf from BabyBearx16 data. fn from(data: BabyBearx16) -> Self { Self { data } } diff --git a/tree/src/lib.rs b/tree/src/lib.rs index 95dfd173..af3d6363 100644 --- a/tree/src/lib.rs +++ b/tree/src/lib.rs @@ -1,3 +1,6 @@ +//! 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::*; diff --git a/tree/src/node.rs b/tree/src/node.rs index 0887d93c..f1d8e2a9 100644 --- a/tree/src/node.rs +++ b/tree/src/node.rs @@ -3,7 +3,7 @@ use std::fmt::Display; use sha2::{Digest, Sha512}; -/// A node is a blob of 32 bytes of data +/// A node in the Merkle tree, representing 32 bytes of data. #[derive(Debug, Copy, Clone, PartialEq, Default)] pub struct Node { pub(crate) data: [u8; 32], @@ -11,18 +11,31 @@ pub struct Node { 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; 32]) -> 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 { - // use sha2-512 to hash the two nodes let mut hasher = Sha512::new(); hasher.update(left.data); hasher.update(right.data); diff --git a/tree/src/path.rs b/tree/src/path.rs index 336a5a00..eaa3c8dc 100644 --- a/tree/src/path.rs +++ b/tree/src/path.rs @@ -6,6 +6,7 @@ use poseidon::PoseidonBabyBearParams; use crate::{Leaf, Node}; +/// Represents a path in the Merkle tree, used for proving membership. #[derive(Clone, Debug, PartialEq)] pub struct Path { pub(crate) path_nodes: Vec, @@ -20,7 +21,7 @@ impl Display for Path { for ((i, node), is_right_node) in self.path_nodes.iter().enumerate().zip(position_list) { writeln!( f, - "{}-th node, is right nide {}, sibling: {}", + "{}-th node, is right node {}, sibling: {}", i, is_right_node, node )?; } @@ -30,16 +31,26 @@ impl Display for Path { } impl Path { - /// The position of on_path node in `leaf_and_sibling_hash` and `non_leaf_and_sibling_hash_path`. - /// `position[i]` is 0 (false) iff `i`th on-path node from top to bottom is on the left. + /// Computes the position of on-path nodes in the Merkle tree. /// - /// This function simply converts `self.leaf_index` to boolean array in big endian form. + /// This function converts the leaf index to a boolean array in big-endian form, + /// where `true` indicates a right child and `false` indicates a left child. #[inline] fn position_list(&'_ self) -> impl '_ + Iterator { (0..self.path_nodes.len() + 1).map(move |i| ((self.index >> i) & 1) != 0) } - /// verifies the path against a root + /// Verifies the path against a given root and leaf. + /// + /// # Arguments + /// + /// * `root` - The root node of the Merkle tree + /// * `leaf` - The leaf node to verify + /// * `hasher` - The Poseidon hash parameters + /// + /// # Returns + /// + /// `true` if the path is valid, `false` otherwise. #[inline] pub fn verify(&self, root: &Node, leaf: &Leaf, hasher: &PoseidonBabyBearParams) -> bool { let timer = start_timer!(|| "path verify"); @@ -48,6 +59,7 @@ impl Path { let leaf_node = leaf.leaf_hash(hasher); let mut current_node = leaf_node; + // Traverse the path from leaf to root for (i, node) in self.path_nodes.iter().rev().enumerate() { if position_list[i] { current_node = Node::node_hash(node, ¤t_node) diff --git a/tree/src/tests.rs b/tree/src/tests.rs index 56edb484..3b82c9a5 100644 --- a/tree/src/tests.rs +++ b/tree/src/tests.rs @@ -7,26 +7,44 @@ use crate::{Leaf, Tree}; #[test] fn test_tree() { + // Initialize a random number generator for the test let mut rng = test_rng(); - let hasher = PoseidonBabyBearParams::new(&mut rng); + // Create a new instance of PoseidonBabyBearParams for hashing + let leaf_hasher = PoseidonBabyBearParams::new(&mut rng); + + // Test trees of different heights, from 4 to 14 for height in 4..15 { + // Generate random leaves for the tree + // The number of leaves is 2^(height-1) let leaves: Vec = (0..(1 << (height - 1))) .map(|_| BabyBearx16::random_unsafe(&mut rng).into()) .collect(); - let tree = Tree::new_with_leaves(&hasher, leaves, height); + // Create a new tree with the generated leaves + let tree = Tree::new_with_leaves(&leaf_hasher, leaves, height); + + // Perform 100 random verifications for each tree for _ in 0..100 { + // Select a random leaf index let index = rng.next_u32() % (1 << (height - 1)); + + // Generate a proof for the selected leaf let proof = tree.gen_proof(index as usize, height); + + // Get the root of the tree let root = tree.root(); + // Print debug information println!("index: {}\n", index); println!("root: {}\n", root); println!("tree {}\n", tree); println!("path {}\n", proof); - assert!(proof.verify(&root, &tree.leaves[index as usize], &hasher)); + // Verify the proof + // This checks that the leaf at the given index is indeed part of the tree + // with the given root, using the generated proof + assert!(proof.verify(&root, &tree.leaves[index as usize], &leaf_hasher)); } } } diff --git a/tree/src/tree.rs b/tree/src/tree.rs index 144e74ce..c4aba7a4 100644 --- a/tree/src/tree.rs +++ b/tree/src/tree.rs @@ -8,10 +8,11 @@ use rayon::iter::{ use crate::{Leaf, Node, Path}; +/// Represents a Merkle tree structure. #[derive(Clone, Debug, PartialEq)] pub struct Tree { pub nodes: Vec, - pub leaves: Vec, // todo: avoid cloning the data here + pub leaves: Vec, } impl Display for Tree { @@ -29,24 +30,24 @@ impl Display for Tree { } impl Tree { - /// create an empty tree + /// Creates an empty tree with default leaves. #[inline] pub fn init(hasher: &PoseidonBabyBearParams, tree_height: usize) -> Self { let leaves = vec![Leaf::default(); 1 << (tree_height - 1)]; Self::new_with_leaves(hasher, leaves, tree_height) } - /// build a tree with leaves + /// Builds a tree with the given leaves. #[inline] pub fn new_with_leaves( - hasher: &PoseidonBabyBearParams, + leaf_hasher: &PoseidonBabyBearParams, leaves: Vec, tree_height: usize, ) -> Self { let leaf_nodes = leaves .as_slice() .into_par_iter() - .map(|leaf| leaf.leaf_hash(hasher)) + .map(|leaf| leaf.leaf_hash(leaf_hasher)) .collect::>(); let nodes = Self::new_with_leaf_nodes(leaf_nodes, tree_height); Self { @@ -55,9 +56,16 @@ impl Tree { } } - /// build a tree with leaves - /// assume the leaves are already hashed via leaf hash - /// returns the leaf nodes and the tree nodes + /// Builds a tree with pre-hashed leaf nodes. + /// + /// # Arguments + /// + /// * `leaf_nodes` - Vector of pre-hashed leaf nodes + /// * `tree_height` - Height of the tree + /// + /// # Returns + /// + /// A tuple containing vectors of non-leaf nodes and leaf nodes. pub fn new_with_leaf_nodes( leaf_nodes: Vec, tree_height: usize, @@ -77,7 +85,7 @@ impl Tree { index = left_child_index(index); } - // compute the hash values for the non-leaf bottom layer + // Compute the hash values for the non-leaf bottom layer { let start_index = level_indices.pop().unwrap(); let upper_bound = left_child_index(start_index); @@ -88,12 +96,8 @@ impl Tree { .take(upper_bound) .skip(start_index) .for_each(|(current_index, e)| { - // `left_child_index(current_index)` and `right_child_index(current_index) returns the position of - // leaf in the whole tree (represented as a list in level order). We need to shift it - // by `-upper_bound` to get the index in `leaf_nodes` list. let left_leaf_index = left_child_index(current_index) - upper_bound; let right_leaf_index = right_child_index(current_index) - upper_bound; - // compute hash *e = Node::node_hash( &leaf_nodes[left_leaf_index], &leaf_nodes[right_leaf_index], @@ -101,11 +105,10 @@ impl Tree { }); } - // compute the hash values for nodes in every other layer in the tree + // Compute the hash values for nodes in every other layer in the tree level_indices.reverse(); for &start_index in &level_indices { - // The layer beginning `start_index` ends at `upper_bound` (exclusive). let upper_bound = left_child_index(start_index); let mut buf = non_leaf_nodes[start_index..upper_bound].to_vec(); buf.par_iter_mut().enumerate().for_each(|(index, node)| { @@ -121,34 +124,30 @@ impl Tree { (non_leaf_nodes, leaf_nodes.to_vec()) } + /// Returns the root node of the tree. #[inline] pub fn root(&self) -> Node { self.nodes[0] } - // generate a membership proof for the given index + /// Generates a membership proof for the given index. #[inline] pub fn gen_proof(&self, index: usize, tree_height: usize) -> Path { let timer = start_timer!(|| "generate membership proof"); - // Get Leaf hash, and leaf sibling hash, let leaf_index_in_tree = convert_index_to_last_level(index, tree_height); let sibling_index_in_tree = sibling_index(leaf_index_in_tree).unwrap(); - // path.len() = `tree height - 1`, the missing elements being the root let mut path_nodes = Vec::with_capacity(tree_height - 1); path_nodes.push(self.nodes[sibling_index_in_tree]); - // Iterate from the bottom layer after the leaves, to the top, storing all nodes and their siblings. + // Iterate from the bottom layer after the leaves to the top let mut current_node = parent_index(leaf_index_in_tree).unwrap(); while current_node != 0 { let sibling_node = sibling_index(current_node).unwrap(); - path_nodes.push(self.nodes[sibling_node]); - current_node = parent_index(current_node).unwrap(); } - // we want to make path from root to bottom path_nodes.reverse(); end_timer!(timer); Path { index, path_nodes } @@ -189,12 +188,13 @@ fn right_child_index(index: usize) -> usize { 2 * index + 2 } +/// Converts a leaf index to its position in the last level of the tree. #[inline] fn convert_index_to_last_level(index: usize, tree_height: usize) -> usize { index + (1 << (tree_height - 1)) - 1 } -/// Returns true iff the given index represents a left child. +/// Returns true if the given index represents a left child. #[inline] fn is_left_child(index: usize) -> bool { index % 2 == 1