From ef6dd7e674f3cb1b64afd799b48e2a0f7919ce14 Mon Sep 17 00:00:00 2001
From: Faye Amacker <33205765+fxamacker@users.noreply.github.com>
Date: Fri, 26 Apr 2024 10:38:06 -0500
Subject: [PATCH] Fix SlabIterator to include nested storage ID

SlabIterator is used by storage health check, which is used
for tests and for new migration features like filtering out
old unreferenced slabs, etc.

Currently, SlabIterator doesn't traverse deep enough to handle
nested storage ID inside another storable such as Cadence SomeValue.

This commit fixes this by handling nested storage ID in element
such as Cadence SomeValue.
---
 storage.go      |  4 ++++
 storage_test.go | 25 ++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/storage.go b/storage.go
index 71b30bf..ee9adca 100644
--- a/storage.go
+++ b/storage.go
@@ -558,6 +558,10 @@ func (s *PersistentSlabStorage) SlabIterator() (SlabIterator, error) {
 
 				storageIDStorable, ok := childStorable.(StorageIDStorable)
 				if !ok {
+					nextChildStorables = append(
+						nextChildStorables,
+						childStorable.ChildStorables()...,
+					)
 					continue
 				}
 
diff --git a/storage_test.go b/storage_test.go
index 8a06ceb..d3dadcf 100644
--- a/storage_test.go
+++ b/storage_test.go
@@ -893,6 +893,7 @@ func TestPersistentStorageSlabIterator(t *testing.T) {
 		id2 := StorageID{Address: address, Index: StorageIndex{0, 0, 0, 0, 0, 0, 0, 2}}
 		id3 := StorageID{Address: address, Index: StorageIndex{0, 0, 0, 0, 0, 0, 0, 3}}
 		id4 := StorageID{Address: address, Index: StorageIndex{0, 0, 0, 0, 0, 0, 0, 4}}
+		id5 := StorageID{Address: address, Index: StorageIndex{0, 0, 0, 0, 0, 0, 0, 5}}
 
 		data := map[StorageID][]byte{
 			// (metadata slab) headers: [{id:2 size:228 count:9} {id:3 size:270 count:11} ]
@@ -969,7 +970,7 @@ func TestPersistentStorageSlabIterator(t *testing.T) {
 				0xd8, 0xff, 0x50, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04,
 			},
 
-			// (data slab) next: 0, data: [0]
+			// (data slab) next: 0, data: [SomeValue(StorageID(...))]
 			id4: {
 				// extra data
 				// version
@@ -981,6 +982,28 @@ func TestPersistentStorageSlabIterator(t *testing.T) {
 				// type info
 				0x18, 0x2b,
 
+				// version
+				0x00,
+				// array data slab flag
+				0x80,
+				// CBOR encoded array head (fixed size 3 byte)
+				0x99, 0x00, 0x01,
+				// CBOR encoded array elements
+				0xd8, cborTagSomeValue, 0xd8, 0xff, 0x50, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05,
+			},
+
+			// (data slab) next: 0, data: [0]
+			id5: {
+				// extra data
+				// version
+				0x00,
+				// extra data flag
+				0x80,
+				// array of extra data
+				0x81,
+				// type info
+				0x18, 0x2b,
+
 				// version
 				0x00,
 				// array data slab flag