gather dictionary keys during migration and use them when migrating d…

…ictionary values avoid re-iterating, as e.g. the iteration using the inlining version of atree will fail for old keys
onflow · May 29, 2024 · 943b9b3 · 943b9b3
1 parent e7f5ee9
commit 943b9b3
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 55 deletions.
diff --git a/migrations/entitlements/migration_test.go b/migrations/entitlements/migration_test.go
@@ -3564,17 +3564,6 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 		)
 	}
 
-	entitlementSetAuthorization := sema.NewEntitlementSetAccess(
-		[]*sema.EntitlementType{
-			sema.NewEntitlementType(
-				nil,
-				fooAddressLocation,
-				"E",
-			),
-		},
-		sema.Conjunction,
-	)
-
 	// Prepare
 	(func() {
 
@@ -3619,8 +3608,8 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 
 		// Note: ID is in the old format
 		assert.Equal(t,
-			common.TypeID("auth&A.4200000000000000.Foo.Bar"),
-			legacyRefType.ID(),
+			common.TypeID("auth&A.4200000000000000.Foo.Bar?"),
+			legacyOptType.ID(),
 		)
 
 		storageMap := storage.GetStorageMap(
@@ -3638,7 +3627,6 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 				false,
 				nil,
 				nil,
-				true, // dictValue is standalone
 			),
 		)
 
@@ -3681,10 +3669,9 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 		err = storage.CheckHealth()
 		require.NoError(t, err)
 
-		require.Len(t, reporter.errors, 2)
+		require.Len(t, reporter.errors, 1)
 
 		assert.ErrorContains(t, reporter.errors[0], importErrorMessage)
-		assert.ErrorContains(t, reporter.errors[1], "key (Type<&A.4200000000000000.Foo.Bar?>()) not found")
 
 		require.Empty(t, reporter.migrated)
 	})()
@@ -3710,9 +3697,11 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 
 		refType := interpreter.NewReferenceStaticType(
 			nil,
-			interpreter.ConvertSemaAccessToStaticAuthorization(nil, entitlementSetAuthorization),
+			interpreter.UnauthorizedAccess,
 			newCompositeType(),
 		)
+		refType.HasLegacyIsAuthorized = true
+		refType.LegacyIsAuthorized = true
 
 		optType := interpreter.NewOptionalStaticType(
 			nil,
@@ -3723,14 +3712,21 @@ func TestUseAfterMigrationFailure(t *testing.T) {
 
 		// Note: ID is in the new format
 		assert.Equal(t,
-			common.TypeID("(auth(A.4200000000000000.E)&A.4200000000000000.Foo.Bar)?"),
+			common.TypeID("(&A.4200000000000000.Foo.Bar)?"),
 			optType.ID(),
 		)
 
 		assert.Equal(t, 1, dictValue.Count())
 
-		// Key did not get migrated, so is inaccessible
+		// 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)
 	})()
 }
diff --git a/migrations/migration.go b/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)

diff --git a/migrations/migration_test.go b/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)