3 files changed, 58 insertions, 14 deletions

diff --git a/src/merkle/merkletree.rs b/src/merkle/merkletree.rs
index fc2879a..4a6a214 100644
--- a/src/merkle/merkletree.rs
+++ b/src/merkle/merkletree.rs
@@ -34,14 +34,21 @@ impl MerkleTree {
     ///
     /// If the number of leaf nodes is odd, the last node is duplicated to ensure all internal
     /// nodes have exactly two children.
-    pub fn new<T: ToString>(hasher: &dyn Hasher, data: Vec<T>) -> Self {
+    pub fn new<I, T>(hasher: &dyn Hasher, data: I) -> Self
+    where
+        I: IntoIterator<Item = T>,
+        T: AsRef<[u8]>,
+    {
+        let owned_data: Vec<T> = data.into_iter().collect();
+        let data_slices: Vec<&[u8]> = owned_data.iter().map(|item| item.as_ref()).collect();
+
         assert!(
-            !data.is_empty(),
+            !data_slices.is_empty(),
             "Merkle Tree requires at least one element"
         );
 
-        let mut leaves: Vec<Node> = data
-            .into_iter()
+        let mut leaves: Vec<Node> = data_slices
+            .iter()
             .map(|x| Node::new_leaf(hasher, x))
             .collect();
 
@@ -58,9 +65,15 @@ impl MerkleTree {
         let mut height = 0;
 
         while nodes.len() > 1 {
+            if nodes.len() % 2 != 0 {
+                // duplicate last node to make the count even
+                nodes.push(nodes[nodes.len() - 1].clone());
+            }
+
             let mut next_level = Vec::new();
             for pair in nodes.chunks(2) {
                 let (left, right) = (pair[0].clone(), pair[1].clone());
+
                 next_level.push(Node::new_internal(hasher, left, right));
             }
             nodes = next_level;
diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs
index 4a6e46f..d00e350 100644
--- a/src/merkle/mod.rs
+++ b/src/merkle/mod.rs
@@ -13,30 +13,30 @@ mod tests {
 
     #[test]
     fn test_merkle_tree_with_default_hasher() {
-        let data = vec!["a", "b", "c", "d"];
+        let data = &["hello".as_bytes(), "world".as_bytes()];
         let tree = merkletree::MerkleTree::new(&DefaultHasher, data);
 
-        assert_eq!(tree.height(), 3);
+        assert_eq!(tree.height(), 2);
         assert_eq!(tree.root().hash(), "0xc0ff3");
     }
 
     #[test]
     #[cfg(feature = "sha256")]
     fn test_merkle_tree_hashing() {
-        let data = vec!["a", "b", "c", "d"];
+        let data = &["hello".as_bytes(), "world".as_bytes()];
         let tree = merkletree::MerkleTree::new(&SHA256Hasher, data);
 
-        assert_eq!(tree.height(), 3);
+        assert_eq!(tree.height(), 2);
         assert_eq!(
             tree.root().hash(),
-            "58c89d709329eb37285837b042ab6ff72c7c8f74de0446b091b6a0131c102cfd"
+            "15e178b71fae8849ee562c9cc0d7ea322fba6cd495411329d47234479167cc8b"
         );
     }
 
     #[test]
     #[cfg(feature = "sha256")]
     fn test_merkle_tree_single_leaf() {
-        let data = vec!["hello"];
+        let data = &["hello".as_bytes()];
         let tree = merkletree::MerkleTree::new(&SHA256Hasher, data);
 
         assert_eq!(tree.height(), 2);
@@ -46,4 +46,24 @@ mod tests {
             "286d189fda11bf4e906b6973a173009f47ede16532f1bae726223f8ee155d73b"
         );
     }
+
+    #[test]
+    #[cfg(feature = "sha256")]
+    fn test_merkle_tree_with_10_elements() {
+        let inputs = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"];
+        let data: Vec<&[u8]> = inputs.iter().map(|s| s.as_bytes()).collect();
+
+        let tree = merkletree::MerkleTree::new(&SHA256Hasher, &data);
+
+        assert_eq!(tree.height(), 5); // 10 elements padded to 16 → log2(16) + 1 = 5
+
+        // You can print the root hash if you're unsure what it should be:
+        println!("Merkle root hash: {}", tree.root().hash());
+
+        // If you know the expected hash, use:
+        assert_eq!(
+            tree.root().hash(),
+            "9da1ff0dfa79217bdbea9ec96407b1e693646cc493f64059fa27182a37cadf94"
+        );
+    }
 }
diff --git a/src/merkle/node.rs b/src/merkle/node.rs
index 0e5e904..021b1c3 100644
--- a/src/merkle/node.rs
+++ b/src/merkle/node.rs
@@ -36,6 +36,8 @@ pub struct Node {
     hash: String,
     /// Type of the node: leaf or internal.
     kind: NodeType,
+    /// Data in bytes.
+    data: Vec<u8>,
 }
 
 impl Node {
@@ -45,10 +47,11 @@ impl Node {
     ///
     /// * `hasher` - A reference to a hashing strategy.
     /// * `data` - The data to be hashed and stored as a leaf.
-    pub fn new_leaf<T: ToString>(hasher: &dyn Hasher, data: T) -> Self {
-        let hash = hasher.hash(&data.to_string());
+    pub fn new_leaf(hasher: &dyn Hasher, data: &[u8]) -> Self {
+        let hash = hasher.hash(&data);
         Self {
             hash,
+            data: data.to_vec(),
             kind: NodeType::Leaf,
         }
     }
@@ -65,10 +68,13 @@ impl Node {
     ///
     /// The internal node hash is computed as the hash of the concatenated children's hashes.
     pub fn new_internal(hasher: &dyn Hasher, left: Node, right: Node) -> Self {
-        let combined = format!("{}{}", left.hash, right.hash);
-        let hash = hasher.hash(&combined);
+        let mut buffer = Vec::<u8>::new();
+        buffer.extend_from_slice(left.hash().as_bytes());
+        buffer.extend_from_slice(right.hash().as_bytes());
+        let hash = hasher.hash(&buffer);
         Self {
             hash,
+            data: buffer,
             kind: NodeType::Internal(Box::new(left), Box::new(right)),
         }
     }
@@ -78,6 +84,11 @@ impl Node {
         &self.hash
     }
 
+    /// Returns the data value in bytes format.
+    pub fn data(&self) -> &[u8] {
+        &self.data
+    }
+
     /// Returns a reference to the node's type (leaf or internal).
     pub fn kind(&self) -> &NodeType {
         &self.kind