trie: move _batch back to trie.batch

ethereumjs · Mar 15, 2024 · 964b0ec · 964b0ec
1 parent e72eeb6
commit 964b0ec
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 115 deletions.
diff --git a/packages/trie/src/node/batch.ts b/packages/trie/src/node/batch.ts
@@ -1,12 +1,7 @@
 import { type BatchDBOp } from '@ethereumjs/util'
 
-import { type TrieNode } from '../types.js'
 import { bytesToNibbles, nibblesCompare } from '../util/nibbles.js'
 
-import { BranchNode } from './branch.js'
-
-import type { Trie } from '../trie.js'
-
 export function orderBatch(
   ops: BatchDBOp[],
   keyTransform: (msg: Uint8Array) => Uint8Array = (msg: Uint8Array) => {
@@ -23,63 +18,3 @@ export function orderBatch(
   })
   return keyNibbles.map(([i, _]) => ops[i])
 }
-
-export async function _batch(
-  trie: Trie,
-  ops: BatchDBOp[],
-  skipKeyTransform?: boolean
-): Promise<void> {
-  const keyTransform =
-    skipKeyTransform === true
-      ? undefined
-      : (msg: Uint8Array) => {
-          return trie['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 : trie['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') {
-      stack = await trie.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()))
-        }
-      }
-    } else if (op.type === 'del') {
-      stack = await trie.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()))
-        }
-      }
-    }
-  }
-}
diff --git a/packages/trie/src/trie.ts b/packages/trie/src/trie.ts
@@ -18,7 +18,7 @@ import debug from 'debug'
 import { keccak256 } from 'ethereum-cryptography/keccak.js'
 
 import { CheckpointDB } from './db/index.js'
-import { _batch } from './node/batch.js'
+import { orderBatch } from './node/batch.js'
 import {
   BranchNode,
   ExtensionNode,
@@ -1149,7 +1149,59 @@ export class Trie {
    * @param ops
    */
   async batch(ops: BatchDBOp[], skipKeyTransform?: boolean): Promise<void> {
-    await _batch(this, ops, skipKeyTransform)
+    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') {
+        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()))
+          }
+        }
+      } else if (op.type === 'del') {
+        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()))
+          }
+        }
+      }
+    }
   }
 
   // This method verifies if all keys in the trie (except the root) are reachable

