Optimise merge registers for migrations - Revert of Revert

cmd/util/ledger/migrations/account_storage_migration.go

@@ -8,6 +8,7 @@ import (
 	"github.com/onflow/cadence/runtime/interpreter"
 	"github.com/rs/zerolog"
 
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/model/flow"
 )
@@ -24,6 +25,7 @@ func NewAccountStorageMigration(
 			payloads,
 			chainID,
 			MigratorRuntimeConfig{},
+			snapshot.SmallChangeSetSnapshot,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create migrator runtime: %w", err)
@@ -58,11 +60,9 @@ func NewAccountStorageMigration(
 			flow.Address(address): {},
 		}
 
-		newPayloads, err := MergeRegisterChanges(
-			migrationRuntime.Snapshot.Payloads,
+		newPayloads, err := migrationRuntime.Snapshot.ApplyChangesAndGetNewPayloads(
 			result.WriteSet,
 			expectedAddresses,
-			nil,
 			log,
 		)
 		if err != nil {

cmd/util/ledger/migrations/atree_register_migration.go

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

cmd/util/ledger/migrations/atree_register_migrator_runtime.go

@@ -8,6 +8,7 @@ import (
 	"github.com/onflow/cadence/runtime/interpreter"
 
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	migrationSnapshot "github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/fvm/storage/state"
 	"github.com/onflow/flow-go/ledger"
@@ -21,7 +22,7 @@ func NewAtreeRegisterMigratorRuntime(
 	*AtreeRegisterMigratorRuntime,
 	error,
 ) {
-	snapshot, err := util.NewPayloadSnapshot(payloads)
+	snapshot, err := migrationSnapshot.NewPayloadSnapshot(payloads, migrationSnapshot.LargeChangeSetOrReadonlySnapshot)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create payload snapshot: %w", err)
 	}
@@ -54,7 +55,7 @@ func NewAtreeRegisterMigratorRuntime(
 }
 
 type AtreeRegisterMigratorRuntime struct {
-	Snapshot            *util.PayloadSnapshot
+	Snapshot            migrationSnapshot.MigrationSnapshot
 	TransactionState    state.NestedTransactionPreparer
 	Interpreter         *interpreter.Interpreter
 	Storage             *runtime.Storage

cmd/util/ledger/migrations/cadence_value_diff_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/onflow/cadence/runtime/common"
 	"github.com/onflow/cadence/runtime/interpreter"
 
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/convert"
@@ -175,6 +176,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				},
 				flow.Emulator,
 				MigratorRuntimeConfig{},
+				snapshot.LargeChangeSetOrReadonlySnapshot,
 			)
 			require.NoError(t, err)
 
@@ -277,6 +279,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				},
 				flow.Emulator,
 				MigratorRuntimeConfig{},
+				snapshot.LargeChangeSetOrReadonlySnapshot,
 			)
 			require.NoError(t, err)
 
@@ -386,6 +389,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				},
 				flow.Emulator,
 				MigratorRuntimeConfig{},
+				snapshot.LargeChangeSetOrReadonlySnapshot,
 			)
 			require.NoError(t, err)
 
@@ -506,6 +510,7 @@ func TestDiffCadenceValues(t *testing.T) {
 				},
 				flow.Emulator,
 				MigratorRuntimeConfig{},
+				snapshot.LargeChangeSetOrReadonlySnapshot,
 			)
 			require.NoError(t, err)
 
@@ -633,6 +638,7 @@ func createStorageMapPayloads(t *testing.T, address common.Address, domain strin
 		},
 		flow.Emulator,
 		MigratorRuntimeConfig{},
+		snapshot.LargeChangeSetOrReadonlySnapshot,
 	)
 	require.NoError(t, err)
 

cmd/util/ledger/migrations/cadence_value_validation.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	migrationSnapshot "github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/ledger"
 )
 
@@ -395,7 +396,7 @@ func newReadonlyStorageRuntime(payloads []*ledger.Payload) (
 	*readonlyStorageRuntime,
 	error,
 ) {
-	snapshot, err := util.NewPayloadSnapshot(payloads)
+	snapshot, err := migrationSnapshot.NewPayloadSnapshot(payloads, migrationSnapshot.LargeChangeSetOrReadonlySnapshot)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create payload snapshot: %w", err)
 	}

cmd/util/ledger/migrations/cadence_values_migration.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/fvm/tracing"
 	"github.com/onflow/flow-go/ledger"
