Move tree key computation to verkle and simplify

packages/evm/package.json

@@ -57,6 +57,7 @@
     "@ethereumjs/statemanager": "^2.3.0",
     "@ethereumjs/tx": "^5.3.0",
     "@ethereumjs/util": "^9.0.3",
+    "@ethereumjs/verkle": "^0.0.2",
     "@types/debug": "^4.1.9",
     "debug": "^4.3.3",
     "ethereum-cryptography": "^2.1.3",

packages/evm/src/opcodes/functions.ts

@@ -1,4 +1,3 @@
-import { getTreeIndexesForStorageSlot } from '@ethereumjs/statemanager'
 import {
   Account,
   Address,
@@ -32,6 +31,7 @@ import {
   setLengthLeft,
   setLengthRight,
 } from '@ethereumjs/util'
+import { getTreeIndexesForStorageSlot } from '@ethereumjs/verkle'
 import { keccak256 } from 'ethereum-cryptography/keccak.js'
 
 import { ERROR } from '../exceptions.js'

packages/evm/src/opcodes/gas.ts

@@ -1,11 +1,4 @@
 import { Hardfork } from '@ethereumjs/common'
-import {
-  BALANCE_LEAF_KEY,
-  CODE_KECCAK_LEAF_KEY,
-  CODE_SIZE_LEAF_KEY,
-  VERSION_LEAF_KEY,
-  getTreeIndexesForStorageSlot,
-} from '@ethereumjs/statemanager'
 import {
   Account,
   Address,
@@ -17,6 +10,13 @@ import {
   bigIntToBytes,
   setLengthLeft,
 } from '@ethereumjs/util'
+import {
+  BALANCE_LEAF_KEY,
+  CODE_KECCAK_LEAF_KEY,
+  CODE_SIZE_LEAF_KEY,
+  VERSION_LEAF_KEY,
+  getTreeIndexesForStorageSlot,
+} from '@ethereumjs/verkle'
 
 import { ERROR } from '../exceptions.js'
 

packages/statemanager/src/accessWitness.ts

@@ -1,5 +1,18 @@
-import { BIGINT_0, bytesToBigInt, bytesToHex, hexToBytes, toBytes } from '@ethereumjs/util'
-import { getKey, getStem } from '@ethereumjs/verkle'
+import { BIGINT_0, bytesToBigInt, bytesToHex, hexToBytes, intToBytes } from '@ethereumjs/util'
+import {
+  BALANCE_LEAF_KEY,
+  CODE_KECCAK_LEAF_KEY,
+  CODE_OFFSET,
+  CODE_SIZE_LEAF_KEY,
+  HEADER_STORAGE_OFFSET,
+  MAIN_STORAGE_OFFSET,
+  NONCE_LEAF_KEY,
+  VERKLE_NODE_WIDTH,
+  VERSION_LEAF_KEY,
+  getKey,
+  getStem,
+  getTreeIndicesForCodeChunk,
+} from '@ethereumjs/verkle'
 import debugDefault from 'debug'
 
 import type { Address, PrefixedHexString } from '@ethereumjs/util'
@@ -11,17 +24,6 @@
 /**
  * Tree key constants.
  */
-export const VERSION_LEAF_KEY = toBytes(0)
-export const BALANCE_LEAF_KEY = toBytes(1)
-export const NONCE_LEAF_KEY = toBytes(2)
-export const CODE_KECCAK_LEAF_KEY = toBytes(3)
-export const CODE_SIZE_LEAF_KEY = toBytes(4)
-
-export const HEADER_STORAGE_OFFSET = 64
-export const CODE_OFFSET = 128
-export const VERKLE_NODE_WIDTH = 256
-export const MAIN_STORAGE_OFFSET = BigInt(256) ** BigInt(31)
-
 const WitnessBranchReadCost = BigInt(1900)
 const WitnessChunkReadCost = BigInt(200)
 const WitnessBranchWriteCost = BigInt(3000)
@@ -270,7 +272,10 @@
       this.stems.set(accessedStemHex, accessedStem)
     }
 
-    const accessedChunkKey = getKey(accessedStemKey, toBytes(subIndex))
+    const accessedChunkKey = getKey(
+      accessedStemKey,
+      typeof subIndex === 'number' ? intToBytes(subIndex) : subIndex
+    )
     const accessedChunkKeyHex = bytesToHex(accessedChunkKey)
     let accessedChunk = this.chunks.get(accessedChunkKeyHex)
     if (accessedChunk === undefined) {
@@ -353,29 +358,6 @@
   }
 }
 
