feat: compatible Merkle tree

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "merkletreers"
-version = "0.4.1"
+version = "1.0.0"
 edition = "2021"
 description = "🌳 The simple and easy implementation of Merkle Tree"
 authors = ["Lucas Oliveira <[email protected]>"]

src/lib.rs

@@ -1,2 +1,15 @@
-pub mod merkletree;
+use node::Node;
+
+pub mod merkle_proof;
+pub mod merkle_proof_check;
+pub mod merkle_root;
+pub mod node;
+pub mod tree;
+pub mod merkle_proof_mixed;
 pub mod utils;
+
+
+pub type Proof = Vec<Node>;
+pub type Hash = [u8; 32];
+pub type Leaf = [u8; 32];
+pub type Root = [u8; 32];

src/merkle_proof.rs

@@ -0,0 +1,84 @@
+use crate::merkle_proof_mixed::merkle_proof_mixed_tree;
+use crate::merkle_root::merkle_root;
+use crate::node::Node;
+use crate::utils::is_power_of_two;
+use crate::{Leaf, Proof};
+
+pub fn merkle_proof(leaves: &[Leaf], leaf: Leaf) -> Proof {
+    let mut proof: Proof = Vec::new();
+
+    if is_power_of_two(leaves.len() as u32) == false {
+        return merkle_proof_mixed_tree(leaves, leaf);
+    }
+
+    let mut current_leaves = leaves;
+    while current_leaves.len() > 1 {
+        let index = match current_leaves.iter().position(|x| x == &leaf) {
+            Some(i) => i,
+            None => panic!("Leaf does not exist in the tree"),
+        };
+        if current_leaves.len() == 2 {
+            if index == 1 {
+                // proof.push(Node(data = leaves[0], side = Side.LEFT))
+                let left: [u8; 32] = current_leaves[0]
+                    .as_slice()
+                    .try_into()
+                    .expect("Failed to convert LEFT to array");
+                proof.push(Node {
+                    data: left,
+                    side: 0.into(),
+                });
+                break;
+            } else {
+                // proof.append(Node(data = leaves[1], side = Side.RIGHT))
+                let right: [u8; 32] = current_leaves[1]
+                    .as_slice()
+                    .try_into()
+                    .expect("Failed to convert RIGHT to array");
+                proof.push(Node {
+                    data: right,
+                    side: 1.into(),
+                });
+                break;
+            }
+        }
+
+        let half_size = current_leaves.len() / 2;
+
+        // divide a lista em 2, left e right
+        let (left, right) = current_leaves.split_at(half_size);
+
+        // se o index estiver em left
+        if index < half_size {
+            // faça a merkle root de right
+            // TODO: make root must be a async (using other thread)
+            let right_root = merkle_root(right);
+
+            // faça o node passando right e 1 (direita)
+            // adicione o node na lista de prova
+            proof.push(Node {
+                data: right_root,
+                side: 1.into(),
+            });
+
+            current_leaves = left;
+        } else {
+            // se o index estiver em right
+            // faça a merkle root de left
+            // TODO: make root must be a async (using other thread)
+            let left_root = merkle_root(left);
+            // faça o node passando left e 0 (esquerda)
+
+            // adicione o node na lista de prova
+            proof.push(Node {
+                data: left_root,
+                side: 0.into(),
+            });
+
+            current_leaves = right;
+        }
+    }
+
+    proof.reverse();
+    proof
+}

src/merkle_proof_check.rs

@@ -0,0 +1,20 @@
+use crate::utils::hash_function;
+use crate::Proof;
+
+pub fn merkle_proof_check(proof: Proof, leaf: [u8; 32]) -> [u8; 32] {
+    let mut current_hash = leaf;
+
+    for node in proof {
+        let mut buffer = [0u8; 32];
+
+        if node.side == 1.into() {
+            hash_function(&current_hash, &node.data, &mut buffer);
+        } else {
+            hash_function(&node.data, &current_hash, &mut buffer);
+        }
+
+        current_hash = buffer;
+    }
+
+    current_hash
+}

