core, trie: cache the code chunking in contract

core/state/statedb.go

@@ -523,12 +523,9 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 			}
 
 			if obj.dirtyCode {
-				if chunks, err := trie.ChunkifyCode(obj.code); err == nil {
-					for i := 0; i < len(chunks); i += 32 {
-						s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(obj.pointEval, uint256.NewInt(uint64(i)/32)), chunks[i:i+32])
-					}
-				} else {
-					s.setError(err)
+				chunks := trie.ChunkifyCode(obj.code)
+				for i := 0; i < len(chunks); i += 32 {
+					s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(obj.pointEval, uint256.NewInt(uint64(i)/32)), chunks[i:i+32])
 				}
 			}
 		} else {
@@ -1019,12 +1016,9 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 			// Write any contract code associated with the state object
 			if obj.code != nil && obj.dirtyCode {
 				if s.trie.IsVerkle() {
-					if chunks, err := trie.ChunkifyCode(obj.code); err == nil {
-						for i := 0; i < len(chunks); i += 32 {
-							s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(obj.pointEval, uint256.NewInt(uint64(i)/32)), chunks[i:32+i])
-						}
-					} else {
-						s.setError(err)
+					chunks := trie.ChunkifyCode(obj.code)
+					for i := 0; i < len(chunks); i += 32 {
+						s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(obj.pointEval, uint256.NewInt(uint64(i)/32)), chunks[i:32+i])
 					}
 				}
 				rawdb.WriteCode(codeWriter, common.BytesToHash(obj.CodeHash()), obj.code)

core/state_processor_test.go

@@ -440,7 +440,7 @@ func TestProcessVerkle(t *testing.T) {
 			t.Fatalf("expected block %d to be present in chain", i+1)
 		}
 		if b.GasUsed() != blockGasUsagesExpected[i] {
-			t.Fatalf("expected block txs to use %d, got %d\n", blockGasUsagesExpected[i], b.GasUsed())
+			t.Fatalf("expected block txs to use %d, got %d in block %d\n", blockGasUsagesExpected[i], b.GasUsed(), i)
 		}
 	}
 }

core/types/access_witness.go

@@ -56,14 +56,20 @@ type AccessWitness struct {
 	// Caches all the points that correspond to an address,
 	// so they are not recalculated.
 	addrToPoint map[string]*verkle.Point
+
+	// Caches which code chunks have been accessed, in order
+	// to reduce the number of times that GetTreeKeyCodeChunk
+	// is called.
+	CodeLocations map[string]map[uint64]struct{}
 }
 
 func NewAccessWitness() *AccessWitness {
 	return &AccessWitness{
-		Branches:     make(map[VerkleStem]Mode),
-		Chunks:       make(map[common.Hash]Mode),
-		InitialValue: make(map[string][]byte),
-		addrToPoint:  make(map[string]*verkle.Point),
+		Branches:      make(map[VerkleStem]Mode),
+		Chunks:        make(map[common.Hash]Mode),
+		InitialValue:  make(map[string][]byte),
+		addrToPoint:   make(map[string]*verkle.Point),
+		CodeLocations: make(map[string]map[uint64]struct{}),
 	}
 }
 
@@ -84,6 +90,31 @@ func (aw *AccessWitness) SetLeafValue(addr []byte, value []byte) {
 	}
 }
 
