Improve dictionary value migration

migrations/entitlements/migration_test.go

@@ -3516,3 +3516,217 @@ func TestIntersectionTypeWithIntersectionLegacyType(t *testing.T) {
 		require.Equal(t, expectedType, typeValue.Type)
 	})()
 }
+
+func TestUseAfterMigrationFailure(t *testing.T) {
+
+	t.Parallel()
+
+	locationRange := interpreter.EmptyLocationRange
+
+	ledger := NewTestLedger(nil, nil)
+
+	storageMapKey := interpreter.StringStorageMapKey("dict")
+	newTestValue := func() interpreter.Value {
+		return interpreter.NewUnmeteredStringValue("test")
+	}
+
+	const fooBarQualifiedIdentifier = "Foo.Bar"
+	testAddress := common.Address{0x42}
+	fooAddressLocation := common.NewAddressLocation(nil, testAddress, "Foo")
+
+	newStorageAndInterpreter := func(t *testing.T) (*runtime.Storage, *interpreter.Interpreter) {
+		storage := runtime.NewStorage(ledger, nil)
+		inter, err := interpreter.NewInterpreter(
+			nil,
+			utils.TestLocation,
+			&interpreter.Config{
+				Storage: storage,
+				// NOTE: disabled, because encoded and decoded values are expected to not match
+				AtreeValueValidationEnabled:   false,
+				AtreeStorageValidationEnabled: true,
+			},
+		)
+		require.NoError(t, err)
+
+		return storage, inter
+	}
+
+	newCompositeType := func() *interpreter.CompositeStaticType {
+		return interpreter.NewCompositeStaticType(
+			nil,
+			fooAddressLocation,
+			fooBarQualifiedIdentifier,
+			common.NewTypeIDFromQualifiedName(
+				nil,
+				fooAddressLocation,
+				fooBarQualifiedIdentifier,
+			),
+		)
+	}
+
+	// Prepare
+	(func() {
+
+		storage, inter := newStorageAndInterpreter(t)
+
+		dictionaryStaticType := interpreter.NewDictionaryStaticType(
+			nil,
+			interpreter.PrimitiveStaticTypeMetaType,
+			interpreter.PrimitiveStaticTypeString,
+		)
+		dictValue := interpreter.NewDictionaryValue(inter, locationRange, dictionaryStaticType)
+
+		refType := interpreter.NewReferenceStaticType(
+			nil,
+			interpreter.UnauthorizedAccess,
+			newCompositeType(),
+		)
+		refType.HasLegacyIsAuthorized = true
+		refType.LegacyIsAuthorized = true
+
+		legacyRefType := &migrations.LegacyReferenceType{
+			ReferenceStaticType: refType,
+		}
+
+		optType := interpreter.NewOptionalStaticType(
+			nil,
+			legacyRefType,
+		)
+
+		legacyOptType := &migrations.LegacyOptionalType{
+			OptionalStaticType: optType,
+		}
+
+		typeValue := interpreter.NewUnmeteredTypeValue(legacyOptType)
+
+		dictValue.Insert(
+			inter,
+			locationRange,
+			typeValue,
+			newTestValue(),
+		)
+
+		// Note: ID is in the old format
+		assert.Equal(t,
+			common.TypeID("auth&A.4200000000000000.Foo.Bar?"),
+			legacyOptType.ID(),
+		)
+
+		storageMap := storage.GetStorageMap(
+			testAddress,
+			common.PathDomainStorage.Identifier(),
+			true,
+		)
+
+		storageMap.SetValue(inter,
+			storageMapKey,
+			dictValue.Transfer(
+				inter,
+				locationRange,
+				atree.Address(testAddress),
+				false,
+				nil,
+				nil,
+			),
+		)
+
+		err := storage.Commit(inter, false)
+		require.NoError(t, err)
+
+		err = storage.CheckHealth()
+		require.NoError(t, err)
+	})()
+
+	// Migrate
+	(func() {
+
+		storage, inter := newStorageAndInterpreter(t)
+
+		const importErrorMessage = "cannot import"
+
+		inter.SharedState.Config.ImportLocationHandler =
+			func(inter *interpreter.Interpreter, location common.Location) interpreter.Import {
+				panic(importErrorMessage)
+			}
+
+		migration, err := migrations.NewStorageMigration(inter, storage, "test", testAddress)
+		require.NoError(t, err)
+
+		reporter := newTestReporter()
+
+		migration.Migrate(
+			migration.NewValueMigrationsPathMigrator(
+				reporter,
+				NewEntitlementsMigration(inter),
+			),
+		)
+
+		err = migration.Commit()
+		require.NoError(t, err)
+
+		// Assert
+
+		err = storage.CheckHealth()
+		require.NoError(t, err)
+
+		require.Len(t, reporter.errors, 1)
+
+		assert.ErrorContains(t, reporter.errors[0], importErrorMessage)
+
+		require.Empty(t, reporter.migrated)
+	})()
+
+	// Load
+	(func() {
+
+		storage, inter := newStorageAndInterpreter(t)
+
+		err := storage.CheckHealth()
+		require.NoError(t, err)
+
+		storageMap := storage.GetStorageMap(
+			testAddress,
+			common.PathDomainStorage.Identifier(),
+			false,
+		)
+		storedValue := storageMap.ReadValue(inter, storageMapKey)
+
+		require.IsType(t, &interpreter.DictionaryValue{}, storedValue)
+
+		dictValue := storedValue.(*interpreter.DictionaryValue)
+
+		refType := interpreter.NewReferenceStaticType(
+			nil,
+			interpreter.UnauthorizedAccess,
+			newCompositeType(),
+		)
+		refType.HasLegacyIsAuthorized = true
+		refType.LegacyIsAuthorized = true
+
+		optType := interpreter.NewOptionalStaticType(
+			nil,
+			refType,
+		)
+
+		typeValue := interpreter.NewUnmeteredTypeValue(optType)
+
+		// Note: ID is in the new format
+		assert.Equal(t,
+			common.TypeID("(&A.4200000000000000.Foo.Bar)?"),
+			optType.ID(),
+		)
+
+		assert.Equal(t, 1, dictValue.Count())
+
+		// Key did not get migrated, so is inaccessible using the "new" type value
+		_, 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

@@ -363,7 +363,7 @@ func (m *StorageMigration) MigrateNestedValue(
 		// The mutating iterator is only able to read new keys,
 		// as it recalculates the stored values' hashes.
 
-		m.migrateDictionaryKeys(
+		keys := m.migrateDictionaryKeys(
 			storageKey,
 			storageMapKey,
 			dictionary,
@@ -379,6 +379,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			valueMigrations,
 			reporter,
 			allowMutation,
+			keys,
 		)
 
 	case *interpreter.PublishedValue:
@@ -468,14 +469,19 @@ 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
@@ -504,6 +510,11 @@ func (m *StorageMigration) migrateDictionaryKeys(
 		)
 
 		if newKey == nil {
+			migratedKeys = append(migratedKeys, migratedDictionaryKey{
+				key:      existingKey,
+				migrated: false,
+			})
+
 			continue
 		}
 
@@ -522,7 +533,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 
 		// Remove the old key-value pair
 
-		existingKey = legacyKey(existingKey)
+		existingKey = LegacyKey(existingKey)
 		existingKeyStorable, existingValueStorable := dictionary.RemoveWithoutTransfer(
 			inter,
 			emptyLocationRange,
@@ -644,8 +655,15 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				newKey,
 				existingValue,
 			)
+
+			migratedKeys = append(migratedKeys, migratedDictionaryKey{
+				key:      newKey,
+				migrated: true,
+			})
 		}
 	}
+
+	return
 }
 
 func (m *StorageMigration) migrateDictionaryValues(
@@ -655,37 +673,21 @@ func (m *StorageMigration) migrateDictionaryValues(
 	valueMigrations []ValueMigration,
 	reporter Reporter,
 	allowMutation bool,
+	migratedDictionaryKeys []migratedDictionaryKey,
 ) {
-
 	inter := m.interpreter
 
-	type keyValuePair struct {
-		key, value interpreter.Value
-	}
+	for _, migratedDictionaryKey := range migratedDictionaryKeys {
 
-	var existingKeysAndValues []keyValuePair
-
-	dictionary.Iterate(
-		inter,
-		func(key, value interpreter.Value) (resume bool) {
-
-			existingKeysAndValues = append(
-				existingKeysAndValues,
-				keyValuePair{
-					key:   key,
-					value: value,
-				},
-			)
-
-			// Continue iteration
-			return true
-		},
-		emptyLocationRange,
-	)
+		existingKey := migratedDictionaryKey.key
+		if !migratedDictionaryKey.migrated {
+			existingKey = LegacyKey(existingKey)
+		}
 
-	for _, existingKeyAndValue := range existingKeysAndValues {
-		existingKey := existingKeyAndValue.key
-		existingValue := existingKeyAndValue.value
+		existingValue, ok := dictionary.Get(inter, emptyLocationRange, existingKey)
+		if !ok {
+			panic(errors.NewUnexpectedError("failed to get existing value for key: %s", existingKey))
+		}
 
 		newValue := m.MigrateNestedValue(
 			storageKey,
@@ -810,8 +812,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)

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)