Test multiple node widths of merkle tree using poseidon

Brechtpd · Jul 10, 2019 · 7f3f1df · 7f3f1df
1 parent 26b0fb9
commit 7f3f1df
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 221 deletions.
diff --git a/ethsnarks/poseidon.py b/ethsnarks/poseidon.py
diff --git a/ethsnarks/shamirspoly.py b/ethsnarks/shamirspoly.py
@@ -40,3 +40,13 @@ def g(i, n):
         coefficient = g(i, n)
         total = total + (yi * coefficient)
     return total
+
+
+def inverse_lagrange(points, y):
+    x = 0
+    for i, (x_i, y_i) in enumerate(points):
+        for j, (x_j, y_j) in enumerate(points):
+            if j != i:
+                x_i = x_i * (y - y_j) / (y_i - y_j)
+        x += x_i 
+    return x
diff --git a/test/test_merkle.py b/test/test_merkle.py
@@ -47,30 +47,37 @@ def test_known1(self):
 
     def test_update(self):
         # Verify that items in the tree can be updated
-        for hasher in [DEFAULT_HASHER, MerkleHasher_Poseidon.factory()]:
-            tree = MerkleTree(2, hasher=hasher)
-            tree.append(FQ.random())
-            tree.append(FQ.random())
-            proof_0_before = tree.proof(0)
-            proof_1_before = tree.proof(1)
-            root_before = tree.root
-            self.assertTrue(proof_0_before.verify(tree.root))
-            self.assertTrue(proof_1_before.verify(tree.root))
-
-            leaf_0_after = FQ.random()
-            tree.update(0, leaf_0_after)
-            root_after_0 = tree.root
-            proof_0_after = tree.proof(0)
-            self.assertTrue(proof_0_after.verify(tree.root))
-            self.assertNotEqual(root_before, root_after_0)
-
-            leaf_1_after = FQ.random()
-            tree.update(1, leaf_1_after)
-            root_after_1 = tree.root
-            proof_1_after = tree.proof(1)
-            self.assertTrue(proof_1_after.verify(tree.root))
-            self.assertNotEqual(root_before, root_after_1)
-            self.assertNotEqual(root_after_0, root_after_1)
+        poseidon_factory = MerkleHasher_Poseidon.factory()
+        for hasher in [DEFAULT_HASHER, poseidon_factory]:
+            widths = [2]
+            if hasher == poseidon_factory:
+                widths += [3, 4]
+            for width in widths:
+                sizes = [width, width**2, width**3]
+                for n_items in sizes:
+                    tree = MerkleTree(n_items, width=width, hasher=hasher)
+                    tree.append(FQ.random())
+                    tree.append(FQ.random())
+                    proof_0_before = tree.proof(0)
+                    proof_1_before = tree.proof(1)
+                    root_before = tree.root
+                    self.assertTrue(proof_0_before.verify(tree.root))
+                    self.assertTrue(proof_1_before.verify(tree.root))
+
+                    leaf_0_after = FQ.random()
+                    tree.update(0, leaf_0_after)
+                    root_after_0 = tree.root
+                    proof_0_after = tree.proof(0)
+                    self.assertTrue(proof_0_after.verify(tree.root))
+                    self.assertNotEqual(root_before, root_after_0)
+
+                    leaf_1_after = FQ.random()
+                    tree.update(1, leaf_1_after)
+                    root_after_1 = tree.root
+                    proof_1_after = tree.proof(1)
+                    self.assertTrue(proof_1_after.verify(tree.root))
+                    self.assertNotEqual(root_before, root_after_1)
+                    self.assertNotEqual(root_after_0, root_after_1)
 
     def test_known_2pow28(self):
         tree = MerkleTree(2<<28)