core/state: rewrite a new optimized keyValueMigrator (#256)

* trie/utils: add helper to calculate code tree indices

* core/state: rewrite optimized version of keyValueMigrator

Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>

* trie/verkle: remove uint256 allocs (#257)

Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>

---------

Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>

core/state_processor.go

@@ -25,6 +25,8 @@ import (
 	"io"
 	"math/big"
 	"os"
+	"runtime"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -170,7 +172,7 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg
 		// mkv will be assiting in the collection of up to maxMovedCount key values to be migrated to the VKT.
 		// It has internal caches to do efficient MPT->VKT key calculations, which will be discarded after
 		// this function.
-		mkv := &keyValueMigrator{vktLeafData: make(map[string]*verkle.BatchNewLeafNodeData)}
+		mkv := newKeyValueMigrator()
 		// move maxCount accounts into the verkle tree, starting with the
 		// slots from the previous account.
 		count := 0
@@ -297,8 +299,17 @@ func (p *StateProcessor) Process(block *types.Block, statedb *state.StateDB, cfg
 		}
 		migrdb.SetCurrentPreimageOffset(preimageSeek)
 
-		log.Info("Collected and prepared key values from base tree", "count", count, "duration", time.Since(now), "last account", statedb.Database().GetCurrentAccountHash())
+		log.Info("Collected key values from base tree", "count", count, "duration", time.Since(now), "last account", statedb.Database().GetCurrentAccountHash())
 
+		// Take all the collected key-values and prepare the new leaf values.
+		// This fires a background routine that will start doing the work that
+		// migrateCollectedKeyValues() will use to insert into the tree.
+		//
+		// TODO: Now both prepare() and migrateCollectedKeyValues() are next to each other, but
+		//       after we fix an existing bug, we can call prepare() before the block execution and
+		//       let it do the work in the background. After the block execution and finalization
+		//       finish, we can call migrateCollectedKeyValues() which should already find everything ready.
+		mkv.prepare()
 		now = time.Now()
 		if err := mkv.migrateCollectedKeyValues(tt.Overlay()); err != nil {
 			return nil, nil, 0, fmt.Errorf("could not migrate key values: %w", err)
@@ -380,30 +391,60 @@ func ApplyTransaction(config *params.ChainConfig, bc ChainContext, author *commo
 	return applyTransaction(msg, config, gp, statedb, header.Number, header.Hash(), tx, usedGas, vmenv)
 }
 
-// keyValueMigrator is a helper struct that collects key-values from the base tree.
+var zeroTreeIndex uint256.Int
-// The walk is done in account order, so **we assume** the APIs hold this invariant. This is
-// useful to be smart about caching banderwagon.Points to make VKT key calculations faster.
-type keyValueMigrator struct {
-	currAddr      []byte
-	currAddrPoint *verkle.Point
 
-	vktLeafData map[string]*verkle.BatchNewLeafNodeData
+// keyValueMigrator is a helper module that collects key-values from the overlay-tree migration for Verkle Trees.
+// It assumes that the walk of the base tree is done in address-order, so it exploit that fact to
+// collect the key-values in a way that is efficient.
+type keyValueMigrator struct {
+	// leafData contains the values for the future leaf for a particular VKT branch.
+	leafData []migratedKeyValue
+
+	// When prepare() is called, it will start a background routine that will process the leafData
+	// saving the result in newLeaves to be used by migrateCollectedKeyValues(). The background
+	// routine signals that it is done by closing processingReady.
+	processingReady chan struct{}
+	newLeaves       []verkle.LeafNode
+	prepareErr      error
+}
+
+func newKeyValueMigrator() *keyValueMigrator {
+	// We do initialize the VKT config since prepare() might indirectly make multiple GetConfig() calls
+	// in different goroutines when we never called GetConfig() before, causing a race considering the way
+	// that `config` is designed in go-verkle.
+	// TODO: jsign as a fix for this in the PR where we move to a file-less precomp, since it allows safe
+	//       concurrent calls to GetConfig(). When that gets merged, we can remove this line.
+	_ = verkle.GetConfig()
+	return &keyValueMigrator{
+		processingReady: make(chan struct{}),
+		leafData:        make([]migratedKeyValue, 0, 10_000),
+	}
+}
+
+type migratedKeyValue struct {
+	branchKey    branchKey
+	leafNodeData verkle.BatchNewLeafNodeData
+}
+type branchKey struct {
+	addr      common.Address
+	treeIndex uint256.Int
+}
+
+func newBranchKey(addr []byte, treeIndex *uint256.Int) branchKey {
+	var sk branchKey
+	copy(sk.addr[:], addr)
+	sk.treeIndex = *treeIndex
+	return sk
 }
 
 func (kvm *keyValueMigrator) addStorageSlot(addr []byte, slotNumber []byte, slotValue []byte) {
-	addrPoint := kvm.getAddrPoint(addr)
+	treeIndex, subIndex := tutils.GetTreeKeyStorageSlotTreeIndexes(slotNumber)
-
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex))
-	vktKey := tutils.GetTreeKeyStorageSlotWithEvaluatedAddress(addrPoint, slotNumber)
+	leafNodeData.Values[subIndex] = slotValue
-	leafNodeData := kvm.getOrInitLeafNodeData(vktKey)
-
-	leafNodeData.Values[vktKey[verkle.StemSize]] = slotValue
 }
 
 func (kvm *keyValueMigrator) addAccount(addr []byte, acc *types.StateAccount) {
-	addrPoint := kvm.getAddrPoint(addr)
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex))
-
-	vktKey := tutils.GetTreeKeyVersionWithEvaluatedAddress(addrPoint)
-	leafNodeData := kvm.getOrInitLeafNodeData(vktKey)
 
 	var version [verkle.LeafValueSize]byte
 	leafNodeData.Values[tutils.VersionLeafKey] = version[:]
@@ -419,16 +460,10 @@ func (kvm *keyValueMigrator) addAccount(addr []byte, acc *types.StateAccount) {
 	leafNodeData.Values[tutils.NonceLeafKey] = nonce[:]
 
 	leafNodeData.Values[tutils.CodeKeccakLeafKey] = acc.CodeHash[:]
-
-	// Code size is ignored here. If this isn't an EOA, the tree-walk will call
-	// addAccountCode with this information.
 }
 
 func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks []byte) {
-	addrPoint := kvm.getAddrPoint(addr)
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex))
-
-	vktKey := tutils.GetTreeKeyVersionWithEvaluatedAddress(addrPoint)
-	leafNodeData := kvm.getOrInitLeafNodeData(vktKey)
 
 	// Save the code size.
 	var codeSizeBytes [verkle.LeafValueSize]byte
