From a97d108f7c1b4c8dd8df26e31eb163abaa33f4f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20M=C3=BCller?= <bastian@turbolent.com>
Date: Tue, 11 Jun 2024 14:36:14 -0700
Subject: [PATCH] also migrate account capability storage maps

---
 migrations/migration.go      |  6 ++++
 migrations/migration_test.go | 70 +++++++++++++++++++++++++++++++-----
 2 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/migrations/migration.go b/migrations/migration.go
index 5139b2d030..ab1b88426b 100644
--- a/migrations/migration.go
+++ b/migrations/migration.go
@@ -133,6 +133,12 @@ func (m *StorageMigration) Migrate(migrator StorageMapKeyMigrator) {
 		stdlib.PathCapabilityStorageDomain,
 		migrator,
 	)
+
+	accountStorage.MigrateUint64Keys(
+		m.interpreter,
+		stdlib.AccountCapabilityStorageDomain,
+		migrator,
+	)
 }
 
 func (m *StorageMigration) NewValueMigrationsPathMigrator(
diff --git a/migrations/migration_test.go b/migrations/migration_test.go
index 236a059c2a..8cf592af51 100644
--- a/migrations/migration_test.go
+++ b/migrations/migration_test.go
@@ -245,6 +245,9 @@ func (testCapConMigration) Migrate(
 
 	case interpreter.UInt64Value:
 		return value + 10, nil
+
+	case interpreter.NilValue:
+		return value, nil
 	}
 
 	return nil, nil
@@ -829,12 +832,21 @@ func TestCapConMigration(t *testing.T) {
 		false,
 	)
 
-	capSet := pathCapStorageMap.ReadValue(nil, interpreter.StringStorageMapKey("foo"))
-	require.IsType(t, &interpreter.DictionaryValue{}, capSet)
-	capSetDict := capSet.(*interpreter.DictionaryValue)
+	pathCapSet := pathCapStorageMap.ReadValue(nil, interpreter.StringStorageMapKey("foo"))
+	require.IsType(t, &interpreter.DictionaryValue{}, pathCapSet)
+	capSetDict := pathCapSet.(*interpreter.DictionaryValue)
 	assert.Equal(t, 1, capSetDict.Count())
 	assert.True(t, bool(capSetDict.ContainsKey(inter, emptyLocationRange, interpreter.UInt64Value(1))))
 
+	accountCapStorageMap := storage.GetStorageMap(
+		testAddress,
+		stdlib.AccountCapabilityStorageDomain,
+		false,
+	)
+
+	accountCapEntry := accountCapStorageMap.ReadValue(nil, interpreter.Uint64StorageMapKey(2))
+	assert.NotNil(t, accountCapEntry)
+
 	// Migrate
 
 	migration, err := NewStorageMigration(inter, storage, "test", testAddress)
@@ -854,9 +866,43 @@ func TestCapConMigration(t *testing.T) {
 
 	// Assert
 
-	assert.Len(t, reporter.migrated, 3)
-
 	require.Empty(t, reporter.errors)
+	assert.Equal(t,
+		map[struct {
+			interpreter.StorageKey
+			interpreter.StorageMapKey
+		}][]string{
+			{
+				StorageKey: interpreter.StorageKey{
+					Address: testAddress,
+					Key:     stdlib.CapabilityControllerStorageDomain,
+				},
+				StorageMapKey: interpreter.Uint64StorageMapKey(1),
+			}: {"testCapConMigration"},
+			{
+				StorageKey: interpreter.StorageKey{
+					Address: testAddress,
+					Key:     stdlib.CapabilityControllerStorageDomain,
+				},
+				StorageMapKey: interpreter.Uint64StorageMapKey(2),
+			}: {"testCapConMigration"},
+			{
+				StorageKey: interpreter.StorageKey{
+					Address: testAddress,
+					Key:     stdlib.PathCapabilityStorageDomain,
+				},
+				StorageMapKey: interpreter.StringStorageMapKey("foo"),
+			}: {"testCapConMigration", "testCapConMigration"},
+			{
+				StorageKey: interpreter.StorageKey{
+					Address: testAddress,
+					Key:     stdlib.AccountCapabilityStorageDomain,
+				},
+				StorageMapKey: interpreter.Uint64StorageMapKey(2),
+			}: {"testCapConMigration"},
+		},
+		reporter.migrated,
+	)
 
 	err = storage.CheckHealth()
 	require.NoError(t, err)
@@ -889,12 +935,20 @@ func TestCapConMigration(t *testing.T) {
 		false,
 	)
 
-	capSet = pathCapStorageMap.ReadValue(nil, interpreter.StringStorageMapKey("foo"))
-	require.IsType(t, &interpreter.DictionaryValue{}, capSet)
-	capSetDict = capSet.(*interpreter.DictionaryValue)
+	pathCapSet = pathCapStorageMap.ReadValue(nil, interpreter.StringStorageMapKey("foo"))
+	require.IsType(t, &interpreter.DictionaryValue{}, pathCapSet)
+	capSetDict = pathCapSet.(*interpreter.DictionaryValue)
 	assert.Equal(t, 1, capSetDict.Count())
 	assert.True(t, bool(capSetDict.ContainsKey(inter, emptyLocationRange, interpreter.UInt64Value(11))))
 
+	accountCapStorageMap = storage.GetStorageMap(
+		testAddress,
+		stdlib.AccountCapabilityStorageDomain,
+		false,
+	)
+
+	accountCapEntry = accountCapStorageMap.ReadValue(nil, interpreter.Uint64StorageMapKey(2))
+	assert.NotNil(t, accountCapEntry)
 }
 
 func TestContractMigration(t *testing.T) {