✨ putWthNodeBag to insert into CachedMerklePatriciaTree (#92)

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rainblock/merkle-patricia-tree",
-  "version": "4.5.0",
+  "version": "4.6.0",
   "description": "An implementation of the modified merkle patricia tree used in Ethereum, optimized for in-memory usage",
   "main": "build/src/index.js",
   "types": "build/src/index.d.js",
@@ -72,8 +72,8 @@
     "ts-loader": "^4.5.0",
     "ts-node": "^7.0.1",
     "typedoc": "^0.14.2",
-    "typescript": "^3.3.4000",
-    "webpack": "^4.29.6",
+    "typescript": "^3.4.3",
+    "webpack": "^4.30.0",
     "webpack-cli": "^3.3.0"
   },
   "publishConfig": {

src/index.spec.ts

@@ -795,6 +795,15 @@ describe('test cached merkle tree', async () => {
   });
 });
 
+describe('Test put in pruned CachedMerklePatriciaTree', async () => {
+  const cache = new CachedMerklePatriciaTree({putCanDelete: false}, 1);
+  cache.put(Buffer.from('abcd'), Buffer.from('abcd'));
+  cache.put(Buffer.from('abxx'), Buffer.from('abxx'));
+  cache.put(Buffer.from('xxxx'), Buffer.from('xxxx'));
+  cache.pruneStateCache();
+  should.throw(() => cache.put(Buffer.from('abcd'), Buffer.from('1234')));
+});
+
 describe('Test getFromCache and rlpToMerkleNode', async () => {
   const cache =
       new CachedMerklePatriciaTree<Buffer, Buffer>({putCanDelete: false}, 1);
@@ -881,10 +890,11 @@ describe('Test getFromCache and rlpToMerkleNode', async () => {
     v3!.should.deep.equal(Buffer.from('xxxx'));
   });
 
+  const nodeMap = new Map();
+  let extensionNodeHash: bigint|undefined;
+
   it('test getFromCache with non-empty nodeMap', async () => {
-    const nodeMap = new Map();
     const bagNodesUsed = new Set<bigint>();
-    let extensionNodeHash: bigint|undefined;
     if (cache.rootNode instanceof BranchNode) {
       for (const branch of (cache.rootNode).branches) {
         if (branch) {
@@ -919,4 +929,38 @@ describe('Test getFromCache and rlpToMerkleNode', async () => {
     bagNodesUsed.size.should.equal(1);
     bagNodesUsed.has(extensionNodeHash!).should.equal(true);
   });
+
+  it('test putWithNodeBag root hashes', async () => {
+    const initRoot = cache.root;
+    const bag1 = new Set<bigint>(), bag2 = new Set<bigint>(),
+          bag3 = new Set<bigint>();
+
+    cache.putWithNodeBag(
+        Buffer.from('abcd'), Buffer.from('abcd'), bag1, nodeMap);
+    cache.putWithNodeBag(
+        Buffer.from('abcx'), Buffer.from('abcx'), bag2, nodeMap);
+    cache.putWithNodeBag(
+        Buffer.from('xxxx'), Buffer.from('xxxx'), bag3, nodeMap);
+
+    bag1.size.should.equal(1);
+    bag1.has(extensionNodeHash!).should.equal(true);
+  });
+
+  it('test putWithNodeBag insertions', async () => {
+    const updatedValue = Buffer.from('1234');
+    cache.putWithNodeBag(Buffer.from('abcd'), updatedValue, undefined, nodeMap);
+    cache.putWithNodeBag(Buffer.from('abcx'), updatedValue, undefined, nodeMap);
+    cache.putWithNodeBag(Buffer.from('xxxx'), updatedValue, undefined, nodeMap);
+
+    const v1 = cache.get(Buffer.from('abcd')).value;
+    const v2 = cache.get(Buffer.from('abcx')).value;
+    const v3 = cache.get(Buffer.from('xxxx')).value;
+
+    should.exist(v1);
+    should.exist(v2);
+    should.exist(v3);
+    v1!.should.deep.equal(updatedValue);
+    v2!.should.deep.equal(updatedValue);
+    v3!.should.deep.equal(updatedValue);
+  });
 });

src/index.ts

@@ -1057,6 +1057,8 @@ export class MerklePatriciaTree<K = Buffer, V = Buffer> implements
             throw new Error('Unexpected non-branch node in update');
           }
         }
+      } else {
+        throw new Error('Unexpected node while inserting into tree');
       }
     }
   }
@@ -1782,6 +1784,114 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
     }
   }
 
