refactor(data): first steps in rv-829

danielpapp-trilitech · danielpapp-trilitech · commit eba7d3331729 · 2026-01-12T17:08:11.000Z
- Replaced MerkleTree::Leaf tuple with a struct
- Removed the impl_modify_map_collect function
- Replaced 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,9 +5,8 @@ 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::hash::Hasher;
 use crate::tree::Tree;
 
 /// Merkle proof tree structure.
@@ -77,6 +76,50 @@ pub enum MerkleProofLeaf {
     Read(Vec<u8>),
 }
 
+enum NodeLeaf {
+    Node,
+    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 get_hash(&self) -> Hash {
+        match self.node_leaf {
+            NodeLeaf::Node => {
+                let mut hasher = Hasher::default();
+
+                for hash in self.hashes.iter() {
+                    hasher.update_with_hash(hash);
+                }
+
+                hasher.to_hash()
+            }
+            NodeLeaf::Leaf => self.hashes[0],
+        }
+    }
+
+    fn get_parent_index(&self) -> usize {
+        self.parent_index
+    }
+}
+
 impl MerkleProof {
     /// Create a new Merkle proof as a read leaf.
     pub fn leaf_read(data: Vec<u8>) -> Self {
@@ -90,18 +133,39 @@ 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),
-        )
+        let mut nodes: Vec<(&MerkleProof, usize)> = vec![(self, 0)];
+        let mut hash_states: Vec<HashState> = vec![];
+
+        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));
+                    }
+                }
+            }
+        }
+
+        while let Some(hash_state) = hash_states.pop() {
+            if hash_states.is_empty() {
+                return hash_state.get_hash();
+            }
+            hash_states[hash_state.get_parent_index()].push(hash_state.get_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
@@ -1,89 +1,3 @@
-// SPDX-FileCopyrightText: 2024 TriliTech <contact@trili.tech>
-//
+// SPDX-FileCopyrightText: 2025 TriliTech <contact@trili.tech>
 // SPDX-License-Identifier: MIT
 
-//! Module for tree utils like types, traversals.
-//!
-//! All the traversals implemented in this module should be the same to maintain consistency,
-//! which is required for serialisation / deserialisation
-
-/// Intermediary either-like type for implementing [`impl_modify_map_collect`]
-#[derive(Clone)]
-pub enum ModifyResult<D, N, L> {
-    /// Current subtree should be replaced with a node containing the given children,
-    /// and an auxiliary data for extra context if needed.
-    /// Traversal should continue recursively.
-    NodeContinue(D, Vec<N>),
-    /// Current subtree is replaced by a leaf containing the given data.
-    LeafStop(L),
-}
-
-/// Perform generic modify_map_collect
-///
-/// This is done in 3 steps while traversing the tree in a pre-order DFS traversal:
-/// 1. Apply `modify` on current subtree: This operation changes the structure of the current
-///    subtree before traversing its children.
-/// 2. When encountering leaves, `map` is called to transform a leaf from `A` to `B` type.
-///    This is done on children of subtrees which have been traversed after `modify` was called.
-/// 3. After modifying & mapping the children of a node, the `collect` method gathers the newly
-///    modified & mapped subtrees to create the new subtree.
-pub fn impl_modify_map_collect<
-    InterimLeafData, // InterimLeafData -> FinalLeafData  when applying `map`
-    FinalLeafData,   // [FinalLeafData] -> FinalLeafData when applying `collect`
-    AuxTreeData,     // Type of auxiliary data held for a subtree
-    InputTree,
-    OutputTree: From<FinalLeafData>,
-    TreeModifier: FnMut(InputTree) -> ModifyResult<AuxTreeData, InputTree, InterimLeafData>,
-    LeafMapper: FnMut(InterimLeafData) -> FinalLeafData,
-    Collector: FnMut(AuxTreeData, Vec<OutputTree>) -> OutputTree,
->(
-    root: InputTree,
-    mut modify: TreeModifier,
-    mut map: LeafMapper,
-    mut collect: Collector,
-) -> OutputTree {
-    enum ProcessEvents<ProcessEvent, CollectAuxTreeData> {
-        Node(ProcessEvent),
-        Collect(CollectAuxTreeData, usize),
-    }
-
-    let mut process = vec![ProcessEvents::Node(root)];
-    let mut done: Vec<OutputTree> = vec![];
-
-    while let Some(event) = process.pop() {
-        match event {
-            ProcessEvents::Node(subtree) => match modify(subtree) {
-                ModifyResult::LeafStop(data) => {
-                    // Instead of pushing a single leaf process on the Process-queue,
-                    // map the data and append it directly to the Done-queue
-                    done.push(OutputTree::from(map(data)));
-                }
-                ModifyResult::NodeContinue(node_data, children) => {
-                    // the only case where we push further process events in the process queue
-                    // We have to first push a collect event to know how many children should be collected when forming back the current subtree
-                    process.push(ProcessEvents::Collect(node_data, children.len()));
-
-                    process.extend(
-                        children
-                            .into_iter()
-                            .rev()
-                            .map(|child| ProcessEvents::Node(child)),
-                    );
-                }
-            },
-            ProcessEvents::Collect(node_data, count) => {
-                // We need to reconstruct a subtree which is made of `count` children
-                // No panic: We are guaranteed count < done.len() since every Collect(size)
-                // corresponds to size nodes pushed to Done-queue
-                let children = done.split_off(done.len() - count);
-                done.push(collect(node_data, children));
-            }
-        }
-    }
-
-    // No Panic: We only add a single node as root at the beginning of the algorithm
-    // which corresponds to this last node in the Done-queue
-    done.pop()
-        .filter(|_| done.is_empty())
-        .expect("Unexpected number of results")
-}
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 [`crate::merkle_tree::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
@@ -51,7 +51,8 @@ use crate::state_backend;
 use crate::state_backend::ManagerBase;
 use crate::state_backend::ManagerClone;
 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;
 
@@ -375,7 +376,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
@@ -134,7 +134,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
@@ -192,7 +192,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
