ensure dictionary keys with legacy representation are always re-stored

even if migration is not necessary or migration failed, ensure the key is restored,
so the key-value pair stays accessible, and inlining atree mutating iterator does not fail

migrations/entitlements/migration_test.go

@@ -3718,15 +3718,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

@@ -503,6 +503,22 @@ func (m *StorageMigration) migrateDictionaryKeys(
 			false,
 		)
 
+		// 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
 		}
@@ -522,7 +538,9 @@ func (m *StorageMigration) migrateDictionaryKeys(
 
 		// Remove the old key-value pair
 
-		existingKey = legacyKey(existingKey)
+		if legacyKey != nil {
+			existingKey = legacyKey
+		}
 		existingKeyStorable, existingValueStorable := dictionary.RemoveWithoutTransfer(
 			inter,
 			emptyLocationRange,
@@ -810,8 +828,8 @@ func (m *StorageMigration) migrate(
 	)
 }
 
-// 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 +848,7 @@ func legacyKey(key interpreter.Value) interpreter.Value {
 		}
 	}
 
-	return key
+	return nil
 }
 
 func legacyType(staticType interpreter.StaticType) interpreter.StaticType {

migrations/migration_test.go

@@ -2480,12 +2480,12 @@ func TestDictionaryKeyConflict(t *testing.T) {
 				dictionaryValue,
 			)
 
-			// NOTE: use legacyKey to ensure the key is encoded in old format
+			// NOTE: use LegacyKey to ensure the key is encoded in old format
 
 			dictionaryValue.InsertWithoutTransfer(
 				inter,
 				emptyLocationRange,
-				legacyKey(dictionaryKey1),
+				LegacyKey(dictionaryKey1),
 				interpreter.NewArrayValue(
 					inter,
 					emptyLocationRange,
@@ -2498,7 +2498,7 @@ func TestDictionaryKeyConflict(t *testing.T) {
 			dictionaryValue.InsertWithoutTransfer(
 				inter,
 				emptyLocationRange,
-				legacyKey(dictionaryKey2),
+				LegacyKey(dictionaryKey2),
 				interpreter.NewArrayValue(
 					inter,
 					emptyLocationRange,
@@ -2511,7 +2511,7 @@ func TestDictionaryKeyConflict(t *testing.T) {
 			oldValue1, ok := dictionaryValue.Get(
 				inter,
 				emptyLocationRange,
-				legacyKey(dictionaryKey1),
+				LegacyKey(dictionaryKey1),
 			)
 			require.True(t, ok)
 
@@ -2530,7 +2530,7 @@ func TestDictionaryKeyConflict(t *testing.T) {
 			oldValue2, ok := dictionaryValue.Get(
 				inter,
 				emptyLocationRange,
-				legacyKey(dictionaryKey2),
+				LegacyKey(dictionaryKey2),
 			)
 			require.True(t, ok)