-export function getTreeIndexesForStorageSlot(storageKey: bigint): {
-  treeIndex: bigint
-  subIndex: number
-} {
-  let position: bigint
-  if (storageKey < CODE_OFFSET - HEADER_STORAGE_OFFSET) {
-    position = BigInt(HEADER_STORAGE_OFFSET) + storageKey
-  } else {
-    position = MAIN_STORAGE_OFFSET + storageKey
-  }
-
-  const treeIndex = position / BigInt(VERKLE_NODE_WIDTH)
-  const subIndex = Number(position % BigInt(VERKLE_NODE_WIDTH))
-
-  return { treeIndex, subIndex }
-}
-
-export function getTreeIndicesForCodeChunk(chunkId: number) {
-  const treeIndex = Math.floor((CODE_OFFSET + chunkId) / VERKLE_NODE_WIDTH)
-  const subIndex = (CODE_OFFSET + chunkId) % VERKLE_NODE_WIDTH
-  return { treeIndex, subIndex }
-}
-
 export function decodeAccessedState(treeIndex: number | bigint, chunkIndex: number): AccessedState {
   const position = BigInt(treeIndex) * BigInt(VERKLE_NODE_WIDTH) + BigInt(chunkIndex)
   switch (position) {

packages/statemanager/src/statelessVerkleStateManager.ts

@@ -13,23 +13,18 @@
   short,
   toBytes,
 } from '@ethereumjs/util'
-import { getKey, getStem } from '@ethereumjs/verkle'
+import {
+  LeafType,
+  getKey,
+  getStem,
+  getTreeKeyForCodeChunk,
+  getTreeKeyForStorageSlot,
+} from '@ethereumjs/verkle'
 import debugDefault from 'debug'
 import { keccak256 } from 'ethereum-cryptography/keccak.js'
 import { loadVerkleCrypto } from 'verkle-cryptography-wasm'
 
-import {
-  AccessWitness,
-  AccessedStateType,
-  BALANCE_LEAF_KEY,
-  CODE_KECCAK_LEAF_KEY,
-  CODE_SIZE_LEAF_KEY,
-  NONCE_LEAF_KEY,
-  VERSION_LEAF_KEY,
-  decodeValue,
-  getTreeIndexesForStorageSlot,
-  getTreeIndicesForCodeChunk,
-} from './accessWitness.js'
+import { AccessWitness, AccessedStateType, decodeValue } from './accessWitness.js'
 import { AccountCache, CacheType, CodeCache, StorageCache } from './cache/index.js'
 import { OriginalStorageCache } from './cache/originalStorageCache.js'
 
@@ -318,50 +313,9 @@
     debug('initVerkleExecutionWitness postState', this._postState)
   }
 
-  getTreeKeyForVersion(stem: Uint8Array) {
-    return getKey(stem, VERSION_LEAF_KEY)
-  }
-
-  getTreeKeyForBalance(stem: Uint8Array) {
-    return getKey(stem, BALANCE_LEAF_KEY)
-  }
-
-  getTreeKeyForNonce(stem: Uint8Array) {
-    return getKey(stem, NONCE_LEAF_KEY)
-  }
-
-  getTreeKeyForCodeHash(stem: Uint8Array) {
-    return getKey(stem, CODE_KECCAK_LEAF_KEY)
-  }
-
-  getTreeKeyForCodeSize(stem: Uint8Array) {
-    return getKey(stem, CODE_SIZE_LEAF_KEY)
-  }
-
-  async getTreeKeyForCodeChunk(address: Address, chunkId: number) {
-    const { treeIndex, subIndex } = getTreeIndicesForCodeChunk(chunkId)
-    return getKey(getStem(this.verkleCrypto, address, treeIndex), toBytes(subIndex))
-  }
-
-  chunkifyCode(code: Uint8Array) {
-    // Pad code to multiple of 31 bytes
-    if (code.length % 31 !== 0) {
-      const paddingLength = 31 - (code.length % 31)
-      code = setLengthRight(code, code.length + paddingLength)
-    }
-
-    throw new Error('Not implemented')
-  }
-
-  async getTreeKeyForStorageSlot(address: Address, storageKey: bigint) {
-    const { treeIndex, subIndex } = getTreeIndexesForStorageSlot(storageKey)
-
-    return getKey(getStem(this.verkleCrypto, address, treeIndex), toBytes(subIndex))
-  }
-
   async checkChunkWitnessPresent(address: Address, codeOffset: number) {
     const chunkId = codeOffset / 31
-    const chunkKey = bytesToHex(await this.getTreeKeyForCodeChunk(address, chunkId))
+    const chunkKey = bytesToHex(await getTreeKeyForCodeChunk(address, chunkId, this.verkleCrypto))
     return this._state[chunkKey] !== undefined
   }
 
