refactor: enabling the node variant of merkle tree to contain data

danielpapp-trilitech · danielpapp-trilitech · commit df56110640cc · 2026-01-16T17:10:03.000Z
- Removes MerkleTree::Leaf tuple with a struct
- Removes the impl_modify_map_collect function
- Replaces the uses of the mentioned function
  with a non-recursive direct implementation.
diff --git a/data/src/components/atom/tests.rs b/data/src/components/atom/tests.rs
@@ -12,6 +12,7 @@ use crate::components::atom::Atom;
 use crate::hash::Hash;
 use crate::hash::PartialHash;
 use crate::merkle_tree::MerkleTree;
+use crate::merkle_tree::MerkleTreeLeafData;
 use crate::mode::Mode;
 use crate::mode::Normal;
 use crate::mode::Prove;
@@ -219,7 +220,11 @@ fn proof_gen() {
         let merkle_tree = MerkleTree::from_foldable(&proof_atoms);
         merkle_tree.check_root_hash();
         match merkle_tree {
-            MerkleTree::Leaf(hash, access_info, _) => {
+            MerkleTree::Leaf(MerkleTreeLeafData {
+                hash,
+                access_info,
+                ..
+            }) => {
                 prop_assert_eq!(hash, initial_root_hash);
                 prop_assert!(access_info);
             }
diff --git a/data/src/components/data_space/tests.rs b/data/src/components/data_space/tests.rs
@@ -20,6 +20,7 @@ use super::PageId;
 use crate::components::data_space::DataSpace;
 use crate::hash::Hash;
 use crate::merkle_tree::MerkleTree;
+use crate::merkle_tree::MerkleTreeLeafData;
 use crate::mode::Normal;
 use crate::mode::Prove;
 use crate::mode::utils::assert_eq_found;
@@ -261,7 +262,7 @@ fn generate_proof() {
         let mut queue = VecDeque::with_capacity(LEAVES + 1);
 
         let pages_tree = match merkle_tree {
-            MerkleTree::Leaf(_, _, _) => panic!("Did not expect leaf"),
+            MerkleTree::Leaf(_) => panic!("Did not expect leaf"),
             MerkleTree::Node(_, mut children) => {
                 // The node for the pages is the second child.
                 children.remove(1)
@@ -273,7 +274,10 @@ fn generate_proof() {
         while let Some(node) = queue.pop_front() {
             match node {
                 MerkleTree::Node(_, children) => queue.extend(children),
-                MerkleTree::Leaf(_, access_info, _) => {
+                MerkleTree::Leaf(MerkleTreeLeafData {
+                    access_info,
+                    ..
+                }) => {
                     prop_assert_eq!(
                         access_info,
                         read_leaves.contains(&leaf) ||
diff --git a/data/src/merkle_proof.rs b/data/src/merkle_proof.rs
@@ -7,7 +7,6 @@
 
 pub mod proof_tree;
 pub mod tag;
-pub mod transform;
 
 use std::error;
 
diff --git a/data/src/merkle_proof/proof_tree.rs b/data/src/merkle_proof/proof_tree.rs
@@ -5,8 +5,6 @@ use bincode::enc::write::Writer;
 
 use super::tag::LeafTag;
 use super::tag::Tag;
-use super::transform::ModifyResult;
-use super::transform::impl_modify_map_collect;
 use crate::hash::Hash;
 use crate::tree::Tree;
 
@@ -77,6 +75,42 @@ pub enum MerkleProofLeaf {
     Read(Vec<u8>),
 }
 
+enum NodeLeaf {
+    Node,
+    Leaf,
+}
+
+/// [`struct@HashState`] is associated with the state of hashing a [`MerkleProof`].
+/// We record whether the node is a leaf or an internal node, the index of its parent(
+/// see [`MerkleProof::root_hash`] for more details) and the hashes of its children
+/// if it's a node and its own hash if its a leaf.
+struct HashState {
+    node_leaf: NodeLeaf,
+    parent_index: usize,
+    hashes: Vec<Hash>,
+}
+
+impl HashState {
+    fn new(node_leaf: NodeLeaf, parent_index: usize, hashes: Vec<Hash>) -> Self {
+        Self {
+            node_leaf,
+            parent_index,
+            hashes,
+        }
+    }
+
+    fn push(&mut self, hash: Hash) {
+        self.hashes.push(hash);
+    }
+
+    fn hash(&self) -> Hash {
+        match self.node_leaf {
+            NodeLeaf::Node => Hash::combine_hashes(&self.hashes),
+            NodeLeaf::Leaf => self.hashes[0],
+        }
+    }
+}
+
 impl MerkleProof {
     /// Create a new Merkle proof as a read leaf.
     pub fn leaf_read(data: Vec<u8>) -> Self {
@@ -90,18 +124,46 @@ impl MerkleProof {
 
     /// Compute the root hash of the Merkle proof.
     pub fn root_hash(&self) -> Hash {
-        impl_modify_map_collect(
-            self,
-            |subtree| match subtree {
-                Tree::Node(vec) => ModifyResult::NodeContinue((), vec.iter().collect()),
-                Tree::Leaf(data) => ModifyResult::LeafStop(data),
-            },
-            |leaf| match leaf {
-                MerkleProofLeaf::Blind(hash) => *hash,
-                MerkleProofLeaf::Read(data) => Hash::hash_bytes(data.as_slice()),
-            },
-            |(), leaves| Hash::combine_hashes(leaves),
-        )
+        // The nodes is essentially used as a stack of pairs where the pair is made up
+        // of a [`MerkleProof`] node and the index of its parent in hash_states.
+        let mut nodes: Vec<(&MerkleProof, usize)> = vec![(self, 0)];
+        // hash_states contains the hashing state associated with a node.
+        // See [`HashState`] for more details.
+        let mut hash_states: Vec<HashState> = vec![];
+
+        // Going through all the nodes of the [`MerkleProof`] in a DFS fashion.
+        while let Some((node, parent_index)) = nodes.pop() {
+            match node {
+                Tree::Leaf(MerkleProofLeaf::Blind(hash)) => {
+                    hash_states.push(HashState::new(NodeLeaf::Leaf, parent_index, vec![*hash]));
+                }
+                Tree::Leaf(MerkleProofLeaf::Read(data)) => {
+                    hash_states.push(HashState::new(
+                        NodeLeaf::Leaf,
+                        parent_index,
+                        vec![Hash::hash_bytes(data.as_slice())],
+                    ));
+                }
+                Tree::Node(children) => {
+                    hash_states.push(HashState::new(NodeLeaf::Node, parent_index, vec![]));
+                    let new_parent_index = hash_states.len() - 1;
+                    for child in children.iter() {
+                        nodes.push((child, new_parent_index));
+                    }
+                }
+            }
+        }
+
+        // Pops the hash states and adds them to their parents.
+        while let Some(hash_state) = hash_states.pop() {
+            if hash_states.is_empty() {
+                // Returns the root hash
+                return hash_state.hash();
+            }
+            hash_states[hash_state.parent_index].push(hash_state.hash());
+        }
+
+        unreachable!("hash_states can't be an empty vector");
     }
 }
 
diff --git a/data/src/merkle_proof/transform.rs b/data/src/merkle_proof/transform.rs
diff --git a/data/src/merkle_tree.rs b/data/src/merkle_tree.rs
@@ -8,14 +8,22 @@ use crate::foldable::Foldable;
 use crate::foldable::NodeFold;
 use crate::hash::Hash;
 
+/// Struct which holds data for the leaves of a [`MerkleTree`]
+#[derive(Debug, Clone)]
+pub struct MerkleTreeLeafData {
+    pub hash: Hash,
+    pub access_info: bool,
+    pub data: Vec<u8>,
+}
+
 /// A variable-width Merkle tree with access metadata for leaves.
 ///
 /// Values of this type are produced by the proof-generating backend to capture
 /// a snapshot of the machine state along with access information for leaves
 /// which hold data that was used in a particular evaluation step.
 #[derive(Debug, Clone)]
 pub enum MerkleTree {
-    Leaf(Hash, bool, Vec<u8>),
+    Leaf(MerkleTreeLeafData),
     Node(Hash, Vec<Self>),
 }
 
@@ -38,14 +46,18 @@ impl MerkleTree {
     pub fn root_hash(&self) -> Hash {
         match self {
             Self::Node(hash, _) => *hash,
-            Self::Leaf(hash, _, _) => *hash,
+            Self::Leaf(MerkleTreeLeafData { hash, .. }) => *hash,
         }
     }
 
     /// Creates a merkle tree which is a single leaf
     pub fn make_merkle_leaf(data: Vec<u8>, access_info: bool) -> Self {
         let hash = Hash::hash_bytes(&data);
-        MerkleTree::Leaf(hash, access_info, data)
+        MerkleTree::Leaf(MerkleTreeLeafData {
+            hash,
+            access_info,
+            data,
+        })
     }
 
     /// Takes a list of children nodes and creates a
@@ -65,7 +77,9 @@ impl MerkleTree {
 
         while let Some(node) = deque.pop_front() {
             let is_valid_hash = match node {
-                Self::Leaf(hash, _, data) => &Hash::hash_bytes(data) == hash,
+                Self::Leaf(MerkleTreeLeafData { hash, data, .. }) => {
+                    &Hash::hash_bytes(data) == hash
+                }
                 Self::Node(hash, children) => {
                     let children_hashes: Vec<Hash> = children
                         .iter()
diff --git a/data/src/tree.rs b/data/src/tree.rs
@@ -8,10 +8,3 @@ pub enum Tree<A> {
     Node(Vec<Self>),
     Leaf(A),
 }
-
-/// Used in [`crate::merkle_proof::transform::impl_modify_map_collect`]
-impl<D> From<D> for Tree<D> {
-    fn from(value: D) -> Self {
-        Tree::Leaf(value)
-    }
-}
diff --git a/src/riscv/lib/src/pvm/common.rs b/src/riscv/lib/src/pvm/common.rs
@@ -54,7 +54,8 @@ use crate::pvm::hooks::PvmHooks;
 use crate::pvm::tezos;
 use crate::range_utils::less_than_bound;
 use crate::state_backend::ProofTree;
-use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof;
+use crate::state_backend::proof_backend::merkle::CompressedMerkleTree;
+use crate::state_backend::proof_backend::merkle::merkle_tree_to_compressed_merkle_tree;
 use crate::state_backend::proof_backend::proof::Proof;
 use crate::state_backend::proof_backend::proof::deserialise_owned;
 
@@ -370,7 +371,9 @@ where
         let _ = self.input_request();
 
         let merkle_tree = MerkleTree::from_foldable(self);
-        let merkle_proof = merkle_tree_to_merkle_proof(merkle_tree);
+        let compressed_merkle_tree: CompressedMerkleTree =
+            merkle_tree_to_compressed_merkle_tree(merkle_tree);
+        let merkle_proof: MerkleProof = compressed_merkle_tree.to_proof();
 
         let final_hash = Hash::from_foldable(self);
         let proof = Proof::new(merkle_proof, final_hash);
diff --git a/src/riscv/lib/src/state_backend.rs b/src/riscv/lib/src/state_backend.rs
@@ -48,7 +48,7 @@ mod tests {
     use rand::RngCore;
 
     use super::*;
-    use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof;
+    use crate::state_backend::proof_backend::merkle::merkle_tree_to_compressed_merkle_tree;
     use crate::state_backend::proof_backend::proof::deserialise_owned::ProofTreeDeserialiser;
 
     /// Data structure whose [`Elem`] implementation only writes to part of the given space
@@ -106,7 +106,7 @@ mod tests {
 
         // The Verify mode needs a proof, so we generate it from the Prove mode
         let merkle_tree = MerkleTree::from_foldable(&mem_prove);
-        let proof_tree = merkle_tree_to_merkle_proof(merkle_tree);
+        let proof_tree = merkle_tree_to_compressed_merkle_tree(merkle_tree).to_proof();
         let proof_deser = ProofTreeDeserialiser::from(ProofTree::Present(&proof_tree));
         let mut mem_verify = DataSpace::from_proof(proof_deser).unwrap().into_result();
 
diff --git a/src/riscv/lib/src/state_backend/proof_backend/merkle.rs b/src/riscv/lib/src/state_backend/proof_backend/merkle.rs
diff --git a/src/riscv/lib/src/state_backend/proof_layout.rs b/src/riscv/lib/src/state_backend/proof_layout.rs
diff --git a/src/riscv/lib/src/state_backend/region.rs b/src/riscv/lib/src/state_backend/region.rs
