refactor: set witness account values in getStateObject() (ethereum#83)

* enforce 32-byte alignment

* save current state

* write account values to witness from getObject

* code cleanup + sanity checks

* fix RLP serialization of missing keys

* remove code redundancy for GetTreeKeyCodeChunk

* fix stem calculation issue for code and storage

* remove redundant SetLeafValue calls

* Add a contract creation tx to the test

* fix botched module version update

* detail gas calculation cost in verkle contract deployment

core/chain_makers.go

@@ -343,14 +343,24 @@ func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine
 			kvs := statedb.Witness().KeyVals()
 			keys := statedb.Witness().Keys()
 			for _, key := range keys {
-				// XXX workaround - there is a problem in the witness creation
-				// so fix the witness creation as well.
-				v, err := vtr.TryGet(key)
+				_, err := vtr.TryGet(key)
 				if err != nil {
 					panic(err)
 				}
-				kvs[string(key)] = v
+
+				// Sanity check: ensure all flagged addresses have an associated
+				// value: keys is built from Chunks and kvs from InitialValue.
+				if _, exists := kvs[string(key)]; !exists {
+					panic(fmt.Sprintf("address not in access witness: %x", key))
+				}
 			}
+
+			// sanity check: ensure all values correspond to a flagged key by
+			// comparing the lengths of both structures: they should be equal
+			if len(kvs) != len(keys) {
+				panic("keys without a value in witness")
+			}
+
 			vtr.Hash()
 			p, k, err := vtr.ProveAndSerialize(keys, kvs)
 			block.SetVerkleProof(p, k)

core/state/state_object.go

@@ -18,6 +18,7 @@ package state
 
 import (
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math/big"
@@ -102,6 +103,24 @@ func (s *stateObject) empty() bool {
 
 // newObject creates a state object.
 func newObject(db *StateDB, address common.Address, data types.StateAccount) *stateObject {
+	if db.trie.IsVerkle() {
+		var nonce, balance, version []byte
+
+		// preserve nil as a balance value, it means it's not in the tree
+		// use is as a heuristic for the nonce being null as well
+		if data.Balance != nil {
+			nonce = make([]byte, 32)
+			balance = make([]byte, 32)
+			version = make([]byte, 32)
+			for i, b := range data.Balance.Bytes() {
+				balance[len(data.Balance.Bytes())-1-i] = b
+			}
+
+			binary.LittleEndian.PutUint64(nonce[:8], data.Nonce)
+		}
+		db.witness.SetGetObjectTouchedLeaves(address.Bytes(), version, balance[:], nonce[:], data.CodeHash)
+	}
+
 	if data.Balance == nil {
 		data.Balance = new(big.Int)
 	}
@@ -497,12 +516,21 @@ func (s *stateObject) Code(db Database) []byte {
 		return s.code
 	}
 	if bytes.Equal(s.CodeHash(), emptyCodeHash) {
+		if s.db.GetTrie().IsVerkle() {
+			// Mark the code size and code hash as empty
+			s.db.witness.SetObjectCodeTouchedLeaves(s.address.Bytes(), nil, nil)
+		}
 		return nil
 	}
 	code, err := db.ContractCode(s.addrHash, common.BytesToHash(s.CodeHash()))
 	if err != nil {
 		s.setError(fmt.Errorf("can't load code hash %x: %v", s.CodeHash(), err))
 	}
+	if s.db.GetTrie().IsVerkle() {
+		var cs [32]byte
+		binary.LittleEndian.PutUint64(cs[:8], uint64(len(code)))
+		s.db.witness.SetObjectCodeTouchedLeaves(s.address.Bytes(), cs[:], s.CodeHash())
+	}
 	s.code = code
 	return code
 }
@@ -515,12 +543,21 @@ func (s *stateObject) CodeSize(db Database) int {
 		return len(s.code)
 	}
 	if bytes.Equal(s.CodeHash(), emptyCodeHash) {
+		if s.db.trie.IsVerkle() {
+			var sz [32]byte
+			s.db.witness.SetLeafValuesMessageCall(s.address.Bytes(), sz[:])
+		}
 		return 0
 	}
 	size, err := db.ContractCodeSize(s.addrHash, common.BytesToHash(s.CodeHash()))
 	if err != nil {
 		s.setError(fmt.Errorf("can't load code size %x: %v", s.CodeHash(), err))
 	}
+	if s.db.trie.IsVerkle() {
+		var sz [32]byte
+		binary.LittleEndian.PutUint64(sz[:8], uint64(size))
+		s.db.witness.SetLeafValuesMessageCall(s.address.Bytes(), sz[:])
+	}
 	return size
 }
 
@@ -560,10 +597,25 @@ func (s *stateObject) Balance() *big.Int {
 	return s.data.Balance
 }
 
+func (s *stateObject) BalanceLE() []byte {
+	var out [32]byte
+	for i, b := range s.data.Balance.Bytes() {
+		out[len(s.data.Balance.Bytes())-1-i] = b
+	}
+
+	return out[:]
+}
+
 func (s *stateObject) Nonce() uint64 {
 	return s.data.Nonce
 }
 
+func (s *stateObject) NonceLE() []byte {
+	var out [32]byte
+	binary.LittleEndian.PutUint64(out[:8], s.data.Nonce)
+	return out[:]
+}
+
 // Never called, but must be present to allow stateObject to be used
 // as a vm.Account interface that also satisfies the vm.ContractRef
 // interface. Interfaces are awesome.

core/state/statedb.go

@@ -607,13 +607,6 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject {
 			}
 		}
 
-		// NOTE: Do not touch the addresses here, kick the can down the
-		// road. That is because I don't want to change the interface
-		// to getDeletedStateObject at this stage, as the PR would then
-		// have a huge footprint.
-		// The alternative is to make accesses available via the state
-		// db instead of the evm. This requires a significant rewrite,
-		// that isn't currently warranted.
 	}
 	// If snapshot unavailable or reading from it failed, load from the database
 	if s.snap == nil || err != nil {

core/state_processor_test.go

@@ -381,17 +381,27 @@ func TestProcessStateless(t *testing.T) {
 	genesis := gspec.MustCommit(db)
 	blockchain, _ := NewBlockChain(db, nil, gspec.Config, ethash.NewFaker(), vm.Config{}, nil, nil)
 	defer blockchain.Stop()
+
+	code := common.FromHex(`6060604052600a8060106000396000f360606040526008565b00`)
 	txCost1 := params.WitnessBranchWriteCost*2 + params.WitnessBranchReadCost*2 + params.WitnessChunkWriteCost*3 + params.WitnessChunkReadCost*10 + params.TxGas
 	txCost2 := params.WitnessBranchWriteCost + params.WitnessBranchReadCost*2 + params.WitnessChunkWriteCost*2 + params.WitnessChunkReadCost*10 + params.TxGas
-	blockGasUsedExpected := txCost1*2 + txCost2
+	intrinsic, _ := IntrinsicGas(code, nil, true, true, true)
+	contractCreationCost := intrinsic + (2*params.WitnessChunkReadCost+params.WitnessChunkWriteCost)*10 + params.WitnessBranchReadCost + params.WitnessBranchWriteCost + uint64(3639) // standard cost of executing the contract
+	blockGasUsagesExpected := []uint64{txCost1*2 + txCost2, txCost1*2 + txCost2 + contractCreationCost}
 	chain, _ := GenerateVerkleChain(gspec.Config, genesis, ethash.NewFaker(), db, 2, func(i int, gen *BlockGen) {
 		// TODO need to check that the tx cost provided is the exact amount used (no remaining left-over)
-		tx, _ := types.SignTx(types.NewTransaction(uint64(i)*3, common.Address{1, 2, 3}, big.NewInt(999), txCost1, big.NewInt(875000000), nil), signer, testKey)
+		tx, _ := types.SignTx(types.NewTransaction(uint64(i)*3, common.Address{byte(i), 2, 3}, big.NewInt(999), txCost1, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
 		tx, _ = types.SignTx(types.NewTransaction(uint64(i)*3+1, common.Address{}, big.NewInt(999), txCost1, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
 		tx, _ = types.SignTx(types.NewTransaction(uint64(i)*3+2, common.Address{}, big.NewInt(0), txCost2, big.NewInt(875000000), nil), signer, testKey)
 		gen.AddTx(tx)
+
+		// Add a contract creation in block #2
+		if i == 1 {
+			tx, _ = types.SignTx(types.NewContractCreation(6, big.NewInt(16), 3000000, big.NewInt(875000000), code), signer, testKey)
+			gen.AddTx(tx)
+		}
 	})
 
 	// Uncomment to extract block #2
@@ -412,8 +422,8 @@ func TestProcessStateless(t *testing.T) {
 		if b == nil {
 			t.Fatalf("expected block %d to be present in chain", i+1)
 		}
-		if b.GasUsed() != blockGasUsedExpected {
-			t.Fatalf("expected block txs to use %d, got %d\n", blockGasUsedExpected, b.GasUsed())
+		if b.GasUsed() != blockGasUsagesExpected[i] {
+			t.Fatalf("expected block txs to use %d, got %d\n", blockGasUsagesExpected[i], b.GasUsed())
 		}
 	}
 }

core/state_transition.go

@@ -17,7 +17,6 @@
 package core
 
 import (
-	"encoding/binary"
 	"fmt"
 	"math"
 	"math/big"
@@ -317,7 +316,7 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 		return nil, fmt.Errorf("%w: have %d, want %d", ErrIntrinsicGas, st.gas, gas)
 	}
 	if st.evm.ChainConfig().IsCancun(st.evm.Context.BlockNumber) {
-		var targetBalance, targetNonce, targetCodeSize, targetCodeKeccak, originBalance, originNonceBytes []byte
+		var originBalance, originNonceBytes []byte
 
 		targetAddr := msg.To()
 		originAddr := msg.From()
@@ -329,21 +328,16 @@ func (st *StateTransition) TransitionDb() (*ExecutionResult, error) {
 		originBalance = st.evm.StateDB.GetBalanceLittleEndian(originAddr)
 		originNonce := st.evm.StateDB.GetNonce(originAddr)
 		originNonceBytes = st.evm.StateDB.GetNonceLittleEndian(originAddr)
-		st.evm.Accesses.SetTxOriginTouchedLeaves(originAddr.Bytes(), originBalance, originNonceBytes)
+		st.evm.Accesses.SetTxOriginTouchedLeaves(originAddr.Bytes(), originBalance, originNonceBytes, st.evm.StateDB.GetCodeSize(originAddr))
 
 		if msg.To() != nil {
 			statelessGasDest := st.evm.Accesses.TouchTxExistingAndComputeGas(targetAddr.Bytes(), msg.Value().Sign() != 0)
 			if !tryConsumeGas(&st.gas, statelessGasDest) {
 				return nil, fmt.Errorf("%w: Insufficient funds to cover witness access costs for transaction: have %d, want %d", ErrInsufficientBalanceWitness, st.gas, gas)
 			}
-			targetBalance = st.evm.StateDB.GetBalanceLittleEndian(*targetAddr)
-			targetNonce = st.evm.StateDB.GetNonceLittleEndian(*targetAddr)
-			targetCodeKeccak = st.evm.StateDB.GetCodeHash(*targetAddr).Bytes()
-
-			codeSize := uint64(st.evm.StateDB.GetCodeSize(*targetAddr))
-			var codeSizeBytes [32]byte
-			binary.LittleEndian.PutUint64(codeSizeBytes[:8], codeSize)
-			st.evm.Accesses.SetTxExistingTouchedLeaves(targetAddr.Bytes(), targetBalance, targetNonce, targetCodeSize, targetCodeKeccak)
+
+			// ensure the code size ends up in the access witness
+			st.evm.StateDB.GetCodeSize(*targetAddr)
 		} else {
 			contractAddr := crypto.CreateAddress(originAddr, originNonce)
 			if !tryConsumeGas(&st.gas, st.evm.Accesses.TouchAndChargeContractCreateInit(contractAddr.Bytes(), msg.Value().Sign() != 0)) {