Implement setHashObjectFn (#62)

* Implement setHashObjectFn * Fix lint * Add missing JSDoc
ChainSafe · Aug 26, 2021 · 35bad6a · 35bad6a
1 parent 2782f3b
commit 35bad6a
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 43 deletions.
diff --git a/src/tree.ts b/src/tree.ts
@@ -6,6 +6,7 @@ import {createSingleProof} from "./proof/single";
 import {zeroNode} from "./zeroNode";
 
 export type Hook = (v: Tree) => void;
+export type HashObjectFn = (hashObject: HashObject) => HashObject;
 
 const ERR_INVALID_TREE = "Invalid tree operation";
 const ERR_PARAM_LT_ZERO = "Param must be >= 0";
@@ -74,8 +75,6 @@ export class Tree {
   }
 
   setNode(gindex: Gindex | GindexBitstring, n: Node, expand = false): void {
-    let node = this.rootNode;
-
     // Pre-compute entire bitstring instead of using an iterator (25% faster)
     let bitstring;
     if (typeof gindex === "string") {
@@ -86,46 +85,8 @@ export class Tree {
       }
       bitstring = gindex.toString(2);
     }
-
-    // Keep a list of all parent nodes of node at gindex `index`. Then walk the list
-    // backwards to rebind them "recursively" with the new nodes without using functions
-    const parentNodes: Node[] = [this.rootNode];
-
-    // Ignore the first bit, left right directions are at bits [1,..]
-    // Ignore the last bit, no need to push the target node to the parentNodes array
-    for (let i = 1; i < bitstring.length - 1; i++) {
-      if (node.isLeaf()) {
-        if (!expand) {
-          throw new Error(ERR_INVALID_TREE);
-        } else {
-          node = zeroNode(bitstring.length - i);
-        }
-      }
-
-      // Compare to string directly to prevent unnecessary type conversions
-      if (bitstring[i] === "1") {
-        node = node.right;
-      } else {
-        node = node.left;
-      }
-
-      parentNodes.push(node);
-    }
-
-    node = n;
-
-    // Ignore the first bit, left right directions are at bits [1,..]
-    // Iterate the list backwards including the last bit, but offset the parentNodes array
-    // by one since the first bit in bitstring was ignored in the previous loop
-    for (let i = bitstring.length - 1; i >= 1; i--) {
-      if (bitstring[i] === "1") {
-        node = parentNodes[i - 1].rebindRight(node);
-      } else {
-        node = parentNodes[i - 1].rebindLeft(node);
-      }
-    }
-
-    this.rootNode = node;
+    const parentNodes = this.getParentNodes(bitstring, expand);
+    this.rebindNodeToRoot(bitstring, parentNodes, n);
   }
 
   getRoot(index: Gindex | GindexBitstring): Uint8Array {
@@ -144,6 +105,30 @@ export class Tree {
     this.setNode(index, new LeafNode(hashObject), expand);
   }
 
+  /**
+   * Traverse from root node to node, get hash object, then apply the function to get new node
+   * and set the new node. This is a convenient method to avoid traversing the tree 2 times to
+   * get and set.
+   */
+  setHashObjectFn(gindex: Gindex | GindexBitstring, hashObjectFn: HashObjectFn, expand = false): void {
+    // Pre-compute entire bitstring instead of using an iterator (25% faster)
+    let bitstring;
+    if (typeof gindex === "string") {
+      bitstring = gindex;
+    } else {
+      if (gindex < 1) {
+        throw new Error("Invalid gindex < 1");
+      }
+      bitstring = gindex.toString(2);
+    }
+    const parentNodes = this.getParentNodes(bitstring, expand);
+    const lastParentNode = parentNodes[parentNodes.length - 1];
+    const lastBit = bitstring[bitstring.length - 1];
+    const oldNode = lastBit === "1" ? lastParentNode.right : lastParentNode.left;
+    const newNode = new LeafNode(hashObjectFn(oldNode));
+    this.rebindNodeToRoot(bitstring, parentNodes, newNode);
+  }
+
   getSubtree(index: Gindex): Tree {
     return new Tree(this.getNode(index), (v: Tree): void => this.setNode(index, v.rootNode));
   }
@@ -328,4 +313,59 @@ export class Tree {
   getProof(input: ProofInput): Proof {
     return createProof(this.rootNode, input);
   }
+
+  /**
+   * Traverse the tree from root node, ignore the last bit to get all parent nodes
+   * of the specified bitstring.
+   */
+  private getParentNodes(bitstring: GindexBitstring, expand = false): Node[] {
+    let node = this.rootNode;
+
+    // Keep a list of all parent nodes of node at gindex `index`. Then walk the list
+    // backwards to rebind them "recursively" with the new nodes without using functions
+    const parentNodes: Node[] = [this.rootNode];
+
+    // Ignore the first bit, left right directions are at bits [1,..]
+    // Ignore the last bit, no need to push the target node to the parentNodes array
+    for (let i = 1; i < bitstring.length - 1; i++) {
+      if (node.isLeaf()) {
+        if (!expand) {
+          throw new Error(ERR_INVALID_TREE);
+        } else {
+          node = zeroNode(bitstring.length - i);
+        }
+      }
+
+      // Compare to string directly to prevent unnecessary type conversions
+      if (bitstring[i] === "1") {
+        node = node.right;
+      } else {
+        node = node.left;
+      }
+
+      parentNodes.push(node);
+    }
+
+    return parentNodes;
+  }
+
+  /**
+   * Build a new tree structure from bitstring, parentNodes and a new node.
+   * Note: keep the same Tree, just mutate the root node.
+   */
+  private rebindNodeToRoot(bitstring: GindexBitstring, parentNodes: Node[], newNode: Node): void {
+    let node = newNode;
+    // Ignore the first bit, left right directions are at bits [1,..]
+    // Iterate the list backwards including the last bit, but offset the parentNodes array
+    // by one since the first bit in bitstring was ignored in the previous loop
+    for (let i = bitstring.length - 1; i >= 1; i--) {
+      if (bitstring[i] === "1") {
+        node = parentNodes[i - 1].rebindRight(node);
+      } else {
+        node = parentNodes[i - 1].rebindLeft(node);
+      }
+    }
+
+    this.rootNode = node;
+  }
 }
diff --git a/test/perf/tree.perf.ts b/test/perf/tree.perf.ts
@@ -1,3 +1,4 @@
+import { HashObject } from "@chainsafe/as-sha256";
 import { itBench, setBenchOpts } from "@dapplion/benchmark";
 import { LeafNode, subtreeFillToContents, Node, countToDepth, Tree, toGindex, uint8ArrayToHashObject, toGindexBitstring } from "../../src";
 
@@ -62,6 +63,21 @@ describe("Track the performance of different Tree methods", () => {
     }
   });
 
+  itBench("getHashObject then setHashObject", () => {
+    for (let i = 0; i < numLoop; i++) {
+      tree.getHashObject(gindex);
+      tree.setHashObject(gindex, newHashObject);
+    }
+  });
+
+  /* Double the speed compared to getHashObject then setHashObject */
+  itBench("hashObjectFn", () => {
+    const hashObjectFn = (hashObject: HashObject) => newHashObject;
+    for (let i = 0; i < numLoop; i++) {
+      tree.setHashObjectFn(gindex, hashObjectFn);
+    }
+  });
+
 });
 
 function createBalanceList(count: number, depth: number): Node {

diff --git a/test/unit/tree.test.ts b/test/unit/tree.test.ts
@@ -1,3 +1,4 @@
+import {byteArrayToHashObject, HashObject} from "@chainsafe/as-sha256";
 import {expect} from "chai";
 
 import {Tree, zeroNode, LeafNode, subtreeFillToContents, uint8ArrayToHashObject} from "../../src";
@@ -71,12 +72,17 @@ describe("subtree mutation", () => {
   });
 });
 
