Refactor witness-accumulation in EVM (#42)

* make push dynamically-charged.  charge witness gas costs for push.  refactor evm witness gas charging to move logic for touching a range of bytecode into a helper method 'touchEachChunksAndChargeGas'

* add witness gas calculation for CodeCopy, ExtCodeCopy, SLoad back to gas_table.go

* witness gas charging for CALL

* remove explicit reference to evm.TxContext

* core/vm: make touchEachChunksAndCharge gas handle nil code value

* core/vm: call implementation, separate out witnesses into touch/set

* some fixes

* remove witness touching from opCall: this will go in evm.go

* remove witness touching for call from gas_table.go

* (hopefully) fix tests

* add SSTORE witness charging that was removed mistakenly

* charge witness gas for call

* clean up and comment touchEachChunksAndChargeGas

* make suggested changes

* address remaining points

* fix build issues

* remove double-charging for contract creation witness gas charging

consensus/ethash/consensus.go

@@ -661,14 +661,20 @@ func accumulateRewards(config *params.ChainConfig, state *state.StateDB, header
 		r.Sub(r, header.Number)
 		r.Mul(r, blockReward)
 		r.Div(r, big8)
+
+		if state.Witness() != nil {
 			uncleCoinbase := utils.GetTreeKeyBalance(uncle.Coinbase.Bytes())
 			state.Witness().TouchAddress(uncleCoinbase, state.GetBalance(uncle.Coinbase).Bytes())
+		}
 		state.AddBalance(uncle.Coinbase, r)
 
 		r.Div(blockReward, big32)
 		reward.Add(reward, r)
 	}
 	coinbase := utils.GetTreeKeyBalance(header.Coinbase.Bytes())
+
+	if state.Witness() != nil {
 		state.Witness().TouchAddress(coinbase, state.GetBalance(header.Coinbase).Bytes())
+	}
 	state.AddBalance(header.Coinbase, reward)
 }

core/state_processor.go

@@ -128,7 +128,9 @@ func applyTransaction(msg types.Message, config *params.ChainConfig, bc ChainCon
 		receipt.ContractAddress = crypto.CreateAddress(evm.TxContext.Origin, tx.Nonce())
 	}
 
+	if config.IsCancun(blockNumber) {
 		statedb.Witness().Merge(txContext.Accesses)
+	}
 
 	// Set the receipt logs and create the bloom filter.
 	receipt.Logs = statedb.GetLogs(tx.Hash(), blockHash)

core/state_transition.go

@@ -304,26 +304,26 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 	if st.gas < gas {
 		return nil, fmt.Errorf("%w: have %d, want %d", ErrIntrinsicGas, st.gas, gas)
 	}
