From 667e1505ec22748b4debe2b94de19aa08aab0acc Mon Sep 17 00:00:00 2001
From: zero <zero@zero>
Date: Sun, 10 Mar 2024 17:14:09 +0100
Subject: [PATCH] smt.py: modify impl to match the impl we have in rust

---
 script/research/smt.py | 78 +++++++++++++++++++++++++++++-------------
 1 file changed, 54 insertions(+), 24 deletions(-)

diff --git a/script/research/smt.py b/script/research/smt.py
index dab0dc6b9..516e2de9a 100644
--- a/script/research/smt.py
+++ b/script/research/smt.py
@@ -3,7 +3,8 @@ import hashlib
 # SMT non-membership check
 # https://hackmd.io/@aztec-network/ryJ8wxfKK
 
-NULL = bytearray(32)
+# We set this to all 
+NULL = bytearray([0xff] * 32)
 
 def from_value(value):
     return value.to_bytes(32, 'little')
@@ -13,20 +14,28 @@ def hash_node(left, right):
 
 #               R
 #             /   \
-#         o           o
+#         o           o             i = 2
 #      /     \      /   \
-#     o       o    o     o
+#     o       o    o     o          i = 1
 #    / \    /  \  /  \  /  \
-#  110  4  77  5  -  -  6  -
+#  110  4  77  5  -  -  6  -        i = 0
 #    0  1   2  3  4  5  6  7
 
-pos = {
-    110: 0,
-    4: 1,
-    77: 2,
-    5: 3,
-    6: 6
-}
+# root    0
+# (2, 0)  1
+# (2, 1)  2
+# (1, 0)  3
+# (1, 1)  4
+# (1, 2)  5
+# (1, 3)  6
+# (0, 0)  7
+# (0, 1)  8
+# (0, 2)  9
+# (0, 3)  10
+# (0, 4)  11
+# (0, 5)  12
+# (0, 6)  13
+# (0, 7)  14
 
 empties = [
     NULL,
@@ -65,33 +74,54 @@ root = hash_node(layer_2[0], layer_2[1])
 
 table = {
     # root
-    0: (layer_2[0], layer_2[1]),
+    0: root,
 
-    # [1, 0]
-    1: (layer_1[0], layer_1[1]),
-    # This subtree contains a single value so just store that
-    # [1, 1]
-    2: (layer_1[2], 2, True),
+    # (2, 0)
+    1: layer_2[0],
+    # This subtree contains a single value
+    # (2, 1)
+    2: layer_2[1],
 
-    # [2, 0]
-    layer_1[0]: (layer_0[0], layer_0[1]),
-    # [2, 1]
-    layer_1[1]: (layer_0[2], layer_0[3]),
-    # ... [2, 2] and [2, 3] are not needed
+    # (1, 0)
+    3: layer_1[0],
+    # (1, 1)
+    4: layer_1[1],
+    # ... (1, 2) (idx=5) is not needed
+    # (1, 3)
+    6: layer_1[3],
+
+    # Don't add (0, 4), (0, 5) and (0, 7)
+    7:  layer_0[0],
+    8:  layer_0[1],
+    9:  layer_0[2],
+    10: layer_0[3],
+    13: layer_0[6],
 }
 
 # Prove that leaf 5 is set to None
 leaf = NULL
-path = [
+pos = 0b101
+
+path = []
+# Sibling of 5 is pos=4, with loc=(0, 4), which is idx=11
+sibling_idx = 11
+assert sibling_idx not in table
+path.append(empties[0])
+# Next node is (1, 2) with sibling (1, 3) and idx=6
+path.append(table[6])
+# Next node is (2, 1) with sibling (2, 0) and idx=1
+path.append(table[1])
+
+assert path == [
     NULL,
     layer_1[3],
     layer_2[0]
 ]
 assert hash_node(path[2], hash_node(hash_node(path[0], leaf), path[1])) == root
 # Starting from the bottom here
-pos = 0b101
 
 # Now do verification
+assert leaf == NULL
 bits = [bool(pos & (1<<n)) for n in range(3)]
 node = leaf
 for bit, other_node in zip(bits, path):