Better SparseMerkleTree Support for Empty-Leaf Operations (#485)

* Making SparseMerkleTree.verifyAndStorePartiallyEmptyPath(...) public and making it so that SparseMerkleTree.getMerkleProof(...) will still get a merkle proof for a disconnected leaf
plasma-group · Oct 2, 2019 · 986cf83 · 986cf83
1 parent e8e77cb
commit 986cf83
Show file tree

Hide file tree

Showing 3 changed files with 155 additions and 35 deletions.
diff --git a/packages/core/src/app/block-production/merkle-tree.ts b/packages/core/src/app/block-production/merkle-tree.ts
@@ -270,48 +270,65 @@ export class SparseMerkleTreeImpl implements SparseMerkleTree {
     leafKey: BigNumber,
     leafValue: Buffer
   ): Promise<MerkleTreeInclusionProof> {
-    return this.treeLock.acquire(SparseMerkleTreeImpl.lockKey, async () => {
-      if (!this.root || !this.root.hash || !this.root.value) {
-        return undefined
-      }
+    const result: MerkleTreeInclusionProof = await this.treeLock.acquire(
+      SparseMerkleTreeImpl.lockKey,
+      async () => {
+        if (!this.root || !this.root.hash) {
+          return undefined
+        }
 
-      let node: MerkleTreeNode = this.root
-      const siblings: Buffer[] = []
-      for (
-        let depth = 0;
-        depth < this.height && node && node.value.length === 64;
-        depth++
-      ) {
-        siblings.push(this.getChildSiblingHash(node, depth, leafKey))
-        node = await this.getChild(node, depth, leafKey)
-      }
+        let node: MerkleTreeNode = this.root
+        const siblings: Buffer[] = []
+        for (
+          let depth = 0;
+          depth < this.height &&
+          !!node &&
+          !!node.value &&
+          node.value.length === 64;
+          depth++
+        ) {
+          siblings.push(this.getChildSiblingHash(node, depth, leafKey))
+          node = await this.getChild(node, depth, leafKey)
+        }
 
-      if (siblings.length !== this.height - 1) {
-        return undefined
-      }
+        if (siblings.length !== this.height - 1) {
+          // TODO: A much better way of indicating this
+          return {
+            rootHash: undefined,
+            key: undefined,
+            value: undefined,
+            siblings: undefined,
+          }
+        }
 
-      if (!node.hash.equals(this.hashFunction(leafValue))) {
-        // Provided leaf doesn't match stored leaf
-        return undefined
+        if (!node.hash.equals(this.hashFunction(leafValue))) {
+          // Provided leaf doesn't match stored leaf
+          return undefined
+        }
+
+        return {
+          rootHash: this.root.hash,
+          key: leafKey,
+          value: leafValue,
+          siblings: siblings.reverse(),
+        }
       }
+    )
+
+    if (!result || !!result.rootHash) {
+      return result
+    }
 
-      return {
-        rootHash: this.root.hash,
-        key: leafKey,
-        value: leafValue,
-        siblings: siblings.reverse(),
+    // If this is for an empty leaf, we can store it and create a MerkleProof
+    if (leafValue.equals(SparseMerkleTreeImpl.emptyBuffer)) {
+      if (await this.verifyAndStorePartiallyEmptyPath(leafKey)) {
+        return this.getMerkleProof(leafKey, leafValue)
       }
-    })
+    }
+    return undefined
   }
 
-  /**
-   * Verifies and stores an empty leaf from a partially non-existent path.
-   *
-   * @param leafKey The leaf to store
-   * @param numExistingNodes The number of existing nodes, if known
-   * @returns True if verified, false otherwise
-   */
-  private async verifyAndStorePartiallyEmptyPath(
+  public async verifyAndStorePartiallyEmptyPath(
     leafKey: BigNumber,
     numExistingNodes?: number
   ): Promise<boolean> {

diff --git a/packages/core/src/types/block-production/merkle-tree.interface.ts b/packages/core/src/types/block-production/merkle-tree.interface.ts
@@ -89,4 +89,16 @@ export interface SparseMerkleTree extends MerkleTree {
    * @return true if the proof was valid (and thus stored), false otherwise
    */
   verifyAndStore(inclusionProof: MerkleTreeInclusionProof): Promise<boolean>
+
+  /**
+   * Verifies and stores an empty leaf from a partially non-existent path.
+   *
+   * @param leafKey The leaf to store
+   * @param numExistingNodes The number of existing nodes, if known
+   * @returns True if verified, false otherwise
+   */
+  verifyAndStorePartiallyEmptyPath(
+    leafKey: BigNumber,
+    numExistingNodes?: number
+  ): Promise<boolean>
 }
diff --git a/packages/core/test/app/block-production/merkle-tree.spec.ts b/packages/core/test/app/block-production/merkle-tree.spec.ts
@@ -266,6 +266,40 @@ describe('SparseMerkleTreeImpl', () => {
     })
   })
 
+  describe('verifyAndStorePartiallyEmptyPath', () => {
+    it('verifies and stores partially empty path sharing 0 populated ancestor other than root', async () => {
+      const tree: SparseMerkleTree = await createAndVerifyEmptyTreeDepthWithDepth(
+        db,
+        ZERO,
+        3
+      )
+
+      const value: Buffer = Buffer.from('zero leaf value')
+      assert(await tree.update(ZERO, value), 'update should have succeeded!')
+
+      assert(
+        await tree.verifyAndStorePartiallyEmptyPath(TWO),
+        'Verify and store should have worked for partially empty path 2'
+      )
+    })
+
+    it('verifies and stores partially empty path sharing 1 populated ancestor other than root', async () => {
+      const tree: SparseMerkleTree = await createAndVerifyEmptyTreeDepthWithDepth(
+        db,
+        ZERO,
+        4
+      )
+
+      const value: Buffer = Buffer.from('zero leaf value')
+      assert(await tree.update(ZERO, value), 'update should have succeeded!')
+
+      assert(
+        await tree.verifyAndStorePartiallyEmptyPath(TWO),
+        'Verify and store should have worked for partially empty path 2'
+      )
+    })
+  })
+
   describe('update', () => {
     it('updates empty tree', async () => {
       const tree: SparseMerkleTree = await SparseMerkleTreeImpl.create(
@@ -778,7 +812,7 @@ describe('SparseMerkleTreeImpl', () => {
   })
 
   describe('getMerkleProof', () => {
-    it('gets empty merkle proof', async () => {
+    it('gets empty merkle proof for existing empty leaf', async () => {
       const tree: MerkleTree = await createAndVerifyEmptyTreeDepthWithDepth(
         db,
         ZERO,
@@ -801,6 +835,32 @@ describe('SparseMerkleTreeImpl', () => {
       assert(proof.rootHash.equals(await tree.getRootHash()))
     })
 
+    it('gets empty merkle proof for disconnected empty leaf', async () => {
+      const tree: SparseMerkleTree = await SparseMerkleTreeImpl.create(
+        db,
+        undefined,
+        3,
+        hashFunction
+      )
+
+      const proof: MerkleTreeInclusionProof = await tree.getMerkleProof(
+        ZERO,
+        SparseMerkleTreeImpl.emptyBuffer
+      )
+
+      assert(!!proof, 'Proof should not be undefined!')
+      assert(proof.value.equals(SparseMerkleTreeImpl.emptyBuffer))
+      assert(proof.key.equals(ZERO))
+      assert(proof.siblings.length === 2)
+
+      let hash: Buffer = zeroHash
+      assert(proof.siblings[0].equals(hash))
+      hash = bufferHashFunction(hashBuffer.fill(hash, 0, 32).fill(hash, 32))
+      assert(proof.siblings[1].equals(hash))
+
+      assert(proof.rootHash.equals(await tree.getRootHash()))
+    })
+
     it('gets merkle proof for non-empty tree', async () => {
       const tree: SparseMerkleTree = await SparseMerkleTreeImpl.create(
         db,
@@ -828,6 +888,37 @@ describe('SparseMerkleTreeImpl', () => {
       assert(proof.rootHash.equals(await tree.getRootHash()))
     })
 
+    it('gets merkle proof for disconnected empty leaf in non-empty tree', async () => {
+      const tree: SparseMerkleTree = await SparseMerkleTreeImpl.create(
+        db,
+        undefined,
+        3,
+        hashFunction
+      )
+      const zeroKeyData: Buffer = Buffer.from('really great leaf data')
+      await tree.update(ZERO, zeroKeyData)
+
+      const proof: MerkleTreeInclusionProof = await tree.getMerkleProof(
+        TWO,
+        SparseMerkleTreeImpl.emptyBuffer
+      )
+
+      assert(!!proof, 'Proof should not be undefined!')
+      assert(proof.value.equals(SparseMerkleTreeImpl.emptyBuffer))
+      assert(proof.key.equals(TWO))
+      assert(proof.siblings.length === 2)
+
+      assert(proof.siblings[0].equals(zeroHash))
+
+      const hashData: Buffer = bufferHashFunction(zeroKeyData)
+      const zeroAndOneParent = bufferHashFunction(
+        hashBuffer.fill(hashData, 0, 32).fill(zeroHash, 32)
+      )
+      assert(proof.siblings[1].equals(zeroAndOneParent))
+
+      assert(proof.rootHash.equals(await tree.getRootHash()))
+    })
+
     it('gets merkle proof for non-empty siblings 0 & 1', async () => {
       const tree: SparseMerkleTree = await SparseMerkleTreeImpl.create(
         db,