@@ -109,6 +110,7 @@ func (m *CadenceBaseMigrator) MigrateAccount(
 		oldPayloads,
 		m.chainID,
 		m.migratorRuntimeConfig,
+		snapshot.SmallChangeSetSnapshot,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create migrator runtime: %w", err)
@@ -214,11 +216,9 @@ func (m *CadenceBaseMigrator) MigrateAccount(
 		flow.Address(address): {},
 	}
 
-	newPayloads, err := MergeRegisterChanges(
-		migrationRuntime.Snapshot.Payloads,
+	newPayloads, err := migrationRuntime.Snapshot.ApplyChangesAndGetNewPayloads(
 		result.WriteSet,
 		expectedAddresses,
-		expectedAddresses,
 		m.log,
 	)
 	if err != nil {

cmd/util/ledger/migrations/cadence_values_migration_test.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/fvm/systemcontracts"
 	"github.com/onflow/flow-go/ledger"
@@ -125,6 +126,7 @@ func checkMigratedPayloads(
 		newPayloads,
 		chainID,
 		MigratorRuntimeConfig{},
+		snapshot.SmallChangeSetSnapshot,
 	)
 	require.NoError(t, err)
 
@@ -743,6 +745,7 @@ func TestProgramParsingError(t *testing.T) {
 		payloads,
 		chainID,
 		MigratorRuntimeConfig{},
+		snapshot.SmallChangeSetSnapshot,
 	)
 	require.NoError(t, err)
 
@@ -793,11 +796,9 @@ func TestProgramParsingError(t *testing.T) {
 		flow.Address(testAddress): {},
 	}
 
-	payloads, err = MergeRegisterChanges(
-		runtime.Snapshot.Payloads,
+	payloads, err = runtime.Snapshot.ApplyChangesAndGetNewPayloads(
 		result.WriteSet,
 		expectedAddresses,
-		nil,
 		logger,
 	)
 	require.NoError(t, err)
@@ -888,6 +889,7 @@ func TestCoreContractUsage(t *testing.T) {
 			payloads,
 			chainID,
 			MigratorRuntimeConfig{},
+			snapshot.SmallChangeSetSnapshot,
 		)
 		require.NoError(t, err)
 
@@ -930,11 +932,9 @@ func TestCoreContractUsage(t *testing.T) {
 			flow.Address(testAddress): {},
 		}
 
-		payloads, err = MergeRegisterChanges(
-			runtime.Snapshot.Payloads,
+		payloads, err = runtime.Snapshot.ApplyChangesAndGetNewPayloads(
 			result.WriteSet,
 			expectedAddresses,
-			nil,
 			logger,
 		)
 		require.NoError(t, err)
@@ -978,6 +978,7 @@ func TestCoreContractUsage(t *testing.T) {
 			payloads,
 			chainID,
 			MigratorRuntimeConfig{},
+			snapshot.SmallChangeSetSnapshot,
 		)
 		require.NoError(t, err)
 

cmd/util/ledger/migrations/contract_checking_migration.go

@@ -10,6 +10,7 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/cmd/util/ledger/reporters"
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/ledger"
 	"github.com/onflow/flow-go/ledger/common/convert"
 	"github.com/onflow/flow-go/model/flow"
@@ -43,6 +44,7 @@ func NewContractCheckingMigration(
 			contractPayloads,
 			chainID,
 			MigratorRuntimeConfig{},
+			snapshot.LargeChangeSetOrReadonlySnapshot,
 		)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create migrator runtime: %w", err)

cmd/util/ledger/migrations/migrator_runtime.go

@@ -9,6 +9,7 @@ import (
 	"github.com/onflow/cadence/runtime/stdlib"
 
 	"github.com/onflow/flow-go/cmd/util/ledger/util"
+	"github.com/onflow/flow-go/cmd/util/ledger/util/snapshot"
 	"github.com/onflow/flow-go/fvm/environment"
 	"github.com/onflow/flow-go/fvm/evm"
 	evmStdlib "github.com/onflow/flow-go/fvm/evm/stdlib"
@@ -18,7 +19,7 @@ import (
 )
 
 type MigratorRuntime struct {
-	Snapshot                *util.PayloadSnapshot
+	Snapshot                snapshot.MigrationSnapshot
 	TransactionState        state.NestedTransactionPreparer
 	Interpreter             *interpreter.Interpreter
 	Storage                 *runtime.Storage
@@ -99,11 +100,12 @@ func NewMigratorRuntime(
 	payloads []*ledger.Payload,
 	chainID flow.ChainID,
 	config MigratorRuntimeConfig,
+	snapshotType snapshot.MigrationSnapshotType,
 ) (
 	*MigratorRuntime,
 	error,
 ) {
-	snapshot, err := util.NewPayloadSnapshot(payloads)
+	snapshot, err := snapshot.NewPayloadSnapshot(payloads, snapshotType)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create payload snapshot: %w", err)
 	}