feat: Sparse Merkle Tree with Generic Key Size

CHANGELOG.md

@@ -9,8 +9,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 
+- [#742](https://github.com/FuelLabs/fuel-vm/pull/742): Modified Sparse Merkle Tree to support generic key sizes. By default, the key size is 32 bytes. 
 - [#732](https://github.com/FuelLabs/fuel-vm/pull/732):  Adds `reset` method to VM memory.
 
+
 #### Breaking
 
 - [#732](https://github.com/FuelLabs/fuel-vm/pull/732): Makes the VM generic over the memory type, allowing reuse of relatively expensive-to-allocate VM memories through `VmMemoryPool`. Functions and traits which require VM initalization such as `estimate_predicates` now take either the memory or `VmMemoryPool` as an argument. The `Interpterter::eq` method now only compares accessible memory regions. `Memory` was renamed into `MemoryInstance` and `Memory` is a trait now.

fuel-merkle/src/common.rs

@@ -27,9 +27,10 @@ pub type Bytes4 = [u8; 4];
 pub type Bytes8 = [u8; 8];
 pub type Bytes16 = [u8; 16];
 pub type Bytes32 = [u8; 32];
+pub type Bytes<const N: usize> = [u8; N];
 
 use alloc::vec::Vec;
-pub type ProofSet = Vec<Bytes32>;
+pub type ProofSet<const KEY_SIZE: usize = 32> = Vec<Bytes<KEY_SIZE>>;
 
 pub use hash::{
     sum,

fuel-merkle/src/common/node.rs

@@ -1,7 +1,4 @@
-use crate::common::{
-    Bytes32,
-    Bytes8,
-};
+use crate::common::Bytes;
 
 use alloc::string::String;
 use core::fmt;
@@ -56,15 +53,7 @@ where
     }
 }
 
-impl KeyFormatting for Bytes8 {
-    type PrettyType = u64;
-
-    fn pretty(&self) -> Self::PrettyType {
-        u64::from_be_bytes(*self)
-    }
-}
-
-impl KeyFormatting for Bytes32 {
+impl<const N: usize> KeyFormatting for Bytes<N> {
     type PrettyType = String;
 
     fn pretty(&self) -> Self::PrettyType {

fuel-merkle/src/sparse.rs

@@ -4,17 +4,19 @@ mod primitive;
 
 pub(crate) use hash::zero_sum;
 
-pub use merkle_tree::{
-    MerkleTree,
-    MerkleTreeError,
-    MerkleTreeKey,
-};
-pub use primitive::Primitive;
 pub mod in_memory;
 pub mod proof;
 
 use crate::common::Bytes32;
 
-pub const fn empty_sum() -> &'static Bytes32 {
+// Define default Merkle Tree structures as concrete implementations of generic
+// types, using 32 byte key sizes
+pub type MerkleTree<TableType, StorageType> =
+    merkle_tree::MerkleTree<32, TableType, StorageType>;
+pub type MerkleTreeError<StorageError> = merkle_tree::MerkleTreeError<32, StorageError>;
+pub type MerkleTreeKey = merkle_tree::MerkleTreeKey<32>;
+pub type Primitive = primitive::Primitive<32>;
+
+pub fn empty_sum() -> &'static Bytes32 {
     zero_sum()
 }

fuel-merkle/src/sparse/hash.rs

@@ -1,28 +1,48 @@
 use crate::common::{
-    sum_iter,
-    Bytes32,
+    sum,
+    Bytes,
     Prefix,
 };
+use std::sync::OnceLock;
 
-pub const fn zero_sum() -> &'static Bytes32 {
-    const ZERO_SUM: Bytes32 = [0; 32];
+pub fn zero_sum<const N: usize>() -> &'static [u8; N] {
+    static ZERO: OnceLock<Vec<u8>> = OnceLock::new();
+    ZERO.get_or_init(|| vec![0; N])
+        .as_slice()
+        .try_into()
+        .expect("Expected valid zero sum")
+}
 
-    &ZERO_SUM
+pub fn sum_truncated<const N: usize, T: AsRef<[u8]>>(data: T) -> Bytes<N> {
+    let hash = sum(data);
+    let mut vec = hash.as_slice().to_vec();
+    vec.truncate(N);
+    vec.try_into().unwrap()
 }
 
-pub fn calculate_hash(
+pub fn calculate_hash<const N: usize>(
     prefix: &Prefix,
-    bytes_lo: &Bytes32,
-    bytes_hi: &Bytes32,
-) -> Bytes32 {
-    let input = [prefix.as_ref(), bytes_lo.as_ref(), bytes_hi.as_ref()];
-    sum_iter(input)
+    bytes_lo: &Bytes<N>,
+    bytes_hi: &Bytes<N>,
+) -> Bytes<N> {
+    let input = [prefix.as_ref(), bytes_lo.as_ref(), bytes_hi.as_ref()]
+        .into_iter()
+        .flatten()
+        .cloned()
+        .collect::<Vec<_>>();
+    sum_truncated(input)
 }
 
-pub fn calculate_leaf_hash(leaf_key: &Bytes32, leaf_value: &Bytes32) -> Bytes32 {
+pub fn calculate_leaf_hash<const N: usize>(
+    leaf_key: &Bytes<N>,
+    leaf_value: &Bytes<N>,
+) -> Bytes<N> {
     calculate_hash(&Prefix::Leaf, leaf_key, leaf_value)
 }
 
-pub fn calculate_node_hash(left_child: &Bytes32, right_child: &Bytes32) -> Bytes32 {
+pub fn calculate_node_hash<const N: usize>(
+    left_child: &Bytes<N>,
+    right_child: &Bytes<N>,
+) -> Bytes<N> {
     calculate_hash(&Prefix::Node, left_child, right_child)
 }

fuel-merkle/src/sparse/in_memory.rs

@@ -1,13 +1,14 @@
 use crate::{
     common::{
+        Bytes,
         Bytes32,
         StorageMap,
     },
     sparse::{
         self,
         merkle_tree::MerkleTreeKey,
+        primitive::Primitive,
         proof::Proof,
-        Primitive,
     },
     storage::{
         Mappable,
@@ -27,10 +28,12 @@ use alloc::{
 #[derive(Debug)]
 pub struct NodesTable;
 
+const DEFAULT_KEY_SIZE: usize = 32;
+
 impl Mappable for NodesTable {
     type Key = Self::OwnedKey;
-    type OwnedKey = Bytes32;
-    type OwnedValue = Primitive;
+    type OwnedKey = Bytes<DEFAULT_KEY_SIZE>;
+    type OwnedValue = Primitive<DEFAULT_KEY_SIZE>;
     type Value = Self::OwnedValue;
 }
 
@@ -58,7 +61,7 @@ impl MerkleTree {
     /// data.
     pub fn from_set<I, D>(set: I) -> Self
     where
-        I: Iterator<Item = (MerkleTreeKey, D)>,
+        I: Iterator<Item = (MerkleTreeKey<DEFAULT_KEY_SIZE>, D)>,
         D: AsRef<[u8]>,
     {
         let tree = SparseMerkleTree::from_set(Storage::new(), set)
@@ -74,9 +77,9 @@ impl MerkleTree {
     /// not incur the overhead of storage writes. This can be helpful when we
     /// know all the key-values in the set upfront and we will not need to
     /// update the set in the future.
-    pub fn root_from_set<I, D>(set: I) -> Bytes32
+    pub fn root_from_set<I, D>(set: I) -> Bytes<DEFAULT_KEY_SIZE>
     where
-        I: Iterator<Item = (MerkleTreeKey, D)>,
+        I: Iterator<Item = (MerkleTreeKey<DEFAULT_KEY_SIZE>, D)>,
         D: AsRef<[u8]>,
     {
         #[derive(Default)]
@@ -85,7 +88,11 @@ impl MerkleTree {
         impl StorageInspect<NodesTable> for EmptyStorage {
             type Error = core::convert::Infallible;
 
-            fn get(&self, _: &Bytes32) -> Result<Option<Cow<Primitive>>, Self::Error> {
+            fn get(
+                &self,
+                _: &Bytes32,
+            ) -> Result<Option<Cow<Primitive<DEFAULT_KEY_SIZE>>>, Self::Error>
+            {
                 Ok(None)
             }
 
@@ -97,13 +104,16 @@ impl MerkleTree {
         impl StorageMutate<NodesTable> for EmptyStorage {
             fn insert(
                 &mut self,
-                _: &Bytes32,
-                _: &Primitive,
-            ) -> Result<Option<Primitive>, Self::Error> {
+                _: &Bytes<DEFAULT_KEY_SIZE>,
+                _: &Primitive<DEFAULT_KEY_SIZE>,
+            ) -> Result<Option<Primitive<DEFAULT_KEY_SIZE>>, Self::Error> {
                 Ok(None)
             }
 
-            fn remove(&mut self, _: &Bytes32) -> Result<Option<Primitive>, Self::Error> {
+            fn remove(
+                &mut self,
+                _: &Bytes<DEFAULT_KEY_SIZE>,
+            ) -> Result<Option<Primitive<DEFAULT_KEY_SIZE>>, Self::Error> {
                 Ok(None)
             }
         }
@@ -121,20 +131,29 @@ impl MerkleTree {
     /// This can be helpful when we know all the key-values in the set upfront
     /// and we need to defer storage writes, such as expensive database inserts,
     /// for batch operations later in the process.
-    pub fn nodes_from_set<I, D>(set: I) -> (Bytes32, Vec<(Bytes32, Primitive)>)
+    pub fn nodes_from_set<I, D>(
+        set: I,
+    ) -> (
+        Bytes<DEFAULT_KEY_SIZE>,
+        Vec<(Bytes<DEFAULT_KEY_SIZE>, Primitive<DEFAULT_KEY_SIZE>)>,
+    )
     where
-        I: Iterator<Item = (MerkleTreeKey, D)>,
+        I: Iterator<Item = (MerkleTreeKey<DEFAULT_KEY_SIZE>, D)>,
         D: AsRef<[u8]>,
     {
         #[derive(Default)]
         struct VectorStorage {
-            storage: Vec<(Bytes32, Primitive)>,
+            storage: Vec<(Bytes32, Primitive<DEFAULT_KEY_SIZE>)>,
         }
 
         impl StorageInspect<NodesTable> for VectorStorage {
             type Error = core::convert::Infallible;
 
-            fn get(&self, _: &Bytes32) -> Result<Option<Cow<Primitive>>, Self::Error> {
+            fn get(
+                &self,
+                _: &Bytes32,
+            ) -> Result<Option<Cow<Primitive<DEFAULT_KEY_SIZE>>>, Self::Error>
+            {
                 unimplemented!("Read operation is not supported")
             }
 
@@ -146,14 +165,17 @@ impl MerkleTree {
         impl StorageMutate<NodesTable> for VectorStorage {
             fn insert(
                 &mut self,
-                key: &Bytes32,
-                value: &Primitive,
-            ) -> Result<Option<Primitive>, Self::Error> {
+                key: &Bytes<DEFAULT_KEY_SIZE>,
+                value: &Primitive<DEFAULT_KEY_SIZE>,
+            ) -> Result<Option<Primitive<DEFAULT_KEY_SIZE>>, Self::Error> {
                 self.storage.push((*key, *value));
                 Ok(None)
             }
 
-            fn remove(&mut self, _: &Bytes32) -> Result<Option<Primitive>, Self::Error> {
+            fn remove(
+                &mut self,
+                _: &Bytes<DEFAULT_KEY_SIZE>,
+            ) -> Result<Option<Primitive<DEFAULT_KEY_SIZE>>, Self::Error> {
                 unimplemented!("Remove operation is not supported")
             }
         }
@@ -167,19 +189,22 @@ impl MerkleTree {
         (root, nodes)
     }
 
-    pub fn update(&mut self, key: MerkleTreeKey, data: &[u8]) {
+    pub fn update(&mut self, key: MerkleTreeKey<DEFAULT_KEY_SIZE>, data: &[u8]) {
         let _ = self.tree.update(key, data);
     }
 
-    pub fn delete(&mut self, key: MerkleTreeKey) {
+    pub fn delete(&mut self, key: MerkleTreeKey<DEFAULT_KEY_SIZE>) {
         let _ = self.tree.delete(key);
     }
 
     pub fn root(&self) -> Bytes32 {
         self.tree.root()
     }
 
-    pub fn generate_proof(&self, key: &MerkleTreeKey) -> Option<Proof> {
+    pub fn generate_proof(
+        &self,
+        key: &MerkleTreeKey<DEFAULT_KEY_SIZE>,
+    ) -> Option<Proof<DEFAULT_KEY_SIZE>> {
         self.tree.generate_proof(key).ok()
     }
 }
@@ -195,7 +220,7 @@ mod test {
     use super::*;
     use crate::common::sum;
 
-    fn key(data: &[u8]) -> MerkleTreeKey {
+    fn key(data: &[u8]) -> MerkleTreeKey<DEFAULT_KEY_SIZE> {
         MerkleTreeKey::new_without_hash(sum(data))
     }