diff --git a/packages/trie/test/trie/updateNode.spec.ts b/packages/trie/test/trie/updateNode.spec.ts
@@ -8,7 +8,7 @@ import {
 import { assert, describe, it } from 'vitest'
 
 import { Trie } from '../../src/index.js'
-import { _batch, orderBatch } from '../../src/node/batch.js'
+import { orderBatch } from '../../src/node/batch.js'
 
 import type { BatchDBOp } from '@ethereumjs/util'
 
@@ -126,7 +126,6 @@ describe('Large Batch Test', async () => {
   })
   const trie_0 = new Trie()
   const trie_1 = new Trie()
-  const trie_2 = new Trie()
 
   for (const op of batchOP) {
     if (op.type === 'put') {
@@ -136,13 +135,9 @@ describe('Large Batch Test', async () => {
 
   await trie_1.batch(batchOP)
 
-  await _batch(trie_2, batchOP)
-
   it('batch should work', () => {
     assert.notDeepEqual(trie_0.root(), trie_0.EMPTY_TRIE_ROOT, 'trie is not empty')
     assert.deepEqual(trie_0.root(), trie_1.root(), 'trie roots should match (v1)')
-    assert.deepEqual(trie_0.root(), trie_2.root(), 'trie roots should match (v2)')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'trie roots should match (v3)')
   })
 
   it('should remain same', async () => {
@@ -170,17 +165,13 @@ describe('Large Batch Test', async () => {
       await trie_0.del(k)
     }
     await trie_1.batch(deleteBatch)
-    await _batch(trie_2, deleteBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
 
     for (const k of putNull) {
       await trie_0.put(k, null)
     }
     await trie_1.batch(putNullBatch)
-    await _batch(trie_2, putNullBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
   })
 })
 describe('Large Batch Test (secure)', async () => {
@@ -194,30 +185,18 @@ describe('Large Batch Test (secure)', async () => {
   })
   const trie_0 = new Trie({ useKeyHashing: true })
   const trie_1 = new Trie({ useKeyHashing: true })
-  const trie_2 = new Trie({ useKeyHashing: true })
   const trie_3 = new Trie({})
 
   const toHash = Uint8Array.from([1, 2, 3, 4])
   const hash0 = trie_0['hash'](toHash)
   const hash1 = trie_1['hash'](toHash)
-  const hash2 = trie_2['hash'](toHash)
   const hash3 = trie_3['hash'](toHash)
   it('should have same hash', async () => {
     assert.deepEqual(hash0, hash1, 'hashes should match')
   })
-  it('should have same hash', async () => {
-    assert.deepEqual(hash0, hash2, 'hashes should match')
-  })
   it('should not have same hash', async () => {
     assert.deepEqual(hash0, hash3, 'hashes should not match')
   })
-  it('tries should start empty', async () => {
-    assert.equal(
-      bytesToHex(trie_1.root()),
-      bytesToHex(trie_2.root()),
-      'trie roots should match (v1)'
-    )
-  })
   it('tries should start empty', async () => {
     assert.deepEqual(
       bytesToHex(trie_0.root()),
@@ -247,15 +226,12 @@ describe('Large Batch Test (secure)', async () => {
     const found3 = await trie_3.get(key)
     assert.deepEqual(found3, value, 'value should be found in trie')
     await trie_1.batch(batchOP)
-    await _batch(trie_2, batchOP)
 
     const root0 = await trie_0.lookupNode(trie_0.root())
     assert.deepEqual(root0.value(), value, 'rootnode should be only node')
 
     const found1 = await trie_1.get(key)
     assert.deepEqual(found1, value, 'value should be found in trie')
-    const found2 = await trie_2.get(key)
-    assert.deepEqual(found2, value, 'value should be found in trie')
 
     assert.notEqual(
       bytesToHex(trie_0.root()),
@@ -267,21 +243,11 @@ describe('Large Batch Test (secure)', async () => {
       bytesToHex(trie_1.EMPTY_TRIE_ROOT),
       'trie is not empty'
     )
-    assert.notEqual(
-      bytesToHex(trie_2.root()),
-      bytesToHex(trie_2.EMPTY_TRIE_ROOT),
-      'trie is not empty'
-    )
     assert.notEqual(
       bytesToHex(trie_3.root()),
       bytesToHex(trie_3.EMPTY_TRIE_ROOT),
       'trie is not empty'
     )
-    assert.equal(
-      bytesToHex(trie_1.root()),
-      bytesToHex(trie_2.root()),
-      'trie roots should match (v3)'
-    )
     assert.notDeepEqual(
       bytesToHex(trie_0.root()),
       bytesToHex(trie_3.root()),
@@ -314,17 +280,13 @@ describe('Large Batch Test (secure)', async () => {
       await trie_0.del(k)
     }
     await trie_1.batch(deleteBatch)
-    await _batch(trie_2, deleteBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
 
     for (const k of putNull) {
       await trie_0.put(k, null)
     }
     await trie_1.batch(putNullBatch)
-    await _batch(trie_2, putNullBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
   })
 })
 describe('Large Batch Test (use node pruning)', async () => {
@@ -338,7 +300,6 @@ describe('Large Batch Test (use node pruning)', async () => {
   })
   const trie_0 = new Trie({ useNodePruning: true, useRootPersistence: true })
   const trie_1 = new Trie({ useNodePruning: true, useRootPersistence: true })
-  const trie_2 = new Trie({ useNodePruning: true, useRootPersistence: true })
 
   for (const op of batchOP) {
     if (op.type === 'put') {
@@ -348,13 +309,9 @@ describe('Large Batch Test (use node pruning)', async () => {
 
   await trie_1.batch(batchOP)
 
-  await _batch(trie_2, batchOP)
-
   it('batch should work', () => {
     assert.notDeepEqual(trie_0.root(), trie_0.EMPTY_TRIE_ROOT, 'trie is not empty')
     assert.deepEqual(trie_0.root(), trie_1.root(), 'trie roots should match (v1)')
-    assert.deepEqual(trie_0.root(), trie_2.root(), 'trie roots should match (v2)')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'trie roots should match (v3)')
   })
 
   it('should remain same', async () => {
@@ -382,16 +339,12 @@ describe('Large Batch Test (use node pruning)', async () => {
       await trie_0.del(k)
     }
     await trie_1.batch(deleteBatch)
-    await _batch(trie_2, deleteBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
 
     for (const k of putNull) {
       await trie_0.put(k, null)
     }
     await trie_1.batch(putNullBatch)
-    await _batch(trie_2, putNullBatch)
     assert.deepEqual(trie_0.root(), trie_1.root(), 'roots should match')
-    assert.deepEqual(trie_1.root(), trie_2.root(), 'roots should match')
   })
 })