Improve migration

migrations/account_storage.go

@@ -39,50 +39,72 @@ func NewAccountStorage(storage *runtime.Storage, address common.Address) Account
 	}
 }
 
-type StorageMapKeyMigrator func(
-	inter *interpreter.Interpreter,
-	storageKey interpreter.StorageKey,
-	storageMap *interpreter.StorageMap,
-	storageMapKey interpreter.StorageMapKey,
-)
+type StorageMapKeyMigrator interface {
+	Migrate(
+		inter *interpreter.Interpreter,
+		storageKey interpreter.StorageKey,
+		storageMap *interpreter.StorageMap,
+		storageMapKey interpreter.StorageMapKey,
+	)
+	Domains() map[string]struct{}
+}
 
 type ValueConverter func(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,
 	value interpreter.Value,
 ) interpreter.Value
 
-func NewValueConverterPathMigrator(convertValue ValueConverter) StorageMapKeyMigrator {
-	return func(
-		inter *interpreter.Interpreter,
-		storageKey interpreter.StorageKey,
-		storageMap *interpreter.StorageMap,
-		storageMapKey interpreter.StorageMapKey,
-	) {
-		value := storageMap.ReadValue(nil, storageMapKey)
-
-		newValue := convertValue(storageKey, storageMapKey, value)
-		if newValue != nil {
-			// If the converter returns a new value,
-			// then replace the existing value with the new one.
-			storageMap.SetValue(
-				inter,
-				storageMapKey,
-				newValue,
-			)
-		}
+type ValueConverterPathMigrator struct {
+	domains      map[string]struct{}
+	ConvertValue ValueConverter
+}
+
+var _ StorageMapKeyMigrator = ValueConverterPathMigrator{}
+
+func NewValueConverterPathMigrator(
+	domains map[string]struct{},
+	convertValue ValueConverter,
+) StorageMapKeyMigrator {
+	return ValueConverterPathMigrator{
+		domains:      domains,
+		ConvertValue: convertValue,
+	}
+}
+
+func (m ValueConverterPathMigrator) Migrate(
+	inter *interpreter.Interpreter,
+	storageKey interpreter.StorageKey,
+	storageMap *interpreter.StorageMap,
+	storageMapKey interpreter.StorageMapKey,
+) {
+	value := storageMap.ReadValue(nil, storageMapKey)
+
+	newValue := m.ConvertValue(storageKey, storageMapKey, value)
+	if newValue != nil {
+		// If the converter returns a new value,
+		// then replace the existing value with the new one.
+		storageMap.SetValue(
+			inter,
+			storageMapKey,
+			newValue,
+		)
 	}
 }
 
+func (m ValueConverterPathMigrator) Domains() map[string]struct{} {
+	return m.domains
+}
+
 func (i *AccountStorage) MigrateStringKeys(
 	inter *interpreter.Interpreter,
 	key string,
-	migrate StorageMapKeyMigrator,
+	migrator StorageMapKeyMigrator,
 ) {
 	i.MigrateStorageMap(
 		inter,
 		key,
-		migrate,
+		migrator,
 		func(key atree.Value) interpreter.StorageMapKey {
 			return interpreter.StringStorageMapKey(key.(interpreter.StringAtreeValue))
 		},
@@ -92,12 +114,12 @@ func (i *AccountStorage) MigrateStringKeys(
 func (i *AccountStorage) MigrateUint64Keys(
 	inter *interpreter.Interpreter,
 	key string,
-	migrate StorageMapKeyMigrator,
+	migrator StorageMapKeyMigrator,
 ) {
 	i.MigrateStorageMap(
 		inter,
 		key,
-		migrate,
+		migrator,
 		func(key atree.Value) interpreter.StorageMapKey {
 			return interpreter.Uint64StorageMapKey(key.(interpreter.Uint64AtreeValue))
 		},
@@ -107,11 +129,17 @@ func (i *AccountStorage) MigrateUint64Keys(
 func (i *AccountStorage) MigrateStorageMap(
 	inter *interpreter.Interpreter,
 	domain string,
-	migrate StorageMapKeyMigrator,
+	migrator StorageMapKeyMigrator,
 	atreeKeyToStorageMapKey func(atree.Value) interpreter.StorageMapKey,
 ) {
 	address := i.address
 
+	if domains := migrator.Domains(); domains != nil {
+		if _, ok := domains[domain]; !ok {
+			return
+		}
+	}
+
 	storageMap := i.storage.GetStorageMap(address, domain, false)
 	if storageMap == nil || storageMap.Count() == 0 {
 		return
@@ -131,7 +159,7 @@ func (i *AccountStorage) MigrateStorageMap(
 
 	for _, storageMapKey := range keys {
 
-		migrate(
+		migrator.Migrate(
 			inter,
 			storageKey,
 			storageMap,

migrations/capcons/capabilitymigration.go

@@ -51,6 +51,10 @@ func (*CapabilityValueMigration) Name() string {
 	return "CapabilityValueMigration"
 }
 
+func (*CapabilityValueMigration) Domains() map[string]struct{} {
+	return nil
+}
+
 var fullyEntitledAccountReferenceStaticType = interpreter.ConvertSemaReferenceTypeToStaticReferenceType(
 	nil,
 	sema.FullyEntitledAccountReferenceType,

migrations/capcons/linkmigration.go

@@ -56,6 +56,15 @@ func (m *LinkValueMigration) CanSkip(valueType interpreter.StaticType) bool {
 	return CanSkipCapabilityValueMigration(valueType)
 }
 
+var linkValueMigrationDomains = map[string]struct{}{
+	common.PathDomainPublic.Identifier():  {},
+	common.PathDomainPrivate.Identifier(): {},
+}
+
+func (m *LinkValueMigration) Domains() map[string]struct{} {
+	return linkValueMigrationDomains
+}
+
 func (m *LinkValueMigration) Migrate(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,

migrations/capcons/migration_test.go

@@ -460,12 +460,8 @@ func testPathCapabilityValueMigration(
 
 	reporter := &testMigrationReporter{}
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&LinkValueMigration{
@@ -476,12 +472,8 @@ func testPathCapabilityValueMigration(
 		),
 	)
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&CapabilityValueMigration{
@@ -1312,12 +1304,8 @@ func testLinkMigration(
 
 	reporter := &testMigrationReporter{}
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&LinkValueMigration{
@@ -2025,12 +2013,8 @@ func TestPublishedPathCapabilityValueMigration(t *testing.T) {
 
 	reporter := &testMigrationReporter{}
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&LinkValueMigration{
@@ -2041,12 +2025,8 @@ func TestPublishedPathCapabilityValueMigration(t *testing.T) {
 		),
 	)
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&CapabilityValueMigration{
@@ -2274,12 +2254,8 @@ func TestUntypedPathCapabilityValueMigration(t *testing.T) {
 
 	reporter := &testMigrationReporter{}
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&LinkValueMigration{
@@ -2290,12 +2266,8 @@ func TestUntypedPathCapabilityValueMigration(t *testing.T) {
 		),
 	)
 
-	migration.Migrate(
-		&migrations.AddressSliceIterator{
-			Addresses: []common.Address{
-				testAddress,
-			},
-		},
+	migration.MigrateAccount(
+		testAddress,
 		migration.NewValueMigrationsPathMigrator(
 			reporter,
 			&CapabilityValueMigration{

migrations/entitlements/migration.go

@@ -41,6 +41,10 @@ func (EntitlementsMigration) Name() string {
 	return "EntitlementsMigration"
 }
 
+func (EntitlementsMigration) Domains() map[string]struct{} {
+	return nil
+}
+
 // ConvertToEntitledType converts the given type to an entitled type according to the following rules:
 //   - ConvertToEntitledType(&T)            --> auth(Entitlements(T)) &T
 //   - ConvertToEntitledType(Capability<T>) --> Capability<ConvertToEntitledType(T)>