refactor(data): first steps in rv-829

danielpapp-trilitech · danielpapp-trilitech · commit 93c06bd76b30 · 2025-12-01T12:44:00.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/hash.rs b/data/src/hash.rs
@@ -57,6 +57,11 @@ pub struct Hash {
 }
 
 impl Hash {
+    /// Creates a new [`crate::hash::Hash`] object from a digest.
+    pub fn new(digest: [u8; DIGEST_SIZE]) -> Self {
+        Hash { digest }
+    }
+
     /// Hash a slice of bytes
     pub fn blake3_hash_bytes(bytes: &[u8]) -> Self {
         let digest = blake3::hash(bytes).into();
@@ -108,6 +113,11 @@ impl Hash {
     pub fn from_foldable(foldable: &impl Foldable<HashFold>) -> Self {
         foldable.fold(HashFold)
     }
+
+    /// Creates a new digest from a [`crate::hash::Hash`] object.
+    pub fn to_digest(self) -> [u8; 32] {
+        self.digest
+    }
 }
 
 impl std::fmt::Display for Hash {
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,50 @@ pub enum MerkleProofLeaf {
     Read(Vec<u8>),
 }
 
+enum NodeLeaf {
+    Node,
+    Leaf,
+}
+
+struct HashState {
+    node_leaf: NodeLeaf,
+    parent_index: usize,
+    digests: Vec<[u8; 32]>,
+}
+
+impl HashState {
+    fn new(node_leaf: NodeLeaf, parent_index: usize, digests: Vec<[u8; 32]>) -> Self {
+        Self {
+            node_leaf,
+            parent_index,
+            digests,
+        }
+    }
+
+    fn push(&mut self, digest: [u8; 32]) {
+        self.digests.push(digest);
+    }
+
+    fn get_digest(&self) -> [u8; 32] {
+        match self.node_leaf {
+            NodeLeaf::Node => {
+                let mut hasher = blake3::Hasher::new();
+
+                for digest in self.digests.iter() {
+                    hasher.update(digest);
+                }
+
+                hasher.finalize().into()
+            }
+            NodeLeaf::Leaf => self.digests[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 +132,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::blake3_hash_bytes(data.as_slice()),
-            },
-            |(), leaves| Hash::combine(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.to_digest(),
+                    ]));
+                }
+                Tree::Leaf(MerkleProofLeaf::Read(data)) => {
+                    hash_states.push(HashState::new(NodeLeaf::Leaf, parent_index, vec![
+                        Hash::blake3_hash_bytes(data.as_slice()).to_digest(),
+                    ]));
+                }
+                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::new(hash_state.get_digest());
+            }
+            hash_states[hash_state.get_parent_index()].push(hash_state.get_digest());
+        }
+
+        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::blake3_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::blake3_hash_bytes(data) == hash,
+                Self::Leaf(MerkleTreeLeafData { hash, data, .. }) => {
+                    &Hash::blake3_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::ManagerBase;
 use crate::state_backend::ManagerClone;
 use crate::state_backend::ProofTree;
 use crate::state_backend::proof_backend::ProofWrapper;
-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;
 use crate::struct_layout;
@@ -400,7 +401,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/layout.rs b/src/riscv/lib/src/state_backend/layout.rs
@@ -287,7 +287,8 @@ mod tests {
     use crate::state_backend::ProofLayout;
     use crate::state_backend::ProofPart;
     use crate::state_backend::proof_backend::ProofWrapper;
-    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::deserialise_owned;
     use crate::state_backend::verify_backend::handle_stepper_panics;
 
@@ -365,7 +366,9 @@ mod tests {
             assert_eq!(hash, tree_root_hash);
 
             // Produce a proof
-            let proof = merkle_tree_to_merkle_proof(tree);
+            let compressed_merkle_tree: CompressedMerkleTree =
+                merkle_tree_to_compressed_merkle_tree(tree);
+            let proof = compressed_merkle_tree.to_proof();
             let proof_hash = proof.root_hash();
             assert_eq!(hash, proof_hash);
 
diff --git a/src/riscv/lib/src/state_backend/proof_backend.rs b/src/riscv/lib/src/state_backend/proof_backend.rs
@@ -469,6 +469,7 @@ mod tests {
 
     use octez_riscv_data::hash::Hash;
     use octez_riscv_data::merkle_tree::MerkleTree;
+    use octez_riscv_data::merkle_tree::MerkleTreeLeafData;
     use octez_riscv_data::mode::Normal;
     use proptest::array;
     use proptest::prop_assert;
@@ -547,7 +548,7 @@ mod tests {
             let merkle_tree = MerkleTree::from_foldable(&proof_cells);
             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);
                 }
@@ -678,7 +679,7 @@ mod tests {
             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)
@@ -690,7 +691,7 @@ mod tests {
             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/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
