add migration which re-stores type values used as dictionary keys

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()
@@ -2858,12 +2861,7 @@ func TestConvertMigratedAccountTypes(t *testing.T) {
 			)
 
 			newValue, err := statictypes.NewStaticTypeMigration().
-				Migrate(
-					interpreter.StorageKey{},
-					nil,
-					value,
-					inter,
-				)
+				Migrate(interpreter.StorageKey{}, nil, value, inter, migrations.ValueMigrationPositionOther)
 			require.NoError(t, err)
 			require.NotNil(t, newValue)
 
@@ -3658,6 +3656,7 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 			migration.NewValueMigrationsPathMigrator(
 				reporter,
 				NewEntitlementsMigration(inter),
+				type_keys.NewTypeKeyMigration(),
 			),
 		)
 
@@ -3673,7 +3672,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

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,24 +518,9 @@ func (m *StorageMigration) migrateDictionaryKeys(
 			reporter,
 			// NOTE: Mutation of keys is not allowed.
 			false,
+			ValueMigrationPositionDictionaryKey,
 		)
 
-		// When we remove the existing key, we need to use the legacy key for the lookup,
-		// as the key might be stored with the old hash function.
-
-		legacyKey := LegacyKey(existingKey)
-
-		// If the key did not get migrated, either if it was not necessary, or because the migration failed,
-		// then still re-store the existing key, so the key-value pair stays accessible.
-		//
-		// This is also important for the inlining version of atree,
-		// as its mutating iterator does not support iterating over values with an old hash,
-		// it re-computes the hash of the stored value with the new hash function.
-
-		if newKey == nil && legacyKey != nil {
-			newKey = existingKey
-		}
-
 		if newKey == nil {
 			continue
 		}
@@ -536,16 +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.
+
+		var existingKeyStorable, existingValueStorable atree.Storable
 
+		legacyKey := LegacyKey(existingKey)
 		if legacyKey != nil {
-			existingKey = legacyKey
+			existingKeyStorable, existingValueStorable = dictionary.RemoveWithoutTransfer(
+				inter,
+				emptyLocationRange,
+				legacyKey,
+			)
+		}
+		if existingKeyStorable == nil {
+			existingKeyStorable, existingValueStorable = dictionary.RemoveWithoutTransfer(
+				inter,
+				emptyLocationRange,
+				existingKey,
+			)
 		}
-		existingKeyStorable, existingValueStorable := dictionary.RemoveWithoutTransfer(
-			inter,
-			emptyLocationRange,
-			existingKey,
-		)
 		if existingKeyStorable == nil {
 			panic(errors.NewUnexpectedError(
 				"failed to remove old value for migrated key: %s",
@@ -596,6 +607,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				valueMigrations,
 				reporter,
 				allowMutation,
+				ValueMigrationPositionOther,
 			)
 
 			var valueToSet interpreter.Value
@@ -712,6 +724,7 @@ func (m *StorageMigration) migrateDictionaryValues(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			ValueMigrationPositionOther,
 		)
 
 		if newValue == nil {
@@ -787,6 +800,7 @@ func (m *StorageMigration) migrate(
 	storageKey interpreter.StorageKey,
 	storageMapKey interpreter.StorageMapKey,
 	value interpreter.Value,
+	position ValueMigrationPosition,
 ) (
 	converted interpreter.Value,
 	err error,
@@ -825,6 +839,7 @@ func (m *StorageMigration) migrate(
 		storageMapKey,
 		value,
 		m.interpreter,
+		position,
 	)
 }