-	if st.evm.TxContext.Accesses != nil {
+	if st.evm.Accesses != nil {
 		if msg.To() != nil {
 			toBalance := trieUtils.GetTreeKeyBalance(msg.To().Bytes())
 			pre := st.state.GetBalance(*msg.To())
-			gas += st.evm.TxContext.Accesses.TouchAddressAndChargeGas(toBalance, pre.Bytes())
+			gas += st.evm.Accesses.TouchAddressAndChargeGas(toBalance, pre.Bytes())
 
 			// NOTE: Nonce also needs to be charged, because it is needed for execution
 			// on the statless side.
 			var preTN [8]byte
 			fromNonce := trieUtils.GetTreeKeyNonce(msg.To().Bytes())
 			binary.BigEndian.PutUint64(preTN[:], st.state.GetNonce(*msg.To()))
-			gas += st.evm.TxContext.Accesses.TouchAddressAndChargeGas(fromNonce, preTN[:])
+			gas += st.evm.Accesses.TouchAddressAndChargeGas(fromNonce, preTN[:])
 		}
 		fromBalance := trieUtils.GetTreeKeyBalance(msg.From().Bytes())
 		preFB := st.state.GetBalance(msg.From()).Bytes()
 		fromNonce := trieUtils.GetTreeKeyNonce(msg.From().Bytes())
 		var preFN [8]byte
 		binary.BigEndian.PutUint64(preFN[:], st.state.GetNonce(msg.From()))
-		gas += st.evm.TxContext.Accesses.TouchAddressAndChargeGas(fromNonce, preFN[:])
+		gas += st.evm.Accesses.TouchAddressAndChargeGas(fromNonce, preFN[:])
-		gas += st.evm.TxContext.Accesses.TouchAddressAndChargeGas(fromBalance, preFB[:])
+		gas += st.evm.Accesses.TouchAddressAndChargeGas(fromBalance, preFB[:])
 	}
 	st.gas -= gas
 

core/vm/common.go

@@ -63,6 +63,18 @@ func getData(data []byte, start uint64, size uint64) []byte {
 	return common.RightPadBytes(data[start:end], int(size))
 }
 
+func getDataAndAdjustedBounds(data []byte, start uint64, size uint64) (codeCopyPadded []byte, actualStart uint64, sizeNonPadded uint64) {
+	length := uint64(len(data))
+	if start > length {
+		start = length
+	}
+	end := start + size
+	if end > length {
+		end = length
+	}
+	return common.RightPadBytes(data[start:end], int(size)), start, end
+}
+
 // toWordSize returns the ceiled word size required for memory expansion.
 func toWordSize(size uint64) uint64 {
 	if size > math.MaxUint64-31 {

core/vm/evm.go

@@ -125,8 +125,6 @@ type EVM struct {
 	// available gas is calculated in gasCall* according to the 63/64 rule and later
 	// applied in opCall*.
 	callGasTemp uint64
-
-	accesses map[common.Hash]common.Hash
 }
 
 // NewEVM returns a new EVM. The returned EVM is not thread safe and should
@@ -170,6 +168,19 @@ func (evm *EVM) Interpreter() *EVMInterpreter {
 	return evm.interpreter
 }
 
+// tryConsumeGas tries to subtract gas from gasPool, setting the result in gasPool
+// if subtracting more gas than remains in gasPool, set gasPool = 0 and return false
+// otherwise, do the subtraction setting the result in gasPool and return true
+func tryConsumeGas(gasPool *uint64, gas uint64) bool {
+	if *gasPool < gas {
+		*gasPool = 0
+		return false
+	}
+
+	*gasPool -= gas
+	return true
+}
+
 // Call executes the contract associated with the addr with the given input as
 // parameters. It also handles any necessary value transfer required and takes
 // the necessary steps to create accounts and reverses the state in case of an
@@ -232,6 +243,7 @@ func (evm *EVM) Call(caller ContractRef, addr common.Address, input []byte, gas
 		if len(code) == 0 {
 			ret, err = nil, nil // gas is unchanged
 		} else {
+			if evm.Accesses != nil {
 				// Touch the account data
 				var data [32]byte
 				evm.Accesses.TouchAddress(utils.GetTreeKeyVersion(addr.Bytes()), data[:])
@@ -241,6 +253,7 @@ func (evm *EVM) Call(caller ContractRef, addr common.Address, input []byte, gas
 				binary.BigEndian.PutUint64(data[:], uint64(len(code)))
 				evm.Accesses.TouchAddress(utils.GetTreeKeyCodeSize(addr[:]), data[:])
 				evm.Accesses.TouchAddress(utils.GetTreeKeyCodeKeccak(addr[:]), evm.StateDB.GetCodeHash(addr).Bytes())
+			}
 
 			addrCopy := addr
 			// If the account has no code, we can abort here

core/vm/gas_table.go

@@ -23,7 +23,6 @@ import (
 	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/params"
 	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
-	"github.com/holiman/uint256"
 )
 
 // memoryGasCost calculates the quadratic gas for memory expansion. It does so
@@ -88,17 +87,6 @@ func memoryCopierGas(stackpos int) gasFunc {
 	}
 }
 
-func gasExtCodeSize(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
-	usedGas := uint64(0)
-	slot := stack.Back(0)
-	if evm.accesses != nil {
-		index := trieUtils.GetTreeKeyCodeSize(slot.Bytes())
-		usedGas += evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
-	}
-
-	return usedGas, nil
-}
-
 var (
 	gasCallDataCopy        = memoryCopierGas(2)
 	gasCodeCopyStateful    = memoryCopierGas(2)
@@ -106,9 +94,20 @@ var (
 	gasReturnDataCopy      = memoryCopierGas(2)
 )
 
+func gasExtCodeSize(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
+	usedGas := uint64(0)
+	slot := stack.Back(0)
+	if evm.Accesses != nil {
+		index := trieUtils.GetTreeKeyCodeSize(slot.Bytes())
+		usedGas += evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
+	}
+
+	return usedGas, nil
+}
+
 func gasCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
 	var statelessGas uint64
-	if evm.accesses != nil {
+	if evm.Accesses != nil {
 		var (
 			codeOffset = stack.Back(1)
 			length     = stack.Back(2)
@@ -117,21 +116,12 @@ func gasCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memory
 		if overflow {
 			uint64CodeOffset = 0xffffffffffffffff
 		}
-		uint64CodeEnd, overflow := new(uint256.Int).Add(codeOffset, length).Uint64WithOverflow()
+		uint64Length, overflow := length.Uint64WithOverflow()
 		if overflow {
-			uint64CodeEnd = 0xffffffffffffffff
+			uint64Length = 0xffffffffffffffff
 		}
-		addr := contract.Address()
+		_, offset, nonPaddedSize := getDataAndAdjustedBounds(contract.Code, uint64CodeOffset, uint64Length)
-		chunk := uint64CodeOffset / 31
+		statelessGas = touchEachChunksAndChargeGas(offset, nonPaddedSize, contract.Address().Bytes()[:], nil, evm.Accesses)
-		endChunk := uint64CodeEnd / 31
-		// XXX uint64 overflow in condition check
-		for ; chunk < endChunk; chunk++ {
-
-			// TODO make a version of GetTreeKeyCodeChunk without the bigint
-			index := trieUtils.GetTreeKeyCodeChunk(addr[:], uint256.NewInt(chunk))
-			statelessGas += evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
-		}
-
 	}
 	usedGas, err := gasCodeCopyStateful(evm, contract, stack, mem, memorySize)
 	return usedGas + statelessGas, err
@@ -139,9 +129,8 @@ func gasCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memory
 
 func gasExtCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
 	var statelessGas uint64
-	if evm.accesses != nil {
+	if evm.Accesses != nil {
 		var (
-			a          = stack.Back(0)
 			codeOffset = stack.Back(2)
 			length     = stack.Back(3)
 		)
@@ -149,20 +138,17 @@ func gasExtCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, mem
 		if overflow {
 			uint64CodeOffset = 0xffffffffffffffff
 		}
-		uint64CodeEnd, overflow := new(uint256.Int).Add(codeOffset, length).Uint64WithOverflow()
+		uint64Length, overflow := length.Uint64WithOverflow()
 		if overflow {
-			uint64CodeEnd = 0xffffffffffffffff
+			uint64Length = 0xffffffffffffffff
 		}
-		addr := common.Address(a.Bytes20())
+		// note:  we must charge witness costs for the specified range regardless of whether it
-		chunk := uint64CodeOffset / 31
+		// is in-bounds of the actual target account code.  This is because we must charge the cost
-		endChunk := uint64CodeEnd / 31
+		// before hitting the db to be able to now what the actual code size is.  This is different
-		// XXX uint64 overflow in condition check
+		// behavior from CODECOPY which only charges witness access costs for the part of the range
-		for ; chunk < endChunk; chunk++ {
+		// which overlaps in the account code.  TODO: clarify this is desired behavior and amend the
-			// TODO(@gballet) make a version of GetTreeKeyCodeChunk without the bigint
+		// spec.
-			index := trieUtils.GetTreeKeyCodeChunk(addr[:], uint256.NewInt(chunk))
+		statelessGas = touchEachChunksAndChargeGas(uint64CodeOffset, uint64Length, nil, nil, evm.Accesses)
-			statelessGas += evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
-		}
-
 	}
 	usedGas, err := gasExtCodeCopyStateful(evm, contract, stack, mem, memorySize)
 	return usedGas + statelessGas, err
@@ -171,11 +157,11 @@ func gasExtCodeCopy(evm *EVM, contract *Contract, stack *Stack, mem *Memory, mem
 func gasSLoad(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize uint64) (uint64, error) {
 	usedGas := uint64(0)
 
-	if evm.accesses != nil {
+	if evm.Accesses != nil {
 		where := stack.Back(0)
 		addr := contract.Address()
 		index := trieUtils.GetTreeKeyStorageSlot(addr[:], where)
-		usedGas += evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
+		usedGas += evm.Accesses.TouchAddressAndChargeGas(index, nil)
 	}
 
 	return usedGas, nil
@@ -207,7 +193,6 @@ func gasSStore(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySi
 			return params.SstoreResetGas + accessGas, nil
 		}
 	}
-
 	// The new gas metering is based on net gas costs (EIP-1283):
 	//
 	// 1. If current value equals new value (this is a no-op), 200 gas is deducted.
@@ -422,7 +407,7 @@ func gasCall(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize
 		transfersValue = !stack.Back(2).IsZero()
 		address        = common.Address(stack.Back(1).Bytes20())
 	)
-	if evm.accesses != nil {
+	if evm.Accesses != nil {
 		// Charge witness costs
 		for i := trieUtils.VersionLeafKey; i <= trieUtils.CodeSizeLeafKey; i++ {
 			index := trieUtils.GetTreeKeyAccountLeaf(address[:], byte(i))
@@ -456,6 +441,7 @@ func gasCall(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySize
 	if gas, overflow = math.SafeAdd(gas, evm.callGasTemp); overflow {
 		return 0, ErrGasUintOverflow
 	}
+
 	return gas, nil
 }
 

core/vm/instructions.go

@@ -19,6 +19,7 @@ package vm
 import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/params"
 	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/holiman/uint256"
@@ -343,9 +344,10 @@ func opReturnDataCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeConte
 func opExtCodeSize(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
 	slot := scope.Stack.peek()
 	cs := uint64(interpreter.evm.StateDB.GetCodeSize(slot.Bytes20()))
-	if interpreter.evm.accesses != nil {
+	if interpreter.evm.Accesses != nil {
 		index := trieUtils.GetTreeKeyCodeSize(slot.Bytes())
-		interpreter.evm.TxContext.Accesses.TouchAddress(index, uint256.NewInt(cs).Bytes())
+		statelessGas := interpreter.evm.Accesses.TouchAddressAndChargeGas(index, uint256.NewInt(cs).Bytes())
+		scope.Contract.UseGas(statelessGas)
 	}
 	slot.SetUint64(cs)
 	return nil, nil
@@ -368,63 +370,92 @@ func opCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([
 	if overflow {
 		uint64CodeOffset = 0xffffffffffffffff
 	}
-	uint64CodeEnd, overflow := new(uint256.Int).Add(&codeOffset, &length).Uint64WithOverflow()
-	if overflow {
-		uint64CodeEnd = 0xffffffffffffffff
-	}
-	if interpreter.evm.accesses != nil {
-		copyCodeFromAccesses(scope.Contract.Address(), uint64CodeOffset, uint64CodeEnd, memOffset.Uint64(), interpreter, scope)
-	} else {
-		codeCopy := getData(scope.Contract.Code, uint64CodeOffset, length.Uint64())
-		scope.Memory.Set(memOffset.Uint64(), length.Uint64(), codeCopy)
 
-		touchEachChunks(uint64CodeOffset, uint64CodeEnd, codeCopy, scope.Contract, interpreter.evm)
+	paddedCodeCopy, copyOffset, nonPaddedCopyLength := getDataAndAdjustedBounds(scope.Contract.Code, uint64CodeOffset, length.Uint64())
+	if interpreter.evm.Accesses != nil {
+		touchEachChunksAndChargeGas(copyOffset, nonPaddedCopyLength, scope.Contract.Address().Bytes()[:], scope.Contract, interpreter.evm.Accesses)
 	}
-
+	scope.Memory.Set(memOffset.Uint64(), uint64(len(paddedCodeCopy)), paddedCodeCopy)
 	return nil, nil
 }
 
-// Helper function to touch every chunk in a code range
+// touchEachChunksAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
-func touchEachChunks(start, end uint64, code []byte, contract *Contract, evm *EVM) {
+func touchEachChunksAndChargeGas(offset, size uint64, address []byte, contract *Contract, accesses *types.AccessWitness) uint64 {
-	for chunk := start / 31; chunk <= end/31 && chunk <= uint64(len(code))/31; chunk++ {
+	// note that in the case where the copied code is outside the range of the
-		index := trieUtils.GetTreeKeyCodeChunk(contract.Address().Bytes(), uint256.NewInt(chunk))
+	// contract code but touches the last leaf with contract code in it,
-		count := uint64(0)
+	// we don't include the last leaf of code in the AccessWitness.  The
-		end := (chunk + 1) * 31
+	// reason that we do not need the last leaf is the account's code size
+	// is already in the AccessWitness so a stateless verifier can see that
+	// the code from the last leaf is not needed.
+	if contract != nil && (size == 0 || offset > uint64(len(contract.Code))) {
+		return 0
+	}
+	var (
+		statelessGasCharged uint64
+		startLeafOffset     uint64
+		endLeafOffset       uint64
+		startOffset         uint64
+		endOffset           uint64
+		numLeaves           uint64
+		code                []byte
+		index               [32]byte
+	)
+	if contract != nil {
+		code = contract.Code[:]
+	}
+	// startLeafOffset, endLeafOffset is the evm code offset of the first byte in the first leaf touched
+	// and the evm code offset of the last byte in the last leaf touched
+	startOffset = offset - (offset % 31)
+	if contract != nil && startOffset+size > uint64(len(contract.Code)) {
+		endOffset = uint64(len(contract.Code))
+	} else {
+		endOffset = startOffset + size
+	}
+	endLeafOffset = endOffset + (endOffset % 31)
+	numLeaves = (endLeafOffset - startLeafOffset) / 31
+	chunkOffset := new(uint256.Int)
+	treeIndex := new(uint256.Int)
 
+	for i := 0; i < int(numLeaves); i++ {
+		chunkOffset.Add(trieUtils.CodeOffset, uint256.NewInt(uint64(i)))
+		treeIndex.Div(chunkOffset, trieUtils.VerkleNodeWidth)
+		var subIndex byte
+		subIndexMod := chunkOffset.Mod(chunkOffset, trieUtils.VerkleNodeWidth).Bytes()
+		if len(subIndexMod) == 0 {
+			subIndex = 0
+		} else {
+			subIndex = subIndexMod[0]
+		}
+		treeKey := trieUtils.GetTreeKey(address, treeIndex, subIndex)
+		copy(index[0:31], treeKey)
+		index[31] = subIndex
+
+		var value []byte
+		if contract != nil {
+			// the offset into the leaf that the first PUSH occurs
+			var firstPushOffset uint64 = 0
 			// Look for the first code byte (i.e. no pushdata)
-		for ; count < 31 && end+count < uint64(len(contract.Code)) && !contract.IsCode(chunk*31+count); count++ {
+			for ; firstPushOffset < 31 && firstPushOffset+uint64(i)*31 < uint64(len(contract.Code)) && !contract.IsCode(uint64(i)*31+firstPushOffset); firstPushOffset++ {
 			}
-		var value [32]byte
+			curEnd := (uint64(i) + 1) * 31
-		value[0] = byte(count)
+			if curEnd > endOffset {
-		if end > uint64(len(code)) {
+				curEnd = endOffset
-			end = uint64(len(code))
 			}
-		copy(value[1:], code[chunk*31:end])
+			valueSize := curEnd - (uint64(i) * 31)
-		evm.Accesses.TouchAddress(index, value[:])
+			value = make([]byte, 32, 32)
+			value[0] = byte(firstPushOffset)
+
+			copy(value[1:valueSize+1], code[i*31:curEnd])
+			if valueSize < 31 {
+				padding := make([]byte, 31-valueSize, 31-valueSize)
+				copy(value[valueSize+1:], padding)
 			}
 		}
 
-// copyCodeFromAccesses perform codecopy from the witness, not from the db.
+		statelessGasCharged += accesses.TouchAddressAndChargeGas(index[:], value)
-func copyCodeFromAccesses(addr common.Address, codeOffset, codeEnd, memOffset uint64, in *EVMInterpreter, scope *ScopeContext) {
-	chunk := codeOffset / 31
-	endChunk := codeEnd / 31
-	start := codeOffset % 31 // start inside the first code chunk
-	offset := uint64(0)      // memory offset to write to
-	// XXX uint64 overflow in condition check
-	for end := uint64(31); chunk < endChunk; chunk, start = chunk+1, 0 {
-		// case of the last chunk: figure out how many bytes need to
-		// be extracted from the last chunk.
-		if chunk+1 == endChunk {
-			end = codeEnd % 31
 	}
 
-		// TODO make a version of GetTreeKeyCodeChunk without the bigint
+	return statelessGasCharged
-		index := common.BytesToHash(trieUtils.GetTreeKeyCodeChunk(addr[:], uint256.NewInt(chunk)))
-		h := in.evm.accesses[index]
-		//in.evm.Accesses.TouchAddress(index.Bytes(), h[1+start:1+end])
-		scope.Memory.Set(memOffset+offset, end-start, h[1+start:end])
-		offset += 31 - start
-	}
 }
 
 func opExtCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
@@ -439,18 +470,12 @@ func opExtCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext)
 	if overflow {
 		uint64CodeOffset = 0xffffffffffffffff
 	}
-	uint64CodeEnd, overflow := new(uint256.Int).Add(&codeOffset, &length).Uint64WithOverflow()
-	if overflow {
-		uint64CodeEnd = 0xffffffffffffffff
-	}
 	addr := common.Address(a.Bytes20())
-	if interpreter.evm.accesses != nil {
+	if interpreter.evm.Accesses != nil {
-		copyCodeFromAccesses(addr, uint64CodeOffset, uint64CodeEnd, memOffset.Uint64(), interpreter, scope)
+		log.Warn("setting witness values for extcodecopy is not currently implemented")
 	} else {
 		codeCopy := getData(interpreter.evm.StateDB.GetCode(addr), uint64CodeOffset, length.Uint64())
 		scope.Memory.Set(memOffset.Uint64(), length.Uint64(), codeCopy)
-
-		touchEachChunks(uint64CodeOffset, uint64CodeEnd, codeCopy, scope.Contract, interpreter.evm)
 	}
 
 	return nil, nil
@@ -579,10 +604,11 @@ func opSload(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]by
 	hash := common.Hash(loc.Bytes32())
 	val := interpreter.evm.StateDB.GetState(scope.Contract.Address(), hash)
 	loc.SetBytes(val.Bytes())
-	// Get the initial value as it might not be present
 
+	if interpreter.evm.Accesses != nil {
 		index := trieUtils.GetTreeKeyStorageSlot(scope.Contract.Address().Bytes(), loc)
-	interpreter.evm.TxContext.Accesses.TouchAddress(index, val.Bytes())
+		interpreter.evm.Accesses.TouchAddressAndChargeGas(index, val.Bytes())
+	}
 	return nil, nil
 }
 
@@ -907,9 +933,11 @@ func opPush1(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]by
 	*pc += 1
 	if *pc < codeLen {
 		scope.Stack.push(integer.SetUint64(uint64(scope.Contract.Code[*pc])))
+
+		if interpreter.evm.Accesses != nil && *pc%31 == 0 {
 			// touch next chunk if PUSH1 is at the boundary. if so, *pc has
 			// advanced past this boundary.
-		if *pc%31 == 0 {
+
 			// touch push data by adding the last byte of the pushdata
 			var value [32]byte
 			chunk := *pc / 31
@@ -924,7 +952,8 @@ func opPush1(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]by
 			}
 			copy(value[1:], scope.Contract.Code[chunk*31:endMin])
 			index := trieUtils.GetTreeKeyCodeChunk(scope.Contract.Address().Bytes(), uint256.NewInt(chunk))
-			interpreter.evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
+			statelessGas := interpreter.evm.Accesses.TouchAddressAndChargeGas(index, nil)
+			scope.Contract.UseGas(statelessGas)
 		}
 	} else {
 		scope.Stack.push(integer.Clear())
@@ -947,43 +976,15 @@ func makePush(size uint64, pushByteSize int) executionFunc {
 			endMin = startMin + pushByteSize
 		}
 
+		if interpreter.evm.Accesses != nil {
+			statelessGas := touchEachChunksAndChargeGas(uint64(startMin), uint64(pushByteSize), scope.Contract.Address().Bytes()[:], scope.Contract, interpreter.evm.Accesses)
+			scope.Contract.UseGas(statelessGas)
+		}
+
 		integer := new(uint256.Int)
 		scope.Stack.push(integer.SetBytes(common.RightPadBytes(
 			scope.Contract.Code[startMin:endMin], pushByteSize)))
 
-		// touch push data by adding the last byte of the pushdata
-		var value [32]byte
-		chunk := uint64(endMin-1) / 31
-		count := uint64(0)
-		// Look for the first code byte (i.e. no pushdata)
-		for ; count < 31 && !scope.Contract.IsCode(chunk*31+count); count++ {
-		}
-		value[0] = byte(count)
-		copy(value[1:], scope.Contract.Code[chunk*31:endMin])
-		index := trieUtils.GetTreeKeyCodeChunk(scope.Contract.Address().Bytes(), uint256.NewInt(chunk))
-		interpreter.evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
-
-		// in the case of PUSH32, the end data might be two chunks away,
-		// so also get the middle chunk. There is a boundary condition
-		// check (endMin > 2) in the case the code is a single PUSH32
-		// insctruction, whose immediate are just 0s.
-		if pushByteSize == 32 && endMin > 2 {
-			chunk = uint64(endMin-2) / 31
-			count = uint64(0)
-			// Look for the first code byte (i.e. no pushdata)
-			for ; count < 31 && !scope.Contract.IsCode(chunk*31+count); count++ {
-			}
-			value[0] = byte(count)
-			end := (chunk + 1) * 31
-			if end > uint64(len(scope.Contract.Code)) {
-				end = uint64(len(scope.Contract.Code))
-			}
-			copy(value[1:], scope.Contract.Code[chunk*31:end])
-			index := trieUtils.GetTreeKeyCodeChunk(scope.Contract.Address().Bytes(), uint256.NewInt(chunk))
-			interpreter.evm.TxContext.Accesses.TouchAddressAndChargeGas(index, nil)
-
-		}
-
 		*pc += size
 		return nil, nil
 	}

core/vm/interpreter.go

@@ -17,15 +17,12 @@
 package vm
 
 import (
-	"errors"
 	"hash"
 	"sync/atomic"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/log"
-	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
-	"github.com/holiman/uint256"
 )
 
 // Config are the configuration options for the Interpreter
@@ -194,51 +191,16 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 			logged, pcCopy, gasCopy = false, pc, contract.Gas
 		}
 
+		if in.evm.TxContext.Accesses != nil {
 			// if the PC ends up in a new "page" of verkleized code, charge the
 			// associated witness costs.
-		inWitness := false
+			contract.Gas -= touchEachChunksAndChargeGas(pc, 1, contract.Address().Bytes()[:], contract, in.evm.TxContext.Accesses)
-		var codePage common.Hash
-		if in.evm.chainRules.IsCancun {
-			index := trieUtils.GetTreeKeyCodeChunk(contract.Address().Bytes(), uint256.NewInt(pc/31))
-
-			var value [32]byte
-			if in.evm.accesses != nil {
-				codePage, inWitness = in.evm.accesses[common.BytesToHash(index)]
-				// Return an error if we're in stateless mode
-				// and the code isn't in the witness. It means
-				// that if code is read beyond the actual code
-				// size, pages of 0s need to be added to the
-				// witness.
-				if !inWitness {
-					return nil, errors.New("code chunk missing from proof")
-				}
-				copy(value[:], codePage[:])
-			} else {
-				// Calculate the chunk
-				chunk := pc / 31
-				end := (chunk + 1) * 31
-				if end >= uint64(len(contract.Code)) {
-					end = uint64(len(contract.Code))
-				}
-				count := uint64(0)
-				// Look for the first code byte (i.e. no pushdata)
-				for ; chunk*31+count < end && count < 31 && !contract.IsCode(chunk*31+count); count++ {
-				}
-				value[0] = byte(count)
-				copy(value[1:], contract.Code[chunk*31:end])
-			}
-			contract.Gas -= in.evm.TxContext.Accesses.TouchAddressAndChargeGas(index, value[:])
 		}
 
-		if inWitness {
+		// TODO how can we tell if we are in stateless mode here and need to get the op from the witness
-			// Get the op from the tree, skipping the header byte
-			op = OpCode(codePage[1+pc%31])
-		} else {
 		// If we are in witness mode, then raise an error
 		op = contract.GetOp(pc)
 
-		}
-
 		// Get the operation from the jump table and validate the stack to ensure there are
 		// enough stack items available to perform the operation.
 		operation := in.cfg.JumpTable[op]