New access-witness module (#235)

Use plain addresses/slot numbers instead of hashing them in the witness

consensus/beacon/consensus.go

@@ -30,6 +30,8 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/rpc"
 	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/holiman/uint256"
 )
 
 // Proof-of-stake protocol constants.
@@ -354,7 +356,7 @@ func (beacon *Beacon) Finalize(chain consensus.ChainHeaderReader, header *types.
 		state.AddBalance(w.Address, amount)
 
 		// The returned gas is not charged
-		state.Witness().TouchAddressOnWriteAndComputeGas(w.Address[:])
+		state.Witness().TouchAddressOnWriteAndComputeGas(w.Address[:], uint256.Int{}, utils.BalanceLeafKey)
 	}
 	// No block reward which is issued by consensus layer instead.
 }

consensus/ethash/consensus.go

@@ -34,6 +34,7 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/holiman/uint256"
 	"golang.org/x/crypto/sha3"
 )
 
@@ -568,8 +569,7 @@ func accumulateRewards(config *params.ChainConfig, state *state.StateDB, header
 
 		// This should not happen, but it's useful for replay tests
 		if config.IsVerkle(header.Number, header.Time) {
-			uncleCoinbase := utils.GetTreeKeyBalance(uncle.Coinbase.Bytes())
+			state.Witness().TouchAddressOnReadAndComputeGas(uncle.Coinbase.Bytes(), uint256.Int{}, utils.BalanceLeafKey)
-			state.Witness().TouchAddressOnReadAndComputeGas(uncleCoinbase)
 		}
 		state.AddBalance(uncle.Coinbase, r)
 
@@ -577,14 +577,10 @@ func accumulateRewards(config *params.ChainConfig, state *state.StateDB, header
 		reward.Add(reward, r)
 	}
 	if config.IsVerkle(header.Number, header.Time) {
-		coinbase := utils.GetTreeKeyBalance(header.Coinbase.Bytes())
+		state.Witness().TouchAddressOnReadAndComputeGas(header.Coinbase.Bytes(), uint256.Int{}, utils.BalanceLeafKey)
-		state.Witness().TouchAddressOnReadAndComputeGas(coinbase)
+		state.Witness().TouchAddressOnReadAndComputeGas(header.Coinbase.Bytes(), uint256.Int{}, utils.VersionLeafKey)
-		coinbase[31] = utils.VersionLeafKey // mark version
+		state.Witness().TouchAddressOnReadAndComputeGas(header.Coinbase.Bytes(), uint256.Int{}, utils.NonceLeafKey)
-		state.Witness().TouchAddressOnReadAndComputeGas(coinbase)
+		state.Witness().TouchAddressOnReadAndComputeGas(header.Coinbase.Bytes(), uint256.Int{}, utils.CodeKeccakLeafKey)
-		coinbase[31] = utils.NonceLeafKey // mark nonce
-		state.Witness().TouchAddressOnReadAndComputeGas(coinbase)
-		coinbase[31] = utils.CodeKeccakLeafKey // mark code keccak
-		state.Witness().TouchAddressOnReadAndComputeGas(coinbase)
 	}
 	state.AddBalance(header.Coinbase, reward)
 }

core/state/access_witness.go

@@ -20,147 +20,160 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/params"
 	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/holiman/uint256"
 )
 
-type VerkleStem [31]byte
+// mode specifies how a tree location has been accessed
-
-// Mode specifies how a tree location has been accessed
 // 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
-type Mode byte
+type mode byte
 
 const (
-	AccessWitnessReadFlag  = Mode(1)
+	AccessWitnessReadFlag  = mode(1)
-	AccessWitnessWriteFlag = Mode(2)
+	AccessWitnessWriteFlag = mode(2)
 )
 
+var zeroTreeIndex uint256.Int
+
 // AccessWitness lists the locations of the state that are being accessed
 // during the production of a block.
 type AccessWitness struct {
-	// Branches flags if a given branch has been loaded
+	branches map[branchAccessKey]mode
-	Branches map[VerkleStem]Mode
+	chunks   map[chunkAccessKey]mode
 
-	// Chunks contains the initial value of each address
+	pointCache *utils.PointCache
-	Chunks map[common.Hash]Mode
-
-	// InitialValue contains either `nil` if the location
-	// didn't exist before it was accessed, or the value
-	// that a location had before the execution of this
-	// block.
-	InitialValue map[string][]byte
-
-	// 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{}
-
-	statedb *StateDB
 }
 
