Merge pull request #3386 from onflow/bastian/improve-dictionary-key-m…

…igration

Improve dictionary key migration

migrations/capcons/capabilitymigration.go

@@ -68,7 +68,12 @@ func (m *CapabilityValueMigration) Migrate(
 	_ interpreter.StorageMapKey,
 	value interpreter.Value,
 	_ *interpreter.Interpreter,
-) (interpreter.Value, error) {
+	_ migrations.ValueMigrationPosition,
+) (
+	interpreter.Value,
+	error,
+) {
+
 	reporter := m.Reporter
 
 	switch value := value.(type) {

migrations/capcons/linkmigration.go

@@ -70,7 +70,11 @@ func (m *LinkValueMigration) Migrate(
 	storageMapKey interpreter.StorageMapKey,
 	value interpreter.Value,
 	inter *interpreter.Interpreter,
-) (interpreter.Value, error) {
+	_ migrations.ValueMigrationPosition,
+) (
+	interpreter.Value,
+	error,
+) {
 
 	pathValue, ok := storageKeyToPathValue(storageKey, storageMapKey)
 	if !ok {

migrations/entitlements/migration.go

@@ -371,6 +371,7 @@ func (m EntitlementsMigration) Migrate(
 	_ interpreter.StorageMapKey,
 	value interpreter.Value,
 	_ *interpreter.Interpreter,
+	_ migrations.ValueMigrationPosition,
 ) (
 	interpreter.Value,
 	error,

migrations/entitlements/migration_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/onflow/cadence"
 	"github.com/onflow/cadence/migrations"
 	"github.com/onflow/cadence/migrations/statictypes"
+	"github.com/onflow/cadence/migrations/type_keys"
 	"github.com/onflow/cadence/runtime"
 	"github.com/onflow/cadence/runtime/ast"
 	"github.com/onflow/cadence/runtime/common"
@@ -690,6 +691,7 @@ func (m testEntitlementsMigration) Migrate(
 	_ interpreter.StorageMapKey,
 	value interpreter.Value,
 	_ *interpreter.Interpreter,
+	_ migrations.ValueMigrationPosition,
 ) (
 	interpreter.Value,
 	error,
@@ -731,6 +733,7 @@ func convertEntireTestValue(
 		},
 		reporter,
 		true,
+		migrations.ValueMigrationPositionOther,
 	)
 
 	err = migration.Commit()
@@ -2863,6 +2866,7 @@ func TestConvertMigratedAccountTypes(t *testing.T) {
 					nil,
 					value,
 					inter,
+					migrations.ValueMigrationPositionOther,
 				)
 			require.NoError(t, err)
 			require.NotNil(t, newValue)
@@ -3658,6 +3662,7 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 			migration.NewValueMigrationsPathMigrator(
 				reporter,
 				NewEntitlementsMigration(inter),
+				type_keys.NewTypeKeyMigration(),
 			),
 		)
 
@@ -3673,7 +3678,21 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 
 		assert.ErrorContains(t, reporter.errors[0], importErrorMessage)
 
-		require.Empty(t, reporter.migrated)
+		assert.Equal(t,
+			map[struct {
+				interpreter.StorageKey
+				interpreter.StorageMapKey
+			}]struct{}{
+				{
+					StorageKey: interpreter.StorageKey{
+						Address: testAddress,
+						Key:     common.PathDomainStorage.Identifier(),
+					},
+					StorageMapKey: interpreter.StringStorageMapKey("dict"),
+				}: {},
+			},
+			reporter.migrated,
+		)
 	})()
 
 	// Load
@@ -3718,15 +3737,9 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 
 		assert.Equal(t, 1, dictValue.Count())
 
-		// Key did not get migrated, so is inaccessible using the "new" type value
+		// Key did not get migrated, but got still re-stored in new format,
+		// so it can be loaded and used after the migration failure
 		_, ok := dictValue.Get(inter, locationRange, typeValue)
-		require.False(t, ok)
-
-		// But the key is still accessible using the "old" type value
-		legacyKey := migrations.LegacyKey(typeValue)
-
-		value, ok := dictValue.Get(inter, locationRange, legacyKey)
 		require.True(t, ok)
-		require.Equal(t, newTestValue(), value)
 	})()
 }