+  /**
+   * putWithNodeBag inserts k,v pairs and rebuilds the pruned
+   * CachedMerklePatriciaTree if needed
+   * @param key : key to insert or modify
+   * @param value : value corresponding to the key
+   * @param usedNodes : A list of nodes used from the nodeBag for rebuilding the
+   * cached tree The usedNodes hashes change due to the insertion of the k, v
+   * pair into the tree
+   * @param nodeBag : A Map of nodes indexed by their hashes
+   */
+  putWithNodeBag(
+      key: K, value: V, usedNodes: Set<bigint>|undefined,
+      ...nodeBag: Array<Map<bigint, MerklePatriciaTreeNode<V>>>) {
+    const result = this.search(key);
+
+    if (result.remainder.length === 0 && result.node !== null) {
+      // Search matches; update the value in the path
+      result.node!.value = value;
+
+    } else if (result.stack[result.stack.length - 1] instanceof HashNode) {
+      // Partial path ending with a HashNode; (search stack depth >= pruneDepth)
+      // Rebuild path in the CachedMerklePatriciaTree with the nodes in the
+      // nodeMap
+      let currNode = result.stack[result.stack.length - 2];
+      const convKey = this.options.keyConverter!(key);
+      const keyNibbles = MerklePatriciaTreeNode.bufferToNibbles(convKey);
+      const lastNibbles = (currNode instanceof BranchNode) ?
+          [keyNibbles[(keyNibbles.length - result.remainder.length) - 1]] :
+          currNode.nibbles;
+      const remainder = lastNibbles.concat(result.remainder);
+      while (remainder && remainder.length !== 0) {
+        if (currNode instanceof BranchNode) {
+          const branch = currNode.branches[remainder[0]];
+          if (branch instanceof HashNode) {
+            // Get the hash of the correct branch
+            const replaceHash = branch.nodeHash;
+            // Read the node form the nodeBag
+            let replaceNode: MerklePatriciaTreeNode<V>|undefined = undefined;
+            for (const nodeMap of nodeBag) {
+              if (nodeMap.has(replaceHash)) {
+                replaceNode = nodeMap.get(replaceHash);
+                if (usedNodes) {
+                  usedNodes.add(replaceHash);
+                }
+                break;
+              }
+            }
+            if (!replaceNode) {
+              throw new Error('Couldn\'t service put using nodeBag');
+            }
+            // Replace HashNode with node from nodeBag
+            currNode.branches[remainder[0]] = replaceNode;
+            currNode = replaceNode;
+            remainder.shift();
+          } else {
+            // Traverse down the tree
+            currNode = branch;
+            remainder.shift();
+          }
+        } else if (currNode instanceof ExtensionNode) {
+          const next = currNode.nextNode;
+          if (next instanceof HashNode) {
+            // Get the hash of the correct branch
+            const replaceHash = next.nodeHash;
+            // Read the node from the nodeBag
+            let replaceNode: MerklePatriciaTreeNode<V>|undefined = undefined;
+            for (const nodeMap of nodeBag) {
+              if (nodeMap.has(replaceHash)) {
+                replaceNode = nodeMap.get(replaceHash);
+                if (usedNodes) {
+                  usedNodes.add(replaceHash);
+                }
+                break;
+              }
+            }
+            if (!replaceNode) {
+              throw new Error('Couldn\'t service put using nodeBag');
+            }
+            // Replace HashNode with node from nodeBag
+            currNode.nextNode = replaceNode;
+            remainder.splice(0, currNode.nibbles.length);
+            currNode = replaceNode;
+          } else {
+            // Traverse down the tree
+            remainder.splice(0, currNode.nibbles.length);
+            currNode = next;
+          }
+        } else if (currNode instanceof LeafNode) {
+          // Update the value in the tree
+          currNode.value = value;
+          break;
+        } else {
+          throw new Error('Unknown nodetype');
+        }
+      }
+
+    } else {
+      // Tree path for the key has depth < 6; perform insertion
+      this.insert(result.stack, result.remainder, value!);
+    }
+
+    // Clear all memoized hashes in the path, they will be reset.
+    const updatedResult = this.search(key);
+    for (const node of updatedResult.stack) {
+      node.clearMemoizedHash();
+    }
+  }
+
   verifyAndAddWitness(root: Buffer, key: K, witness: Witness<V>) {
     // Verify witness
     const convKey = this.options.keyConverter!(key);
@@ -1809,7 +1919,7 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
           currNode.nextNode = witness.proof[result.stack.length];
         }
       } else {
-        // Tree path for the key has depth < 6; perform insertion
+        // Tree path for the key has depth < pruneDepth; perform insertion
         this.insert(result.stack, result.remainder, witness.value!);
       }
       // Clear all memoized hashes in the path, they will be reset.