@@ -432,7 +386,7 @@
 
     const chunks = Math.floor(codeSize / 31) + 1
     for (let chunkId = 0; chunkId < chunks; chunkId++) {
-      const chunkKey = bytesToHex(await this.getTreeKeyForCodeChunk(address, chunkId))
+      const chunkKey = bytesToHex(await getTreeKeyForCodeChunk(address, chunkId, this.verkleCrypto))
       const codeChunk = this._state[chunkKey]
       if (codeChunk === null) {
         const errorMsg = `Invalid access to a non existent code chunk with chunkKey=${chunkKey}`
@@ -504,7 +458,11 @@
       }
     }
 
-    const storageKey = await this.getTreeKeyForStorageSlot(address, BigInt(bytesToHex(key)))
+    const storageKey = await getTreeKeyForStorageSlot(
+      address,
+      BigInt(bytesToHex(key)),
+      this.verkleCrypto
+    )
     const storageValue = toBytes(this._state[bytesToHex(storageKey)])
 
     if (!this._storageCacheSettings.deactivate) {
@@ -526,7 +484,11 @@
       this._storageCache!.put(address, key, value)
     } else {
       // TODO: Consider refactoring this in a writeContractStorage function? Like in stateManager.ts
-      const storageKey = await this.getTreeKeyForStorageSlot(address, BigInt(bytesToHex(key)))
+      const storageKey = await getTreeKeyForStorageSlot(
+        address,
+        BigInt(bytesToHex(key)),
+        this.verkleCrypto
+      )
       this._state[bytesToHex(storageKey)] = bytesToHex(setLengthRight(value, 32))
     }
   }
@@ -539,7 +501,7 @@
    */
   async clearContractStorage(address: Address): Promise<void> {
     const stem = getStem(this.verkleCrypto, address, 0)
-    const codeHashKey = this.getTreeKeyForCodeHash(stem)
+    const codeHashKey = getKey(stem, LeafType.CodeKeccak)
     this._storageCache?.clearContractStorage(address)
     // Update codeHash to `c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470`
     this._state[bytesToHex(codeHashKey)] = KECCAK256_NULL_S
@@ -556,11 +518,11 @@
     }
 
     const stem = getStem(this.verkleCrypto, address, 0)
-    const versionKey = this.getTreeKeyForVersion(stem)
-    const balanceKey = this.getTreeKeyForBalance(stem)
-    const nonceKey = this.getTreeKeyForNonce(stem)
-    const codeHashKey = this.getTreeKeyForCodeHash(stem)
-    const codeSizeKey = this.getTreeKeyForCodeSize(stem)
+    const versionKey = getKey(stem, LeafType.Version)
+    const balanceKey = getKey(stem, LeafType.Balance)
+    const nonceKey = getKey(stem, LeafType.Nonce)
+    const codeHashKey = getKey(stem, LeafType.CodeKeccak)
+    const codeSizeKey = getKey(stem, LeafType.CodeSize)
 
     const versionRaw = this._state[bytesToHex(versionKey)]
     const balanceRaw = this._state[bytesToHex(balanceKey)]
