refactor(data): make the merkle tree use the generic one

Making the MerkleTree be a type synonym for
a Tree<LeafData, NodeData>

Author: danielpapp-trilitech

data/src/compressed_merkle_tree.rs

@@ -7,6 +7,8 @@ use crate::hash::Hash;
 use crate::merkle_proof::proof_tree::MerkleProof;
 use crate::merkle_proof::proof_tree::MerkleProofLeaf;
 use crate::merkle_tree::MerkleTree;
+use crate::merkle_tree::MerkleTreeLeafData;
+use crate::merkle_tree::MerkleTreeNodeData;
 use crate::tree::Tree;
 
 // TODO RV-322: Choose optimal Merkleisation parameters for main memory.
@@ -83,9 +85,12 @@ impl CompressedMerkleTree {
 pub fn merkle_tree_to_compressed_merkle_tree(merkle_tree: MerkleTree) -> CompressedMerkleTree {
     match merkle_tree {
         MerkleTree::Leaf {
-            hash,
-            access_info,
-            data,
+            data:
+                MerkleTreeLeafData {
+                    hash,
+                    access_info,
+                    data,
+                },
         } => {
             if access_info {
                 CompressedMerkleTree::Leaf {
@@ -103,7 +108,10 @@ pub fn merkle_tree_to_compressed_merkle_tree(merkle_tree: MerkleTree) -> Compres
                 }
             }
         }
-        MerkleTree::Node(hash, children) => {
+        MerkleTree::Node {
+            data: MerkleTreeNodeData { hash },
+            children,
+        } => {
             let compressed_children: Vec<CompressedMerkleTree> = children
                 .into_iter()
                 .map(merkle_tree_to_compressed_merkle_tree)
@@ -209,6 +217,7 @@ mod tests {
     use crate::merkle_proof::proof_tree::MerkleProof;
     use crate::merkle_proof::proof_tree::MerkleProofLeaf;
     use crate::merkle_tree::MerkleTree;
+    use crate::merkle_tree::MerkleTreeLeafData;
 
     impl CompressedMerkleTree {
         /// Get the root hash of a compressed Merkle tree
@@ -265,9 +274,11 @@ mod tests {
     fn m_l(data: &[u8], access: bool) -> Result<MerkleTree, HashError> {
         let hash = Hash::blake3_hash_bytes(data);
         Ok(MerkleTree::Leaf {
-            hash,
-            access_info: access,
-            data: data.to_vec(),
+            data: MerkleTreeLeafData {
+                hash,
+                access_info: access,
+                data: data.to_vec(),
+            },
         })
     }
 

data/src/lib.rs

@@ -4,8 +4,8 @@
 // SPDX-License-Identifier: MIT
 
 pub mod clone;
-pub mod foldable;
 pub mod compressed_merkle_tree;
+pub mod foldable;
 pub mod hash;
 pub mod merkle_proof;
 pub mod merkle_tree;

data/src/merkle_tree.rs

@@ -7,21 +7,26 @@ use crate::foldable::Fold;
 use crate::foldable::Foldable;
 use crate::foldable::NodeFold;
 use crate::hash::Hash;
+use crate::tree::Tree;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct MerkleTreeLeafData {
+    pub hash: Hash,
+    pub access_info: bool,
+    pub data: Vec<u8>,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct MerkleTreeNodeData {
+    pub hash: Hash,
+}
 
 /// 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: Hash,
-        access_info: bool,
-        data: Vec<u8>,
-    },
-    Node(Hash, Vec<Self>),
-}
+pub type MerkleTree = Tree<MerkleTreeLeafData, MerkleTreeNodeData>;
 
 impl MerkleTree {
     /// Returns the precalculated root hash of the node.
@@ -41,27 +46,37 @@ impl MerkleTree {
     /// ```
     pub fn root_hash(&self) -> Hash {
         match self {
-            Self::Node(hash, _) => *hash,
-            Self::Leaf { hash, .. } => *hash,
+            Self::Node {
+                data: MerkleTreeNodeData { hash },
+                ..
+            } => *hash,
+            Self::Leaf {
+                data: 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,
+            data: MerkleTreeLeafData {
+                hash,
+                access_info,
+                data,
+            },
         }
     }
 
     /// Takes a list of children nodes and creates a
     /// new parent node on top of them.
     pub fn make_merkle_node(children: Vec<Self>) -> Self {
         let children_hashes = children.iter().map(|t| t.root_hash());
-        let node_hash = Hash::combine(children_hashes);
-        MerkleTree::Node(node_hash, children)
+        let hash = Hash::combine(children_hashes);
+        MerkleTree::Node {
+            data: MerkleTreeNodeData { hash },
+            children,
+        }
     }
 
     /// Recomputes the hashes for the whole tree
@@ -73,8 +88,18 @@ 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::Node(hash, children) => {
+                Self::Leaf {
+                    data:
+                        MerkleTreeLeafData {
+                            hash,
+                            data: node_data,
+                            ..
+                        },
+                } => &Hash::blake3_hash_bytes(node_data) == hash,
+                Self::Node {
+                    data: MerkleTreeNodeData { hash },
+                    children,
+                } => {
                     let children_hashes: Vec<Hash> = children
                         .iter()
                         .map(|child| {

src/riscv/lib/src/pvm/common.rs

@@ -10,11 +10,11 @@ use std::ops::ControlFlow;
 use bincode::Decode;
 use bincode::Encode;
 use octez_riscv_data::clone::CloneState;
+use octez_riscv_data::compressed_merkle_tree::CompressedMerkleTree;
+use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_compressed_merkle_tree;
 use octez_riscv_data::foldable::Fold;
 use octez_riscv_data::foldable::Foldable;
 use octez_riscv_data::foldable::NodeFold;
-use octez_riscv_data::compressed_merkle_tree::CompressedMerkleTree;
-use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_compressed_merkle_tree;
 use octez_riscv_data::hash::Hash;
 use octez_riscv_data::hash::HashError;
 use octez_riscv_data::hash::HashFold;

src/riscv/lib/src/state_backend/proof_backend.rs

@@ -470,6 +470,7 @@ mod tests {
     use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE;
     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;
@@ -548,9 +549,11 @@ mod tests {
             merkle_tree.check_root_hash();
             match merkle_tree {
                 MerkleTree::Leaf{
-                    hash,
-                    access_info,
-                    ..
+                    data: MerkleTreeLeafData {
+                        hash,
+                        access_info,
+                        ..
+                    }
                 } => {
                     prop_assert_eq!(hash, initial_root_hash);
                     prop_assert!(access_info);
@@ -683,7 +686,7 @@ mod tests {
 
             let pages_tree = match merkle_tree {
                 MerkleTree::Leaf{ .. } => panic!("Did not expect leaf"),
-                MerkleTree::Node(_, mut children) => {
+                MerkleTree::Node{ mut children, ..} => {
                     // The node for the pages is the second child.
                     children.remove(1)
                 },
@@ -693,10 +696,9 @@ mod tests {
             let mut leaf: usize = 0;
             while let Some(node) = queue.pop_front() {
                 match node {
-                    MerkleTree::Node(_, children) => queue.extend(children),
+                    MerkleTree::Node{ children, ..} => queue.extend(children),
                     MerkleTree::Leaf{
-                        access_info,
-                        ..
+                        data: MerkleTreeLeafData { access_info, .. }
                     } => {
                         prop_assert_eq!(
                             access_info,

src/riscv/lib/src/state_backend/proof_backend/merkle.rs

@@ -4,7 +4,6 @@
 
 //! Merkle trees used for proof generation by the PVM
 
-
 use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE;
 use octez_riscv_data::hash::HashError;
 use octez_riscv_data::merkle_tree::MerkleTree;

src/riscv/lib/src/state_backend/region.rs

@@ -13,6 +13,8 @@ use bincode::enc::Encoder;
 use bincode::error::DecodeError;
 use bincode::error::EncodeError;
 use octez_riscv_data::clone::CloneState;
+use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE;
+use octez_riscv_data::compressed_merkle_tree::chunks_to_writer;
 use octez_riscv_data::foldable::Fold;
 use octez_riscv_data::foldable::Foldable;
 use octez_riscv_data::hash::Hash;
@@ -25,6 +27,7 @@ use octez_riscv_data::mode::Normal;
 use octez_riscv_data::mode::Prove;
 use octez_riscv_data::serialisation::serialise;
 use perfect_derive::perfect_derive;
+
 use super::FnManager;
 use super::ManagerAlloc;
 use super::ManagerBase;
@@ -43,8 +46,6 @@ use crate::state_backend::RegionProj;
 use crate::state_backend::normal_backend::region_elem_offset;
 use crate::state_backend::proof_backend::merkle::MERKLE_ARITY;
 use crate::state_backend::proof_backend::merkle::MerkleWriter;
-use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE;
-use octez_riscv_data::compressed_merkle_tree::chunks_to_writer;
 use crate::state_context::projection::ApplyCons;
 use crate::state_context::projection::CellCons;
 use crate::state_context::projection::CellsCons;