Generalize the migrations and make common codes re-usable

migrations/account_storage.go

@@ -39,13 +39,16 @@ func NewAccountStorage(storage *runtime.Storage, address common.Address) Account
 
 // ForEachValue iterates over the values in the account.
 // The `valueConverter takes a function to be applied to each value.
-// It returns the `newValue`, if a new value was created during conversion,
-// or a flag, indicating whether the old value was updated in-place.
+// It returns the converted, if a new value was created during conversion.
 func (i *AccountStorage) ForEachValue(
 	inter *interpreter.Interpreter,
 	domains []common.PathDomain,
-	valueConverter func(interpreter.Value) (newValue interpreter.Value, updated bool),
-	reporter Reporter,
+	valueConverter func(
+		value interpreter.Value,
+		address common.Address,
+		domain common.PathDomain,
+		key string,
+	) interpreter.Value,
 ) {
 	for _, domain := range domains {
 		storageMap := i.storage.GetStorageMap(i.address, domain.Identifier(), false)
@@ -68,23 +71,17 @@ func (i *AccountStorage) ForEachValue(
 
 			value := storageMap.ReadValue(nil, storageKey)
 
-			newValue, updated := valueConverter(value)
-			if newValue == nil && !updated {
+			newValue := valueConverter(value, i.address, domain, key)
+			if newValue == nil {
 				continue
 			}
 
-			if newValue != nil {
-				// If the converter returns a new value, then replace the existing value with the new one.
-				storageMap.SetValue(
-					inter,
-					storageKey,
-					newValue,
-				)
-			}
-
-			if reporter != nil {
-				reporter.Report(i.address, domain, key)
-			}
+			// If the converter returns a new value, then replace the existing value with the new one.
+			storageMap.SetValue(
+				inter,
+				storageKey,
+				newValue,
+			)
 		}
 	}
 }

migrations/account_type/migration.go

@@ -20,108 +20,63 @@ package account_type
 
 import (
 	"github.com/onflow/cadence/migrations"
-	"github.com/onflow/cadence/runtime"
 	"github.com/onflow/cadence/runtime/common"
 	"github.com/onflow/cadence/runtime/interpreter"
 	"github.com/onflow/cadence/runtime/sema"
 )
 
-type AccountTypeMigration struct {
-	storage     *runtime.Storage
-	interpreter *interpreter.Interpreter
-}
-
-func NewAccountTypeMigration(
-	interpreter *interpreter.Interpreter,
-	storage *runtime.Storage,
-) *AccountTypeMigration {
-	return &AccountTypeMigration{
-		storage:     storage,
-		interpreter: interpreter,
-	}
-}
+type AccountTypeMigration struct{}
 
-func (m *AccountTypeMigration) Migrate(
-	addressIterator migrations.AddressIterator,
-	reporter migrations.Reporter,
-) {
-	for {
-		address := addressIterator.NextAddress()
-		if address == common.ZeroAddress {
-			break
-		}
-
-		m.migrateTypeValuesInAccount(
-			address,
-			reporter,
-		)
-	}
+var _ migrations.Migration = AccountTypeMigration{}
 
-	err := m.storage.Commit(m.interpreter, false)
-	if err != nil {
-		panic(err)
-	}
-}
-
-// migrateTypeValuesInAccount migrates `AuthAccount` and `PublicAccount` types in a given account
-// to the account reference type (&Account).
-func (m *AccountTypeMigration) migrateTypeValuesInAccount(
-	address common.Address,
-	reporter migrations.Reporter,
-) {
-
-	accountStorage := migrations.NewAccountStorage(m.storage, address)
-
-	accountStorage.ForEachValue(
-		m.interpreter,
-		common.AllPathDomains,
-		m.migrateValue,
-		reporter,
-	)
+func NewAccountTypeMigration() AccountTypeMigration {
+	return AccountTypeMigration{}
 }
 
-func (m *AccountTypeMigration) migrateValue(value interpreter.Value) (newValue interpreter.Value, updatedInPlace bool) {
-	return migrations.MigrateNestedValue(m.interpreter, value, m.migrateTypeValue)
+func (AccountTypeMigration) Name() string {
+	return "AccountTypeMigration"
 }
 
-func (m *AccountTypeMigration) migrateTypeValue(value interpreter.Value) (newValue interpreter.Value, updatedInPlace bool) {
+// Migrate migrates `AuthAccount` and `PublicAccount` types inside `TypeValue`s,
+// to the account reference type (&Account).
+func (AccountTypeMigration) Migrate(value interpreter.Value) (newValue interpreter.Value) {
 	switch value := value.(type) {
 	case interpreter.TypeValue:
-		convertedType := m.maybeConvertAccountType(value.Type)
+		convertedType := maybeConvertAccountType(value.Type)
 		if convertedType == nil {
 			return
 		}
-		return interpreter.NewTypeValue(nil, convertedType), true
+		return interpreter.NewTypeValue(nil, convertedType)
 
 	case *interpreter.CapabilityValue:
-		convertedBorrowType := m.maybeConvertAccountType(value.BorrowType)
+		convertedBorrowType := maybeConvertAccountType(value.BorrowType)
 		if convertedBorrowType == nil {
 			return
 		}
-		return interpreter.NewUnmeteredCapabilityValue(value.ID, value.Address, convertedBorrowType), true
+		return interpreter.NewUnmeteredCapabilityValue(value.ID, value.Address, convertedBorrowType)
 
 	default:
-		return nil, false
+		return nil
 	}
 }
 
-func (m *AccountTypeMigration) maybeConvertAccountType(staticType interpreter.StaticType) interpreter.StaticType {
+func maybeConvertAccountType(staticType interpreter.StaticType) interpreter.StaticType {
 	switch staticType := staticType.(type) {
 	case *interpreter.ConstantSizedStaticType:
-		convertedType := m.maybeConvertAccountType(staticType.Type)
+		convertedType := maybeConvertAccountType(staticType.Type)
 		if convertedType != nil {
 			return interpreter.NewConstantSizedStaticType(nil, convertedType, staticType.Size)
 		}
 
 	case *interpreter.VariableSizedStaticType:
-		convertedType := m.maybeConvertAccountType(staticType.Type)
+		convertedType := maybeConvertAccountType(staticType.Type)
 		if convertedType != nil {
 			return interpreter.NewVariableSizedStaticType(nil, convertedType)
 		}
 
 	case *interpreter.DictionaryStaticType:
-		convertedKeyType := m.maybeConvertAccountType(staticType.KeyType)
-		convertedValueType := m.maybeConvertAccountType(staticType.ValueType)
+		convertedKeyType := maybeConvertAccountType(staticType.KeyType)
+		convertedValueType := maybeConvertAccountType(staticType.ValueType)
 		if convertedKeyType != nil && convertedValueType != nil {
 			return interpreter.NewDictionaryStaticType(nil, convertedKeyType, convertedValueType)
 		}
@@ -133,7 +88,7 @@ func (m *AccountTypeMigration) maybeConvertAccountType(staticType interpreter.St
 		}
 
 	case *interpreter.CapabilityStaticType:
-		convertedBorrowType := m.maybeConvertAccountType(staticType.BorrowType)
+		convertedBorrowType := maybeConvertAccountType(staticType.BorrowType)
 		if convertedBorrowType != nil {
 			return interpreter.NewCapabilityStaticType(nil, convertedBorrowType)
 		}
@@ -142,14 +97,14 @@ func (m *AccountTypeMigration) maybeConvertAccountType(staticType interpreter.St
 		// Nothing to do. Inner types can only be interfaces.
 
 	case *interpreter.OptionalStaticType:
-		convertedInnerType := m.maybeConvertAccountType(staticType.Type)
+		convertedInnerType := maybeConvertAccountType(staticType.Type)
 		if convertedInnerType != nil {
 			return interpreter.NewOptionalStaticType(nil, convertedInnerType)
 		}
 
 	case *interpreter.ReferenceStaticType:
 		// TODO: Reference of references must not be allowed?
-		convertedReferencedType := m.maybeConvertAccountType(staticType.ReferencedType)
+		convertedReferencedType := maybeConvertAccountType(staticType.ReferencedType)
 		if convertedReferencedType != nil {
 			switch convertedReferencedType {
 
@@ -175,7 +130,7 @@ func (m *AccountTypeMigration) maybeConvertAccountType(staticType interpreter.St
 		// This is for testing the migration.
 		// i.e: wrapper was only to make it possible to use as a dictionary-key.
 		// Ignore the wrapper, and continue with the inner type.
-		return m.maybeConvertAccountType(staticType.PrimitiveStaticType)
+		return maybeConvertAccountType(staticType.PrimitiveStaticType)
 
 	default:
 		// Is it safe to do so?

migrations/account_type/migration_test.go

@@ -50,6 +50,7 @@ func (t *testReporter) Report(
 	address common.Address,
 	domain common.PathDomain,
 	identifier string,
+	_ string,
 ) {
 	migratedPathsInAddress, ok := t.migratedPaths[address]
 	if !ok {
@@ -66,7 +67,7 @@ func (t *testReporter) Report(
 	migratedPathsInDomain[identifier] = struct{}{}
 }
 
-func TestMigration(t *testing.T) {
+func TestTypeValueMigration(t *testing.T) {
 	t.Parallel()
 
 	account := common.Address{0x42}
@@ -341,7 +342,7 @@ func TestMigration(t *testing.T) {
 
 	// Migrate
 
-	migration := NewAccountTypeMigration(inter, storage)
+	migration := migrations.NewStorageMigration(inter, storage)
 
 	reporter := newTestReporter()
 
@@ -352,6 +353,7 @@ func TestMigration(t *testing.T) {
 			},
 		},
 		reporter,
+		NewAccountTypeMigration(),
 	)
 
 	// Check reported migrated paths
@@ -651,7 +653,7 @@ func TestNestedTypeValueMigration(t *testing.T) {
 
 	// Migrate
 
-	migration := NewAccountTypeMigration(inter, storage)
+	migration := migrations.NewStorageMigration(inter, storage)
 
 	migration.Migrate(
 		&migrations.AddressSliceIterator{
@@ -660,6 +662,7 @@ func TestNestedTypeValueMigration(t *testing.T) {
 			},
 		},
 		nil,
+		NewAccountTypeMigration(),
 	)
 
 	// Assert: Traverse through the storage and see if the values are updated now.
@@ -761,7 +764,7 @@ func TestValuesWithStaticTypeMigration(t *testing.T) {
 
 	// Migrate
 
-	migration := NewAccountTypeMigration(inter, storage)
+	migration := migrations.NewStorageMigration(inter, storage)
 
 	migration.Migrate(
 		&migrations.AddressSliceIterator{
@@ -770,6 +773,7 @@ func TestValuesWithStaticTypeMigration(t *testing.T) {
 			},
 		},
 		nil,
+		NewAccountTypeMigration(),
 	)
 
 	// Assert: Traverse through the storage and see if the values are updated now.