@@ -442,8 +477,8 @@ func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks
 
 	// Potential further chunks, have their own leaf nodes.
 	for i := 128; i < len(chunks)/32; {
-		vktKey := tutils.GetTreeKeyCodeChunkWithEvaluatedAddress(addrPoint, uint256.NewInt(uint64(i)))
+		treeIndex, _ := tutils.GetTreeKeyCodeChunkIndices(uint256.NewInt(uint64(i)))
-		leafNodeData := kvm.getOrInitLeafNodeData(vktKey)
+		leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex))
 
 		j := i
 		for ; (j-i) < 256 && j < len(chunks)/32; j++ {
@@ -453,41 +488,79 @@ func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks
 	}
 }
 
-func (kvm *keyValueMigrator) getAddrPoint(addr []byte) *verkle.Point {
+func (kvm *keyValueMigrator) getOrInitLeafNodeData(bk branchKey) *verkle.BatchNewLeafNodeData {
-	if bytes.Equal(addr, kvm.currAddr) {
+	// Remember that keyValueMigration receives actions ordered by (address, subtreeIndex).
-		return kvm.currAddrPoint
+	// This means that we can assume that the last element of leafData is the one that we
-	}
+	// are looking for, or that we need to create a new one.
-	kvm.currAddr = addr
+	if len(kvm.leafData) == 0 || kvm.leafData[len(kvm.leafData)-1].branchKey != bk {
-	kvm.currAddrPoint = tutils.EvaluateAddressPoint(addr)
+		kvm.leafData = append(kvm.leafData, migratedKeyValue{
-	return kvm.currAddrPoint
+			branchKey: bk,
-}
+			leafNodeData: verkle.BatchNewLeafNodeData{
-
+				Stem:   nil, // It will be calculated in the prepare() phase, since it's CPU heavy.
-func (kvm *keyValueMigrator) getOrInitLeafNodeData(stem []byte) *verkle.BatchNewLeafNodeData {
-	stemStr := string(stem)
-	if _, ok := kvm.vktLeafData[stemStr]; !ok {
-		kvm.vktLeafData[stemStr] = &verkle.BatchNewLeafNodeData{
-			Stem:   stem[:verkle.StemSize],
 				Values: make(map[byte][]byte),
+			},
+		})
 	}
-	}
+	return &kvm.leafData[len(kvm.leafData)-1].leafNodeData
-	return kvm.vktLeafData[stemStr]
 }
 
-func (kvm *keyValueMigrator) migrateCollectedKeyValues(tree *trie.VerkleTrie) error {
+func (kvm *keyValueMigrator) prepare() {
-	// Transform the map into a slice.
+	// We fire a background routine to process the leafData and save the result in newLeaves.
-	nodeValues := make([]verkle.BatchNewLeafNodeData, 0, len(kvm.vktLeafData))
+	// The background routine signals that it is done by closing processingReady.
-	for _, vld := range kvm.vktLeafData {
+	go func() {
-		nodeValues = append(nodeValues, *vld)
+		// Step 1: We split kvm.leafData in numBatches batches, and we process each batch in a separate goroutine.
+		//         This fills each leafNodeData.Stem with the correct value.
+		var wg sync.WaitGroup
+		batchNum := runtime.NumCPU()
+		batchSize := (len(kvm.leafData) + batchNum - 1) / batchNum
+		for i := 0; i < len(kvm.leafData); i += batchSize {
+			start := i
+			end := i + batchSize
+			if end > len(kvm.leafData) {
+				end = len(kvm.leafData)
+			}
+			wg.Add(1)
+
+			batch := kvm.leafData[start:end]
+			go func() {
+				defer wg.Done()
+				var currAddr common.Address
+				var currPoint *verkle.Point
+				for i := range batch {
+					if batch[i].branchKey.addr != currAddr {
+						currAddr = batch[i].branchKey.addr
+						currPoint = tutils.EvaluateAddressPoint(currAddr[:])
+					}
+					stem := tutils.GetTreeKeyWithEvaluatedAddess(currPoint, &batch[i].branchKey.treeIndex, 0)
+					stem = stem[:verkle.StemSize]
+					batch[i].leafNodeData.Stem = stem
+				}
+			}()
+		}
+		wg.Wait()
+
+		// Step 2: Now that we have all stems (i.e: tree keys) calcualted, we can create the new leaves.
+		nodeValues := make([]verkle.BatchNewLeafNodeData, len(kvm.leafData))
+		for i := range kvm.leafData {
+			nodeValues[i] = kvm.leafData[i].leafNodeData
 		}
 
 		// Create all leaves in batch mode so we can optimize cryptography operations.
-	newLeaves, err := verkle.BatchNewLeafNode(nodeValues)
+		kvm.newLeaves, kvm.prepareErr = verkle.BatchNewLeafNode(nodeValues)
-	if err != nil {
+		close(kvm.processingReady)
-		return fmt.Errorf("failed to batch-create new leaf nodes")
+	}()
+}
+
+func (kvm *keyValueMigrator) migrateCollectedKeyValues(tree *trie.VerkleTrie) error {
+	now := time.Now()
+	<-kvm.processingReady
+	if kvm.prepareErr != nil {
+		return fmt.Errorf("failed to prepare key values: %w", kvm.prepareErr)
 	}
+	log.Info("Prepared key values from base tree", "duration", time.Since(now))
 
 	// Insert into the tree.
-	if err := tree.InsertMigratedLeaves(newLeaves); err != nil {
+	if err := tree.InsertMigratedLeaves(kvm.newLeaves); err != nil {
 		return fmt.Errorf("failed to insert migrated leaves: %w", err)
 	}
 

trie/utils/verkle.go

@@ -17,7 +17,7 @@
 package utils
 
 import (
-	"math/big"
+	"encoding/binary"
 	"sync"
 
 	"github.com/crate-crypto/go-ipa/bandersnatch/fr"
@@ -35,18 +35,14 @@ const (
 
 var (
 	zero                                = uint256.NewInt(0)
+	VerkleNodeWidthLog2                 = 8
 	HeaderStorageOffset                 = uint256.NewInt(64)
+	mainStorageOffsetLshVerkleNodeWidth = new(uint256.Int).Lsh(uint256.NewInt(256), 31-uint(VerkleNodeWidthLog2))
 	CodeOffset                          = uint256.NewInt(128)
 	MainStorageOffset                   = new(uint256.Int).Lsh(uint256.NewInt(256), 31)
 	VerkleNodeWidth                     = uint256.NewInt(256)
 	codeStorageDelta                    = uint256.NewInt(0).Sub(CodeOffset, HeaderStorageOffset)
 
-	// BigInt versions of the above.
-	headerStorageOffsetBig = HeaderStorageOffset.ToBig()
-	mainStorageOffsetBig   = MainStorageOffset.ToBig()
-	verkleNodeWidthBig     = VerkleNodeWidth.ToBig()
-	codeStorageDeltaBig    = codeStorageDelta.ToBig()
-
 	getTreePolyIndex0Point *verkle.Point
 )
 
@@ -164,6 +160,11 @@ func GetTreeKeyCodeSize(address []byte) []byte {
 }
 
 func GetTreeKeyCodeChunk(address []byte, chunk *uint256.Int) []byte {
+	treeIndex, subIndex := GetTreeKeyCodeChunkIndices(chunk)
+	return GetTreeKey(address, treeIndex, subIndex)
+}
+
+func GetTreeKeyCodeChunkIndices(chunk *uint256.Int) (*uint256.Int, byte) {
 	chunkOffset := new(uint256.Int).Add(CodeOffset, chunk)
 	treeIndex := new(uint256.Int).Div(chunkOffset, VerkleNodeWidth)
 	subIndexMod := new(uint256.Int).Mod(chunkOffset, VerkleNodeWidth)
@@ -171,7 +172,7 @@ func GetTreeKeyCodeChunk(address []byte, chunk *uint256.Int) []byte {
 	if len(subIndexMod) != 0 {
 		subIndex = byte(subIndexMod[0])
 	}
-	return GetTreeKey(address, treeIndex, subIndex)
+	return treeIndex, subIndex
 }
 
 func GetTreeKeyCodeChunkWithEvaluatedAddress(addressPoint *verkle.Point, chunk *uint256.Int) []byte {
@@ -230,8 +231,8 @@ func GetTreeKeyWithEvaluatedAddess(evaluated *verkle.Point, treeIndex *uint256.I
 
 	// little-endian, 32-byte aligned treeIndex
 	var index [32]byte
-	for i, b := range treeIndex.Bytes() {
+	for i := 0; i < len(treeIndex); i++ {
-		index[len(treeIndex.Bytes())-1-i] = b
+		binary.LittleEndian.PutUint64(index[i*8:(i+1)*8], treeIndex[i])
 	}
 	verkle.FromLEBytes(&poly[3], index[:16])
 	verkle.FromLEBytes(&poly[4], index[16:])
@@ -274,22 +275,27 @@ func GetTreeKeyStorageSlotWithEvaluatedAddress(evaluated *verkle.Point, storageK
 }
 
 func GetTreeKeyStorageSlotTreeIndexes(storageKey []byte) (*uint256.Int, byte) {
-	// Note that `pos` must be a big.Int and not a uint256.Int, because the subsequent
+	var pos uint256.Int
-	// arithmetics operations could overflow. (e.g: imagine if storageKey is 2^256-1)
+	pos.SetBytes(storageKey)
-	pos := new(big.Int).SetBytes(storageKey)
-	if pos.Cmp(codeStorageDeltaBig) < 0 {
-		pos.Add(headerStorageOffsetBig, pos)
-	} else {
-		pos.Add(mainStorageOffsetBig, pos)
-	}
-	treeIndex, overflow := uint256.FromBig(big.NewInt(0).Div(pos, verkleNodeWidthBig))
-	if overflow { // Must never happen considering the EIP definition.
-		panic("tree index overflow")
-	}
-	// calculate the sub_index, i.e. the index in the stem tree.
-	// Because the modulus is 256, it's the last byte of treeIndex
-	posBytes := pos.Bytes()
-	subIndex := posBytes[len(posBytes)-1]
 
-	return treeIndex, subIndex
+	// If the storage slot is in the header, we need to add the header offset.
+	if pos.Cmp(codeStorageDelta) < 0 {
+		// This addition is always safe; it can't ever overflow since pos<codeStorageDelta.
+		pos.Add(HeaderStorageOffset, &pos)
+
+		// In this branch, the tree-index is zero since we're in the account header,
+		// and the sub-index is the LSB of the modified storage key.
+		return zero, byte(pos[0] & 0xFF)
+
+	}
+	// If the storage slot is in the main storage, we need to add the main storage offset.
+
+	// We first divide by VerkleNodeWidth to create room to avoid an overflow next.
+	pos.Rsh(&pos, uint(VerkleNodeWidthLog2))
+	// We add mainStorageOffset/VerkleNodeWidth which can't overflow.
+	pos.Add(&pos, mainStorageOffsetLshVerkleNodeWidth)
+
+	// The sub-index is the LSB of the original storage key, since mainStorageOffset
+	// doesn't affect this byte, so we can avoid masks or shifts.
+	return &pos, storageKey[len(storageKey)-1]
 }