From d787f40b63a76084fc6d08f6484a67cfa4e5560f Mon Sep 17 00:00:00 2001
From: Gabriel Rocheleau <contact@rockwaterweb.com>
Date: Tue, 16 Jul 2024 14:57:24 -0400
Subject: [PATCH] verkle: use verkle code chunk size constant

---
 .../src/statelessVerkleStateManager.ts        | 19 ++++++++++++-------
 packages/util/src/verkle.ts                   |  9 +++++----
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/packages/statemanager/src/statelessVerkleStateManager.ts b/packages/statemanager/src/statelessVerkleStateManager.ts
index 7c2eaebf5f..9525d915fe 100644
--- a/packages/statemanager/src/statelessVerkleStateManager.ts
+++ b/packages/statemanager/src/statelessVerkleStateManager.ts
@@ -2,6 +2,7 @@ import {
   Account,
   KECCAK256_NULL,
   KECCAK256_NULL_S,
+  VERKLE_CODE_CHUNK_SIZE,
   VerkleLeafType,
   bigIntToBytes,
   bytesToHex,
@@ -309,7 +310,7 @@ export class StatelessVerkleStateManager implements EVMStateManagerInterface {
   }
 
   async checkChunkWitnessPresent(address: Address, codeOffset: number) {
-    const chunkId = codeOffset / 31
+    const chunkId = codeOffset / VERKLE_CODE_CHUNK_SIZE
     const chunkKey = bytesToHex(
       await getVerkleTreeKeyForCodeChunk(address, chunkId, this.verkleCrypto)
     )
@@ -379,9 +380,9 @@ export class StatelessVerkleStateManager implements EVMStateManagerInterface {
     // allocate the code and copy onto it from the available witness chunks
     const codeSize = account.codeSize
     // allocate enough to fit the last chunk
-    const accessedCode = new Uint8Array(codeSize + 31)
+    const accessedCode = new Uint8Array(codeSize + VERKLE_CODE_CHUNK_SIZE)
 
-    const chunks = Math.floor(codeSize / 31) + 1
+    const chunks = Math.floor(codeSize / VERKLE_CODE_CHUNK_SIZE) + 1
     for (let chunkId = 0; chunkId < chunks; chunkId++) {
       const chunkKey = bytesToHex(
         await getVerkleTreeKeyForCodeChunk(address, chunkId, this.verkleCrypto)
@@ -393,7 +394,7 @@ export class StatelessVerkleStateManager implements EVMStateManagerInterface {
         throw Error(errorMsg)
       }
 
-      const codeOffset = chunkId * 31
+      const codeOffset = chunkId * VERKLE_CODE_CHUNK_SIZE
       // if code chunk was accessed as per the provided witnesses copy it over
       if (codeChunk !== undefined) {
         // actual code starts from index 1 in chunk, 0th index is if there are any push data bytes
@@ -402,7 +403,7 @@ export class StatelessVerkleStateManager implements EVMStateManagerInterface {
       } else {
         // else fill this unaccessed segment with invalid opcode since the evm execution shouldn't
         // end up here
-        accessedCode.fill(0xfe, codeOffset, 31)
+        accessedCode.fill(0xfe, codeOffset, VERKLE_CODE_CHUNK_SIZE)
       }
     }
 
@@ -808,8 +809,12 @@ export class StatelessVerkleStateManager implements EVMStateManagerInterface {
         // we can only compare the actual code because to compare the first byte would
         // be very tricky and impossible in certain scenarios like when the previous code chunk
         // was not accessed and hence not even provided in the witness
-        const chunkSize = 31
-        return bytesToHex(setLengthRight(code.slice(codeOffset, codeOffset + chunkSize), chunkSize))
+        return bytesToHex(
+          setLengthRight(
+            code.slice(codeOffset, codeOffset + VERKLE_CODE_CHUNK_SIZE),
+            VERKLE_CODE_CHUNK_SIZE
+          )
+        )
       }
 
       case AccessedStateType.Storage: {
diff --git a/packages/util/src/verkle.ts b/packages/util/src/verkle.ts
index eeff673509..9f1ce2b569 100644
--- a/packages/util/src/verkle.ts
+++ b/packages/util/src/verkle.ts
@@ -146,10 +146,11 @@ export const VERKLE_BALANCE_BYTES_LENGTH = 16
 export const VERKLE_BASIC_DATA_LEAF_KEY = intToBytes(VerkleLeafType.BasicData)
 export const VERKLE_CODE_HASH_LEAF_KEY = intToBytes(VerkleLeafType.CodeHash)
 
+export const VERKLE_CODE_CHUNK_SIZE = 31
 export const VERKLE_HEADER_STORAGE_OFFSET = 64
 export const VERKLE_CODE_OFFSET = 128
 export const VERKLE_NODE_WIDTH = 256
-export const VERKLE_MAIN_STORAGE_OFFSET = BigInt(256) ** BigInt(31)
+export const VERKLE_MAIN_STORAGE_OFFSET = BigInt(256) ** BigInt(VERKLE_CODE_CHUNK_SIZE)
 
 /**
  * @dev Returns the tree key for a given verkle tree stem, and sub index.
@@ -203,9 +204,9 @@ export const getVerkleTreeKeyForCodeChunk = async (
 }
 
 export const chunkifyCode = (code: Uint8Array) => {
-  // Pad code to multiple of 31 bytes
-  if (code.length % 31 !== 0) {
-    const paddingLength = 31 - (code.length % 31)
+  // Pad code to multiple of VERKLE_CODE_CHUNK_SIZE bytes
+  if (code.length % VERKLE_CODE_CHUNK_SIZE !== 0) {
+    const paddingLength = VERKLE_CODE_CHUNK_SIZE - (code.length % VERKLE_CODE_CHUNK_SIZE)
     code = setLengthRight(code, code.length + paddingLength)
   }