From 86040b30171f26c9ec2749c7760d59eb8ac73b3a Mon Sep 17 00:00:00 2001
From: Faye Amacker <33205765+fxamacker@users.noreply.github.com>
Date: Tue, 11 Jul 2023 09:31:09 -0500
Subject: [PATCH] Add more tests for loaded value iterator

Added tests that remove slabs from storage during iteration
over loaded values.

The scenario tested shouldn't happen in practice but
we want to test that it would be handled gracefully.
---
 array_test.go | 36 ++++++++++++++++++++++++++++++++++++
 map_test.go   | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/array_test.go b/array_test.go
index 6794abab..c2a733b7 100644
--- a/array_test.go
+++ b/array_test.go
@@ -2727,6 +2727,42 @@ func TestArrayLoadedValueIterator(t *testing.T) {
 		verifyArrayLoadedElements(t, array, values)
 	})
 
+	t.Run("root data slab with composite values, unload composite elements during iteration", func(t *testing.T) {
+		storage := newTestPersistentStorage(t)
+
+		const arraySize = 3
+		array, values := createArrayWithCompositeValues(t, storage, address, typeInfo, arraySize)
+
+		// parent array: 1 root data slab
+		// nested composite elements: 1 root data slab for each
+		require.Equal(t, 1+arraySize, len(storage.deltas))
+		require.Equal(t, 0, getArrayMetaDataSlabCount(storage))
+
+		verifyArrayLoadedElements(t, array, values)
+
+		i := 0
+		err := array.IterateLoadedValues(func(v Value) (bool, error) {
+			// At this point, iterator returned first element (v).
+
+			// Remove all other nested composite elements (except first element) from storage.
+			for _, value := range values[1:] {
+				nestedArray, ok := value.(*Array)
+				require.True(t, ok)
+
+				err := storage.Remove(nestedArray.StorageID())
+				require.NoError(t, err)
+			}
+
+			require.Equal(t, 0, i)
+			valueEqual(t, typeInfoComparator, values[0], v)
+			i++
+			return true, nil
+		})
+
+		require.NoError(t, err)
+		require.Equal(t, 1, i) // Only first element is iterated because other elements are remove during iteration.
+	})
+
 	t.Run("root data slab with simple and composite values, unload composite element", func(t *testing.T) {
 		const arraySize = 3
 
diff --git a/map_test.go b/map_test.go
index 28f8e4cd..fca734fc 100644
--- a/map_test.go
+++ b/map_test.go
@@ -4800,6 +4800,51 @@ func TestMapLoadedValueIterator(t *testing.T) {
 		verifyMapLoadedElements(t, m, values)
 	})
 
+	t.Run("root data slab with composite values, unload composite elements during iteration", func(t *testing.T) {
+		storage := newTestPersistentStorage(t)
+
+		const mapSize = 3
+		m, values := createMapWithCompositeValues(
+			t,
+			storage,
+			address,
+			typeInfo,
+			mapSize,
+			func(i int) []Digest { return []Digest{Digest(i)} },
+		)
+
+		// parent map: 1 root data slab
+		// nested composite elements: 1 root data slab for each
+		require.Equal(t, 1+mapSize, len(storage.deltas))
+		require.Equal(t, 0, getMapMetaDataSlabCount(storage))
+
+		verifyMapLoadedElements(t, m, values)
+
+		i := 0
+		err := m.IterateLoadedValues(func(k Value, v Value) (bool, error) {
+			// At this point, iterator returned first element (v).
+
+			// Remove all other nested composite elements (except first element) from storage.
+			for _, element := range values[1:] {
+				value := element[1]
+				nestedArray, ok := value.(*Array)
+				require.True(t, ok)
+
+				err := storage.Remove(nestedArray.StorageID())
+				require.NoError(t, err)
+			}
+
+			require.Equal(t, 0, i)
+			valueEqual(t, typeInfoComparator, values[0][0], k)
+			valueEqual(t, typeInfoComparator, values[0][1], v)
+			i++
+			return true, nil
+		})
+
+		require.NoError(t, err)
+		require.Equal(t, 1, i) // Only first element is iterated because other elements are remove during iteration.
+	})
+
 	t.Run("root data slab with simple and composite values, unloading composite value", func(t *testing.T) {
 		const mapSize = 3