Revert "Optimise merge registers for migrations"

cmd/util/ledger/migrations/account_storage_migration.go

@@ -58,9 +58,11 @@ func NewAccountStorageMigration(
 			flow.Address(address): {},
 		}
 
-		newPayloads, err := migrationRuntime.Snapshot.ApplyChangesAndGetNewPayloads(
+		newPayloads, err := MergeRegisterChanges(
+			migrationRuntime.Snapshot.Payloads,
 			result.WriteSet,
 			expectedAddresses,
+			nil,
 			log,
 		)
 		if err != nil {

cmd/util/ledger/migrations/atree_register_migration.go

@@ -129,7 +129,7 @@ func (m *AtreeRegisterMigrator) MigrateAccount(
 		m.rw.Write(migrationProblem{
 			Address: address.Hex(),
 			Key:     "",
-			Size:    mr.Snapshot.Len(),
+			Size:    len(mr.Snapshot.Payloads),
 			Kind:    "more_registers_after_migration",
 			Msg:     fmt.Sprintf("original: %d, new: %d", originalLen, newLen),
 		})
@@ -227,7 +227,7 @@ func (m *AtreeRegisterMigrator) convertStorageDomain(
 
 			m.rw.Write(migrationProblem{
 				Address: mr.Address.Hex(),
-				Size:    mr.Snapshot.Len(),
+				Size:    len(mr.Snapshot.Payloads),
 				Key:     string(key),
 				Kind:    "migration_failure",
 				Msg:     err.Error(),
@@ -245,7 +245,7 @@ func (m *AtreeRegisterMigrator) validateChangesAndCreateNewRegisters(
 	storageMapIds map[string]struct{},
 ) ([]*ledger.Payload, error) {
 	originalPayloadsSnapshot := mr.Snapshot
-	originalPayloads := originalPayloadsSnapshot.PayloadMap()
+	originalPayloads := originalPayloadsSnapshot.Payloads
 	newPayloads := make([]*ledger.Payload, 0, len(originalPayloads))
 
 	// store state payload so that it can be updated
@@ -331,7 +331,7 @@ func (m *AtreeRegisterMigrator) validateChangesAndCreateNewRegisters(
 			m.rw.Write(migrationProblem{
 				Address: mr.Address.Hex(),
 				Key:     id.String(),
-				Size:    mr.Snapshot.Len(),
+				Size:    len(mr.Snapshot.Payloads),
 				Kind:    "not_migrated",
 				Msg:     fmt.Sprintf("%x", value.Value()),
 			})

cmd/util/ledger/migrations/cadence_value_validation.go

@@ -363,7 +363,7 @@ func newReadonlyStorageRuntime(payloads []*ledger.Payload) (
 	*readonlyStorageRuntime,
 	error,
 ) {
-	snapshot, err := util.NewMapBasedPayloadSnapshot(payloads)
+	snapshot, err := util.NewPayloadSnapshot(payloads)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create payload snapshot: %w", err)
 	}

cmd/util/ledger/migrations/cadence_values_migration.go

@@ -145,9 +145,11 @@ func (m *CadenceBaseMigrator) MigrateAccount(
 		flow.Address(address): {},
 	}
 
-	newPayloads, err := migrationRuntime.Snapshot.ApplyChangesAndGetNewPayloads(
+	newPayloads, err := MergeRegisterChanges(
+		migrationRuntime.Snapshot.Payloads,
 		result.WriteSet,
 		expectedAddresses,
+		expectedAddresses,
 		m.log,
 	)
 	if err != nil {

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -796,9 +796,11 @@ func TestProgramParsingError(t *testing.T) {
 		flow.Address(testAddress): {},
 	}
 
-	payloads, err = runtime.Snapshot.ApplyChangesAndGetNewPayloads(
+	payloads, err = MergeRegisterChanges(
+		runtime.Snapshot.Payloads,
 		result.WriteSet,
 		expectedAddresses,
+		nil,
 		logger,
 	)
 	require.NoError(t, err)
@@ -932,9 +934,11 @@ func TestCoreContractUsage(t *testing.T) {
 			flow.Address(testAddress): {},
 		}
 
-		payloads, err = runtime.Snapshot.ApplyChangesAndGetNewPayloads(
+		payloads, err = MergeRegisterChanges(
+			runtime.Snapshot.Payloads,
 			result.WriteSet,
 			expectedAddresses,
+			nil,
 			logger,
 		)
 		require.NoError(t, err)

cmd/util/ledger/migrations/merge.go

@@ -0,0 +1,72 @@
+package migrations
+
+import (
+	"github.com/rs/zerolog"
+
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/convert"
+	"github.com/onflow/flow-go/model/flow"
+)
+
+func MergeRegisterChanges(
+	originalPayloads map[flow.RegisterID]*ledger.Payload,
+	changes map[flow.RegisterID]flow.RegisterValue,
+	expectedChangeAddresses map[flow.Address]struct{},
+	expectedOriginalAddresses map[flow.Address]struct{},
+	logger zerolog.Logger,
+) ([]*ledger.Payload, error) {
+
+	newPayloads := make([]*ledger.Payload, 0, len(originalPayloads))
+
+	// Add all new payloads.
+	for id, value := range changes {
+		delete(originalPayloads, id)
+		if len(value) == 0 {
+			continue
+		}
+
+		if expectedChangeAddresses != nil {
+			ownerAddress := flow.BytesToAddress([]byte(id.Owner))
+
+			if _, ok := expectedChangeAddresses[ownerAddress]; !ok {
+				// something was changed that does not belong to this account. Log it.
+				logger.Error().
+					Str("key", id.String()).
+					Str("actual_address", ownerAddress.Hex()).
+					Interface("expected_addresses", expectedChangeAddresses).
+					Hex("value", value).
+					Msg("key is part of the change set, but is for a different account")
+			}
+		}
+
+		key := convert.RegisterIDToLedgerKey(id)
+		newPayloads = append(newPayloads, ledger.NewPayload(key, value))
+	}
+
+	// Add any old payload that wasn't updated.
+	for id, value := range originalPayloads {
+		if len(value.Value()) == 0 {
+			// This is strange, but we don't want to add empty values. Log it.
+			logger.Warn().Msgf("empty value for key %s", id)
+			continue
+		}
+
+		if expectedOriginalAddresses != nil {
+			ownerAddress := flow.BytesToAddress([]byte(id.Owner))
+
+			if _, ok := expectedOriginalAddresses[ownerAddress]; !ok {
+				// something was changed that does not belong to this account. Log it.
+				logger.Error().
+					Str("key", id.String()).
+					Str("actual_address", ownerAddress.Hex()).
+					Interface("expected_addresses", expectedOriginalAddresses).
+					Hex("value", value.Value()).
+					Msg("key is part of the original set, but is for a different account")
+			}
+		}
+
+		newPayloads = append(newPayloads, value)
+	}
+
+	return newPayloads, nil
+}

cmd/util/ledger/migrations/migrator_runtime.go

@@ -34,7 +34,7 @@ func NewMigratorRuntime(
 	*migratorRuntime,
 	error,
 ) {
-	snapshot, err := util.NewMapBasedPayloadSnapshot(payloads)
+	snapshot, err := util.NewPayloadSnapshot(payloads)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create payload snapshot: %w", err)
 	}
@@ -107,7 +107,7 @@ func NewMigratorRuntime(
 }
 
 type migratorRuntime struct {
-	Snapshot                util.MigrationStorageSnapshot
+	Snapshot                *util.PayloadSnapshot
 	TransactionState        state.NestedTransactionPreparer
 	Interpreter             *interpreter.Interpreter
 	Storage                 *runtime.Storage

cmd/util/ledger/migrations/transaction_migration.go

@@ -27,7 +27,7 @@ func NewTransactionBasedMigration(
 			fvm.WithTransactionFeesEnabled(false))
 		ctx := fvm.NewContext(options...)
 
-		snapshot, err := util.NewMapBasedPayloadSnapshot(payloads)
+		snapshot, err := util.NewPayloadSnapshot(payloads)
 		if err != nil {
 			return nil, err
 		}
@@ -48,9 +48,11 @@ func NewTransactionBasedMigration(
 			return nil, res.Err
 		}
 
-		return snapshot.ApplyChangesAndGetNewPayloads(
+		return MergeRegisterChanges(
+			snapshot.Payloads,
 			executionSnapshot.WriteSet,
 			expectedWriteAddresses,
+			nil,
 			logger,
 		)
 	}