@@ -643,9 +605,9 @@
 
     if (this._accountCacheSettings.deactivate) {
       const stem = getStem(this.verkleCrypto, address, 0)
-      const balanceKey = this.getTreeKeyForBalance(stem)
-      const nonceKey = this.getTreeKeyForNonce(stem)
-      const codeHashKey = this.getTreeKeyForCodeHash(stem)
+      const balanceKey = getKey(stem, LeafType.Balance)
+      const nonceKey = getKey(stem, LeafType.Nonce)
+      const codeHashKey = getKey(stem, LeafType.CodeKeccak)
 
       const balanceBuf = setLengthRight(bigIntToBytes(account.balance, true), 32)
       const nonceBuf = setLengthRight(bigIntToBytes(account.nonce, true), 32)

packages/statemanager/test/statelessVerkleStateManager.spec.ts

@@ -9,7 +9,7 @@ import {
   hexToBytes,
   randomBytes,
 } from '@ethereumjs/util'
-import { getStem } from '@ethereumjs/verkle'
+import { LeafType, getKey, getStem } from '@ethereumjs/verkle'
 import { loadVerkleCrypto } from 'verkle-cryptography-wasm'
 import { assert, beforeAll, describe, it, test } from 'vitest'
 
@@ -109,16 +109,16 @@ describe('StatelessVerkleStateManager: Kaustinen Verkle Block', () => {
     )
   })
 
-  it('getTreeKeyFor* functions', async () => {
+  it('getKey function', async () => {
     const stateManager = await StatelessVerkleStateManager.create({ common, verkleCrypto })
     stateManager.initVerkleExecutionWitness(block.header.number, block.executionWitness)
 
     const address = Address.fromString('0x6177843db3138ae69679a54b95cf345ed759450d')
     const stem = getStem(stateManager.verkleCrypto, address, 0n)
 
-    const balanceKey = stateManager.getTreeKeyForBalance(stem)
-    const nonceKey = stateManager.getTreeKeyForNonce(stem)
-    const codeHashKey = stateManager.getTreeKeyForCodeHash(stem)
+    const balanceKey = getKey(stem, LeafType.Balance)
+    const nonceKey = getKey(stem, LeafType.Nonce)
+    const codeHashKey = getKey(stem, LeafType.CodeKeccak)
 
     const balanceRaw = stateManager['_state'][bytesToHex(balanceKey)]
     const nonceRaw = stateManager['_state'][bytesToHex(nonceKey)]

packages/verkle/src/index.ts

@@ -2,4 +2,5 @@ export * from './db/index.js'
 export * from './node/index.js'
 export * from './types.js'
 export * from './util/index.js'
+export * from './util/keys.js'
 export * from './verkleTree.js'

packages/verkle/src/types.ts

@@ -1,4 +1,4 @@
-import { utf8ToBytes } from '@ethereumjs/util'
+import { intToBytes, utf8ToBytes } from '@ethereumjs/util'
 
 import type { VerkleNode } from './node/index.js'
 import type { WalkController } from './util/walkController.js'
@@ -123,3 +123,22 @@ export type FoundNodeFunction = (
 export const ROOT_DB_KEY = utf8ToBytes('__root__')
 
 export type VerkleCrypto = VerkleFFI
+
+export enum LeafType {
+  Version = 0,
+  Balance = 1,
+  Nonce = 2,
+  CodeKeccak = 3,
+  CodeSize = 4,
+}
+
+export const VERSION_LEAF_KEY = intToBytes(LeafType.Version)
+export const BALANCE_LEAF_KEY = intToBytes(LeafType.Balance)
+export const NONCE_LEAF_KEY = intToBytes(LeafType.Nonce)
+export const CODE_KECCAK_LEAF_KEY = intToBytes(LeafType.CodeKeccak)
+export const CODE_SIZE_LEAF_KEY = intToBytes(LeafType.CodeSize)
+
+export const HEADER_STORAGE_OFFSET = 64
+export const CODE_OFFSET = 128
+export const VERKLE_NODE_WIDTH = 256
+export const MAIN_STORAGE_OFFSET = BigInt(256) ** BigInt(31)

packages/verkle/src/util/crypto.ts

@@ -1,7 +1,6 @@
 import {
   type Address,
   bigIntToBytes,
-  concatBytes,
   int32ToBytes,
   setLengthLeft,
   setLengthRight,
@@ -36,17 +35,4 @@ export function getStem(
   return treeStem
 }
 
-/**
- * @dev Returns the tree key for a given verkle tree stem, and sub index.
- * @dev Assumes that the verkle node width = 256
- * @param stem The 31-bytes verkle tree stem as a Uint8Array.
- * @param subIndex The sub index of the tree to generate the key for as a Uint8Array.
- * @return The tree key as a Uint8Array.
- */
-
-export function getKey(stem: Uint8Array, subIndex: Uint8Array): Uint8Array {
-  const treeKey = concatBytes(stem, subIndex)
-  return treeKey
-}
-
 export const POINT_IDENTITY = new Uint8Array(0)