Batch-optimization

packages/trie/src/trie.ts

@@ -29,7 +29,7 @@
 import { verifyRangeProof } from './proof/range.js'
 import { ROOT_DB_KEY } from './types.js'
 import { _walkTrie } from './util/asyncWalk.js'
-import { bytesToNibbles, matchingNibbleLength } from './util/nibbles.js'
+import { bytesToNibbles, matchingNibbleLength, orderBatch } from './util/nibbles.js'
 import { TrieReadStream as ReadStream } from './util/readStream.js'
 import { WalkController } from './util/walkController.js'
 
@@ -495,14 +495,15 @@
   async put(
     key: Uint8Array,
     value: Uint8Array | null,
-    skipKeyTransform: boolean = false
-  ): Promise<void> {
+    skipKeyTransform: boolean = false,
+    path?: { stack: TrieNode[]; remaining: number[] }
+  ): Promise<TrieNode[]> {
     this.DEBUG && this.debug(`Key: ${bytesToHex(key)}`, ['PUT'])
     this.DEBUG && this.debug(`Value: ${value === null ? 'null' : bytesToHex(key)}`, ['PUT'])
     if (this._opts.useRootPersistence && equalsBytes(key, ROOT_DB_KEY) === true) {
       throw new Error(`Attempted to set '${bytesToUtf8(ROOT_DB_KEY)}' key but it is not allowed.`)
     }
-
+    let returnStack: TrieNode[] = []
     // If value is empty, delete
     if (value === null || value.length === 0) {
       return this.del(key)
@@ -512,10 +513,10 @@
     const appliedKey = skipKeyTransform ? key : this.appliedKey(key)
     if (equalsBytes(this.root(), this.EMPTY_TRIE_ROOT) === true) {
       // If no root, initialize this trie
-      await this._createInitialNode(appliedKey, value)
+      returnStack = [await this._createInitialNode(appliedKey, value)]
     } else {
       // First try to find the given key or its nearest node
-      const { remaining, stack } = await this.findPath(appliedKey)
+      const { remaining, stack } = path ?? (await this.findPath(appliedKey))
       let ops: BatchDBOp[] = []
       if (this._opts.useNodePruning) {
         const val = await this.get(key)
@@ -540,14 +541,15 @@
         }
       }
       // then update
-      await this._updateNode(appliedKey, value, remaining, stack)
+      returnStack = await this._updateNode(appliedKey, value, remaining, stack)
       if (this._opts.useNodePruning) {
         // Only after updating the node we can delete the keyhashes
         await this._db.batch(ops)
       }
     }
     await this.persistRoot()
     this._lock.release()
+    return returnStack
   }
 
   /**
@@ -556,11 +558,16 @@
    * @param key
    * @returns A Promise that resolves once value is deleted.
    */
-  async del(key: Uint8Array, skipKeyTransform: boolean = false): Promise<void> {
+  async del(
+    key: Uint8Array,
+    skipKeyTransform: boolean = false,
+    path?: { stack: TrieNode[]; remaining: number[] }
+  ): Promise<TrieNode[]> {
     this.DEBUG && this.debug(`Key: ${bytesToHex(key)}`, ['DEL'])
     await this._lock.acquire()
     const appliedKey = skipKeyTransform ? key : this.appliedKey(key)
-    const { node, stack } = await this.findPath(appliedKey)
+    const partialPath = path && { stack: path.stack }
+    const { node, stack } = await this.findPath(appliedKey, false, partialPath)
 
     let ops: BatchDBOp[] = []
     // Only delete if the `key` currently has any value
@@ -588,6 +595,7 @@
     }
     await this.persistRoot()
     this._lock.release()
+    return stack
   }
 
   /**
@@ -770,7 +778,7 @@
    * Creates the initial node from an empty tree.
    * @private
    */
