docs

tree/src/leaf.rs

@@ -7,25 +7,38 @@ 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,
 }
 
 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::<BabyBearx16, [u8; 64]>(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<BabyBearx16> for Leaf {
+    /// Implements the From trait to allow creation of a Leaf from BabyBearx16 data.
     fn from(data: BabyBearx16) -> Self {
         Self { data }
     }

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::*;
 

tree/src/node.rs

@@ -3,26 +3,39 @@ 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],
 }
 
 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);

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<Node>,
@@ -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<Item = bool> {
         (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, &current_node)

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<Leaf> = (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));
         }
     }
 }

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<Node>,
-    pub leaves: Vec<Leaf>, // todo: avoid cloning the data here
+    pub leaves: Vec<Leaf>,
 }
 
 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<Leaf>,
         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::<Vec<Node>>();
         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<Node>,
         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,24 +96,19 @@ 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],
                     );
                 });
         }
 
-        // 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