From c3301a472afe35bf4b19a5dd78a5999f71a43692 Mon Sep 17 00:00:00 2001
From: Scotty <66335769+ScottyPoi@users.noreply.github.com>
Date: Wed, 26 Feb 2025 17:57:22 -0700
Subject: [PATCH] binarytree: implement proof methods (#3888)

* binarytree: implement create proof method

* binarytree: implement fromProof method

* binarytree: implement verify proof method

* test proof methods

* test proof of non-existence
---
 packages/binarytree/src/binaryTree.ts  | 33 ++++++++--
 packages/binarytree/test/proof.spec.ts | 83 ++++++++++++++++++++++++++
 2 files changed, 110 insertions(+), 6 deletions(-)
 create mode 100644 packages/binarytree/test/proof.spec.ts

diff --git a/packages/binarytree/src/binaryTree.ts b/packages/binarytree/src/binaryTree.ts
index ceb148003b..a45d0364c4 100644
--- a/packages/binarytree/src/binaryTree.ts
+++ b/packages/binarytree/src/binaryTree.ts
@@ -11,6 +11,7 @@ import {
 } from '@ethereumjs/util'
 import debug from 'debug'
 
+import { createBinaryTree } from './constructors.js'
 import { CheckpointDB } from './db/index.js'
 import { InternalBinaryNode } from './node/internalNode.js'
 import { StemBinaryNode } from './node/stemNode.js'
@@ -576,16 +577,25 @@ export class BinaryTree {
    * Saves the nodes from a proof into the tree.
    * @param proof
    */
-  async fromProof(_proof: any): Promise<void> {
-    throw new Error('Not implemented')
+  async fromProof(_proof: Uint8Array[]): Promise<BinaryTree> {
+    const proofTrie = await createBinaryTree()
+    const putStack: [Uint8Array, BinaryNode][] = _proof.map((bytes) => {
+      const node = decodeBinaryNode(bytes)
+      return [this.merkelize(node), node]
+    })
+    await proofTrie.saveStack(putStack)
+    const root = putStack[0][0]
+    proofTrie.root(root)
+    return proofTrie
   }
 
   /**
    * Creates a proof from a tree and key that can be verified using {@link BinaryTree.verifyBinaryProof}.
    * @param key
    */
-  async createBinaryProof(_key: Uint8Array): Promise<any> {
-    throw new Error('Not implemented')
+  async createBinaryProof(_key: Uint8Array): Promise<Uint8Array[]> {
+    const { stack } = await this.findPath(_key)
+    return stack.map(([node, _]) => node.serialize())
   }
 
   /**
@@ -599,9 +609,20 @@ export class BinaryTree {
   async verifyBinaryProof(
     _rootHash: Uint8Array,
     _key: Uint8Array,
-    _proof: any,
+    _proof: Uint8Array[],
   ): Promise<Uint8Array | null> {
-    throw new Error('Not implemented')
+    const proofTrie = await this.fromProof(_proof)
+    const [value] = await proofTrie.get(_key.slice(0, 31), [_key[31]])
+    const valueNode = decodeBinaryNode(_proof[_proof.length - 1]) as StemBinaryNode
+    const expectedValue = valueNode.values[_key[31]]
+    if (!expectedValue) {
+      if (value) {
+        throw new Error('Proof is invalid')
+      }
+    } else if (value && !equalsBytes(value, expectedValue)) {
+      throw new Error('Proof is invalid')
+    }
+    return value
   }
 
   /**
diff --git a/packages/binarytree/test/proof.spec.ts b/packages/binarytree/test/proof.spec.ts
new file mode 100644
index 0000000000..a6a667a69a
--- /dev/null
+++ b/packages/binarytree/test/proof.spec.ts
@@ -0,0 +1,83 @@
+import { blake3 } from '@noble/hashes/blake3'
+import { assert, describe, it } from 'vitest'
+
+import { createBinaryTree } from '../src/constructors.js'
+import { decodeBinaryNode } from '../src/index.js'
+
+import type { StemBinaryNode } from '../src/node/stemNode.js'
+
+// Create an array of 100 random key/value pairs by hashing keys.
+
+const keyValuePairs: { originalKey: Uint8Array; hashedKey: Uint8Array; value: Uint8Array }[] = []
+
+for (let i = 0; i < 100; i++) {
+  const key = new Uint8Array(32).fill(0)
+  key[31] = i // vary the last byte to differentiate keys
+
+  const hashedKey = blake3(key)
+
+  // Create a value also based on i (filled with 0xBB and ending with i)
+  const value = new Uint8Array(32).fill(1)
+  value[31] = i
+
+  keyValuePairs.push({ originalKey: key, hashedKey, value })
+}
+
+describe('binary tree proof', async () => {
+  const tree1 = await createBinaryTree()
+
+  // Insert each key/value pair into the tree.
+  for (const { hashedKey, value } of keyValuePairs) {
+    const stem = hashedKey.slice(0, 31)
+    const index = hashedKey[31]
+    await tree1.put(stem, [index], [value])
+  }
+
+  it('should create and verify a merkle proof for existing key', async () => {
+    // create merkle proof for first key/value pair
+    const proof = await tree1.createBinaryProof(keyValuePairs[0].hashedKey)
+
+    const rootNode = decodeBinaryNode(proof[0])
+    assert.deepEqual(
+      tree1['merkelize'](rootNode),
+      tree1.root(),
+      'first value in proof should be root node',
+    )
+    const valueNode = decodeBinaryNode(proof[proof.length - 1]) as StemBinaryNode
+    assert.deepEqual(
+      keyValuePairs[0].value,
+      valueNode.values[keyValuePairs[0].hashedKey[31]],
+      'last value in proof should be target node',
+    )
+
+    // create sparse tree from proof
+    const tree2 = await tree1.fromProof(proof)
+    assert.deepEqual(
+      tree2.root(),
+      tree1.root(),
+      'tree from proof should be created with correct root node',
+    )
+
+    // get value from sparse tree
+    const [value] = await tree2.get(keyValuePairs[0].hashedKey.slice(0, 31), [
+      keyValuePairs[0].hashedKey[31],
+    ])
+    assert.deepEqual(value, keyValuePairs[0].value)
+
+    // verify proof using verifyBinaryProof
+    const proofValue = await tree1.verifyBinaryProof(
+      tree1.root(),
+      keyValuePairs[0].hashedKey,
+      proof,
+    )
+    assert.deepEqual(keyValuePairs[0].value, proofValue, 'verify proof should return target value')
+  })
+
+  it('should create and verify a proof of non-existence', async () => {
+    const fakeKey = new Uint8Array(keyValuePairs[0].hashedKey.length).fill(5)
+
+    const proof = await tree1.createBinaryProof(fakeKey)
+    const proofValue = await tree1.verifyBinaryProof(tree1.root(), fakeKey, proof)
+    assert.deepEqual(proofValue, undefined, 'verify proof of non-existence should return undefined')
+  })
+})