nodeBag cache (#91)

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rainblock/merkle-patricia-tree",
-  "version": "4.4.0",
+  "version": "4.5.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",

src/index.spec.ts

@@ -766,7 +766,7 @@ describe('test cached merkle tree', async () => {
     should.not.equal(null, witness.value);
 
     // And getting from cache should as well
-    const result = tree.getFromCache(errorAccount, new Map());
+    const result = tree.getFromCache(errorAccount, undefined, new Map());
     should.not.equal(null, result);
     result!.should.deep.equal(value);
   });
@@ -791,7 +791,7 @@ describe('test cached merkle tree', async () => {
     tree.pruneStateCache();
 
     // This should return an error
-    should.throw(() => tree.getFromCache(errorAccount, new Map()));
+    should.throw(() => tree.getFromCache(errorAccount, undefined, new Map()));
   });
 });
 
@@ -873,16 +873,18 @@ describe('Test getFromCache and rlpToMerkleNode', async () => {
 
   it('test getFromCache with empty nodeMap', async () => {
     const nodeMap = new Map();
-    const v1 = cache.getFromCache(Buffer.from('abcd'), nodeMap);
-    const v2 = cache.getFromCache(Buffer.from('abcx'), nodeMap);
-    const v3 = cache.getFromCache(Buffer.from('xxxx'), nodeMap);
+    const v1 = cache.getFromCache(Buffer.from('abcd'), undefined, nodeMap);
+    const v2 = cache.getFromCache(Buffer.from('abcx'), undefined, nodeMap);
+    const v3 = cache.getFromCache(Buffer.from('xxxx'), undefined, nodeMap);
     v1!.should.deep.equal(Buffer.from('abcd'));
     v2!.should.deep.equal(Buffer.from('abcx'));
     v3!.should.deep.equal(Buffer.from('xxxx'));
   });
 
   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) {
@@ -892,6 +894,9 @@ describe('Test getFromCache and rlpToMerkleNode', async () => {
               cache.options as {} as MerklePatriciaTreeOptions<{}, Buffer>);
           const mappedNode = cache.rlpToMerkleNode(node, (val: Buffer) => val);
           nodeMap.set(hash, mappedNode);
+          if (mappedNode instanceof ExtensionNode) {
+            extensionNodeHash = hash;
+          }
         }
       }
       const branch6 = cache.rootNode.branches[6];
@@ -905,11 +910,13 @@ describe('Test getFromCache and rlpToMerkleNode', async () => {
       }
     }
     cache.pruneStateCache();
-    const v1 = cache.getFromCache(Buffer.from('abcd'), nodeMap);
-    const v2 = cache.getFromCache(Buffer.from('abcx'), nodeMap);
-    const v3 = cache.getFromCache(Buffer.from('xxxx'), nodeMap);
+    const v1 = cache.getFromCache(Buffer.from('abcd'), bagNodesUsed, nodeMap);
+    const v2 = cache.getFromCache(Buffer.from('abcx'), bagNodesUsed, nodeMap);
+    const v3 = cache.getFromCache(Buffer.from('xxxx'), bagNodesUsed, nodeMap);
     v1!.should.deep.equal(Buffer.from('abcd'));
     v2!.should.deep.equal(Buffer.from('abcx'));
     v3!.should.deep.equal(Buffer.from('xxxx'));
+    bagNodesUsed.size.should.equal(1);
+    bagNodesUsed.has(extensionNodeHash!).should.equal(true);
   });
 });

src/index.ts

@@ -1629,16 +1629,18 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
    * throws an exception if key is not found in the cache and nodeMap
    */
   private _getRecursive(
-      key: number[], nodeMap: Map<bigint, MerklePatriciaTreeNode<V>>,
-      node: MerklePatriciaTreeNode<V>): V {
+      node: MerklePatriciaTreeNode<V>, key: number[],
+      bagNodesUsed: Set<bigint>|undefined,
+      ...nodeMap: Array<Map<bigint, MerklePatriciaTreeNode<V>>>): V {
     if (node instanceof BranchNode) {
       // If key ends at a BranchNode; return the BranchNode value
       const nib = key.shift();
       if (nib === undefined) {
         return node.value;
       }
       // Search down the appropriate branch of the BranchNode
-      const ret = this._getRecursive(key, nodeMap, node.branches[nib]);
+      const ret =
+          this._getRecursive(node.branches[nib], key, bagNodesUsed, ...nodeMap);
       return ret;
 
     } else if (node instanceof ExtensionNode) {
@@ -1649,7 +1651,8 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
       // Remove the matchingNibbles from the key
       key.splice(0, node.nibbles.length);
       // Search down the nextNode of the ExtensionNode
-      const ret = this._getRecursive(key, nodeMap, node.nextNode);
+      const ret =
+          this._getRecursive(node.nextNode, key, bagNodesUsed, ...nodeMap);
       return ret;
 
     } else if (node instanceof LeafNode) {
@@ -1664,12 +1667,21 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
       // Read the nodeHash of the HashNode
       const hash = node.nodeHash;
       // Get the MerkleNode corresponding to nodeHash from nodeMap
-      const mappedNode = nodeMap.get(hash);
+      let mappedNode = undefined;
+      for (const map of nodeMap) {
+        mappedNode = map.get(hash);
+        if (mappedNode) {
+          if (bagNodesUsed) {
+            bagNodesUsed.add(hash);
+          }
+          break;
+        }
+      }
       if (!mappedNode) {
         throw new MerklePrunedError();
       }
       // Search down the mappedNode
-      const ret = this._getRecursive(key, nodeMap, mappedNode);
+      const ret = this._getRecursive(mappedNode, key, bagNodesUsed, ...nodeMap);
       return ret;
 
     } else if (NullNode) {
@@ -1684,17 +1696,22 @@ export class CachedMerklePatriciaTree<K, V> extends MerklePatriciaTree<K, V> {
   /**
    * getFromCache returns the value corresponding to the key using nodeMap
    * @param key : key to get from the CachedMerklePatriciaTree
+   * @param bagNodesUsed: A list of nodes used from the nodeMap updated by the
+   * Tree
    * @param nodeMap : Bag of recent MerklePatriciaTree nodes from the client
    *
    * @returns value corresponding to the key if present; null if otherwise
    */
-  getFromCache(key: K, nodeMap: Map<bigint, MerklePatriciaTreeNode<V>>): V {
+  getFromCache(
+      key: K, bagNodesUsed: Set<bigint>|undefined,
+      ...nodeMap: Array<Map<bigint, MerklePatriciaTreeNode<V>>>): V {
     const convKey = this.options.keyConverter!(key);
     const keyNibbles = MerklePatriciaTreeNode.bufferToNibbles(convKey);
     // getRecursiveKey throws an exception if key is not searchable
     // in the cache and the nodeBag or if key is not present
-    const ret = this._getRecursive(keyNibbles, nodeMap, this.rootNode);
-    return ret;
+    const value =
+        this._getRecursive(this.rootNode, keyNibbles, bagNodesUsed, ...nodeMap);
+    return value;
   }
 
   /**