start on merkle proofs part

aiken-lang · May 28, 2024 · ce7002c · ce7002c
1 parent e51d219
commit ce7002c
Show file tree

Hide file tree

Showing 2 changed files with 164 additions and 6 deletions.
diff --git a/offchain/src/index.test.ts b/offchain/src/index.test.ts
@@ -1,7 +1,7 @@
 import { expect, test } from "vitest";
 import { Buffer } from "buffer";
 
-import { Leaf, SparseMerkleTree } from "./index.js";
+import { SparseMerkleTree } from "./index.js";
 
 test("Test 1", () => {
   const x = new SparseMerkleTree();
@@ -20,7 +20,7 @@ test("Test 1", () => {
     "ce21ae7b870c1012db2b9d469e95a05540ad74640c236776139e52118d39f2fc",
     "0170b41f8f90f96eb95a0dfc66b959fb4e7060ed738ee162076d03597a0f468f",
     "92cc9d3ed08668c5d71243ccac72b76b46c924e5cee13583665a13920b244e23",
-  ].map((x) => Buffer.from(x, "hex").toString("hex"));
+  ];
 
   x.insert("apple (0)");
 
@@ -74,3 +74,26 @@ test("Test 1", () => {
     expectedList
   );
 });
+
+test("Test 2", () => {
+  const x = new SparseMerkleTree();
+  let rootList: Uint8Array[] = [];
+  const expectedList: string[] = [];
+
+  x.insert("apple (0)");
+  x.insert("apricot (0)");
+  x.insert("banana (328)");
+  x.insert("blackberry (0)");
+  x.insert("blueberry (92383)");
+  x.insert("cherry (0)");
+  x.insert("coconut (0)");
+  x.insert("cranberry (0)");
+  x.insert("durian (0)");
+  x.insert("fig (0)");
+  x.insert("grape (110606)");
+  x.insert("grapefruit (0)");
+
+  const thing = x.merkleProof("grapefruit (0)");
+
+  console.log(thing);
+});
diff --git a/offchain/src/index.ts b/offchain/src/index.ts
@@ -4,6 +4,8 @@ import { Buffer } from "buffer";
 
 const leafBytes = new Uint8Array(Buffer.from("0deeffaad07783", "hex"));
 
+type Side = "left" | "right";
+
 export class Leaf {
   key: BitSet;
   value: Uint8Array;
@@ -28,6 +30,10 @@ export class Leaf {
     );
   }
 
+  getHash() {
+    return this.leafHash;
+  }
+
   doInsert(key: BitSet, value: string | Buffer) {
     let other = new Leaf(value);
 
@@ -52,6 +58,18 @@ export class Leaf {
     }
   }
 
+  doMerkeProof(_key: BitSet): [Uint8Array, number, Side][] {
+    return [[blake2b(this.value, undefined, 32), 0, "left"]];
+  }
+
+  traverseLeft(): [Uint8Array, number, Side][] {
+    return [[blake2b(this.value, undefined, 32), 0, "left"]];
+  }
+
+  traverseRight(): [Uint8Array, number, Side][] {
+    return [[blake2b(this.value, undefined, 32), 0, "right"]];
+  }
+
   static boundaryLeaf(isMin: boolean) {
     if (isMin) {
       let x = new Leaf(Buffer.from("00", "hex"));
@@ -156,6 +174,10 @@ export class Branch {
     }
   }
 
+  getHash() {
+    return this.branchHash;
+  }
+
   doInsert(key: BitSet, value: string | Buffer): Branch | Leaf {
     if (key.slice(this.height + 1).equals(this.key)) {
       let leftHeight = this.height - 1;
@@ -224,9 +246,9 @@ export class Branch {
         }
       }
 
-      if (leftHeight > 0 && rightHeight < 0) {
+      if (leftHeight >= 0 && rightHeight < 0) {
         this.leftChild = this.leftChild.doInsert(key, value);
-      } else if (leftHeight < 0 && rightHeight > 0) {
+      } else if (leftHeight < 0 && rightHeight >= 0) {
         this.rightChild = this.rightChild.doInsert(key, value);
       } else {
         throw new Error("Impossible");
@@ -263,6 +285,106 @@ export class Branch {
       return new Branch(new Leaf(value), this);
     }
   }
+
+  doMerkeProof(key: BitSet): [Uint8Array, number, Side][] {
+    let leftHeight = this.height - 1;
+
+    if (this.leftChild instanceof Leaf) {
+      while (leftHeight > -1) {
+        if (
+          key
+            .slice(leftHeight + 1)
+            .equals(this.leftChild.key.slice(leftHeight + 1))
+        ) {
+          break;
+        }
+
+        leftHeight--;
+      }
+    } else {
+      while (leftHeight >= this.leftChild.height) {
+        if (
+          key
+            .slice(leftHeight + 1)
+            .equals(
+              this.leftChild.key.slice(leftHeight - this.leftChild.height)
+            )
+        ) {
+          break;
+        }
+
+        leftHeight--;
+      }
+      if (leftHeight < this.leftChild.height) {
+        leftHeight = -1;
+      }
+    }
+
+    let rightHeight = this.height - 1;
+
+    if (this.rightChild instanceof Leaf) {
+      while (rightHeight > -1) {
+        if (
+          key
+            .slice(rightHeight + 1)
+            .equals(this.rightChild.key.slice(rightHeight + 1))
+        ) {
+          break;
+        }
+
+        rightHeight--;
+      }
+    } else {
+      while (rightHeight >= this.rightChild.height) {
+        if (
+          key
+            .slice(rightHeight + 1)
+            .equals(
+              this.rightChild.key.slice(rightHeight - this.rightChild.height)
+            )
+        ) {
+          break;
+        }
+
+        rightHeight--;
+      }
+      if (rightHeight < this.rightChild.height) {
+        rightHeight = -1;
+      }
+    }
+
+    if (key.equals(this.leftChild.key)) {
+      return this.rightChild.traverseLeft();
+    } else if (key.equals(this.rightChild.key)) {
+      return this.leftChild.traverseRight();
+    } else if (leftHeight >= 0 && rightHeight < 0) {
+      return [
+        [this.rightChild.getHash(), this.height, "right"],
+        ...this.leftChild.doMerkeProof(key),
+      ];
+    } else if (leftHeight < 0 && rightHeight >= 0) {
+      return [
+        [this.leftChild.getHash(), this.height, "left"],
+        ...this.rightChild.doMerkeProof(key),
+      ];
+    } else {
+      throw new Error("Impossible");
+    }
+  }
+
+  traverseLeft(): [Uint8Array, number, Side][] {
+    return [
+      [this.rightChild.getHash(), this.height, "right"],
+      ...this.leftChild.traverseLeft(),
+    ];
+  }
+
+  traverseRight(): [Uint8Array, number, Side][] {
+    return [
+      [this.leftChild.getHash(), this.height, "left"],
+      ...this.rightChild.traverseRight(),
+    ];
+  }
 }
 
 export class SparseMerkleTree extends Branch {
@@ -281,12 +403,25 @@ export class SparseMerkleTree extends Branch {
         ? new TextEncoder().encode(value)
         : new Uint8Array(value);
 
-    console.log(blake2bHex(bufferValue, undefined, 32));
-
     const initialKey = new BitSet(
       Buffer.from(blake2b(bufferValue, undefined, 32)).reverse()
     );
 
     super.doInsert(initialKey, value);
   }
+
+  merkleProof(value: string | Buffer) {
+    const bufferValue: Uint8Array =
+      typeof value == "string"
+        ? new TextEncoder().encode(value)
+        : new Uint8Array(value);
+
+    const initialKey = new BitSet(
+      Buffer.from(blake2b(bufferValue, undefined, 32)).reverse()
+    );
+
+    let proofArray = super.doMerkeProof(initialKey);
+
+    return proofArray;
+  }
 }