Update witness gas charging to differentiate between access events and write events

consensus/ethash/consensus.go

@@ -669,5 +669,6 @@ func accumulateRewards(config *params.ChainConfig, state *state.StateDB, header
 	}
 	state.AddBalance(header.Coinbase, reward)
 	coinbase := utils.GetTreeKeyBalance(header.Coinbase.Bytes())
-	state.Witness().TouchAddress(coinbase, state.GetBalance(header.Coinbase).Bytes())
+	state.Witness().TouchAddressOnReadAndChargeGas(coinbase)
+	state.Witness().SetLeafValue(coinbase, state.GetBalance(header.Coinbase).Bytes())
 }

core/state_processor_test.go

@@ -378,11 +378,11 @@ func TestProcessStateless(t *testing.T) {
 	blockchain, _ := NewBlockChain(db, nil, gspec.Config, ethash.NewFaker(), vm.Config{}, nil, nil)
 	defer blockchain.Stop()
 	chain, _ := GenerateVerkleChain(gspec.Config, genesis, ethash.NewFaker(), db, 1, func(_ int, gen *BlockGen) {
-		tx, _ := types.SignTx(types.NewTransaction(0, common.Address{1, 2, 3}, big.NewInt(999), params.TxGas, big.NewInt(875000000), nil), signer, testKey)
+		tx, _ := types.SignTx(types.NewTransaction(0, common.Address{1, 2, 3}, big.NewInt(999), 28600, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
-		tx, _ = types.SignTx(types.NewTransaction(1, common.Address{}, big.NewInt(999), params.TxGas, big.NewInt(875000000), nil), signer, testKey)
+		tx, _ = types.SignTx(types.NewTransaction(1, common.Address{}, big.NewInt(999), 28600, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
-		tx, _ = types.SignTx(types.NewTransaction(2, common.Address{}, big.NewInt(0), params.TxGas, big.NewInt(875000000), nil), signer, testKey)
+		tx, _ = types.SignTx(types.NewTransaction(2, common.Address{}, big.NewInt(0), 28600, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
 
 	})

core/state_transition.go

@@ -301,29 +301,36 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 	if err != nil {
 		return nil, err
 	}
-	if st.gas < gas {
-		return nil, fmt.Errorf("%w: have %d, want %d", ErrIntrinsicGas, st.gas, gas)
-	}
+
 	if st.evm.TxContext.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())
+			if msg.Value().Cmp(big.NewInt(0)) != 0 {
+				gas += st.evm.TxContext.Accesses.TouchAddressOnWriteAndChargeGas(toBalance)
+				st.evm.TxContext.Accesses.SetLeafValue(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.TxContext.Accesses.TouchAddressOnWriteAndChargeGas(fromNonce)
+			st.evm.TxContext.Accesses.SetLeafValue(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.TxContext.Accesses.TouchAddressAndChargeGas(fromBalance, preFB[:])
+		gas += st.evm.TxContext.Accesses.TouchAddressOnReadAndChargeGas(fromNonce)
+		st.evm.TxContext.Accesses.SetLeafValue(fromNonce, preFN[:])
+		gas += st.evm.TxContext.Accesses.TouchAddressOnReadAndChargeGas(fromBalance)
+		st.evm.TxContext.Accesses.SetLeafValue(fromBalance, preFB[:])
+	}
+	if st.gas < gas {
+		return nil, fmt.Errorf("%w: have %d, want %d", ErrIntrinsicGas, st.gas, gas)
 	}
 	st.gas -= gas
 

core/types/access_witness.go

@@ -17,94 +17,168 @@
 package types
 
 import (
+	"fmt"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/params"
 )
 
+type VerkleStem [31]byte
+
+type ChunkValue struct {
+	mode byte
+	value []byte
+}
+
 // AccessWitness lists the locations of the state that are being accessed
 // during the production of a block.
 // TODO(@gballet) this doesn't fully support deletions
 type AccessWitness struct {
 	// Branches flags if a given branch has been loaded
-	Branches map[[31]byte]struct{}
+	// for the byte value:
+	//	the first bit is set if the branch has been edited
+	//	the second bit is set if the branch has been read
+	Branches map[VerkleStem]byte
 
 	// Chunks contains the initial value of each address
-	Chunks map[common.Hash][]byte
-
-	// The initial value isn't always available at the time an
-	// address is touched, this map references addresses that
-	// were touched but can not yet be put in Chunks.
-	Undefined map[common.Hash]struct{}
+	Chunks map[common.Hash]ChunkValue
 }
 
 func NewAccessWitness() *AccessWitness {
 	return &AccessWitness{
-		Branches:  make(map[[31]byte]struct{}),
-		Chunks:    make(map[common.Hash][]byte),
-		Undefined: make(map[common.Hash]struct{}),
+		Branches:  make(map[VerkleStem]byte),
+		Chunks:    make(map[common.Hash]ChunkValue),
+	}
+}
+
+const (
+	AccessWitnessReadFlag = 1
+	AccessWitnessWriteFlag = 2
+)
+
+// because of the way Geth's EVM is implemented, the gas cost of an operation
+// may be needed before the value of the leaf-key can be retrieved. Hence, we
+// break witness access (for the purpose of gas accounting), and filling witness
+// values into two methods
+func (aw *AccessWitness) SetLeafValue(addr []byte, value []byte) {
+	var stem [31]byte
+	copy(stem[:], addr[:31])
+
+	if chunk, exists := aw.Chunks[common.BytesToHash(addr)]; exists {
+		chunk.value = value
+	} else {
+		panic(fmt.Sprintf("address not in access witness: %x", addr))
+	}
+}
+
+func (aw *AccessWitness) touchAddressOnWrite(addr []byte) (bool, bool, bool) {
+	var stem        VerkleStem
+	var stemWrite, chunkWrite, chunkFill bool
+	copy(stem[:], addr[:31])
+
+	// NOTE: stem, selector access flags already exist in their
+	// respective maps because this function is called at the end of 
+	// processing a read access event
+
+	if (aw.Branches[stem] & AccessWitnessWriteFlag) == 0 {
+		stemWrite = true
+		aw.Branches[stem] |= AccessWitnessWriteFlag
+	}
+
+	chunkValue := aw.Chunks[common.BytesToHash(addr)]
+	if chunkValue.mode & AccessWitnessWriteFlag != 0 {
+		chunkWrite = true
+		chunkValue.mode |= AccessWitnessWriteFlag
+		aw.Chunks[common.BytesToHash(addr)] = chunkValue
+	}
+
+	// TODO charge chunk filling costs if the leaf was previously empty in the state
+	/*
+	if chunkWrite {
+		if _, err := verkleDb.TryGet(addr); err != nil {
+			chunkFill = true
+		}
 	}
+	*/
+
+	return stemWrite, chunkWrite, chunkFill
 }
 
 // TouchAddress adds any missing addr to the witness and returns respectively
 // true if the stem or the stub weren't arleady present.
-func (aw *AccessWitness) TouchAddress(addr, value []byte) (bool, bool) {
+func (aw *AccessWitness) touchAddress(addr []byte, isWrite bool) (bool, bool, bool, bool, bool) {
 	var (
 		stem        [31]byte
-		newStem     bool
-		newSelector bool
+		stemRead bool
+		selectorRead bool
 	)
 	copy(stem[:], addr[:31])
 
 	// Check for the presence of the stem
-	if _, newStem := aw.Branches[stem]; !newStem {
-		aw.Branches[stem] = struct{}{}
-	}
-
-	// Check for the presence of the selector
-	if _, newSelector := aw.Chunks[common.BytesToHash(addr)]; !newSelector {
-		if value == nil {
-			aw.Undefined[common.BytesToHash(addr)] = struct{}{}
-		} else {
-			if _, ok := aw.Undefined[common.BytesToHash(addr)]; !ok {
-				delete(aw.Undefined, common.BytesToHash(addr))
-			}
-			aw.Chunks[common.BytesToHash(addr)] = value
+	if _, hasStem := aw.Branches[stem]; !hasStem {
+		stemRead = true
+		aw.Branches[stem] = AccessWitnessReadFlag
+	}
+
+	// always charge read cost whether the access event is read/write
+	// literal interpretation of the spec
+	selectorRead = true
+
+	// Check for the presence of the leaf selector
+	if _, hasSelector := aw.Chunks[common.BytesToHash(addr)]; !hasSelector {
+		aw.Chunks[common.BytesToHash(addr)] = ChunkValue{
+			AccessWitnessReadFlag,
+			nil,
 		}
 	}
 
-	return newStem, newSelector
+	var stemWrite, selectorWrite, chunkFill bool
+
+	if isWrite {
+		stemWrite, selectorWrite, chunkFill = aw.touchAddressOnWrite(addr)
+	}
+
+	return stemRead, selectorRead, stemWrite, selectorWrite, chunkFill
 }
 
-// TouchAddressAndChargeGas checks if a location has already been touched in
-// the current witness, and charge extra gas if that isn't the case. This is
-// meant to only be called on a tx-context access witness (i.e. before it is
-// merged), not a block-context witness: witness costs are charged per tx.
-func (aw *AccessWitness) TouchAddressAndChargeGas(addr, value []byte) uint64 {
+func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, isWrite bool) uint64 {
 	var gas uint64
 
-	nstem, nsel := aw.TouchAddress(addr, value)
-	if nstem {
-		gas += params.WitnessBranchCost
+	stemRead, selectorRead, stemWrite, selectorWrite, selectorFill := aw.touchAddress(addr, isWrite)
+	if stemRead {
+		gas += params.WitnessBranchReadCost
+	}
+	if selectorRead {
+		gas += params.WitnessChunkReadCost
+	}
+	if stemWrite {
+		gas += params.WitnessBranchWriteCost
 	}
-	if nsel {
-		gas += params.WitnessChunkCost
+	if selectorWrite {
+		gas += params.WitnessChunkWriteCost
 	}
+	if selectorFill {
+		gas += params.WitnessChunkFillCost
+	}
+
 	return gas
 }
 
+func (aw *AccessWitness) TouchAddressOnWriteAndChargeGas(addr []byte) uint64 {
+	return aw.touchAddressAndChargeGas(addr, true)
+}
+
+func (aw *AccessWitness) TouchAddressOnReadAndChargeGas(addr []byte) uint64 {
+	return aw.touchAddressAndChargeGas(addr, false)
+}
+
 // Merge is used to merge the witness that got generated during the execution
 // of a tx, with the accumulation of witnesses that were generated during the
 // execution of all the txs preceding this one in a given block.
 func (aw *AccessWitness) Merge(other *AccessWitness) {
-	for k := range other.Undefined {
-		if _, ok := aw.Undefined[k]; !ok {
-			aw.Undefined[k] = struct{}{}
-		}
-	}
-
 	for k := range other.Branches {
 		if _, ok := aw.Branches[k]; !ok {
-			aw.Branches[k] = struct{}{}
+			aw.Branches[k] = other.Branches[k]
 		}
 	}
 
@@ -128,14 +202,17 @@ func (aw *AccessWitness) Keys() [][]byte {
 }
 
 func (aw *AccessWitness) KeyVals() map[common.Hash][]byte {
-	return aw.Chunks
+	result := make(map[common.Hash][]byte)
+	for k, v := range aw.Chunks {
+		result[k] = v.value
+	}
+	return result
 }
 
 func (aw *AccessWitness) Copy() *AccessWitness {
 	naw := &AccessWitness{
-		Branches:  make(map[[31]byte]struct{}),
-		Chunks:    make(map[common.Hash][]byte),
-		Undefined: make(map[common.Hash]struct{}),
+		Branches:  make(map[VerkleStem]byte),
+		Chunks:    make(map[common.Hash]ChunkValue),
 	}
 
 	naw.Merge(aw)

core/vm/evm.go

@@ -170,6 +170,20 @@ func (evm *EVM) Interpreter() *EVMInterpreter {
 	return evm.interpreter
 }
 
+func (evm *EVM) tryTouchWitnessAndConsumeGas(key, value []byte, gasLeft *uint64) bool {
+	gasConsumed := evm.Accesses.TouchAddressOnReadAndChargeGas(key)
+
+	if gasConsumed > *gasLeft {
+		*gasLeft = 0
+		return false
+	}
+
+	evm.Accesses.SetLeafValue(key, value)
+
+	*gasLeft -= gasConsumed
+	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
@@ -234,13 +248,28 @@ func (evm *EVM) Call(caller ContractRef, addr common.Address, input []byte, gas
 		} else {
 			// Touch the account data
 			var data [32]byte
-			evm.Accesses.TouchAddress(utils.GetTreeKeyVersion(addr.Bytes()), data[:])
+			if !evm.tryTouchWitnessAndConsumeGas(utils.GetTreeKeyVersion(addr.Bytes()), data[:], &gas) {
+				evm.StateDB.RevertToSnapshot(snapshot)
+				return []byte{}, gas, ErrOutOfGas
+			}
 			binary.BigEndian.PutUint64(data[:], evm.StateDB.GetNonce(addr))
-			evm.Accesses.TouchAddress(utils.GetTreeKeyNonce(addr[:]), data[:])
-			evm.Accesses.TouchAddress(utils.GetTreeKeyBalance(addr[:]), evm.StateDB.GetBalance(addr).Bytes())
+			if !evm.tryTouchWitnessAndConsumeGas(utils.GetTreeKeyNonce(addr.Bytes()), data[:], &gas) {
+				evm.StateDB.RevertToSnapshot(snapshot)
+				return []byte{}, gas, ErrOutOfGas
+			}
+			if !evm.tryTouchWitnessAndConsumeGas(utils.GetTreeKeyBalance(addr.Bytes()), data[:], &gas) {
+				evm.StateDB.RevertToSnapshot(snapshot)
+				return []byte{}, gas, ErrOutOfGas
+			}
 			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())
+			if !evm.tryTouchWitnessAndConsumeGas(utils.GetTreeKeyCodeSize(addr.Bytes()), data[:], &gas) {
+				evm.StateDB.RevertToSnapshot(snapshot)
+				return []byte{}, gas, ErrOutOfGas
+			}
+			if !evm.tryTouchWitnessAndConsumeGas(utils.GetTreeKeyCodeKeccak(addr.Bytes()), data[:], &gas) {
+				evm.StateDB.RevertToSnapshot(snapshot)
+				return []byte{}, gas, ErrOutOfGas
+			}
 
 			addrCopy := addr
 			// If the account has no code, we can abort here