+func (aw *AccessWitness) HasCodeChunk(addr []byte, chunknr uint64) bool {
+	if locs, ok := aw.CodeLocations[string(addr)]; ok {
+		if _, ok = locs[chunknr]; ok {
+			return true
+		}
+	}
+
+	return false
+}
+
+// SetCodeLeafValue does the same thing as SetLeafValue, but for code chunks. It
+// maintains a cache of which (address, chunk) were calculated, in order to avoid
+// calling GetTreeKey more than once per chunk.
+func (aw *AccessWitness) SetCachedCodeChunk(addr []byte, chunknr uint64) {
+	if locs, ok := aw.CodeLocations[string(addr)]; ok {
+		if _, ok = locs[chunknr]; ok {
+			return
+		}
+	} else {
+		aw.CodeLocations[string(addr)] = map[uint64]struct{}{}
+	}
+
+	aw.CodeLocations[string(addr)][chunknr] = struct{}{}
+}
+
 func (aw *AccessWitness) touchAddressOnWrite(addr []byte) (bool, bool, bool) {
 	var stem VerkleStem
 	var stemWrite, chunkWrite, chunkFill bool

core/vm/contract.go

@@ -20,6 +20,7 @@ import (
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/gballet/go-verkle"
 	"github.com/holiman/uint256"
@@ -57,6 +58,7 @@ type Contract struct {
 	analysis  bitvec                 // Locally cached result of JUMPDEST analysis
 
 	Code     []byte
+	Chunks   trie.ChunkedCode
 	CodeHash common.Hash
 	CodeAddr *common.Address
 	Input    []byte

core/vm/gas_table.go

@@ -89,10 +89,10 @@ func memoryCopierGas(stackpos int) gasFunc {
 }
 
 var (
-	gasCallDataCopy        = memoryCopierGas(2)
-	gasCodeCopyStateful    = memoryCopierGas(2)
-	gasExtCodeCopyStateful = memoryCopierGas(3)
-	gasReturnDataCopy      = memoryCopierGas(2)
+	gasCallDataCopy   = memoryCopierGas(2)
+	gasCodeCopy       = memoryCopierGas(2)
+	gasExtCodeCopy    = memoryCopierGas(3)
+	gasReturnDataCopy = memoryCopierGas(2)
 )
 
 func gasExtCodeSize(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
@@ -106,56 +106,6 @@ func gasExtCodeSize(evm *EVM, contract *Contract, stack *Stack, mem *Memory, mem
 	return usedGas, nil
 }
 
-func gasCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
-	var statelessGas uint64
-	if evm.chainConfig.IsCancun(evm.Context.BlockNumber) {
-		var (
-			codeOffset = stack.Back(1)
-			length     = stack.Back(2)
-		)
-		uint64CodeOffset, overflow := codeOffset.Uint64WithOverflow()
-		if overflow {
-			uint64CodeOffset = 0xffffffffffffffff
-		}
-		uint64Length, overflow := length.Uint64WithOverflow()
-		if overflow {
-			uint64Length = 0xffffffffffffffff
-		}
-		_, offset, nonPaddedSize := getDataAndAdjustedBounds(contract.Code, uint64CodeOffset, uint64Length)
-		statelessGas = touchEachChunksOnReadAndChargeGas(offset, nonPaddedSize, contract.AddressPoint(), nil, evm.Accesses, contract.IsDeployment)
-	}
-	usedGas, err := gasCodeCopyStateful(evm, contract, stack, mem, memorySize)
-	return usedGas + statelessGas, err
-}
-
-func gasExtCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
-	var statelessGas uint64
-	if evm.chainConfig.IsCancun(evm.Context.BlockNumber) {
-		var (
-			codeOffset = stack.Back(2)
-			length     = stack.Back(3)
-			targetAddr = stack.Back(0).Bytes20()
-		)
-		uint64CodeOffset, overflow := codeOffset.Uint64WithOverflow()
-		if overflow {
-			uint64CodeOffset = 0xffffffffffffffff
-		}
-		uint64Length, overflow := length.Uint64WithOverflow()
-		if overflow {
-			uint64Length = 0xffffffffffffffff
-		}
-		// note:  we must charge witness costs for the specified range regardless of whether it
-		// is in-bounds of the actual target account code.  This is because we must charge the cost
-		// before hitting the db to be able to now what the actual code size is.  This is different
-		// behavior from CODECOPY which only charges witness access costs for the part of the range
-		// which overlaps in the account code.  TODO: clarify this is desired behavior and amend the
-		// spec.
-		statelessGas = touchEachChunksOnReadAndChargeGasWithAddress(uint64CodeOffset, uint64Length, targetAddr[:], nil, evm.Accesses, contract.IsDeployment)
-	}
-	usedGas, err := gasExtCodeCopyStateful(evm, contract, stack, mem, memorySize)
-	return usedGas + statelessGas, err
-}
-
 func gasSLoad(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
 	usedGas := uint64(0)
 

core/vm/instructions.go

@@ -26,7 +26,6 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/trie"
 	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
-	"github.com/gballet/go-verkle"
 	"github.com/holiman/uint256"
 	"golang.org/x/crypto/sha3"
 )
@@ -378,15 +377,14 @@ func opCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([
 
 	paddedCodeCopy, copyOffset, nonPaddedCopyLength := getDataAndAdjustedBounds(scope.Contract.Code, uint64CodeOffset, length.Uint64())
 	if interpreter.evm.chainConfig.IsCancun(interpreter.evm.Context.BlockNumber) {
-		touchEachChunksOnReadAndChargeGas(copyOffset, nonPaddedCopyLength, scope.Contract.AddressPoint(), scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
+		scope.Contract.UseGas(touchEachChunksOnReadAndChargeGas(copyOffset, nonPaddedCopyLength, scope.Contract, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment))
 	}
 	scope.Memory.Set(memOffset.Uint64(), uint64(len(paddedCodeCopy)), paddedCodeCopy)
 	return nil, nil
 }
 
-func touchEachChunksOnReadAndChargeGasWithAddress(offset, size uint64, address, code []byte, accesses *types.AccessWitness, deployment bool) uint64 {
-	addrPoint := trieUtils.EvaluateAddressPoint(address)
-	return touchEachChunksOnReadAndChargeGas(offset, size, addrPoint, code, accesses, deployment)
+func touchEachChunksOnReadAndChargeGasWithAddress(offset, size uint64, contract *Contract, code []byte, accesses *types.AccessWitness, deployment bool) uint64 {
+	return touchEachChunksOnReadAndChargeGas(offset, size, contract, code, accesses, deployment)
 }
 
 // touchChunkOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
@@ -428,7 +426,7 @@ func touchChunkOnReadAndChargeGas(chunks trie.ChunkedCode, offset uint64, evals
 }
 
 // touchEachChunksOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
-func touchEachChunksOnReadAndChargeGas(offset, size uint64, addrPoint *verkle.Point, code []byte, accesses *types.AccessWitness, deployment bool) uint64 {
+func touchEachChunksOnReadAndChargeGas(offset, size uint64, contract *Contract, code []byte, accesses *types.AccessWitness, deployment bool) uint64 {
 	// note that in the case where the copied code is outside the range of the
 	// contract code but touches the last leaf with contract code in it,
 	// we don't include the last leaf of code in the AccessWitness.  The
@@ -447,28 +445,29 @@ func touchEachChunksOnReadAndChargeGas(offset, size uint64, addrPoint *verkle.Po
 	} else {
 		endOffset = offset + size
 	}
-	chunks, err := trie.ChunkifyCode(code)
-	if err != nil {
-		panic(err)
-	}
 
 	// endOffset - 1 since if the end offset is aligned on a chunk boundary,
 	// the last chunk should not be included.
 	for i := offset / 31; i <= (endOffset-1)/31; i++ {
-		index := trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(addrPoint, uint256.NewInt(i))
-
-		var overflow bool
-		statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, accesses.TouchAddressOnReadAndComputeGas(index))
-		if overflow {
-			panic("overflow when adding gas")
-		}
+		// only charge for+cache the chunk if it isn't already present
+		if !accesses.HasCodeChunk(contract.Address().Bytes(), i) {
+			index := trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(contract.AddressPoint(), uint256.NewInt(i))
+
+			var overflow bool
+			statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, accesses.TouchAddressOnReadAndComputeGas(index))
+			if overflow {
+				panic("overflow when adding gas")
+			}
 
-		if len(code) > 0 {
-			if deployment {
-				accesses.SetLeafValue(index[:], nil)
-			} else {
-				accesses.SetLeafValue(index[:], chunks[32*i:(i+1)*32])
+			if len(code) > 0 {
+				if deployment {
+					accesses.SetLeafValue(index[:], nil)
+				} else {
+					accesses.SetLeafValue(index[:], contract.Chunks[32*i:(i+1)*32])
+				}
 			}
+
+			accesses.SetCachedCodeChunk(contract.Address().Bytes(), i)
 		}
 	}
 
@@ -490,8 +489,13 @@ func opExtCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext)
 	addr := common.Address(a.Bytes20())
 	if interpreter.evm.chainConfig.IsCancun(interpreter.evm.Context.BlockNumber) {
 		code := interpreter.evm.StateDB.GetCode(addr)
+		contract := &Contract{
+			Code:   code,
+			Chunks: trie.ChunkifyCode(code),
+			self:   AccountRef(addr),
+		}
 		paddedCodeCopy, copyOffset, nonPaddedCopyLength := getDataAndAdjustedBounds(code, uint64CodeOffset, length.Uint64())
-		touchEachChunksOnReadAndChargeGasWithAddress(copyOffset, nonPaddedCopyLength, addr[:], code, interpreter.evm.Accesses, false)
+		touchEachChunksOnReadAndChargeGasWithAddress(copyOffset, nonPaddedCopyLength, contract, code, interpreter.evm.Accesses, false)
 		scope.Memory.Set(memOffset.Uint64(), length.Uint64(), paddedCodeCopy)
 	} else {
 		codeCopy := getData(interpreter.evm.StateDB.GetCode(addr), uint64CodeOffset, length.Uint64())
@@ -1009,7 +1013,7 @@ func opPush1(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]by
 		if interpreter.evm.chainConfig.IsCancun(interpreter.evm.Context.BlockNumber) && *pc%31 == 0 {
 			// touch next chunk if PUSH1 is at the boundary. if so, *pc has
 			// advanced past this boundary.
-			statelessGas := touchEachChunksOnReadAndChargeGas(*pc+1, uint64(1), scope.Contract.AddressPoint(), scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
+			statelessGas := touchEachChunksOnReadAndChargeGas(*pc+1, uint64(1), scope.Contract, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
 			scope.Contract.UseGas(statelessGas)
 		}
 	} else {
@@ -1034,7 +1038,7 @@ func makePush(size uint64, pushByteSize int) executionFunc {
 		}
 
 		if interpreter.evm.chainConfig.IsCancun(interpreter.evm.Context.BlockNumber) {
-			statelessGas := touchEachChunksOnReadAndChargeGas(uint64(startMin), uint64(pushByteSize), scope.Contract.AddressPoint(), scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
+			statelessGas := touchEachChunksOnReadAndChargeGas(uint64(startMin), uint64(pushByteSize), scope.Contract, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
 			scope.Contract.UseGas(statelessGas)
 		}
 

core/vm/interpreter.go

@@ -158,7 +158,6 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 		logged  bool   // deferred EVMLogger should ignore already logged steps
 		res     []byte // result of the opcode execution function
 
-		chunks     trie.ChunkedCode
 		chunkEvals [][]byte
 	)
 	// Don't move this deferred function, it's placed before the capturestate-deferred method,
@@ -183,7 +182,7 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 
 	// Evaluate one address per group of 256, 31-byte chunks
 	if in.evm.ChainConfig().IsCancun(in.evm.Context.BlockNumber) && !contract.IsDeployment {
-		chunks, err = trie.ChunkifyCode(contract.Code)
+		contract.Chunks = trie.ChunkifyCode(contract.Code)
 
 		totalEvals := len(contract.Code) / 31 / 256
 		if len(contract.Code)%(256*31) != 0 {
@@ -206,10 +205,10 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 			logged, pcCopy, gasCopy = false, pc, contract.Gas
 		}
 
-		if chunks != nil {
+		if contract.Chunks != nil {
 			// if the PC ends up in a new "chunk" of verkleized code, charge the
 			// associated costs.
-			contract.Gas -= touchChunkOnReadAndChargeGas(chunks, pc, chunkEvals, contract.Code, in.evm.TxContext.Accesses, contract.IsDeployment)
+			contract.Gas -= touchChunkOnReadAndChargeGas(contract.Chunks, pc, chunkEvals, contract.Code, in.evm.TxContext.Accesses, contract.IsDeployment)
 		}
 
 		// Get the operation from the jump table and validate the stack to ensure there are

trie/verkle.go

@@ -284,7 +284,7 @@ const (
 )
 
 // ChunkifyCode generates the chunked version of an array representing EVM bytecode
-func ChunkifyCode(code []byte) (ChunkedCode, error) {
+func ChunkifyCode(code []byte) ChunkedCode {
 	var (
 		chunkOffset = 0 // offset in the chunk
 		chunkCount  = len(code) / 31
@@ -331,5 +331,5 @@ func ChunkifyCode(code []byte) (ChunkedCode, error) {
 		}
 	}
 
-	return chunks, nil
+	return chunks
 }