Conversion and TransitionTrie fixes

core/overlay/conversion.go

@@ -47,7 +47,7 @@ var zeroTreeIndex uint256.Int
 // 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
+	leafData map[branchKey]*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
@@ -66,7 +66,7 @@ func newKeyValueMigrator() *keyValueMigrator {
 	_ = verkle.GetConfig()
 	return &keyValueMigrator{
 		processingReady: make(chan struct{}),
-		leafData:        make([]migratedKeyValue, 0, 10_000),
+		leafData:        make(map[branchKey]*migratedKeyValue, 10_000),
 	}
 }
 
@@ -138,19 +138,17 @@ func (kvm *keyValueMigrator) addAccountCode(addr []byte, codeSize uint64, chunks
 }
 
 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),
-			},
-		})
+	if ld, ok := kvm.leafData[bk]; ok {
+		return &ld.leafNodeData
 	}
-	return &kvm.leafData[len(kvm.leafData)-1].leafNodeData
+	kvm.leafData[bk] = &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, 256),
+		},
+	}
+	return &kvm.leafData[bk].leafNodeData
 }
 
 func (kvm *keyValueMigrator) prepare() {
@@ -159,6 +157,10 @@ func (kvm *keyValueMigrator) prepare() {
 	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.
+		leafData := make([]migratedKeyValue, 0, len(kvm.leafData))
+		for _, v := range kvm.leafData {
+			leafData = append(leafData, *v)
+		}
 		var wg sync.WaitGroup
 		batchNum := runtime.NumCPU()
 		batchSize := (len(kvm.leafData) + batchNum - 1) / batchNum
@@ -170,7 +172,7 @@ func (kvm *keyValueMigrator) prepare() {
 			}
 			wg.Add(1)
 
-			batch := kvm.leafData[start:end]
+			batch := leafData[start:end]
 			go func() {
 				defer wg.Done()
 				var currAddr common.Address
@@ -190,8 +192,8 @@ func (kvm *keyValueMigrator) prepare() {
 
 		// 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
+		for i := range leafData {
+			nodeValues[i] = leafData[i].leafNodeData
 		}
 
 		// Create all leaves in batch mode so we can optimize cryptography operations.

trie/transition.go

@@ -17,6 +17,8 @@
 package trie
 
 import (
+	"fmt"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/ethdb"
@@ -62,7 +64,11 @@ func (t *TransitionTrie) GetKey(key []byte) []byte {
 // not be modified by the caller. If a node was not found in the database, a
 // trie.MissingNodeError is returned.
 func (t *TransitionTrie) GetStorage(addr common.Address, key []byte) ([]byte, error) {
-	if val, err := t.overlay.GetStorage(addr, key); len(val) != 0 || err != nil {
+	val, err := t.overlay.GetStorage(addr, key)
+	if err != nil {
+		return nil, fmt.Errorf("get storage from overlay: %s", err)
+	}
+	if len(val) != 0 {
 		return val, nil
 	}
 	// TODO also insert value into overlay