-func NewAccessWitness(statedb *StateDB) *AccessWitness {
+func NewAccessWitness(pointCache *utils.PointCache) *AccessWitness {
 	return &AccessWitness{
-		Branches:      make(map[VerkleStem]Mode),
+		branches:   make(map[branchAccessKey]mode),
-		Chunks:        make(map[common.Hash]Mode),
+		chunks:     make(map[chunkAccessKey]mode),
-		InitialValue:  make(map[string][]byte),
+		pointCache: pointCache,
-		CodeLocations: make(map[string]map[uint64]struct{}),
-		statedb:       statedb,
 	}
 }
 
-func (aw *AccessWitness) HasCodeChunk(addr []byte, chunknr uint64) bool {
+// Merge is used to merge the witness that got generated during the execution
-	if locs, ok := aw.CodeLocations[string(addr)]; ok {
+// of a tx, with the accumulation of witnesses that were generated during the
-		if _, ok = locs[chunknr]; ok {
+// execution of all the txs preceding this one in a given block.
-			return true
+func (aw *AccessWitness) Merge(other *AccessWitness) {
+	for k := range other.branches {
+		aw.branches[k] |= other.branches[k]
+	}
+	for k, chunk := range other.chunks {
+		aw.chunks[k] |= chunk
 	}
 }
 
-	return false
+// Key returns, predictably, the list of keys that were touched during the
+// buildup of the access witness.
+func (aw *AccessWitness) Keys() [][]byte {
+	// TODO: consider if parallelizing this is worth it, probably depending on len(aw.chunks).
+	keys := make([][]byte, 0, len(aw.chunks))
+	for chunk := range aw.chunks {
+		basePoint := aw.pointCache.GetTreeKeyHeader(chunk.addr[:])
+		key := utils.GetTreeKeyWithEvaluatedAddess(basePoint, &chunk.treeIndex, chunk.leafKey)
+		keys = append(keys, key)
+	}
+	return keys
 }
 
-// SetCodeLeafValue does the same thing as SetLeafValue, but for code chunks. It
+func (aw *AccessWitness) Copy() *AccessWitness {
-// maintains a cache of which (address, chunk) were calculated, in order to avoid
+	naw := &AccessWitness{
-// calling GetTreeKey more than once per chunk.
+		branches:   make(map[branchAccessKey]mode),
-func (aw *AccessWitness) SetCachedCodeChunk(addr []byte, chunknr uint64) {
+		chunks:     make(map[chunkAccessKey]mode),
-	if locs, ok := aw.CodeLocations[string(addr)]; ok {
+		pointCache: aw.pointCache,
-		if _, ok = locs[chunknr]; ok {
-			return
 	}
-	} else {
+	naw.Merge(aw)
-		aw.CodeLocations[string(addr)] = map[uint64]struct{}{}
+	return naw
 }
 
-	aw.CodeLocations[string(addr)][chunknr] = struct{}{}
+func (aw *AccessWitness) TouchAndChargeProofOfAbsence(addr []byte) uint64 {
-}
-
-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 XOR AccessWitnessWriteFlag
-	if ((chunkValue & AccessWitnessWriteFlag) == 0) && ((chunkValue | AccessWitnessWriteFlag) != 0) {
-		chunkWrite = true
-		chunkValue |= 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 []byte, isWrite bool) (bool, bool, bool, bool, bool) {
-	var (
-		stem                                [31]byte
-		stemRead, selectorRead              bool
-		stemWrite, selectorWrite, chunkFill bool
-	)
-	copy(stem[:], addr[:31])
-
-	// Check for the presence of the stem
-	if _, hasStem := aw.Branches[stem]; !hasStem {
-		stemRead = true
-		aw.Branches[stem] = AccessWitnessReadFlag
-	}
-
-	// Check for the presence of the leaf selector
-	if _, hasSelector := aw.Chunks[common.BytesToHash(addr)]; !hasSelector {
-		selectorRead = true
-		aw.Chunks[common.BytesToHash(addr)] = AccessWitnessReadFlag
-	}
-
-	if isWrite {
-		stemWrite, selectorWrite, chunkFill = aw.touchAddressOnWrite(addr)
-	}
-
-	return stemRead, selectorRead, stemWrite, selectorWrite, chunkFill
-}
-
-func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, isWrite bool) uint64 {
 	var gas uint64
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.BalanceLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.CodeSizeLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.CodeKeccakLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.NonceLeafKey)
+	return gas
+}
 
-	stemRead, selectorRead, stemWrite, selectorWrite, selectorFill := aw.touchAddress(addr, isWrite)
+func (aw *AccessWitness) TouchAndChargeMessageCall(addr []byte) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(addr, zeroTreeIndex, utils.CodeSizeLeafKey)
+	return gas
+}
 
+func (aw *AccessWitness) TouchAndChargeValueTransfer(callerAddr, targetAddr []byte) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnWriteAndComputeGas(callerAddr, zeroTreeIndex, utils.BalanceLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(targetAddr, zeroTreeIndex, utils.BalanceLeafKey)
+	return gas
+}
+
+// TouchAndChargeContractCreateInit charges access costs to initiate
+// a contract creation
+func (aw *AccessWitness) TouchAndChargeContractCreateInit(addr []byte, createSendsValue bool) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.NonceLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.CodeKeccakLeafKey)
+	if createSendsValue {
+		gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.BalanceLeafKey)
+	}
+	return gas
+}
+
+// TouchAndChargeContractCreateCompleted charges access access costs after
+// the completion of a contract creation to populate the created account in
+// the tree
+func (aw *AccessWitness) TouchAndChargeContractCreateCompleted(addr []byte) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.BalanceLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.CodeSizeLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.CodeKeccakLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(addr, zeroTreeIndex, utils.NonceLeafKey)
+	return gas
+}
+
+func (aw *AccessWitness) TouchTxOriginAndComputeGas(originAddr []byte) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnReadAndComputeGas(originAddr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(originAddr, zeroTreeIndex, utils.CodeSizeLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(originAddr, zeroTreeIndex, utils.CodeKeccakLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(originAddr, zeroTreeIndex, utils.NonceLeafKey)
+	gas += aw.TouchAddressOnWriteAndComputeGas(originAddr, zeroTreeIndex, utils.BalanceLeafKey)
+	return gas
+}
+
+func (aw *AccessWitness) TouchTxExistingAndComputeGas(targetAddr []byte, sendsValue bool) uint64 {
+	var gas uint64
+	gas += aw.TouchAddressOnReadAndComputeGas(targetAddr, zeroTreeIndex, utils.VersionLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(targetAddr, zeroTreeIndex, utils.CodeSizeLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(targetAddr, zeroTreeIndex, utils.CodeKeccakLeafKey)
+	gas += aw.TouchAddressOnReadAndComputeGas(targetAddr, zeroTreeIndex, utils.NonceLeafKey)
+	if sendsValue {
+		gas += aw.TouchAddressOnWriteAndComputeGas(targetAddr, zeroTreeIndex, utils.BalanceLeafKey)
+	} else {
+		gas += aw.TouchAddressOnReadAndComputeGas(targetAddr, zeroTreeIndex, utils.BalanceLeafKey)
+	}
+	return gas
+}
+
+func (aw *AccessWitness) TouchAddressOnWriteAndComputeGas(addr []byte, treeIndex uint256.Int, subIndex byte) uint64 {
+	return aw.touchAddressAndChargeGas(addr, treeIndex, subIndex, true)
+}
+
+func (aw *AccessWitness) TouchAddressOnReadAndComputeGas(addr []byte, treeIndex uint256.Int, subIndex byte) uint64 {
+	return aw.touchAddressAndChargeGas(addr, treeIndex, subIndex, false)
+}
+
+func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, treeIndex uint256.Int, subIndex byte, isWrite bool) uint64 {
+	stemRead, selectorRead, stemWrite, selectorWrite, selectorFill := aw.touchAddress(addr, treeIndex, subIndex, isWrite)
+
+	var gas uint64
 	if stemRead {
 		gas += params.WitnessBranchReadCost
 	}
@@ -180,231 +193,62 @@ func (aw *AccessWitness) touchAddressAndChargeGas(addr []byte, isWrite bool) uin
 	return gas
 }
 
-func (aw *AccessWitness) TouchAddressOnWriteAndComputeGas(addr []byte) uint64 {
+// touchAddress adds any missing access event to the witness.
-	return aw.touchAddressAndChargeGas(addr, true)
+func (aw *AccessWitness) touchAddress(addr []byte, treeIndex uint256.Int, subIndex byte, isWrite bool) (bool, bool, bool, bool, bool) {
+	branchKey := newBranchAccessKey(addr, treeIndex)
+	chunkKey := newChunkAccessKey(branchKey, subIndex)
+
+	// Read access.
+	var branchRead, chunkRead bool
+	if _, hasStem := aw.branches[branchKey]; !hasStem {
+		branchRead = true
+		aw.branches[branchKey] = AccessWitnessReadFlag
+	}
+	if _, hasSelector := aw.chunks[chunkKey]; !hasSelector {
+		chunkRead = true
+		aw.chunks[chunkKey] = AccessWitnessReadFlag
 	}
 
-func (aw *AccessWitness) TouchAddressOnReadAndComputeGas(addr []byte) uint64 {
+	// Write access.
-	return aw.touchAddressAndChargeGas(addr, false)
+	var branchWrite, chunkWrite, chunkFill bool
+	if isWrite {
+		if (aw.branches[branchKey] & AccessWitnessWriteFlag) == 0 {
+			branchWrite = true
+			aw.branches[branchKey] |= AccessWitnessWriteFlag
 		}
 
-// Merge is used to merge the witness that got generated during the execution
+		chunkValue := aw.chunks[chunkKey]
-// of a tx, with the accumulation of witnesses that were generated during the
+		if (chunkValue & AccessWitnessWriteFlag) == 0 {
-// execution of all the txs preceding this one in a given block.
+			chunkWrite = true
-func (aw *AccessWitness) Merge(other *AccessWitness) {
+			aw.chunks[chunkKey] |= AccessWitnessWriteFlag
-	for k := range other.Branches {
-		if _, ok := aw.Branches[k]; !ok {
-			aw.Branches[k] = other.Branches[k]
-		}
 		}
 
-	for k, chunk := range other.Chunks {
+		// TODO: charge chunk filling costs if the leaf was previously empty in the state
-		if _, ok := aw.Chunks[k]; !ok {
-			aw.Chunks[k] = chunk
-		}
 	}
 
-	for k, v := range other.InitialValue {
+	return branchRead, chunkRead, branchWrite, chunkWrite, chunkFill
-		if _, ok := aw.InitialValue[k]; !ok {
-			aw.InitialValue[k] = v
-		}
 }
 
-	// TODO see if merging improves performance
+type branchAccessKey struct {
-	//for k, v := range other.addrToPoint {
+	addr      common.Address
-	//if _, ok := aw.addrToPoint[k]; !ok {
+	treeIndex uint256.Int
-	//aw.addrToPoint[k] = v
-	//}
-	//}
 }
 
-// Key returns, predictably, the list of keys that were touched during the
+func newBranchAccessKey(addr []byte, treeIndex uint256.Int) branchAccessKey {
-// buildup of the access witness.
+	var sk branchAccessKey
-func (aw *AccessWitness) Keys() [][]byte {
+	copy(sk.addr[:], addr)
-	keys := make([][]byte, 0, len(aw.Chunks))
+	sk.treeIndex = treeIndex
-	for key := range aw.Chunks {
+	return sk
-		var k [32]byte
-		copy(k[:], key[:])
-		keys = append(keys, k[:])
-	}
-	return keys
 }
 
-func (aw *AccessWitness) KeyVals() map[string][]byte {
+type chunkAccessKey struct {
-	result := make(map[string][]byte)
+	branchAccessKey
-	for k, v := range aw.InitialValue {
+	leafKey byte
-		result[k] = v
-	}
-	return result
 }
 
-func (aw *AccessWitness) Copy() *AccessWitness {
+func newChunkAccessKey(branchKey branchAccessKey, leafKey byte) chunkAccessKey {
-	naw := &AccessWitness{
+	var lk chunkAccessKey
-		Branches:     make(map[VerkleStem]Mode),
+	lk.branchAccessKey = branchKey
-		Chunks:       make(map[common.Hash]Mode),
+	lk.leafKey = leafKey
-		InitialValue: make(map[string][]byte),
+	return lk
-	}
-
-	naw.Merge(aw)
-
-	return naw
-}
-
-func (aw *AccessWitness) GetTreeKeyVersionCached(addr []byte) []byte {
-	return aw.statedb.db.(*cachingDB).addrToPoint.GetTreeKeyVersionCached(addr)
-}
-
-func (aw *AccessWitness) TouchAndChargeProofOfAbsence(addr []byte) uint64 {
-	var (
-		balancekey, cskey, ckkey, noncekey [32]byte
-		gas                                uint64
-	)
-
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(addr[:])
-	copy(balancekey[:], versionkey)
-	balancekey[31] = utils.BalanceLeafKey
-	copy(noncekey[:], versionkey)
-	noncekey[31] = utils.NonceLeafKey
-	copy(cskey[:], versionkey)
-	cskey[31] = utils.CodeSizeLeafKey
-	copy(ckkey[:], versionkey)
-	ckkey[31] = utils.CodeKeccakLeafKey
-
-	gas += aw.TouchAddressOnReadAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnReadAndComputeGas(balancekey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(cskey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(ckkey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(noncekey[:])
-	return gas
-}
-
-func (aw *AccessWitness) TouchAndChargeMessageCall(addr []byte) uint64 {
-	var (
-		gas   uint64
-		cskey [32]byte
-	)
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(addr[:])
-	copy(cskey[:], versionkey)
-	cskey[31] = utils.CodeSizeLeafKey
-	gas += aw.TouchAddressOnReadAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnReadAndComputeGas(cskey[:])
-	return gas
-}
-
-func (aw *AccessWitness) TouchAndChargeValueTransfer(callerAddr, targetAddr []byte) uint64 {
-	var gas uint64
-	gas += aw.TouchAddressOnWriteAndComputeGas(utils.GetTreeKeyBalance(callerAddr[:]))
-	gas += aw.TouchAddressOnWriteAndComputeGas(utils.GetTreeKeyBalance(targetAddr[:]))
-	return gas
-}
-
-// TouchAndChargeContractCreateInit charges access costs to initiate
-// a contract creation
-func (aw *AccessWitness) TouchAndChargeContractCreateInit(addr []byte, createSendsValue bool) uint64 {
-	var (
-		balancekey, ckkey, noncekey [32]byte
-		gas                         uint64
-	)
-
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(addr[:])
-	copy(balancekey[:], versionkey)
-	balancekey[31] = utils.BalanceLeafKey
-	copy(noncekey[:], versionkey)
-	noncekey[31] = utils.NonceLeafKey
-	copy(ckkey[:], versionkey)
-	ckkey[31] = utils.CodeKeccakLeafKey
-
-	gas += aw.TouchAddressOnWriteAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnWriteAndComputeGas(noncekey[:])
-	if createSendsValue {
-		gas += aw.TouchAddressOnWriteAndComputeGas(balancekey[:])
-	}
-	gas += aw.TouchAddressOnWriteAndComputeGas(ckkey[:])
-	return gas
-}
-
-// TouchAndChargeContractCreateCompleted charges access access costs after
-// the completion of a contract creation to populate the created account in
-// the tree
-func (aw *AccessWitness) TouchAndChargeContractCreateCompleted(addr []byte, withValue bool) uint64 {
-	var (
-		balancekey, cskey, ckkey, noncekey [32]byte
-		gas                                uint64
-	)
-
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(addr[:])
-	copy(balancekey[:], versionkey)
-	balancekey[31] = utils.BalanceLeafKey
-	copy(noncekey[:], versionkey)
-	noncekey[31] = utils.NonceLeafKey
-	copy(cskey[:], versionkey)
-	cskey[31] = utils.CodeSizeLeafKey
-	copy(ckkey[:], versionkey)
-	ckkey[31] = utils.CodeKeccakLeafKey
-
-	gas += aw.TouchAddressOnWriteAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnWriteAndComputeGas(balancekey[:])
-	gas += aw.TouchAddressOnWriteAndComputeGas(cskey[:])
-	gas += aw.TouchAddressOnWriteAndComputeGas(ckkey[:])
-	gas += aw.TouchAddressOnWriteAndComputeGas(noncekey[:])
-	return gas
-}
-
-func (aw *AccessWitness) TouchTxOriginAndComputeGas(originAddr []byte) uint64 {
-	var (
-		balancekey, cskey, ckkey, noncekey [32]byte
-		gas                                uint64
-	)
-
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(originAddr[:])
-	copy(balancekey[:], versionkey)
-	balancekey[31] = utils.BalanceLeafKey
-	copy(noncekey[:], versionkey)
-	noncekey[31] = utils.NonceLeafKey
-	copy(cskey[:], versionkey)
-	cskey[31] = utils.CodeSizeLeafKey
-	copy(ckkey[:], versionkey)
-	ckkey[31] = utils.CodeKeccakLeafKey
-
-	gas += aw.TouchAddressOnReadAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnReadAndComputeGas(cskey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(ckkey[:])
-	gas += aw.TouchAddressOnWriteAndComputeGas(noncekey[:])
-	gas += aw.TouchAddressOnWriteAndComputeGas(balancekey[:])
-
-	return gas
-}
-
-func (aw *AccessWitness) TouchTxExistingAndComputeGas(targetAddr []byte, sendsValue bool) uint64 {
-	var (
-		balancekey, cskey, ckkey, noncekey [32]byte
-		gas                                uint64
-	)
-
-	// Only evaluate the polynomial once
-	versionkey := aw.GetTreeKeyVersionCached(targetAddr[:])
-	copy(balancekey[:], versionkey)
-	balancekey[31] = utils.BalanceLeafKey
-	copy(noncekey[:], versionkey)
-	noncekey[31] = utils.NonceLeafKey
-	copy(cskey[:], versionkey)
-	cskey[31] = utils.CodeSizeLeafKey
-	copy(ckkey[:], versionkey)
-	ckkey[31] = utils.CodeKeccakLeafKey
-
-	gas += aw.TouchAddressOnReadAndComputeGas(versionkey)
-	gas += aw.TouchAddressOnReadAndComputeGas(cskey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(ckkey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(noncekey[:])
-	gas += aw.TouchAddressOnReadAndComputeGas(balancekey[:])
-
-	if sendsValue {
-		gas += aw.TouchAddressOnWriteAndComputeGas(balancekey[:])
-	}
-	return gas
 }

core/state/statedb.go

@@ -174,7 +174,7 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 		hasher:               crypto.NewKeccakState(),
 	}
 	if tr.IsVerkle() {
-		sdb.witness = NewAccessWitness(sdb)
+		sdb.witness = sdb.NewAccessWitness()
 		// if sdb.snaps == nil {
 		// snapconfig := snapshot.Config{
 		// 	CacheSize:  256,
@@ -206,9 +206,13 @@ func (s *StateDB) Snaps() *snapshot.Tree {
 	return s.snaps
 }
 
+func (s *StateDB) NewAccessWitness() *AccessWitness {
+	return NewAccessWitness(s.db.(*cachingDB).addrToPoint)
+}
+
 func (s *StateDB) Witness() *AccessWitness {
 	if s.witness == nil {
-		s.witness = NewAccessWitness(s)
+		s.witness = s.NewAccessWitness()
 	}
 	return s.witness
 }

core/state_processor.go

@@ -320,7 +320,7 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg
 func applyTransaction(msg *Message, config *params.ChainConfig, gp *GasPool, statedb *state.StateDB, blockNumber *big.Int, blockHash common.Hash, tx *types.Transaction, usedGas *uint64, evm *vm.EVM) (*types.Receipt, error) {
 	// Create a new context to be used in the EVM environment.
 	txContext := NewEVMTxContext(msg)
-	txContext.Accesses = state.NewAccessWitness(statedb)
+	txContext.Accesses = statedb.NewAccessWitness()
 	evm.Reset(txContext, statedb)
 
 	// Apply the transaction to the current state (included in the env).

core/vm/contract.go

@@ -20,9 +20,6 @@ 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"
 )
 
@@ -52,13 +49,11 @@ type Contract struct {
 	CallerAddress common.Address
 	caller        ContractRef
 	self          ContractRef
-	addressPoint  *verkle.Point
 
 	jumpdests map[common.Hash]bitvec // Aggregated result of JUMPDEST analysis.
 	analysis  bitvec                 // Locally cached result of JUMPDEST analysis
 
 	Code     []byte
-	Chunks   trie.ChunkedCode
 	CodeHash common.Hash
 	CodeAddr *common.Address
 	Input    []byte
@@ -180,14 +175,6 @@ func (c *Contract) Address() common.Address {
 	return c.self.Address()
 }
 
-func (c *Contract) AddressPoint() *verkle.Point {
-	if c.addressPoint == nil {
-		c.addressPoint = utils.EvaluateAddressPoint(c.Address().Bytes())
-	}
-
-	return c.addressPoint
-}
-
 // Value returns the contract's value (sent to it from it's caller)
 func (c *Contract) Value() *big.Int {
 	return c.value

core/vm/evm.go

@@ -138,7 +138,7 @@ func NewEVM(blockCtx BlockContext, txCtx TxContext, statedb StateDB, chainConfig
 		chainRules:  chainConfig.Rules(blockCtx.BlockNumber, blockCtx.Random != nil, blockCtx.Time),
 	}
 	if txCtx.Accesses == nil && chainConfig.IsVerkle(blockCtx.BlockNumber, blockCtx.Time) {
-		txCtx.Accesses = state.NewAccessWitness(evm.StateDB.(*state.StateDB))
+		txCtx.Accesses = evm.StateDB.(*state.StateDB).NewAccessWitness()
 	}
 	evm.interpreter = NewEVMInterpreter(evm)
 	return evm
@@ -148,7 +148,7 @@ func NewEVM(blockCtx BlockContext, txCtx TxContext, statedb StateDB, chainConfig
 // This is not threadsafe and should only be done very cautiously.
 func (evm *EVM) Reset(txCtx TxContext, statedb StateDB) {
 	if txCtx.Accesses == nil && evm.chainRules.IsVerkle {
-		txCtx.Accesses = state.NewAccessWitness(evm.StateDB.(*state.StateDB))
+		txCtx.Accesses = evm.StateDB.(*state.StateDB).NewAccessWitness()
 	}
 	evm.TxContext = txCtx
 	evm.StateDB = statedb
@@ -530,7 +530,7 @@ func (evm *EVM) create(caller ContractRef, codeAndHash *codeAndHash, gas uint64,
 	}
 
 	if err == nil && evm.chainRules.IsVerkle {
-		if !contract.UseGas(evm.Accesses.TouchAndChargeContractCreateCompleted(address.Bytes()[:], value.Sign() != 0)) {
+		if !contract.UseGas(evm.Accesses.TouchAndChargeContractCreateCompleted(address.Bytes()[:])) {
 			evm.StateDB.RevertToSnapshot(snapshot)
 			err = ErrOutOfGas
 		}

core/vm/gas_table.go

@@ -24,6 +24,7 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"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
@@ -101,8 +102,7 @@ func gasExtCodeSize(evm *EVM, contract *Contract, stack *Stack, mem *Memory, mem
 	usedGas := uint64(0)
 	slot := stack.Back(0)
 	if evm.chainRules.IsVerkle {
-		index := trieUtils.GetTreeKeyCodeSize(slot.Bytes())
+		usedGas += evm.TxContext.Accesses.TouchAddressOnReadAndComputeGas(slot.Bytes(), uint256.Int{}, trieUtils.CodeSizeLeafKey)
-		usedGas += evm.TxContext.Accesses.TouchAddressOnReadAndComputeGas(index)
 	}
 
 	return usedGas, nil
@@ -113,8 +113,8 @@ func gasSLoad(evm *EVM, contract *Contract, stack *Stack, mem *Memory, memorySiz
 
 	if evm.chainRules.IsVerkle {
 		where := stack.Back(0)
-		index := trieUtils.GetTreeKeyStorageSlotWithEvaluatedAddress(contract.AddressPoint(), where.Bytes())
+		treeIndex, subIndex := trieUtils.GetTreeKeyStorageSlotTreeIndexes(where.Bytes())
-		usedGas += evm.Accesses.TouchAddressOnReadAndComputeGas(index)
+		usedGas += evm.Accesses.TouchAddressOnReadAndComputeGas(contract.Address().Bytes(), *treeIndex, subIndex)
 	}
 
 	return usedGas, nil

core/vm/instructions.go

@@ -23,7 +23,6 @@ import (
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/params"
-	"github.com/ethereum/go-ethereum/trie"
 	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/holiman/uint256"
 )
@@ -347,8 +346,7 @@ func opExtCodeSize(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext)
 	slot := scope.Stack.peek()
 	cs := uint64(interpreter.evm.StateDB.GetCodeSize(slot.Bytes20()))
 	if interpreter.evm.chainRules.IsVerkle {
-		index := trieUtils.GetTreeKeyCodeSize(slot.Bytes())
+		statelessGas := interpreter.evm.Accesses.TouchAddressOnReadAndComputeGas(slot.Bytes(), uint256.Int{}, trieUtils.CodeSizeLeafKey)
-		statelessGas := interpreter.evm.Accesses.TouchAddressOnReadAndComputeGas(index)
 		scope.Contract.UseGas(statelessGas)
 	}
 	slot.SetUint64(cs)
@@ -373,80 +371,43 @@ func opCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([
 		uint64CodeOffset = 0xffffffffffffffff
 	}
 
+	contractAddr := scope.Contract.Address()
 	paddedCodeCopy, copyOffset, nonPaddedCopyLength := getDataAndAdjustedBounds(scope.Contract.Code, uint64CodeOffset, length.Uint64())
 	if interpreter.evm.chainRules.IsVerkle {
-		scope.Contract.UseGas(touchEachChunksOnReadAndChargeGas(copyOffset, nonPaddedCopyLength, scope.Contract, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment))
+		scope.Contract.UseGas(touchCodeChunksRangeOnReadAndChargeGas(contractAddr[:], copyOffset, nonPaddedCopyLength, uint64(len(scope.Contract.Code)), interpreter.evm.Accesses))
 	}
 	scope.Memory.Set(memOffset.Uint64(), uint64(len(paddedCodeCopy)), paddedCodeCopy)
 	return nil, nil
 }
 
-// touchChunkOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
+// touchCodeChunksRangeOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
-func touchChunkOnReadAndChargeGas(chunks trie.ChunkedCode, offset uint64, evals [][]byte, code []byte, accesses *state.AccessWitness, deployment bool) uint64 {
+func touchCodeChunksRangeOnReadAndChargeGas(contractAddr []byte, startPC, size uint64, codeLen uint64, accesses *state.AccessWitness) uint64 {
-	// note that in the case where the executed 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
-	// 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 code != nil && offset > uint64(len(code)) {
-		return 0
-	}
-	var (
-		chunknr             = offset / 31
-		statelessGasCharged uint64
-	)
-
-	// Build the chunk address from the evaluated address of its whole group
-	var index [32]byte
-	copy(index[:], evals[chunknr/256])
-	index[31] = byte((128 + chunknr) % 256)
-
-	var overflow bool
-	statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, accesses.TouchAddressOnReadAndComputeGas(index[:]))
-	if overflow {
-		panic("overflow when adding gas")
-	}
-
-	return statelessGasCharged
-}
-
-// touchEachChunksOnReadAndChargeGas is a helper function to touch every chunk in a code range and charge witness gas costs
-func touchEachChunksOnReadAndChargeGas(offset, size uint64, contract *Contract, code []byte, accesses *state.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
 	// 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 len(code) == 0 && size == 0 || offset > uint64(len(code)) {
+	if (codeLen == 0 && size == 0) || startPC > codeLen {
 		return 0
 	}
-	var (
+
-		statelessGasCharged uint64
+	// endPC is the last PC that must be touched.
-		endOffset           uint64
+	endPC := startPC + size - 1
-	)
+	if startPC+size > codeLen {
-	if code != nil && offset+size > uint64(len(code)) {
+		endPC = codeLen
-		endOffset = uint64(len(code))
-	} else {
-		endOffset = offset + size
 	}
 
-	// endOffset - 1 since if the end offset is aligned on a chunk boundary,
+	var statelessGasCharged uint64
-	// the last chunk should not be included.
+	for chunkNumber := startPC / 31; chunkNumber <= endPC/31; chunkNumber++ {
-	for i := offset / 31; i <= (endOffset-1)/31; i++ {
+		treeIndex := *uint256.NewInt((chunkNumber + 128) / 256)
-		// only charge for+cache the chunk if it isn't already present
+		subIndex := byte((chunkNumber + 128) % 256)
-		if !accesses.HasCodeChunk(contract.Address().Bytes(), i) {
+		gas := accesses.TouchAddressOnReadAndComputeGas(contractAddr, treeIndex, subIndex)
-			index := trieUtils.GetTreeKeyCodeChunkWithEvaluatedAddress(contract.AddressPoint(), uint256.NewInt(i))
-
 		var overflow bool
-			statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, accesses.TouchAddressOnReadAndComputeGas(index))
+		statelessGasCharged, overflow = math.SafeAdd(statelessGasCharged, gas)
 		if overflow {
 			panic("overflow when adding gas")
 		}
-
-			accesses.SetCachedCodeChunk(contract.Address().Bytes(), i)
-		}
 	}
 
 	return statelessGasCharged
@@ -469,11 +430,11 @@ func opExtCodeCopy(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext)
 		code := interpreter.evm.StateDB.GetCode(addr)
 		contract := &Contract{
 			Code: code,
-			Chunks: trie.ChunkedCode(code),
 			self: AccountRef(addr),
 		}
 		paddedCodeCopy, copyOffset, nonPaddedCopyLength := getDataAndAdjustedBounds(code, uint64CodeOffset, length.Uint64())
-		touchEachChunksOnReadAndChargeGas(copyOffset, nonPaddedCopyLength, contract, code, interpreter.evm.Accesses, false)
+		gas := touchCodeChunksRangeOnReadAndChargeGas(addr[:], copyOffset, nonPaddedCopyLength, uint64(len(contract.Code)), interpreter.evm.Accesses)
+		scope.Contract.UseGas(gas)
 		scope.Memory.Set(memOffset.Uint64(), length.Uint64(), paddedCodeCopy)
 	} else {
 		codeCopy := getData(interpreter.evm.StateDB.GetCode(addr), uint64CodeOffset, length.Uint64())
@@ -1001,7 +962,8 @@ func opPush1(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]by
 		if interpreter.evm.chainRules.IsVerkle && *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, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
+			contractAddr := scope.Contract.Address()
+			statelessGas := touchCodeChunksRangeOnReadAndChargeGas(contractAddr[:], *pc+1, uint64(1), uint64(len(scope.Contract.Code)), interpreter.evm.Accesses)
 			scope.Contract.UseGas(statelessGas)
 		}
 	} else {
@@ -1026,7 +988,8 @@ func makePush(size uint64, pushByteSize int) executionFunc {
 		}
 
 		if interpreter.evm.chainRules.IsVerkle {
-			statelessGas := touchEachChunksOnReadAndChargeGas(uint64(startMin), uint64(pushByteSize), scope.Contract, scope.Contract.Code, interpreter.evm.Accesses, scope.Contract.IsDeployment)
+			contractAddr := scope.Contract.Address()
+			statelessGas := touchCodeChunksRangeOnReadAndChargeGas(contractAddr[:], uint64(startMin), uint64(pushByteSize), uint64(len(scope.Contract.Code)), interpreter.evm.Accesses)
 			scope.Contract.UseGas(statelessGas)
 		}
 

core/vm/interpreter.go

@@ -21,10 +21,6 @@ import (
 	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/log"
-	"github.com/ethereum/go-ethereum/trie"
-	"github.com/ethereum/go-ethereum/trie/utils"
-	"github.com/gballet/go-verkle"
-	"github.com/holiman/uint256"
 )
 
 // Config are the configuration options for the Interpreter
@@ -152,8 +148,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
 		debug   = in.evm.Config.Tracer != nil
-
-		chunkEvals [][]byte
 	)
 	// Don't move this deferred function, it's placed before the capturestate-deferred method,
 	// so that it get's executed _after_: the capturestate needs the stacks before
@@ -175,21 +169,6 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 		}()
 	}
 
-	// Evaluate one address per group of 256, 31-byte chunks
-	if in.evm.chainRules.IsCancun && !contract.IsDeployment {
-		contract.Chunks = trie.ChunkifyCode(contract.Code)
-
-		// number of extra stems to evaluate after the header stem
-		extraEvals := (len(contract.Chunks) + 127) / verkle.NodeWidth
-
-		chunkEvals = make([][]byte, extraEvals+1)
-		for i := 1; i < extraEvals+1; i++ {
-			chunkEvals[i] = utils.GetTreeKeyCodeChunkWithEvaluatedAddress(contract.AddressPoint(), uint256.NewInt(uint64(i)*256))
-		}
-		// Header account is already known, it's the header account
-		chunkEvals[0] = utils.GetTreeKeyVersionWithEvaluatedAddress(contract.AddressPoint())
-	}
-
 	// The Interpreter main run loop (contextual). This loop runs until either an
 	// explicit STOP, RETURN or SELFDESTRUCT is executed, an error occurred during
 	// the execution of one of the operations or until the done flag is set by the
@@ -200,10 +179,11 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 			logged, pcCopy, gasCopy = false, pc, contract.Gas
 		}
 
-		if contract.Chunks != nil {
+		if in.evm.chainRules.IsCancun && !contract.IsDeployment {
 			// if the PC ends up in a new "chunk" of verkleized code, charge the
 			// associated costs.
-			contract.Gas -= touchChunkOnReadAndChargeGas(contract.Chunks, pc, chunkEvals, contract.Code, in.evm.TxContext.Accesses, contract.IsDeployment)
+			contractAddr := contract.Address()
+			contract.Gas -= touchCodeChunksRangeOnReadAndChargeGas(contractAddr[:], pc, 1, uint64(len(contract.Code)), in.evm.TxContext.Accesses)
 		}
 
 		// Get the operation from the jump table and validate the stack to ensure there are

core/vm/operations_acl.go

@@ -22,7 +22,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/math"
 	"github.com/ethereum/go-ethereum/params"
-	trieUtils "github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/ethereum/go-ethereum/trie/utils"
 )
 
 func makeGasSStoreFunc(clearingRefund uint64) gasFunc {
@@ -53,8 +53,8 @@ func makeGasSStoreFunc(clearingRefund uint64) gasFunc {
 		value := common.Hash(y.Bytes32())
 
 		if evm.chainRules.IsVerkle {
-			index := trieUtils.GetTreeKeyStorageSlotWithEvaluatedAddress(contract.AddressPoint(), x.Bytes())
+			treeIndex, subIndex := utils.GetTreeKeyStorageSlotTreeIndexes(x.Bytes())
-			cost += evm.Accesses.TouchAddressOnWriteAndComputeGas(index)
+			cost += evm.Accesses.TouchAddressOnWriteAndComputeGas(contract.Address().Bytes(), *treeIndex, subIndex)
 		}
 
 		if current == value { // noop (1)
@@ -113,9 +113,9 @@ func gasSLoadEIP2929(evm *EVM, contract *Contract, stack *Stack, mem *Memory, me
 
 	if evm.chainRules.IsVerkle {
 		where := stack.Back(0)
+		treeIndex, subIndex := utils.GetTreeKeyStorageSlotTreeIndexes(where.Bytes())
 		addr := contract.Address()
-		index := trieUtils.GetTreeKeyStorageSlot(addr[:], where)
+		gasUsed += evm.Accesses.TouchAddressOnReadAndComputeGas(addr.Bytes(), *treeIndex, subIndex)
-		gasUsed += evm.Accesses.TouchAddressOnReadAndComputeGas(index)
 	}
 
 	// Check slot presence in the access list

trie/utils/verkle.go

@@ -144,7 +144,7 @@ func GetTreeKeyVersion(address []byte) []byte {
 }
 
 func GetTreeKeyVersionWithEvaluatedAddress(addrp *verkle.Point) []byte {
-	return getTreeKeyWithEvaluatedAddess(addrp, zero, VersionLeafKey)
+	return GetTreeKeyWithEvaluatedAddess(addrp, zero, VersionLeafKey)
 }
 
 func GetTreeKeyBalance(address []byte) []byte {
@@ -182,7 +182,7 @@ func GetTreeKeyCodeChunkWithEvaluatedAddress(addressPoint *verkle.Point, chunk *
 	if len(subIndexMod) != 0 {
 		subIndex = byte(subIndexMod[0])
 	}
-	return getTreeKeyWithEvaluatedAddess(addressPoint, treeIndex, subIndex)
+	return GetTreeKeyWithEvaluatedAddess(addressPoint, treeIndex, subIndex)
 }
 
 func GetTreeKeyStorageSlot(address []byte, storageKey *uint256.Int) []byte {
@@ -221,7 +221,7 @@ func PointToHash(evaluated *verkle.Point, suffix byte) []byte {
 	return retb[:]
 }
 
-func getTreeKeyWithEvaluatedAddess(evaluated *verkle.Point, treeIndex *uint256.Int, subIndex byte) []byte {
+func GetTreeKeyWithEvaluatedAddess(evaluated *verkle.Point, treeIndex *uint256.Int, subIndex byte) []byte {
 	var poly [5]fr.Element
 
 	poly[0].SetZero()
@@ -269,6 +269,11 @@ func EvaluateAddressPoint(address []byte) *verkle.Point {
 }
 
 func GetTreeKeyStorageSlotWithEvaluatedAddress(evaluated *verkle.Point, storageKey []byte) []byte {
+	treeIndex, subIndex := GetTreeKeyStorageSlotTreeIndexes(storageKey)
+	return GetTreeKeyWithEvaluatedAddess(evaluated, treeIndex, subIndex)
+}
+
+func GetTreeKeyStorageSlotTreeIndexes(storageKey []byte) (*uint256.Int, byte) {
 	// Note that `pos` must be a big.Int and not a uint256.Int, because the subsequent
 	// arithmetics operations could overflow. (e.g: imagine if storageKey is 2^256-1)
 	pos := new(big.Int).SetBytes(storageKey)
@@ -286,5 +291,5 @@ func GetTreeKeyStorageSlotWithEvaluatedAddress(evaluated *verkle.Point, storageK
 	posBytes := pos.Bytes()
 	subIndex := posBytes[len(posBytes)-1]
 
-	return getTreeKeyWithEvaluatedAddess(evaluated, treeIndex, subIndex)
+	return treeIndex, subIndex
 }

trie/verkle_test.go

@@ -65,11 +65,9 @@ func TestReproduceTree(t *testing.T) {
 	}
 
 	root := verkle.New()
-	kv := make(map[string][]byte)
 
 	for i, key := range presentKeys {
 		root.Insert(key, values[i], nil)
-		kv[string(key)] = values[i]
 	}
 
 	proof, Cs, zis, yis, _ := verkle.MakeVerkleMultiProof(root, append(presentKeys, absentKeys...))
@@ -286,11 +284,9 @@ func TestReproduceCondrieuStemAggregationInProofOfAbsence(t *testing.T) {
 	}
 
 	root := verkle.New()
-	kv := make(map[string][]byte)
 
 	for i, key := range presentKeys {
 		root.Insert(key, values[i], nil)
-		kv[string(key)] = values[i]
 	}
 
 	proof, Cs, zis, yis, _ := verkle.MakeVerkleMultiProof(root, append(presentKeys, absentKeys...))
@@ -333,11 +329,9 @@ func TestReproduceCondrieuPoAStemConflictWithAnotherStem(t *testing.T) {
 	}
 
 	root := verkle.New()
-	kv := make(map[string][]byte)
 
 	for i, key := range presentKeys {
 		root.Insert(key, values[i], nil)
-		kv[string(key)] = values[i]
 	}
 
 	proof, Cs, zis, yis, _ := verkle.MakeVerkleMultiProof(root, append(presentKeys, absentKeys...))