From 9bd300745cb5ea51ca923124d6361d04eb8233ac Mon Sep 17 00:00:00 2001
From: Brandon Vrooman <brandon.vrooman@gmail.com>
Date: Fri, 31 May 2024 19:10:41 -0400
Subject: [PATCH 1/4] Sparse Merkle Tree with Generic Key Size

---
 fuel-merkle/src/common.rs                    |   3 +-
 fuel-merkle/src/common/node.rs               |  15 +-
 fuel-merkle/src/sparse.rs                    |   2 +-
 fuel-merkle/src/sparse/hash.rs               |  46 +++--
 fuel-merkle/src/sparse/in_memory.rs          |  86 +++++++---
 fuel-merkle/src/sparse/merkle_tree.rs        | 170 +++++++++++--------
 fuel-merkle/src/sparse/merkle_tree/branch.rs |  34 ++--
 fuel-merkle/src/sparse/merkle_tree/node.rs   | 129 ++++++++------
 fuel-merkle/src/sparse/primitive.rs          |  20 +--
 fuel-merkle/src/sparse/proof.rs              |  67 ++++----
 fuel-merkle/src/tests/sparse.rs              |   4 +-
 11 files changed, 334 insertions(+), 242 deletions(-)

diff --git a/fuel-merkle/src/common.rs b/fuel-merkle/src/common.rs
index 3a11cfa119..58a1973ccf 100644
--- a/fuel-merkle/src/common.rs
+++ b/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,
diff --git a/fuel-merkle/src/common/node.rs b/fuel-merkle/src/common/node.rs
index 8d0845d6f3..a8fd27c5af 100644
--- a/fuel-merkle/src/common/node.rs
+++ b/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 {
diff --git a/fuel-merkle/src/sparse.rs b/fuel-merkle/src/sparse.rs
index c8fd73c506..27563937dd 100644
--- a/fuel-merkle/src/sparse.rs
+++ b/fuel-merkle/src/sparse.rs
@@ -15,6 +15,6 @@ pub mod proof;
 
 use crate::common::Bytes32;
 
-pub const fn empty_sum() -> &'static Bytes32 {
+pub fn empty_sum() -> &'static Bytes32 {
     zero_sum()
 }
diff --git a/fuel-merkle/src/sparse/hash.rs b/fuel-merkle/src/sparse/hash.rs
index 0d9086c193..380b4203c7 100644
--- a/fuel-merkle/src/sparse/hash.rs
+++ b/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)
 }
diff --git a/fuel-merkle/src/sparse/in_memory.rs b/fuel-merkle/src/sparse/in_memory.rs
index ae1c77bec0..f9a84a4bb8 100644
--- a/fuel-merkle/src/sparse/in_memory.rs
+++ b/fuel-merkle/src/sparse/in_memory.rs
@@ -1,5 +1,6 @@
 use crate::{
     common::{
+        Bytes,
         Bytes32,
         StorageMap,
     },
@@ -27,15 +28,17 @@ 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;
 }
 
 type Storage = StorageMap<NodesTable>;
-type SparseMerkleTree = sparse::MerkleTree<NodesTable, Storage>;
+type SparseMerkleTree = sparse::MerkleTree<DEFAULT_KEY_SIZE, NodesTable, Storage>;
 
 #[derive(Debug)]
 pub struct MerkleTree {
@@ -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,19 +104,25 @@ 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)
             }
         }
 
-        let tree = sparse::MerkleTree::<NodesTable, _>::from_set(EmptyStorage, set)
-            .expect("`Storage` can't return error");
+        let tree = sparse::MerkleTree::<DEFAULT_KEY_SIZE, NodesTable, _>::from_set(
+            EmptyStorage,
+            set,
+        )
+        .expect("`Storage` can't return error");
         tree.root()
     }
 
@@ -121,20 +134,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,32 +168,37 @@ 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")
             }
         }
 
-        let tree =
-            sparse::MerkleTree::<NodesTable, _>::from_set(VectorStorage::default(), set)
-                .expect("`Storage` can't return error");
+        let tree = sparse::MerkleTree::<DEFAULT_KEY_SIZE, NodesTable, _>::from_set(
+            VectorStorage::default(),
+            set,
+        )
+        .expect("`Storage` can't return error");
         let root = tree.root();
         let nodes = tree.into_storage().storage;
 
         (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);
     }
 
@@ -179,7 +206,10 @@ impl MerkleTree {
         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 +225,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))
     }
 