-  protected async _createInitialNode(key: Uint8Array, value: Uint8Array): Promise<void> {
+  protected async _createInitialNode(key: Uint8Array, value: Uint8Array): Promise<TrieNode> {
     const newNode = new LeafNode(bytesToNibbles(key), value)
 
     const encoded = newNode.serialize()
@@ -779,6 +787,7 @@
     rootKey = this._opts.keyPrefix ? concatBytes(this._opts.keyPrefix, rootKey) : rootKey
     await this._db.put(rootKey, encoded)
     await this.persistRoot()
+    return newNode
   }
 
   /**
@@ -817,7 +826,7 @@
     value: Uint8Array,
     keyRemainder: Nibbles,
     stack: TrieNode[]
-  ): Promise<void> {
+  ): Promise<TrieNode[]> {
     const toSave: BatchDBOp[] = []
     const lastNode = stack.pop()
     if (!lastNode) {
@@ -908,7 +917,9 @@
       }
     }
 
+    const returnStack = [...stack]
     await this.saveStack(key, stack, toSave)
+    return returnStack
   }
 
   /**
@@ -1137,17 +1148,59 @@
    * @param ops
    */
   async batch(ops: BatchDBOp[], skipKeyTransform?: boolean): Promise<void> {
-    for (const op of ops) {
+    const keyTransform =
+      skipKeyTransform === true
+        ? undefined
+        : (msg: Uint8Array) => {
+            return this.appliedKey(msg)
+          }
+    const sortedOps = orderBatch(ops, keyTransform)
+    let stack: TrieNode[] = []
+    const stackPathCache: Map<string, TrieNode> = new Map()
+    for (const op of sortedOps) {
+      const appliedKey = skipKeyTransform === true ? op.key : this.appliedKey(op.key)
+      const nibbles = bytesToNibbles(appliedKey)
+      stack = []
+      let remaining = nibbles
+      for (let i = 0; i < nibbles.length; i++) {
+        const p: string = JSON.stringify(nibbles.slice(0, i) as number[])
+        if (stackPathCache.has(p)) {
+          const node = stackPathCache.get(p)!
+          stack.push(node)
+          remaining = nibbles.slice(i)
+        }
+      }
+      const _path =
+        stack.length > 0
+          ? {
+              stack,
+              remaining,
+            }
+          : undefined
       if (op.type === 'put') {
-        if (op.value === null || op.value === undefined) {
-          throw new Error('Invalid batch db operation')
+        stack = await this.put(op.key, op.value, skipKeyTransform, _path)
+        const path: number[] = []
+        for (const node of stack) {
+          stackPathCache.set(JSON.stringify([...path]), node)
+          if (node instanceof BranchNode) {
+            path.push(nibbles.shift()!)
+          } else {
+            path.push(...nibbles.splice(0, node.keyLength()))
+          }
         }
-        await this.put(op.key, op.value, skipKeyTransform)
       } else if (op.type === 'del') {
-        await this.del(op.key, skipKeyTransform)
+        stack = await this.put(op.key, null, skipKeyTransform, _path)
+        const path: number[] = []
+        for (const node of stack) {
+          stackPathCache.set(JSON.stringify([...path]), node)
+          if (node instanceof BranchNode) {
+            path.push(nibbles.shift()!)
+          } else {
+            path.push(...nibbles.splice(0, node.keyLength()))
+          }
+        }
       }
     }
-    await this.persistRoot()
   }
 
   // This method verifies if all keys in the trie (except the root) are reachable

packages/trie/src/util/nibbles.ts

@@ -1,6 +1,7 @@
 import { toBytes } from '@ethereumjs/util'
 
 import type { Nibbles } from '../types.js'
+import type { BatchDBOp } from '@ethereumjs/util'
 
 /**
  * Converts a bytes to a nibble array.
@@ -91,3 +92,20 @@ export function doKeysMatch(keyA: Nibbles, keyB: Nibbles): boolean {
   const length = matchingNibbleLength(keyA, keyB)
   return length === keyA.length && length === keyB.length
 }
+
+export function orderBatch(
+  ops: BatchDBOp[],
+  keyTransform: (msg: Uint8Array) => Uint8Array = (msg: Uint8Array) => {
+    return msg
+  }
+): BatchDBOp[] {
+  const keyNibbles: [number, number[]][] = ops.map((o, i) => {
+    const appliedKey = keyTransform(o.key)
+    const nibbles: number[] = bytesToNibbles(appliedKey)
+    return [i, nibbles]
+  })
+  keyNibbles.sort(([_, a], [__, b]) => {
+    return nibblesCompare(a, b)
+  })
+  return keyNibbles.map(([i, _]) => ops[i])
+}