refactor(data): first steps in rv-829

danielpapp-trilitech · danielpapp-trilitech · commit 2f8c160bbb4e · 2025-11-28T13:04:29.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
@@ -56,6 +56,10 @@ pub struct Hash {
 }
 
 impl Hash {
+    pub fn new(digest: [u8; 32]) -> Self {
+        Hash { digest }
+    }
+
     /// Hash a slice of bytes
     pub fn blake3_hash_bytes(bytes: &[u8]) -> Self {
         let digest = blake3::hash(bytes).into();
@@ -107,6 +111,10 @@ impl Hash {
     pub fn from_foldable(foldable: &impl Foldable<HashFold>) -> Self {
         foldable.fold(HashFold)
     }
+
+    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,58 @@ pub enum MerkleProofLeaf {
     Read(Vec<u8>),
 }
 
+enum NodeLeaf {
+    Node,
+    Leaf,
+}
+
+struct HashState {
+    node_leaf: NodeLeaf,
+    parent_index: usize,
+    digests: Vec<[u8; 32]>,
+}
+
+impl HashState {
+    pub fn new(node_leaf: NodeLeaf, parent_index: usize) -> Self {
+        Self {
+            node_leaf,
+            parent_index,
+            digests: vec![],
+        }
+    }
+
+    pub fn new_with_digest(node_leaf: NodeLeaf, parent_index: usize, digest: [u8; 32]) -> Self {
+        Self {
+            node_leaf,
+            parent_index,
+            digests: vec![digest],
+        }
+    }
+
+    pub fn push(&mut self, digest: [u8; 32]) {
+        self.digests.push(digest);
+    }
+
+    pub 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],
+        }
+    }
+
+    pub 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 +140,43 @@ 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 hashes: Vec<HashState> = vec![];
+
+        while let Some((node, parent_index)) = nodes.pop() {
+            match node {
+                Tree::Leaf(MerkleProofLeaf::Blind(hash)) => {
+                    hashes.push(HashState::new_with_digest(
+                        NodeLeaf::Leaf,
+                        parent_index,
+                        hash.to_digest(),
+                    ));
+                }
+                Tree::Leaf(MerkleProofLeaf::Read(data)) => {
+                    hashes.push(HashState::new_with_digest(
+                        NodeLeaf::Leaf,
+                        parent_index,
+                        Hash::blake3_hash_bytes(data.as_slice()).to_digest(),
+                    ));
+                }
+                Tree::Node(children) => {
+                    hashes.push(HashState::new(NodeLeaf::Node, parent_index));
+                    let new_parent_index = hashes.len() - 1;
+                    for child in children.iter() {
+                        nodes.push((child, new_parent_index));
+                    }
+                }
+            }
+        }
+
+        while hashes.len() > 1 {
+            let hash_state = hashes
+                .pop()
+                .expect("This can't panic since we verified it's not empty");
+            hashes[hash_state.get_parent_index()].push(hash_state.get_digest());
+        }
+
+        Hash::new(hashes[0].get_digest())
     }
 }
 
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,21 @@ use crate::foldable::Foldable;
 use crate::foldable::NodeFold;
 use crate::hash::Hash;
 
+#[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 +45,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 +76,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
@@ -47,7 +47,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;
@@ -396,7 +397,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
@@ -309,7 +309,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;
 
@@ -387,7 +388,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

Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,10 @@ pub struct Hash {`
`56`	`56`	`}`
`57`	`57`
`58`	`58`	`impl Hash {`
	`59`	`+ pub fn new(digest: [u8; 32]) -> Self {`
	`60`	`+ Hash { digest }`
	`61`	`+ }`
	`62`	`+`
`59`	`63`	`/// Hash a slice of bytes`
`60`	`64`	`pub fn blake3_hash_bytes(bytes: &[u8]) -> Self {`
`61`	`65`	`let digest = blake3::hash(bytes).into();`
`@@ -107,6 +111,10 @@ impl Hash {`
`107`	`111`	`pub fn from_foldable(foldable: &impl Foldable<HashFold>) -> Self {`
`108`	`112`	`foldable.fold(HashFold)`
`109`	`113`	`}`
	`114`	`+`
	`115`	`+ pub fn to_digest(self) -> [u8; 32] {`
	`116`	`+ self.digest`
	`117`	`+ }`
`110`	`118`	`}`
`111`	`119`
`112`	`120`	`impl std::fmt::Display for Hash {`