Add Array.SetType() to allow updating TypeInfo

array.go

@@ -2371,6 +2371,22 @@ func (a *Array) Type() TypeInfo {
 	return nil
 }
 
+func (a *Array) SetType(typeInfo TypeInfo) error {
+	extraData := a.root.ExtraData()
+	extraData.TypeInfo = typeInfo
+
+	a.root.SetExtraData(extraData)
+
+	// Store modified root slab in storage since typeInfo is part of extraData stored in root slab.
+	err := a.Storage.Store(a.root.Header().id, a.root)
+	if err != nil {
+		// Wrap err as external error (if needed) because err is returned by SlabStorage interface.
+		return wrapErrorfAsExternalErrorIfNeeded(err, fmt.Sprintf("failed to store slab %s", a.root.Header().id))
+	}
+
+	return nil
+}
+
 func (a *Array) String() string {
 	iterator, err := a.Iterator()
 	if err != nil {

array_test.go

@@ -23,6 +23,7 @@ import (
 	"math"
 	"math/rand"
 	"reflect"
+	"runtime"
 	"strings"
 	"testing"
 
@@ -3436,3 +3437,125 @@ func getArrayMetaDataSlabCount(storage *PersistentSlabStorage) int {
 	}
 	return counter
 }
+
+func TestArraySetType(t *testing.T) {
+	typeInfo := testTypeInfo{42}
+	newTypeInfo := testTypeInfo{43}
+	address := Address{1, 2, 3, 4, 5, 6, 7, 8}
+
+	t.Run("empty", func(t *testing.T) {
+		storage := newTestPersistentStorage(t)
+
+		// Create a new array in memory
+		array, err := NewArray(storage, address, typeInfo)
+		require.NoError(t, err)
+		require.Equal(t, uint64(0), array.Count())
+		require.Equal(t, typeInfo, array.Type())
+		require.True(t, array.root.IsData())
+
+		// Modify type info of new array
+		err = array.SetType(newTypeInfo)
+		require.NoError(t, err)
+		require.Equal(t, newTypeInfo, array.Type())
+
+		// Commit new array to storage
+		err = storage.FastCommit(runtime.NumCPU())
+		require.NoError(t, err)
+
+		testExistingArraySetType(t, array.StorageID(), storage.baseStorage, newTypeInfo, array.Count())
+	})
+
+	t.Run("data slab root", func(t *testing.T) {
+		storage := newTestPersistentStorage(t)
+
+		array, err := NewArray(storage, address, typeInfo)
+		require.NoError(t, err)
+
+		arraySize := 10
+		for i := 0; i < arraySize; i++ {
+			v := Uint64Value(i)
+			err := array.Append(v)
+			require.NoError(t, err)
+		}
+
+		require.Equal(t, uint64(arraySize), array.Count())
+		require.Equal(t, typeInfo, array.Type())
+		require.True(t, array.root.IsData())
+
+		err = array.SetType(newTypeInfo)
+		require.NoError(t, err)
+		require.Equal(t, newTypeInfo, array.Type())
+
+		// Commit modified slabs in storage
+		err = storage.FastCommit(runtime.NumCPU())
+		require.NoError(t, err)
+
+		testExistingArraySetType(t, array.StorageID(), storage.baseStorage, newTypeInfo, array.Count())
+	})
+
+	t.Run("metadata slab root", func(t *testing.T) {
+		storage := newTestPersistentStorage(t)
+
+		array, err := NewArray(storage, address, typeInfo)
+		require.NoError(t, err)
+
+		arraySize := 10_000
+		for i := 0; i < arraySize; i++ {
+			v := Uint64Value(i)
+			err := array.Append(v)
+			require.NoError(t, err)
+		}
+
+		require.Equal(t, uint64(arraySize), array.Count())
+		require.Equal(t, typeInfo, array.Type())
+		require.False(t, array.root.IsData())
+
+		err = array.SetType(newTypeInfo)
+		require.NoError(t, err)
+		require.Equal(t, newTypeInfo, array.Type())
+
+		// Commit modified slabs in storage
+		err = storage.FastCommit(runtime.NumCPU())
+		require.NoError(t, err)
+
+		testExistingArraySetType(t, array.StorageID(), storage.baseStorage, newTypeInfo, array.Count())
+	})
+}
+
+func testExistingArraySetType(
+	t *testing.T,
+	id StorageID,
+	baseStorage BaseStorage,
+	expectedTypeInfo testTypeInfo,
+	expectedCount uint64,
+) {
+	newTypeInfo := testTypeInfo{value: expectedTypeInfo.value + 1}
+
+	// Create storage from existing data
+	storage := newTestPersistentStorageWithBaseStorage(t, baseStorage)
+
+	// Load existing array by ID
+	array, err := NewArrayWithRootID(storage, id)
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, array.Count())
+	require.Equal(t, expectedTypeInfo, array.Type())
+
+	// Modify type info of existing array
+	err = array.SetType(newTypeInfo)
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, array.Count())
+	require.Equal(t, newTypeInfo, array.Type())
+
+	// Commit data in storage
+	err = storage.FastCommit(runtime.NumCPU())
+	require.NoError(t, err)
+
+	// Create storage from existing data
+	storage2 := newTestPersistentStorageWithBaseStorage(t, storage.baseStorage)
+
+	// Load existing array again from storage
+	array2, err := NewArrayWithRootID(storage2, id)
+	require.NoError(t, err)
+	require.Equal(t, expectedCount, array2.Count())
+	require.Equal(t, newTypeInfo, array2.Type())
+}