src/merkle_proof_mixed.rs

@@ -0,0 +1,50 @@
+use crate::node::Node;
+use crate::utils::hash_function;
+use crate::{Leaf, Proof};
+
+pub fn merkle_proof_mixed_tree(leaves: &[Leaf], leaf: Leaf) -> Proof {
+    let mut proof: Proof = Vec::new();
+    let mut leaf_index = leaves.iter().position(|x| x == &leaf).unwrap_or_else(|| {
+        panic!("Leaf does not exist in the tree");
+    });
+
+    let mut current_leaves = leaves.to_vec();
+    while current_leaves.len() > 1 {
+        let next_leaves = up_layer(&current_leaves);
+
+        if leaf_index % 2 == 0 {
+            if leaf_index + 1 < current_leaves.len() {
+                proof.push(Node {
+                    data: current_leaves[leaf_index + 1],
+                    side: 1.into(),
+                });
+            }
+        } else {
+            proof.push(Node {
+                data: current_leaves[leaf_index - 1],
+                side: 0.into(),
+            });
+        }
+
+        current_leaves = next_leaves;
+        leaf_index /= 2;
+    }
+
+    proof
+}
+
+fn up_layer(leaves: &[Leaf]) -> Vec<Leaf> {
+    let mut new_layer: Vec<[u8; 32]> = vec![];
+
+    for pair in leaves.chunks(2) {
+        if pair.len() == 1 {
+            new_layer.push(*pair.first().unwrap());
+        } else {
+            let mut buffer = [0u8; 32];
+            hash_function(pair.first().unwrap(), &pair[1], &mut buffer);
+            new_layer.push(buffer);
+        }
+    }
+
+    new_layer
+}

src/merkle_root.rs

@@ -0,0 +1,24 @@
+use crate::utils::hash_function;
+use crate::{Hash, Leaf, Root};
+
+pub fn merkle_root(leaves: &[Leaf]) -> Root {
+    let mut node: Hash = [0u8; 32];
+
+    let mut tmp = leaves.to_vec();
+
+    while tmp.len() > 1 {
+        let mut next_level: Vec<Leaf> = vec![];
+
+        for leaf_pair in tmp.chunks(2) {
+            match leaf_pair {
+                [left, right] => hash_function(left, right, &mut node),
+                [left] => node.copy_from_slice(left),
+                _ => unreachable!(),
+            };
+            next_level.push(node);
+        }
+
+        tmp = next_level;
+    }
+    node
+}

src/node.rs

@@ -0,0 +1,23 @@
+use crate::Hash;
+
+#[derive(PartialEq, Debug, Clone)]
+pub enum Side {
+    LEFT = 0,
+    RIGHT = 1,
+}
+impl From<u8> for Side {
+    fn from(num: u8) -> Self {
+        match num {
+            0 => Side::LEFT,
+            1 => Side::RIGHT,
+            _ => panic!("Invalid value for Side enum, must be either `0` or `1`"),
+        }
+    }
+}
+
+/// # 🍃 Leaf of Tree
+#[derive(PartialEq, Debug, Clone)]
+pub struct Node {
+    pub data: Hash,
+    pub side: Side,
+}

src/tree.rs

@@ -0,0 +1,27 @@
+use crate::merkle_proof::merkle_proof;
+use crate::merkle_root::merkle_root;
+use crate::node::Node;
+use crate::{Leaf, Root};
+
+/// # 🌳 Merkle Tree
+/// - You can pass raw data
+/// - They will hashed by `keccak-256`
+pub struct MerkleTree {
+    pub leaves: Vec<Leaf>,
+    pub root: Root,
+}
+
+impl MerkleTree {
+    pub fn new(leaves: Vec<Leaf>) -> Self {
+        MerkleTree {
+            leaves: leaves.clone(),
+            root: merkle_root(&leaves),
+        }
+    }
+}
+
+impl MerkleTree {
+    pub fn make_proof(self, leaf: Leaf) -> Vec<Node> {
+        merkle_proof(&self.leaves, leaf)
+    }
+}