migrations/migration.go

@@ -22,6 +22,8 @@ import (
 	"fmt"
 	"runtime/debug"
 
+	"github.com/onflow/atree"
+
 	"github.com/onflow/cadence/runtime"
 	"github.com/onflow/cadence/runtime/common"
 	"github.com/onflow/cadence/runtime/errors"
@@ -37,11 +39,19 @@ type ValueMigration interface {
 		storageMapKey interpreter.StorageMapKey,
 		value interpreter.Value,
 		interpreter *interpreter.Interpreter,
+		position ValueMigrationPosition,
 	) (newValue interpreter.Value, err error)
 	CanSkip(valueType interpreter.StaticType) bool
 	Domains() map[string]struct{}
 }
 
+type ValueMigrationPosition uint8
+
+const (
+	ValueMigrationPositionOther ValueMigrationPosition = iota
+	ValueMigrationPositionDictionaryKey
+)
+
 type DomainMigration interface {
 	Name() string
 	Migrate(
@@ -162,6 +172,7 @@ func (m *StorageMigration) NewValueMigrationsPathMigrator(
 				valueMigrations,
 				reporter,
 				true,
+				ValueMigrationPositionOther,
 			)
 		},
 	)
@@ -176,6 +187,7 @@ func (m *StorageMigration) MigrateNestedValue(
 	valueMigrations []ValueMigration,
 	reporter Reporter,
 	allowMutation bool,
+	position ValueMigrationPosition,
 ) (migratedValue interpreter.Value) {
 
 	defer func() {
@@ -240,6 +252,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 		if newInnerValue != nil {
 			migratedValue = interpreter.NewSomeValueNonCopying(inter, newInnerValue)
@@ -264,6 +277,7 @@ func (m *StorageMigration) MigrateNestedValue(
 				valueMigrations,
 				reporter,
 				allowMutation,
+				ValueMigrationPositionOther,
 			)
 
 			if newElement == nil {
@@ -325,6 +339,7 @@ func (m *StorageMigration) MigrateNestedValue(
 				valueMigrations,
 				reporter,
 				allowMutation,
+				ValueMigrationPositionOther,
 			)
 
 			if newValue == nil {
@@ -363,7 +378,7 @@ func (m *StorageMigration) MigrateNestedValue(
 		// The mutating iterator is only able to read new keys,
 		// as it recalculates the stored values' hashes.
 
-		keys := m.migrateDictionaryKeys(
+		m.migrateDictionaryKeys(
 			storageKey,
 			storageMapKey,
 			dictionary,
@@ -379,7 +394,6 @@ func (m *StorageMigration) MigrateNestedValue(
 			valueMigrations,
 			reporter,
 			allowMutation,
-			keys,
 		)
 
 	case *interpreter.PublishedValue:
@@ -391,6 +405,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 		if newInnerValue != nil {
 			newInnerCapability := newInnerValue.(interpreter.CapabilityValue)
@@ -415,6 +430,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			storageKey,
 			storageMapKey,
 			value,
+			position,
 		)
 
 		if err != nil {
@@ -469,19 +485,14 @@ func (m *StorageMigration) MigrateNestedValue(
 
 }
 
-type migratedDictionaryKey struct {
-	key      interpreter.Value
-	migrated bool
-}
-
 func (m *StorageMigration) migrateDictionaryKeys(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,
 	dictionary *interpreter.DictionaryValue,
 	valueMigrations []ValueMigration,
 	reporter Reporter,
 	allowMutation bool,
-) (migratedKeys []migratedDictionaryKey) {
+) {
 	inter := m.interpreter
 
 	var existingKeys []interpreter.Value
@@ -507,14 +518,10 @@ func (m *StorageMigration) migrateDictionaryKeys(
 			reporter,
 			// NOTE: Mutation of keys is not allowed.
 			false,
+			ValueMigrationPositionDictionaryKey,
 		)
 
 		if newKey == nil {
-			migratedKeys = append(migratedKeys, migratedDictionaryKey{
-				key:      existingKey,
-				migrated: false,
-			})
-
 			continue
 		}
 
@@ -531,14 +538,25 @@ func (m *StorageMigration) migrateDictionaryKeys(
 
 		// We only reach here because key needs to be migrated.
 
-		// Remove the old key-value pair
+		// Remove the old key-value pair.
 
-		existingKey = LegacyKey(existingKey)
-		existingKeyStorable, existingValueStorable := dictionary.RemoveWithoutTransfer(
-			inter,
-			emptyLocationRange,
-			existingKey,
-		)
+		var existingKeyStorable, existingValueStorable atree.Storable
+
+		legacyKey := LegacyKey(existingKey)
+		if legacyKey != nil {
+			existingKeyStorable, existingValueStorable = dictionary.RemoveWithoutTransfer(
+				inter,
+				emptyLocationRange,
+				legacyKey,
+			)
+		}
+		if existingKeyStorable == nil {
+			existingKeyStorable, existingValueStorable = dictionary.RemoveWithoutTransfer(
+				inter,
+				emptyLocationRange,
+				existingKey,
+			)
+		}
 		if existingKeyStorable == nil {
 			panic(errors.NewUnexpectedError(
 				"failed to remove old value for migrated key: %s",
@@ -589,6 +607,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				valueMigrations,
 				reporter,
 				allowMutation,
+				ValueMigrationPositionOther,
 			)
 
 			var valueToSet interpreter.Value
@@ -655,15 +674,8 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				newKey,
 				existingValue,
 			)
-
-			migratedKeys = append(migratedKeys, migratedDictionaryKey{
-				key:      newKey,
-				migrated: true,
-			})
 		}
 	}
-
-	return
 }
 
 func (m *StorageMigration) migrateDictionaryValues(
@@ -673,21 +685,37 @@ func (m *StorageMigration) migrateDictionaryValues(
 	valueMigrations []ValueMigration,
 	reporter Reporter,
 	allowMutation bool,
-	migratedDictionaryKeys []migratedDictionaryKey,
 ) {
+
 	inter := m.interpreter
 
-	for _, migratedDictionaryKey := range migratedDictionaryKeys {
+	type keyValuePair struct {
+		key, value interpreter.Value
+	}
 
-		existingKey := migratedDictionaryKey.key
-		if !migratedDictionaryKey.migrated {
-			existingKey = LegacyKey(existingKey)
-		}
+	var existingKeysAndValues []keyValuePair
 
-		existingValue, ok := dictionary.Get(inter, emptyLocationRange, existingKey)
-		if !ok {
-			panic(errors.NewUnexpectedError("failed to get existing value for key: %s", existingKey))
-		}
+	dictionary.Iterate(
+		inter,
+		func(key, value interpreter.Value) (resume bool) {
+
+			existingKeysAndValues = append(
+				existingKeysAndValues,
+				keyValuePair{
+					key:   key,
+					value: value,
+				},
+			)
+
+			// Continue iteration
+			return true
+		},
+		emptyLocationRange,
+	)
+
+	for _, existingKeyAndValue := range existingKeysAndValues {
+		existingKey := existingKeyAndValue.key
+		existingValue := existingKeyAndValue.value
 
 		newValue := m.MigrateNestedValue(
 			storageKey,
@@ -696,6 +724,7 @@ func (m *StorageMigration) migrateDictionaryValues(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 
 		if newValue == nil {
@@ -771,6 +800,7 @@ func (m *StorageMigration) migrate(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,
 	value interpreter.Value,
+	position ValueMigrationPosition,
 ) (
 	converted interpreter.Value,
 	err error,
@@ -809,6 +839,7 @@ func (m *StorageMigration) migrate(
 		storageMapKey,
 		value,
 		m.interpreter,
+		position,
 	)
 }
 
@@ -832,7 +863,7 @@ func LegacyKey(key interpreter.Value) interpreter.Value {
 		}
 	}
 
-	return key
+	return nil
 }
 
 func legacyType(staticType interpreter.StaticType) interpreter.StaticType {