fix: import cycle in conversion refactor

consensus/beacon/consensus.go

@@ -25,7 +25,7 @@ import (
 	"github.com/ethereum/go-ethereum/consensus"
 	"github.com/ethereum/go-ethereum/consensus/misc/eip1559"
 	"github.com/ethereum/go-ethereum/consensus/misc/eip4844"
-	"github.com/ethereum/go-ethereum/core"
+	"github.com/ethereum/go-ethereum/core/overlay"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/params"
@@ -373,7 +373,7 @@ func (beacon *Beacon) Finalize(chain consensus.ChainHeaderReader, header *types.
 	if chain.Config().IsPrague(header.Number, header.Time) {
 		fmt.Println("at block", header.Number, "performing transition?", state.Database().InTransition())
 		parent := chain.GetHeaderByHash(header.ParentHash)
-		core.OverlayVerkleTransition(state, parent.Root)
+		overlay.OverlayVerkleTransition(state, parent.Root)
 	}
 }
 

core/chain_makers.go

@@ -425,13 +425,15 @@ func GenerateVerkleChain(config *params.ChainConfig, parent *types.Block, engine
 		return nil, nil
 	}
 	var snaps *snapshot.Tree
+	triedb := state.NewDatabaseWithConfig(db, nil)
+	triedb.StartVerkleTransition(common.Hash{}, common.Hash{}, config, config.PragueTime, common.Hash{})
+	triedb.EndVerkleTransition()
+	statedb, err := state.New(parent.Root(), triedb, snaps)
+	if err != nil {
+		panic(fmt.Sprintf("could not find state for block %d: err=%v, parent root=%x", parent.NumberU64(), err, parent.Root()))
+	}
+	statedb.Database().SaveTransitionState(parent.Root())
 	for i := 0; i < n; i++ {
-		triedb := state.NewDatabaseWithConfig(db, nil)
-		triedb.EndVerkleTransition()
-		statedb, err := state.New(parent.Root(), triedb, snaps)
-		if err != nil {
-			panic(fmt.Sprintf("could not find state for block %d: err=%v, parent root=%x", i, err, parent.Root()))
-		}
 		block, receipt := genblock(i, parent, statedb)
 		blocks[i] = block
 		receipts[i] = receipt

core/overlay_transition.go → core/overlay/conversion.go

@@ -14,14 +14,17 @@
 // You should have received a copy of the GNU Lesser General Public License
 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
 
-package core
+package overlay
 
 import (
 	"bufio"
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"os"
+	"runtime"
+	"sync"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -32,8 +35,187 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie"
+	"github.com/ethereum/go-ethereum/trie/utils"
+	"github.com/gballet/go-verkle"
+	"github.com/holiman/uint256"
 )
 
+var zeroTreeIndex uint256.Int
+
+// 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) {
+	treeIndex, subIndex := utils.GetTreeKeyStorageSlotTreeIndexes(slotNumber)
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex))
+	leafNodeData.Values[subIndex] = slotValue
+}
+
+func (kvm *keyValueMigrator) addAccount(addr []byte, acc *types.StateAccount) {
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex))
+
+	var version [verkle.LeafValueSize]byte
+	leafNodeData.Values[utils.VersionLeafKey] = version[:]
+
+	var balance [verkle.LeafValueSize]byte
+	for i, b := range acc.Balance.Bytes() {
+		balance[len(acc.Balance.Bytes())-1-i] = b
+	}
+	leafNodeData.Values[utils.BalanceLeafKey] = balance[:]
+
+	var nonce [verkle.LeafValueSize]byte
+	binary.LittleEndian.PutUint64(nonce[:8], acc.Nonce)
+	leafNodeData.Values[utils.NonceLeafKey] = nonce[:]
+
+	leafNodeData.Values[utils.CodeKeccakLeafKey] = acc.CodeHash[:]
+}
+
+func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks []byte) {
+	leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, &zeroTreeIndex))
+
+	// Save the code size.
+	var codeSizeBytes [verkle.LeafValueSize]byte
+	binary.LittleEndian.PutUint64(codeSizeBytes[:8], codeSize)
+	leafNodeData.Values[utils.CodeSizeLeafKey] = codeSizeBytes[:]
+
+	// The first 128 chunks are stored in the account header leaf.
+	for i := 0; i < 128 && i < len(chunks)/32; i++ {
+		leafNodeData.Values[byte(128+i)] = chunks[32*i : 32*(i+1)]
+	}
+
+	// Potential further chunks, have their own leaf nodes.
+	for i := 128; i < len(chunks)/32; {
+		treeIndex, _ := utils.GetTreeKeyCodeChunkIndices(uint256.NewInt(uint64(i)))
+		leafNodeData := kvm.getOrInitLeafNodeData(newBranchKey(addr, treeIndex))
+
+		j := i
+		for ; (j-i) < 256 && j < len(chunks)/32; j++ {
+			leafNodeData.Values[byte((j-128)%256)] = chunks[32*j : 32*(j+1)]
+		}
+		i = j
+	}
+}
+
+func (kvm *keyValueMigrator) getOrInitLeafNodeData(bk branchKey) *verkle.BatchNewLeafNodeData {
+	// Remember that keyValueMigration receives actions ordered by (address, subtreeIndex).
+	// 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.
+	if len(kvm.leafData) == 0 || kvm.leafData[len(kvm.leafData)-1].branchKey != bk {
+		kvm.leafData = append(kvm.leafData, migratedKeyValue{
+			branchKey: bk,
+			leafNodeData: verkle.BatchNewLeafNodeData{
+				Stem:   nil, // It will be calculated in the prepare() phase, since it's CPU heavy.
+				Values: make(map[byte][]byte),
+			},
+		})
+	}
+	return &kvm.leafData[len(kvm.leafData)-1].leafNodeData
+}
+
+func (kvm *keyValueMigrator) prepare() {
+	// We fire a background routine to process the leafData and save the result in newLeaves.
+	// The background routine signals that it is done by closing processingReady.
+	go func() {
+		// 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 == (common.Address{}) {
+						currAddr = batch[i].branchKey.addr
+						currPoint = utils.EvaluateAddressPoint(currAddr[:])
+					}
+					stem := utils.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) calculated, 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.
+		kvm.newLeaves, kvm.prepareErr = verkle.BatchNewLeafNode(nodeValues)
+		close(kvm.processingReady)
+	}()
+}
+
+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(kvm.newLeaves); err != nil {
+		return fmt.Errorf("failed to insert migrated leaves: %w", err)
+	}
+
+	return nil
+}
+
 // OverlayVerkleTransition contains the overlay conversion logic
 func OverlayVerkleTransition(statedb *state.StateDB, root common.Hash) error {
 	migrdb := statedb.Database()