-describe("Tree.setNode", () => {
+describe("Tree.setNode vs Tree.setHashObjectFn", () => {
   it("Should compute root correctly after setting a leaf", () => {
     const depth = 4;
     const tree = new Tree(zeroNode(depth));
     tree.setNode(BigInt(18), new LeafNode(Buffer.alloc(32, 2)));
     expect(toHex(tree.root)).to.equal("3cfd85690fdd88abcf22ca7acf45bb47835326ff3166d3c953d5a23263fea2b2");
+    // setHashObjectFn
+    const hashObjectFn = (hashObject: HashObject): HashObject => byteArrayToHashObject(Buffer.alloc(32, 2));
+    const tree2 = new Tree(zeroNode(depth));
+    tree2.setHashObjectFn(BigInt(18), hashObjectFn);
+    expect(toHex(tree2.root)).to.equal("3cfd85690fdd88abcf22ca7acf45bb47835326ff3166d3c953d5a23263fea2b2");
   });
 
   it("Should compute root correctly after setting 3 leafs", () => {
@@ -86,6 +92,13 @@ describe("Tree.setNode", () => {
     tree.setNode(BigInt(46), new LeafNode(Buffer.alloc(32, 2)));
     tree.setNode(BigInt(60), new LeafNode(Buffer.alloc(32, 2)));
     expect(toHex(tree.root)).to.equal("02607e58782c912e2f96f4ff9daf494d0d115e7c37e8c2b7ddce17213591151b");
+    // setHashObjectFn
+    const hashObjectFn = (hashObject: HashObject): HashObject => byteArrayToHashObject(Buffer.alloc(32, 2));
+    const tree2 = new Tree(zeroNode(depth));
+    tree2.setHashObjectFn(BigInt(18), hashObjectFn);
+    tree2.setHashObjectFn(BigInt(46), hashObjectFn);
+    tree2.setHashObjectFn(BigInt(60), hashObjectFn);
+    expect(toHex(tree2.root)).to.equal("02607e58782c912e2f96f4ff9daf494d0d115e7c37e8c2b7ddce17213591151b");
   });
 
   it("Should throw for gindex 0", () => {