diff --git a/fuel-merkle/src/sparse/merkle_tree.rs b/fuel-merkle/src/sparse/merkle_tree.rs
index 83bb8f5912..08a4c6e6a1 100644
--- a/fuel-merkle/src/sparse/merkle_tree.rs
+++ b/fuel-merkle/src/sparse/merkle_tree.rs
@@ -16,10 +16,10 @@ use crate::{
         error::DeserializeError,
         node::ChildError,
         AsPathIterator,
-        Bytes32,
+        Bytes,
     },
     sparse::{
-        empty_sum,
+        hash::sum_truncated,
         proof::{
             ExclusionLeaf,
             ExclusionLeafData,
@@ -50,12 +50,12 @@ use core::{
 };
 
 #[derive(Debug, Clone, derive_more::Display)]
-pub enum MerkleTreeError<StorageError> {
+pub enum MerkleTreeError<const KEY_SIZE: usize, StorageError> {
     #[display(
         fmt = "cannot load node with key {}; the key is not found in storage",
         "hex::encode(_0)"
     )]
-    LoadError(Bytes32),
+    LoadError(Bytes<KEY_SIZE>),
 
     #[display(fmt = "{}", _0)]
     StorageError(StorageError),
@@ -64,11 +64,13 @@ pub enum MerkleTreeError<StorageError> {
     DeserializeError(DeserializeError),
 
     #[display(fmt = "{}", _0)]
-    ChildError(ChildError<Bytes32, StorageNodeError<StorageError>>),
+    ChildError(ChildError<Bytes<KEY_SIZE>, StorageNodeError<StorageError>>),
 }
 
-impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
-    fn from(err: StorageError) -> MerkleTreeError<StorageError> {
+impl<const KEY_SIZE: usize, StorageError> From<StorageError>
+    for MerkleTreeError<KEY_SIZE, StorageError>
+{
+    fn from(err: StorageError) -> MerkleTreeError<KEY_SIZE, StorageError> {
         MerkleTreeError::StorageError(err)
     }
 }
@@ -77,19 +79,16 @@ impl<StorageError> From<StorageError> for MerkleTreeError<StorageError> {
 /// The type contains only one constructor that hashes the storage key.
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 #[cfg_attr(test, derive(proptest_derive::Arbitrary))]
-pub struct MerkleTreeKey(Bytes32);
+pub struct MerkleTreeKey<const N: usize>(Bytes<N>);
 
-impl MerkleTreeKey {
+impl<const N: usize> MerkleTreeKey<N> {
     /// The safe way to create a `Self`. It hashes the `storage_key`, making
     /// it entirely random and preventing SMT structure manipulation.
     pub fn new<B>(storage_key: B) -> Self
     where
         B: AsRef<[u8]>,
     {
-        use digest::Digest;
-        let mut hash = sha2::Sha256::new();
-        hash.update(storage_key.as_ref());
-        let hash = hash.finalize().into();
+        let hash = sum_truncated(storage_key);
 
         Self(hash)
     }
@@ -103,7 +102,7 @@ impl MerkleTreeKey {
     /// was randomly generated like `ContractId` or `AssetId`.
     pub unsafe fn convert<B>(storage_key: B) -> Self
     where
-        B: Into<Bytes32>,
+        B: Into<Bytes<N>>,
     {
         Self(storage_key.into())
     }
@@ -111,38 +110,38 @@ impl MerkleTreeKey {
     #[cfg(any(test, feature = "test-helpers"))]
     pub fn new_without_hash<B>(storage_key: B) -> Self
     where
-        B: Into<Bytes32>,
+        B: Into<Bytes<N>>,
     {
         unsafe { Self::convert(storage_key) }
     }
 }
 
-impl Debug for MerkleTreeKey {
+impl<const N: usize> Debug for MerkleTreeKey<N> {
     fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
         f.write_str(&format!("MerkleTreeKey({})", hex::encode(self.0)))
     }
 }
 
-impl From<MerkleTreeKey> for Bytes32 {
-    fn from(value: MerkleTreeKey) -> Self {
+impl<const N: usize> From<MerkleTreeKey<N>> for Bytes<N> {
+    fn from(value: MerkleTreeKey<N>) -> Self {
         value.0
     }
 }
 
-impl AsRef<[u8]> for MerkleTreeKey {
+impl<const N: usize> AsRef<[u8]> for MerkleTreeKey<N> {
     fn as_ref(&self) -> &[u8] {
         self.0.as_ref()
     }
 }
 
-impl AsRef<Bytes32> for MerkleTreeKey {
-    fn as_ref(&self) -> &Bytes32 {
+impl<const N: usize> AsRef<Bytes<N>> for MerkleTreeKey<N> {
+    fn as_ref(&self) -> &Bytes<N> {
         &self.0
     }
 }
 
-impl Deref for MerkleTreeKey {
-    type Target = Bytes32;
+impl<const N: usize> Deref for MerkleTreeKey<N> {
+    type Target = Bytes<N>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
@@ -150,25 +149,27 @@ impl Deref for MerkleTreeKey {
 }
 
 #[cfg(any(test, feature = "test-helpers"))]
-impl From<Bytes32> for MerkleTreeKey {
-    fn from(value: Bytes32) -> Self {
+impl<const N: usize> From<Bytes<N>> for MerkleTreeKey<N> {
+    fn from(value: Bytes<N>) -> Self {
         Self::new_without_hash(value)
     }
 }
 
 #[derive(Debug)]
-pub struct MerkleTree<TableType, StorageType> {
-    root_node: Node,
+pub struct MerkleTree<const KEY_SIZE: usize, TableType, StorageType> {
+    root_node: Node<KEY_SIZE>,
     storage: StorageType,
     phantom_table: PhantomData<TableType>,
 }
 
-impl<TableType, StorageType> MerkleTree<TableType, StorageType> {
-    pub const fn empty_root() -> &'static Bytes32 {
-        empty_sum()
+impl<const KEY_SIZE: usize, TableType, StorageType>
+    MerkleTree<KEY_SIZE, TableType, StorageType>
+{
+    pub fn empty_root() -> &'static Bytes<KEY_SIZE> {
+        crate::sparse::hash::zero_sum()
     }
 
-    pub fn root(&self) -> Bytes32 {
+    pub fn root(&self) -> Bytes<KEY_SIZE> {
         *self.root_node().hash()
     }
 
@@ -180,19 +181,24 @@ impl<TableType, StorageType> MerkleTree<TableType, StorageType> {
         &self.storage
     }
 
-    fn root_node(&self) -> &Node {
+    fn root_node(&self) -> &Node<KEY_SIZE> {
         &self.root_node
     }
 
-    fn set_root_node(&mut self, node: Node) {
-        debug_assert!(node.is_leaf() || node.height() == Node::max_height());
+    fn set_root_node(&mut self, node: Node<KEY_SIZE>) {
+        debug_assert!(node.is_leaf() || node.height() == Node::<KEY_SIZE>::max_height());
         self.root_node = node;
     }
 }
 
-impl<TableType, StorageType, StorageError> MerkleTree<TableType, StorageType>
+impl<const KEY_SIZE: usize, TableType, StorageType, StorageError>
+    MerkleTree<KEY_SIZE, TableType, StorageType>
 where
-    TableType: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    TableType: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
     StorageType: StorageInspect<TableType, Error = StorageError>,
 {
     pub fn new(storage: StorageType) -> Self {
@@ -205,8 +211,8 @@ where
 
     pub fn load(
         storage: StorageType,
-        root: &Bytes32,
-    ) -> Result<Self, MerkleTreeError<StorageError>> {
+        root: &Bytes<KEY_SIZE>,
+    ) -> Result<Self, MerkleTreeError<KEY_SIZE, StorageError>> {
         if root == Self::empty_root() {
             let tree = Self::new(storage);
             Ok(tree)
@@ -228,21 +234,25 @@ where
 
     fn path_set(
         &self,
-        leaf_key: &Bytes32,
-    ) -> Result<(Vec<Node>, Vec<Node>), MerkleTreeError<StorageError>> {
+        leaf_key: &Bytes<KEY_SIZE>,
+    ) -> Result<
+        (Vec<Node<KEY_SIZE>>, Vec<Node<KEY_SIZE>>),
+        MerkleTreeError<KEY_SIZE, StorageError>,
+    > {
         let root_node = self.root_node().clone();
         let root_storage_node = StorageNode::new(&self.storage, root_node);
-        let (mut path_nodes, mut side_nodes): (Vec<Node>, Vec<Node>) = root_storage_node
-            .as_path_iter(leaf_key)
-            .map(|(path_node, side_node)| {
-                Ok((
-                    path_node.map_err(MerkleTreeError::ChildError)?.into_node(),
-                    side_node.map_err(MerkleTreeError::ChildError)?.into_node(),
-                ))
-            })
-            .collect::<Result<Vec<_>, MerkleTreeError<StorageError>>>()?
-            .into_iter()
-            .unzip();
+        let (mut path_nodes, mut side_nodes): (Vec<Node<KEY_SIZE>>, Vec<Node<KEY_SIZE>>) =
+            root_storage_node
+                .as_path_iter(leaf_key)
+                .map(|(path_node, side_node)| {
+                    Ok((
+                        path_node.map_err(MerkleTreeError::ChildError)?.into_node(),
+                        side_node.map_err(MerkleTreeError::ChildError)?.into_node(),
+                    ))
+                })
+                .collect::<Result<Vec<_>, MerkleTreeError<KEY_SIZE, StorageError>>>()?
+                .into_iter()
+                .unzip();
         path_nodes.reverse();
         side_nodes.reverse();
         side_nodes.pop(); // The last element in the side nodes list is the
@@ -252,9 +262,14 @@ where
     }
 }
 
-impl<TableType, StorageType, StorageError> MerkleTree<TableType, StorageType>
+impl<const KEY_SIZE: usize, TableType, StorageType, StorageError>
+    MerkleTree<KEY_SIZE, TableType, StorageType>
 where
-    TableType: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    TableType: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
     StorageType: StorageMutate<TableType, Error = StorageError>,
 {
     /// Build a sparse Merkle tree from a set of key-value pairs. This is
@@ -270,18 +285,18 @@ where
     ) -> Result<Self, StorageError>
     where
         I: Iterator<Item = (B, D)>,
-        B: Into<Bytes32>,
+        B: Into<Bytes<KEY_SIZE>>,
         D: AsRef<[u8]>,
     {
         let sorted = set
             .into_iter()
             .map(|(k, v)| (k.into(), v))
-            .collect::<alloc::collections::BTreeMap<Bytes32, D>>();
+            .collect::<alloc::collections::BTreeMap<Bytes<KEY_SIZE>, D>>();
         let mut branches = sorted
             .iter()
             .filter(|(_, value)| !value.as_ref().is_empty())
             .map(|(key, data)| Node::create_leaf(key, data))
-            .map(Into::<Branch>::into)
+            .map(Into::<Branch<KEY_SIZE>>::into)
             .collect::<Vec<_>>();
 
         for branch in branches.iter() {
@@ -301,7 +316,7 @@ where
             return Ok(tree)
         }
 
-        let mut nodes = Vec::<Branch>::with_capacity(branches.len());
+        let mut nodes = Vec::<Branch<KEY_SIZE>>::with_capacity(branches.len());
         let mut proximities = Vec::<u32>::with_capacity(branches.len());
 
         // Building the tree starts by merging all leaf nodes where possible.
@@ -389,7 +404,7 @@ where
         let path = top.bits;
         let height = node.height();
         #[allow(clippy::arithmetic_side_effects)] // height <= max_height
-        let depth = Node::max_height() - height;
+        let depth = Node::<KEY_SIZE>::max_height() - height;
         let placeholders = iter::repeat(Node::create_placeholder()).take(depth as usize);
         for placeholder in placeholders {
             node = Node::create_node_on_path(&path, &node, &placeholder);
@@ -406,9 +421,9 @@ where
 
     pub fn update(
         &mut self,
-        key: MerkleTreeKey,
+        key: MerkleTreeKey<KEY_SIZE>,
         data: &[u8],
-    ) -> Result<(), MerkleTreeError<StorageError>> {
+    ) -> Result<(), MerkleTreeError<KEY_SIZE, StorageError>> {
         if data.is_empty() {
             // If the data is empty, this signifies a delete operation for the
             // given key.
@@ -436,15 +451,15 @@ where
 
     pub fn delete(
         &mut self,
-        key: MerkleTreeKey,
-    ) -> Result<(), MerkleTreeError<StorageError>> {
+        key: MerkleTreeKey<KEY_SIZE>,
+    ) -> Result<(), MerkleTreeError<KEY_SIZE, StorageError>> {
         if self.root() == *Self::empty_root() {
             // The zero root signifies that all leaves are empty, including the
             // given key.
             return Ok(())
         }
 
-        let (path_nodes, side_nodes): (Vec<Node>, Vec<Node>) =
+        let (path_nodes, side_nodes): (Vec<Node<KEY_SIZE>>, Vec<Node<KEY_SIZE>>) =
             self.path_set(key.as_ref())?;
 
         match path_nodes.first() {
@@ -463,9 +478,9 @@ where
 
     fn update_with_path_set(
         &mut self,
-        requested_leaf_node: &Node,
-        path_nodes: &[Node],
-        side_nodes: &[Node],
+        requested_leaf_node: &Node<KEY_SIZE>,
+        path_nodes: &[Node<KEY_SIZE>],
+        side_nodes: &[Node<KEY_SIZE>],
     ) -> Result<(), StorageError> {
         let path = requested_leaf_node.leaf_key();
         let actual_leaf_node = &path_nodes[0];
@@ -541,9 +556,9 @@ where
 
     fn delete_with_path_set(
         &mut self,
-        requested_leaf_key: &Bytes32,
-        path_nodes: &[Node],
-        side_nodes: &[Node],
+        requested_leaf_key: &Bytes<KEY_SIZE>,
+        path_nodes: &[Node<KEY_SIZE>],
+        side_nodes: &[Node<KEY_SIZE>],
     ) -> Result<(), StorageError> {
         for node in path_nodes {
             self.storage.remove(node.hash())?;
@@ -606,15 +621,20 @@ where
     }
 }
 
-impl<TableType, StorageType, StorageError> MerkleTree<TableType, StorageType>
+impl<const KEY_SIZE: usize, TableType, StorageType, StorageError>
+    MerkleTree<KEY_SIZE, TableType, StorageType>
 where
-    TableType: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    TableType: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
     StorageType: StorageInspect<TableType, Error = StorageError>,
 {
     pub fn generate_proof(
         &self,
-        key: &MerkleTreeKey,
-    ) -> Result<Proof, MerkleTreeError<StorageError>> {
+        key: &MerkleTreeKey<KEY_SIZE>,
+    ) -> Result<Proof<KEY_SIZE>, MerkleTreeError<KEY_SIZE, StorageError>> {
         let path = key.as_ref();
         let (path_nodes, side_nodes) = self.path_set(path)?;
         // Identify the closest leaf that is included in the tree to the
@@ -713,11 +733,11 @@ mod test {
     impl Mappable for TestTable {
         type Key = Self::OwnedKey;
         type OwnedKey = Bytes32;
-        type OwnedValue = Primitive;
+        type OwnedValue = Primitive<32>;
         type Value = Self::OwnedValue;
     }
 
-    fn key<B: AsRef<[u8]>>(data: B) -> MerkleTreeKey {
+    fn key<B: AsRef<[u8]>>(data: B) -> MerkleTreeKey<32> {
         MerkleTreeKey::new(data.as_ref())
     }
 
diff --git a/fuel-merkle/src/sparse/merkle_tree/branch.rs b/fuel-merkle/src/sparse/merkle_tree/branch.rs
index 7d597f38f6..cb2f0c8384 100644
--- a/fuel-merkle/src/sparse/merkle_tree/branch.rs
+++ b/fuel-merkle/src/sparse/merkle_tree/branch.rs
@@ -1,9 +1,6 @@
 use super::Node;
 use crate::{
-    common::{
-        path::Path,
-        Bytes32,
-    },
+    common::path::Path,
     sparse::Primitive,
 };
 use fuel_storage::{
@@ -11,15 +8,16 @@ use fuel_storage::{
     StorageMutate,
 };
 
+use crate::common::Bytes;
 use core::iter;
 
-pub(super) struct Branch {
-    pub bits: Bytes32,
-    pub node: Node,
+pub(super) struct Branch<const KEY_SIZE: usize> {
+    pub bits: Bytes<KEY_SIZE>,
+    pub node: Node<KEY_SIZE>,
 }
 
-impl From<Node> for Branch {
-    fn from(leaf: Node) -> Self {
+impl<const KEY_SIZE: usize> From<Node<KEY_SIZE>> for Branch<KEY_SIZE> {
+    fn from(leaf: Node<KEY_SIZE>) -> Self {
         Self {
             bits: *leaf.leaf_key(),
             node: leaf,
@@ -27,26 +25,30 @@ impl From<Node> for Branch {
     }
 }
 
-pub(super) fn merge_branches<Storage, Table>(
+pub(super) fn merge_branches<const KEY_SIZE: usize, Storage, Table>(
     storage: &mut Storage,
-    mut left_branch: Branch,
-    mut right_branch: Branch,
-) -> Result<Branch, Storage::Error>
+    mut left_branch: Branch<KEY_SIZE>,
+    mut right_branch: Branch<KEY_SIZE>,
+) -> Result<Branch<KEY_SIZE>, Storage::Error>
 where
     Storage: StorageMutate<Table>,
-    Table: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    Table: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
 {
     #[allow(clippy::cast_possible_truncation)] // Key is 32 bytes, never truncates
     let ancestor_height = if left_branch.node.is_leaf() && right_branch.node.is_leaf() {
         let parent_depth = left_branch.node.common_path_length(&right_branch.node) as u32;
         #[allow(clippy::arithmetic_side_effects)] // common_path_length <= max_height
-        let parent_height = Node::max_height() - parent_depth;
+        let parent_height = Node::<KEY_SIZE>::max_height() - parent_depth;
         parent_height
     } else {
         let ancestor_depth =
             left_branch.bits.common_path_length(&right_branch.bits) as u32;
         #[allow(clippy::arithmetic_side_effects)] // common_path_length <= max_height
-        let ancestor_height = Node::max_height() - ancestor_depth;
+        let ancestor_height = Node::<KEY_SIZE>::max_height() - ancestor_depth;
 
         for branch in [&mut right_branch, &mut left_branch] {
             if branch.node.is_node() {
diff --git a/fuel-merkle/src/sparse/merkle_tree/node.rs b/fuel-merkle/src/sparse/merkle_tree/node.rs
index 85371beebe..0381c99cf1 100644
--- a/fuel-merkle/src/sparse/merkle_tree/node.rs
+++ b/fuel-merkle/src/sparse/merkle_tree/node.rs
@@ -11,8 +11,6 @@ use crate::{
             Path,
             Side,
         },
-        sum,
-        Bytes32,
         Prefix,
     },
     sparse::{
@@ -33,33 +31,37 @@ use crate::{
     },
 };
 
+use crate::{
+    common::Bytes,
+    sparse::hash::sum_truncated,
+};
 use core::{
     fmt,
     marker::PhantomData,
 };
 
 #[derive(Clone, PartialEq, Eq)]
-pub(super) enum Node {
+pub(super) enum Node<const N: usize> {
     Node {
-        hash: Bytes32,
+        hash: Bytes<N>,
         height: u32,
         prefix: Prefix,
-        bytes_lo: Bytes32,
-        bytes_hi: Bytes32,
+        bytes_lo: Bytes<N>,
+        bytes_hi: Bytes<N>,
     },
     Placeholder,
 }
 
-impl Node {
+impl<const N: usize> Node<N> {
     pub fn max_height() -> u32 {
-        Node::key_size_bits()
+        Node::<N>::key_size_bits()
     }
 
     pub fn new(
         height: u32,
         prefix: Prefix,
-        bytes_lo: Bytes32,
-        bytes_hi: Bytes32,
+        bytes_lo: Bytes<N>,
+        bytes_hi: Bytes<N>,
     ) -> Self {
         Self::Node {
             hash: calculate_hash(&prefix, &bytes_lo, &bytes_hi),
@@ -70,8 +72,11 @@ impl Node {
         }
     }
 
-    pub fn create_leaf<D: AsRef<[u8]>>(key: &Bytes32, data: D) -> Self {
-        let bytes_hi = sum(data);
+    pub fn create_leaf<D>(key: &Bytes<N>, data: D) -> Self
+    where
+        D: AsRef<[u8]>,
+    {
+        let bytes_hi = sum_truncated(data);
         Self::Node {
             hash: calculate_leaf_hash(key, &bytes_hi),
             height: 0u32,
@@ -81,7 +86,7 @@ impl Node {
         }
     }
 
-    pub fn create_node(left_child: &Node, right_child: &Node, height: u32) -> Self {
+    pub fn create_node(left_child: &Node<N>, right_child: &Node<N>, height: u32) -> Self {
         let bytes_lo = *left_child.hash();
         let bytes_hi = *right_child.hash();
         Self::Node {
@@ -95,8 +100,8 @@ impl Node {
 
     pub fn create_node_on_path(
         path: &dyn Path,
-        path_node: &Node,
-        side_node: &Node,
+        path_node: &Node<N>,
+        side_node: &Node<N>,
     ) -> Self {
         if path_node.is_leaf() && side_node.is_leaf() {
             // When joining two leaves, the joined node is found where the paths
@@ -106,7 +111,7 @@ impl Node {
             #[allow(clippy::cast_possible_truncation)] // Key is 32 bytes
             let parent_depth = path_node.common_path_length(side_node) as u32;
             #[allow(clippy::arithmetic_side_effects)] // parent_depth <= max_height
-            let parent_height = Node::max_height() - parent_depth;
+            let parent_height = Node::<N>::max_height() - parent_depth;
             match path.get_instruction(parent_depth).unwrap() {
                 Side::Left => Node::create_node(path_node, side_node, parent_height),
                 Side::Right => Node::create_node(side_node, path_node, parent_height),
@@ -119,7 +124,7 @@ impl Node {
             #[allow(clippy::arithmetic_side_effects)] // Neither node cannot be root
             let parent_height = path_node.height().max(side_node.height()) + 1;
             #[allow(clippy::arithmetic_side_effects)] // parent_height <= max_height
-            let parent_depth = Node::max_height() - parent_height;
+            let parent_depth = Node::<N>::max_height() - parent_height;
             match path.get_instruction(parent_depth).unwrap() {
                 Side::Left => Node::create_node(path_node, side_node, parent_height),
                 Side::Right => Node::create_node(side_node, path_node, parent_height),
@@ -131,7 +136,7 @@ impl Node {
         Self::Placeholder
     }
 
-    pub fn common_path_length(&self, other: &Node) -> u64 {
+    pub fn common_path_length(&self, other: &Node<N>) -> u64 {
         debug_assert!(self.is_leaf());
         debug_assert!(other.is_leaf());
 
@@ -165,7 +170,7 @@ impl Node {
         &Self::Placeholder == self
     }
 
-    pub fn hash(&self) -> &Bytes32 {
+    pub fn hash(&self) -> &Bytes<N> {
         match self {
             Node::Node { hash, .. } => hash,
             Node::Placeholder => zero_sum(),
@@ -179,14 +184,14 @@ impl Node {
         }
     }
 
-    fn bytes_lo(&self) -> &Bytes32 {
+    fn bytes_lo(&self) -> &Bytes<N> {
         match self {
             Node::Node { bytes_lo, .. } => bytes_lo,
             Node::Placeholder => zero_sum(),
         }
     }
 
-    fn bytes_hi(&self) -> &Bytes32 {
+    fn bytes_hi(&self) -> &Bytes<N> {
         match self {
             Node::Node { bytes_hi, .. } => bytes_hi,
             Node::Placeholder => zero_sum(),
@@ -204,7 +209,7 @@ impl Node {
     /// In `debug`, this method will panic if the node is not a leaf node to
     /// indicate to the developer that there is a potential problem in the
     /// tree's implementation.  
-    pub(super) fn leaf_key(&self) -> &Bytes32 {
+    pub(super) fn leaf_key(&self) -> &Bytes<N> {
         debug_assert!(self.is_leaf());
         self.bytes_lo()
     }
@@ -220,7 +225,7 @@ impl Node {
     /// In `debug`, this method will panic if the node is not a leaf node to
     /// indicate to the developer that there is a potential problem in the
     /// tree's implementation.
-    pub(super) fn leaf_data(&self) -> &Bytes32 {
+    pub(super) fn leaf_data(&self) -> &Bytes<N> {
         debug_assert!(self.is_leaf());
         self.bytes_hi()
     }
@@ -236,7 +241,7 @@ impl Node {
     /// In `debug`, this method will panic if the node is not an internal node
     /// to indicate to the developer that there is a potential problem in the
     /// tree's implementation.
-    pub(super) fn left_child_key(&self) -> &Bytes32 {
+    pub(super) fn left_child_key(&self) -> &Bytes<N> {
         debug_assert!(self.is_node());
         self.bytes_lo()
     }
@@ -252,20 +257,20 @@ impl Node {
     /// In `debug`, this method will panic if the node is not an internal node
     /// to indicate to the developer that there is a potential problem in the
     /// tree's implementation.
-    pub(super) fn right_child_key(&self) -> &Bytes32 {
+    pub(super) fn right_child_key(&self) -> &Bytes<N> {
         debug_assert!(self.is_node());
         self.bytes_hi()
     }
 }
 
-impl AsRef<Node> for Node {
-    fn as_ref(&self) -> &Node {
+impl<const N: usize> AsRef<Node<N>> for Node<N> {
+    fn as_ref(&self) -> &Node<N> {
         self
     }
 }
 
-impl NodeTrait for Node {
-    type Key = Bytes32;
+impl<const N: usize> NodeTrait for Node<N> {
+    type Key = Bytes<N>;
 
     fn height(&self) -> u32 {
         Node::height(self)
@@ -289,8 +294,8 @@ impl NodeTrait for Node {
     }
 }
 
-impl From<&Node> for Primitive {
-    fn from(node: &Node) -> Self {
+impl<const N: usize> From<&Node<N>> for Primitive<N> {
+    fn from(node: &Node<N>) -> Self {
         (
             node.height(),
             node.prefix() as u8,
@@ -300,10 +305,10 @@ impl From<&Node> for Primitive {
     }
 }
 
-impl TryFrom<Primitive> for Node {
+impl<const N: usize> TryFrom<Primitive<N>> for Node<N> {
     type Error = DeserializeError;
 
-    fn try_from(primitive: Primitive) -> Result<Self, Self::Error> {
+    fn try_from(primitive: Primitive<N>) -> Result<Self, Self::Error> {
         let height = primitive.height();
         let prefix = primitive.prefix()?;
         let bytes_lo = *primitive.bytes_lo();
@@ -313,7 +318,7 @@ impl TryFrom<Primitive> for Node {
     }
 }
 
-impl fmt::Debug for Node {
+impl<const N: usize> fmt::Debug for Node<N> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         if self.is_node() {
             f.debug_struct("Node (Internal)")
@@ -333,13 +338,15 @@ impl fmt::Debug for Node {
     }
 }
 
-pub(super) struct StorageNode<'storage, TableType, StorageType> {
+pub(super) struct StorageNode<'storage, const KEY_SIZE: usize, TableType, StorageType> {
     storage: &'storage StorageType,
-    node: Node,
+    node: Node<KEY_SIZE>,
     phantom_table: PhantomData<TableType>,
 }
 
-impl<TableType, StorageType> Clone for StorageNode<'_, TableType, StorageType> {
+impl<const KEY_SIZE: usize, TableType, StorageType> Clone
+    for StorageNode<'_, KEY_SIZE, TableType, StorageType>
+{
     fn clone(&self) -> Self {
         Self {
             storage: self.storage,
@@ -349,8 +356,10 @@ impl<TableType, StorageType> Clone for StorageNode<'_, TableType, StorageType> {
     }
 }
 
-impl<'s, TableType, StorageType> StorageNode<'s, TableType, StorageType> {
-    pub fn new(storage: &'s StorageType, node: Node) -> Self {
+impl<'s, const KEY_SIZE: usize, TableType, StorageType>
+    StorageNode<'s, KEY_SIZE, TableType, StorageType>
+{
+    pub fn new(storage: &'s StorageType, node: Node<KEY_SIZE>) -> Self {
         Self {
             node,
             storage,
@@ -359,18 +368,22 @@ impl<'s, TableType, StorageType> StorageNode<'s, TableType, StorageType> {
     }
 }
 
-impl<TableType, StorageType> StorageNode<'_, TableType, StorageType> {
-    pub fn hash(&self) -> &Bytes32 {
+impl<const KEY_SIZE: usize, TableType, StorageType>
+    StorageNode<'_, KEY_SIZE, TableType, StorageType>
+{
+    pub fn hash(&self) -> &Bytes<KEY_SIZE> {
         self.node.hash()
     }
 
-    pub fn into_node(self) -> Node {
+    pub fn into_node(self) -> Node<KEY_SIZE> {
         self.node
     }
 }
 
-impl<TableType, StorageType> NodeTrait for StorageNode<'_, TableType, StorageType> {
-    type Key = Bytes32;
+impl<const KEY_SIZE: usize, TableType, StorageType> NodeTrait
+    for StorageNode<'_, KEY_SIZE, TableType, StorageType>
+{
+    type Key = Bytes<KEY_SIZE>;
 
     fn height(&self) -> u32 {
         self.node.height()
@@ -402,10 +415,15 @@ pub enum StorageNodeError<StorageError> {
     DeserializeError(DeserializeError),
 }
 
-impl<TableType, StorageType> ParentNodeTrait for StorageNode<'_, TableType, StorageType>
+impl<const KEY_SIZE: usize, TableType, StorageType> ParentNodeTrait
+    for StorageNode<'_, KEY_SIZE, TableType, StorageType>
 where
     StorageType: StorageInspect<TableType>,
-    TableType: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    TableType: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
 {
     type Error = StorageNodeError<StorageType::Error>;
 
@@ -450,10 +468,15 @@ where
     }
 }
 
-impl<TableType, StorageType> fmt::Debug for StorageNode<'_, TableType, StorageType>
+impl<const KEY_SIZE: usize, TableType, StorageType> fmt::Debug
+    for StorageNode<'_, KEY_SIZE, TableType, StorageType>
 where
     StorageType: StorageInspect<TableType>,
-    TableType: Mappable<Key = Bytes32, Value = Primitive, OwnedValue = Primitive>,
+    TableType: Mappable<
+        Key = Bytes<KEY_SIZE>,
+        Value = Primitive<KEY_SIZE>,
+        OwnedValue = Primitive<KEY_SIZE>,
+    >,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         if self.is_node() {
@@ -531,7 +554,7 @@ mod test_node {
 
     #[test]
     fn test_create_placeholder_returns_a_placeholder_node() {
-        let node = Node::create_placeholder();
+        let node: Node<32> = Node::create_placeholder();
         assert_eq!(node.is_placeholder(), true);
         assert_eq!(node.hash(), zero_sum());
     }
@@ -540,7 +563,7 @@ mod test_node {
     fn test_create_leaf_from_primitive_returns_a_valid_leaf() {
         let primitive = (0, Prefix::Leaf as u8, [0xff; 32], [0xff; 32]);
 
-        let node: Node = primitive.try_into().unwrap();
+        let node: Node<32> = primitive.try_into().unwrap();
         assert_eq!(node.is_leaf(), true);
         assert_eq!(node.is_node(), false);
         assert_eq!(node.height(), 0);
@@ -553,7 +576,7 @@ mod test_node {
     fn test_create_node_from_primitive_returns_a_valid_node() {
         let primitive = (255, Prefix::Node as u8, [0xff; 32], [0xff; 32]);
 
-        let node: Node = primitive.try_into().unwrap();
+        let node: Node<32> = primitive.try_into().unwrap();
         assert_eq!(node.is_leaf(), false);
         assert_eq!(node.is_node(), true);
         assert_eq!(node.height(), 255);
@@ -566,7 +589,7 @@ mod test_node {
     fn test_create_from_primitive_returns_deserialize_error_if_invalid_prefix() {
         let primitive = (0xff, 0xff, [0xff; 32], [0xff; 32]);
 
-        // Should return Error; prefix 0xff is does not represent a node or leaf
+        // Should return Error; prefix 0xff does not represent a node or leaf
         let err = Node::try_from(primitive)
             .expect_err("Expected try_from() to be Error; got OK");
         assert!(matches!(
@@ -675,7 +698,7 @@ mod test_storage_node {
     impl Mappable for TestTable {
         type Key = Self::OwnedKey;
         type OwnedKey = Bytes32;
-        type OwnedValue = Primitive;
+        type OwnedValue = Primitive<32>;
         type Value = Self::OwnedValue;
     }
 
diff --git a/fuel-merkle/src/sparse/primitive.rs b/fuel-merkle/src/sparse/primitive.rs
index a46a7c20cd..866f67a3a1 100644
--- a/fuel-merkle/src/sparse/primitive.rs
+++ b/fuel-merkle/src/sparse/primitive.rs
@@ -1,10 +1,10 @@
 use crate::common::{
-    Bytes32,
+    Bytes,
     Prefix,
     PrefixError,
 };
 
-/// **Leaf buffer:**
+/// **Leaf buffer (32 byte KEY_SIZE):**
 ///
 /// | Allocation | Data                       |
 /// |------------|----------------------------|
@@ -13,7 +13,7 @@ use crate::common::{
 /// | `05 - 37`  | hash(Key) (32 bytes)       |
 /// | `37 - 69`  | hash(Data) (32 bytes)      |
 ///
-/// **Node buffer:**
+/// **Node buffer (32 byte KEY_SIZE):**
 ///
 /// | Allocation | Data                       |
 /// |------------|----------------------------|
@@ -21,16 +21,16 @@ use crate::common::{
 /// | `04 - 05`  | Prefix (1 byte, `0x01`)    |
 /// | `05 - 37`  | Left child key (32 bytes)  |
 /// | `37 - 69`  | Right child key (32 bytes) |
-pub type Primitive = (u32, u8, Bytes32, Bytes32);
+pub type Primitive<const KEY_SIZE: usize> = (u32, u8, Bytes<KEY_SIZE>, Bytes<KEY_SIZE>);
 
-pub trait PrimitiveView {
+pub trait PrimitiveView<const KEY_SIZE: usize> {
     fn height(&self) -> u32;
     fn prefix(&self) -> Result<Prefix, PrefixError>;
-    fn bytes_lo(&self) -> &Bytes32;
-    fn bytes_hi(&self) -> &Bytes32;
+    fn bytes_lo(&self) -> &Bytes<KEY_SIZE>;
+    fn bytes_hi(&self) -> &Bytes<KEY_SIZE>;
 }
 
-impl PrimitiveView for Primitive {
+impl<const KEY_SIZE: usize> PrimitiveView<KEY_SIZE> for Primitive<KEY_SIZE> {
     fn height(&self) -> u32 {
         self.0
     }
@@ -39,11 +39,11 @@ impl PrimitiveView for Primitive {
         Prefix::try_from(self.1)
     }
 
-    fn bytes_lo(&self) -> &Bytes32 {
+    fn bytes_lo(&self) -> &Bytes<KEY_SIZE> {
         &self.2
     }
 
-    fn bytes_hi(&self) -> &Bytes32 {
+    fn bytes_hi(&self) -> &Bytes<KEY_SIZE> {
         &self.3
     }
 }
diff --git a/fuel-merkle/src/sparse/proof.rs b/fuel-merkle/src/sparse/proof.rs
index 6d9337c1db..b070951cdc 100644
--- a/fuel-merkle/src/sparse/proof.rs
+++ b/fuel-merkle/src/sparse/proof.rs
@@ -4,8 +4,6 @@ use crate::{
             Path,
             Side,
         },
-        sum,
-        Bytes32,
         ProofSet,
     },
     sparse::{
@@ -18,6 +16,10 @@ use crate::{
     },
 };
 
+use crate::{
+    common::Bytes,
+    sparse::hash::sum_truncated,
+};
 use alloc::vec::Vec;
 use core::{
     fmt,
@@ -25,13 +27,13 @@ use core::{
 };
 
 #[derive(Debug, Clone, Eq, PartialEq)]
-pub enum Proof {
-    Inclusion(InclusionProof),
-    Exclusion(ExclusionProof),
+pub enum Proof<const KEY_SIZE: usize> {
+    Inclusion(InclusionProof<KEY_SIZE>),
+    Exclusion(ExclusionProof<KEY_SIZE>),
 }
 
-impl Proof {
-    pub fn proof_set(&self) -> &ProofSet {
+impl<const KEY_SIZE: usize> Proof<KEY_SIZE> {
+    pub fn proof_set(&self) -> &ProofSet<KEY_SIZE> {
         match self {
             Proof::Inclusion(proof) => &proof.proof_set,
             Proof::Exclusion(proof) => &proof.proof_set,
@@ -51,19 +53,24 @@ impl Proof {
 }
 
 #[derive(Clone, Eq, PartialEq)]
-pub struct InclusionProof {
-    pub proof_set: ProofSet,
+pub struct InclusionProof<const KEY_SIZE: usize> {
+    pub proof_set: ProofSet<KEY_SIZE>,
 }
 
-impl InclusionProof {
-    pub fn verify(&self, root: &Bytes32, key: &MerkleTreeKey, value: &[u8]) -> bool {
+impl<const KEY_SIZE: usize> InclusionProof<KEY_SIZE> {
+    pub fn verify(
+        &self,
+        root: &Bytes<KEY_SIZE>,
+        key: &MerkleTreeKey<KEY_SIZE>,
+        value: &[u8],
+    ) -> bool {
         let Self { proof_set } = self;
 
         if proof_set.len() > 256usize {
             return false;
         }
 
-        let mut current = calculate_leaf_hash(key, &sum(value));
+        let mut current = calculate_leaf_hash(key, &sum_truncated(value));
         for (i, side_hash) in proof_set.iter().enumerate() {
             #[allow(clippy::arithmetic_side_effects)] // Cannot underflow
             let index =
@@ -77,7 +84,7 @@ impl InclusionProof {
     }
 }
 
-impl Debug for InclusionProof {
+impl<const KEY_SIZE: usize> Debug for InclusionProof<KEY_SIZE> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let proof_set = self.proof_set.iter().map(hex::encode).collect::<Vec<_>>();
         f.debug_struct("InclusionProof")
@@ -87,20 +94,20 @@ impl Debug for InclusionProof {
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
-pub enum ExclusionLeaf {
-    Leaf(ExclusionLeafData),
+pub enum ExclusionLeaf<const KEY_SIZE: usize> {
+    Leaf(ExclusionLeafData<KEY_SIZE>),
     Placeholder,
 }
 
 #[derive(Clone, Eq, PartialEq)]
-pub struct ExclusionLeafData {
+pub struct ExclusionLeafData<const KEY_SIZE: usize> {
     /// The leaf key.
-    pub leaf_key: Bytes32,
+    pub leaf_key: Bytes<KEY_SIZE>,
     /// Hash of the value of the leaf.
-    pub leaf_value: Bytes32,
+    pub leaf_value: Bytes<KEY_SIZE>,
 }
 
-impl Debug for ExclusionLeafData {
+impl<const KEY_SIZE: usize> Debug for ExclusionLeafData<KEY_SIZE> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("ExclusionLeafData")
             .field("Leaf key", &hex::encode(self.leaf_key))
@@ -109,8 +116,8 @@ impl Debug for ExclusionLeafData {
     }
 }
 
-impl ExclusionLeaf {
-    fn hash(&self) -> Bytes32 {
+impl<const KEY_SIZE: usize> ExclusionLeaf<KEY_SIZE> {
+    fn hash(&self) -> Bytes<KEY_SIZE> {
         match self {
             ExclusionLeaf::Leaf(data) => {
                 calculate_leaf_hash(&data.leaf_key, &data.leaf_value)
@@ -121,13 +128,13 @@ impl ExclusionLeaf {
 }
 
 #[derive(Clone, Eq, PartialEq)]
-pub struct ExclusionProof {
-    pub proof_set: ProofSet,
-    pub leaf: ExclusionLeaf,
+pub struct ExclusionProof<const KEY_SIZE: usize> {
+    pub proof_set: ProofSet<KEY_SIZE>,
+    pub leaf: ExclusionLeaf<KEY_SIZE>,
 }
 
-impl ExclusionProof {
-    pub fn verify(&self, root: &Bytes32, key: &MerkleTreeKey) -> bool {
+impl<const KEY_SIZE: usize> ExclusionProof<KEY_SIZE> {
+    pub fn verify(&self, root: &Bytes<KEY_SIZE>, key: &MerkleTreeKey<KEY_SIZE>) -> bool {
         let Self { proof_set, leaf } = self;
 
         if let ExclusionLeaf::Leaf(data) = leaf {
@@ -154,7 +161,7 @@ impl ExclusionProof {
     }
 }
 
-impl Debug for ExclusionProof {
+impl<const KEY_SIZE: usize> Debug for ExclusionProof<KEY_SIZE> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let proof_set = self.proof_set.iter().map(hex::encode).collect::<Vec<_>>();
         f.debug_struct("ExclusionProof")
@@ -186,7 +193,7 @@ mod test {
     impl Mappable for TestTable {
         type Key = Self::OwnedKey;
         type OwnedKey = Bytes32;
-        type OwnedValue = Primitive;
+        type OwnedValue = Primitive<32>;
         type Value = Self::OwnedValue;
     }
 
@@ -637,7 +644,7 @@ mod test_random {
     impl Mappable for TestTable {
         type Key = Self::OwnedKey;
         type OwnedKey = Bytes32;
-        type OwnedValue = Primitive;
+        type OwnedValue = Primitive<32>;
         type Value = Self::OwnedValue;
     }
 
@@ -764,7 +771,7 @@ mod test_random {
         let root = tree.root();
 
         // Given
-        let key: MerkleTreeKey = random_bytes32(&mut rng).into();
+        let key: MerkleTreeKey<32> = random_bytes32(&mut rng).into();
         let proof = tree.generate_proof(&key).unwrap();
 
         // When
diff --git a/fuel-merkle/src/tests/sparse.rs b/fuel-merkle/src/tests/sparse.rs
index fdf56d66fc..ccc01cd81c 100644
--- a/fuel-merkle/src/tests/sparse.rs
+++ b/fuel-merkle/src/tests/sparse.rs
@@ -44,7 +44,7 @@ struct TestTable;
 impl Mappable for TestTable {
     type Key = Self::OwnedKey;
     type OwnedKey = Bytes32;
-    type OwnedValue = Primitive;
+    type OwnedValue = Primitive<32>;
     type Value = Self::OwnedValue;
 }
 
@@ -105,7 +105,7 @@ prop_compose! {
 }
 
 prop_compose! {
-    fn random_tree(min: usize, max: usize)(kv in key_values(min, max)) -> (Vec<(Key, Value)>, MerkleTree<TestTable, StorageMap<TestTable>>) {
+    fn random_tree(min: usize, max: usize)(kv in key_values(min, max)) -> (Vec<(Key, Value)>, MerkleTree<32, TestTable, StorageMap<TestTable>>) {
         let storage = StorageMap::<TestTable>::new();
         let iter = kv.clone().into_iter().map(|(key, value)| (MerkleTreeKey::new(key), value));
         let tree = MerkleTree::from_set(storage, iter).expect("Unable to create Merkle tree");

From 66fe940674b6d826632bd6cf2ec989e5e2fe886b Mon Sep 17 00:00:00 2001
From: Brandon Vrooman <brandon.vrooman@gmail.com>
Date: Fri, 31 May 2024 19:12:27 -0400
Subject: [PATCH 2/4] Update CHANGELOG.md

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f053c407ea..986e1ab722 100644
--- a/CHANGELOG.md
+++ b/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.

From efca5d3b41af78061673d54c0b7230dae8d85c62 Mon Sep 17 00:00:00 2001
From: Brandon Vrooman <brandon.vrooman@gmail.com>
Date: Fri, 31 May 2024 19:15:32 -0400
Subject: [PATCH 3/4] Default exports

---
 fuel-merkle/src/sparse.rs | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/fuel-merkle/src/sparse.rs b/fuel-merkle/src/sparse.rs
index 27563937dd..728e60b89b 100644
--- a/fuel-merkle/src/sparse.rs
+++ b/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;
 
+// 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()
 }

From 9903bed1c635ee0e41cb42dd7e3e05734161b015 Mon Sep 17 00:00:00 2001
From: Brandon Vrooman <brandon.vrooman@gmail.com>
Date: Fri, 31 May 2024 19:18:20 -0400
Subject: [PATCH 4/4] Update in_memory.rs

---
 fuel-merkle/src/sparse/in_memory.rs | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/fuel-merkle/src/sparse/in_memory.rs b/fuel-merkle/src/sparse/in_memory.rs
index f9a84a4bb8..3b0c97c2b0 100644
--- a/fuel-merkle/src/sparse/in_memory.rs
+++ b/fuel-merkle/src/sparse/in_memory.rs
@@ -7,8 +7,8 @@ use crate::{
     sparse::{
         self,
         merkle_tree::MerkleTreeKey,
+        primitive::Primitive,
         proof::Proof,
-        Primitive,
     },
     storage::{
         Mappable,
@@ -38,7 +38,7 @@ impl Mappable for NodesTable {
 }
 
 type Storage = StorageMap<NodesTable>;
-type SparseMerkleTree = sparse::MerkleTree<DEFAULT_KEY_SIZE, NodesTable, Storage>;
+type SparseMerkleTree = sparse::MerkleTree<NodesTable, Storage>;
 
 #[derive(Debug)]
 pub struct MerkleTree {
@@ -118,11 +118,8 @@ impl MerkleTree {
             }
         }
 
-        let tree = sparse::MerkleTree::<DEFAULT_KEY_SIZE, NodesTable, _>::from_set(
-            EmptyStorage,
-            set,
-        )
-        .expect("`Storage` can't return error");
+        let tree = sparse::MerkleTree::<NodesTable, _>::from_set(EmptyStorage, set)
+            .expect("`Storage` can't return error");
         tree.root()
     }
 
@@ -183,11 +180,9 @@ impl MerkleTree {
             }
         }
 
-        let tree = sparse::MerkleTree::<DEFAULT_KEY_SIZE, NodesTable, _>::from_set(
-            VectorStorage::default(),
-            set,
-        )
-        .expect("`Storage` can't return error");
+        let tree =
+            sparse::MerkleTree::<NodesTable, _>::from_set(VectorStorage::default(), set)
+                .expect("`Storage` can't return error");
         let root = tree.root();
         let nodes = tree.into_storage().storage;