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 {
@@ -390,6 +405,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 		if newInnerValue != nil {
 			newInnerCapability := newInnerValue.(interpreter.CapabilityValue)
@@ -414,6 +430,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			storageKey,
 			storageMapKey,
 			value,
+			position,
 		)
 
 		if err != nil {
@@ -501,6 +518,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 			reporter,
 			// NOTE: Mutation of keys is not allowed.
 			false,
+			ValueMigrationPositionDictionaryKey,
 		)
 
 		if newKey == nil {
@@ -520,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",
@@ -578,6 +607,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				valueMigrations,
 				reporter,
 				allowMutation,
+				ValueMigrationPositionOther,
 			)
 
 			var valueToSet interpreter.Value
@@ -694,6 +724,7 @@ func (m *StorageMigration) migrateDictionaryValues(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 
 		if newValue == nil {
@@ -769,6 +800,7 @@ func (m *StorageMigration) migrate(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,
 	value interpreter.Value,
+	position ValueMigrationPosition,
 ) (
 	converted interpreter.Value,
 	err error,
@@ -807,11 +839,12 @@ func (m *StorageMigration) migrate(
 		storageMapKey,
 		value,
 		m.interpreter,
+		position,
 	)
 }
 
-// legacyKey return the same type with the "old" hash/ID generation function.
-func legacyKey(key interpreter.Value) interpreter.Value {
+// LegacyKey return the same type with the "old" hash/ID generation function.
+func LegacyKey(key interpreter.Value) interpreter.Value {
 	switch key := key.(type) {
 	case interpreter.TypeValue:
 		legacyType := legacyType(key.Type)
@@ -830,7 +863,7 @@ func legacyKey(key interpreter.Value) interpreter.Value {
 		}
 	}
 
-	return key
+	return nil
 }
 
 func legacyType(staticType interpreter.StaticType) interpreter.StaticType {