From 44fe95cb538d7b1fdc03809a66e942290ca35641 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <olivmath@protonmail.com>
Date: Fri, 7 Jul 2023 17:35:19 -0300
Subject: [PATCH] add: any size leaves proof

---
 merkly/mtree.py          | 30 ++++++++++++++++++++++++++++++
 merkly/utils.py          | 17 ++++++++++-------
 test/test_merkle_tree.py | 17 +++++++++--------
 test/utils/test_math.py  |  6 +-----
 4 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/merkly/mtree.py b/merkly/mtree.py
index 22d6c37..8ef3220 100644
--- a/merkly/mtree.py
+++ b/merkly/mtree.py
@@ -8,6 +8,7 @@
 from merkly.node import Node, Side
 from merkly.utils import (
     hash_function_type_checking,
+    is_power_2,
     slice_in_pairs,
     keccak,
     half,
@@ -109,6 +110,9 @@ def make_proof(self, leafs: List[str], proof: List[Node], leaf: str) -> List[Nod
             msg = f"Leaf: {leaf} does not exist in the tree: {leafs}"
             raise ValueError(msg) from err
 
+        if is_power_2(len(leafs)) is False:
+            return self.mix_tree(leafs, [], index)
+
         if len(leafs) == 2:
             if index == 1:
                 proof.append(Node(data=leafs[0], side=Side.LEFT))
@@ -125,3 +129,29 @@ def make_proof(self, leafs: List[str], proof: List[Node], leaf: str) -> List[Nod
         else:
             proof.append(Node(data=self.make_root(left)[0], side=Side.LEFT))
             return self.make_proof(right, proof, leaf)
+
+    def mix_tree(
+        self, leaves: List[str], proof: List[Node], leaf_index: int
+    ) -> List[Node]:
+        if len(leaves) == 1:
+            return proof
+
+        if leaf_index % 2 == 0:
+            if leaf_index + 1 < len(leaves):
+                node = Node(data=leaves[leaf_index + 1], side=Side.RIGHT)
+                proof.append(node)
+        else:
+            node = Node(data=leaves[leaf_index - 1], side=Side.LEFT)
+            proof.append(node)
+
+        return self.mix_tree(self.up_layer(leaves), proof, leaf_index // 2)
+
+    def up_layer(self, leaves: List[str]) -> List[str]:
+        new_layer = []
+        for pair in slice_in_pairs(leaves):
+            if len(pair) == 1:
+                new_layer.append(pair[0])
+            else:
+                data = self.hash_function(pair[0], pair[1])
+                new_layer.append(data)
+        return new_layer
diff --git a/merkly/utils.py b/merkly/utils.py
index 824699e..359624a 100644
--- a/merkly/utils.py
+++ b/merkly/utils.py
@@ -85,12 +85,15 @@ def slice_in_pairs(list_item: list):
 
 
 def hash_function_type_checking(hash_function: Callable[[str], str]) -> bool:
-    is_valid = (
-        isinstance(hash_function, types.FunctionType)
-        and callable(hash_function)
-        and isinstance(hash_function(str()), str)
-    )
-    if hash_function is not None and not is_valid:
+    a = isinstance(hash_function, types.FunctionType)
+    b = callable(hash_function)
+    try:
+        c = isinstance(hash_function(str(), str()), str)
+    except TypeError:
+        c = False
+
+    valid = a and b and c
+    if not valid:
         raise InvalidHashFunctionError()
 
 
@@ -113,4 +116,4 @@ def is_power_2(number: int) -> bool:
     if left and right:
         return True
     else:
-        raise PowerOfTwoError(number)
+        return False
\ No newline at end of file
diff --git a/test/test_merkle_tree.py b/test/test_merkle_tree.py
index a8a0d15..8a1f613 100644
--- a/test/test_merkle_tree.py
+++ b/test/test_merkle_tree.py
@@ -70,14 +70,13 @@ def test_proof_simple_odd_merkle():
     Instantiated a simple Merkle Tree
     """
     leafs = ["a", "b", "c", "d", "e"]
-    tree = MerkleTree(leafs)
+    tree = MerkleTree(leafs, lambda x, y: x + y)
     proof = [
-        Node(right="b5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510"),
-        Node(right="ed3a2f2a068b98ea5eb600912326df7b62037603d2633eba4ccc9a3845674b90"),
+        Node(data="abcd", side=Side.LEFT),
     ]
 
-    assert tree.proof("a") == proof
-    assert tree.verify(proof, "a") == True
+    assert tree.proof("e") == proof, "Proofs dont's match"
+    assert tree.verify(proof, "e"), "Proof dont's right"
 
 
 def test_proof_simple_merkle():
@@ -86,8 +85,7 @@ def test_proof_simple_merkle():
     """
     leafs = ["a", "b", "c", "d"]
     tree = MerkleTree(leafs)
-
-    assert tree.proof("a") == [
+    proof = [
         Node(
             side=Side.RIGHT,
             data="b5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510",
@@ -150,4 +148,7 @@ def invalid_hash_function_that_returns_an_integer_instead_of_a_string(data):
         return 123
 
     with raises(InvalidHashFunctionError):
-        MerkleTree(["a", "b", "c", "d"], invalid_hash_function_that_returns_an_integer_instead_of_a_string)
+        MerkleTree(
+            ["a", "b", "c", "d"],
+            invalid_hash_function_that_returns_an_integer_instead_of_a_string,
+        )
diff --git a/test/utils/test_math.py b/test/utils/test_math.py
index cacb25a..c5ef4c5 100644
--- a/test/utils/test_math.py
+++ b/test/utils/test_math.py
@@ -18,8 +18,4 @@
     ],
 )
 def test_of_is_power_2(number: int, ok: bool):
-    if ok:
-        assert ok == is_power_2(number)
-    else:
-        with raises(PowerOfTwoError):
-            is_power_2(number)
+    assert ok == is_power_2(number)