From 5fcf093668e033675dd1ce06d898b8176efe1215 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Thu, 14 Nov 2024 15:39:06 +0100
Subject: [PATCH] core/state/snapshot: fix binary iterator

---
 core/state/snapshot/iterator_binary.go | 56 ++++++++++++++++----------
 core/state/snapshot/iterator_test.go   |  9 ++---
 2 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/core/state/snapshot/iterator_binary.go b/core/state/snapshot/iterator_binary.go
index d6f84f2e34..6d17d37759 100644
--- a/core/state/snapshot/iterator_binary.go
+++ b/core/state/snapshot/iterator_binary.go
@@ -117,33 +117,47 @@ func (dl *diffLayer) initBinaryStorageIterator(account, seek common.Hash) Iterat
 // or an error if iteration failed for some reason (e.g. root being iterated
 // becomes stale and garbage collected).
 func (it *binaryIterator) Next() bool {
+	for {
+		ok := it.next()
+		if !ok {
+			return ok
+		}
+		if len(it.Account()) == 0 && len(it.Slot()) == 0 {
+			continue
+		}
+		return ok
+	}
+}
+
+func (it *binaryIterator) next() bool {
 	if it.aDone && it.bDone {
 		return false
 	}
-first:
-	if it.aDone {
-		it.k = it.b.Hash()
+	for {
+		if it.aDone {
+			it.k = it.b.Hash()
+			it.bDone = !it.b.Next()
+			return true
+		}
+		if it.bDone {
+			it.k = it.a.Hash()
+			it.aDone = !it.a.Next()
+			return true
+		}
+		nextA, nextB := it.a.Hash(), it.b.Hash()
+		if diff := bytes.Compare(nextA[:], nextB[:]); diff < 0 {
+			it.aDone = !it.a.Next()
+			it.k = nextA
+			return true
+		} else if diff == 0 {
+			// Now we need to advance one of them
+			it.aDone = !it.a.Next()
+			continue
+		}
 		it.bDone = !it.b.Next()
+		it.k = nextB
 		return true
 	}
-	if it.bDone {
-		it.k = it.a.Hash()
-		it.aDone = !it.a.Next()
-		return true
-	}
-	nextA, nextB := it.a.Hash(), it.b.Hash()
-	if diff := bytes.Compare(nextA[:], nextB[:]); diff < 0 {
-		it.aDone = !it.a.Next()
-		it.k = nextA
-		return true
-	} else if diff == 0 {
-		// Now we need to advance one of them
-		it.aDone = !it.a.Next()
-		goto first
-	}
-	it.bDone = !it.b.Next()
-	it.k = nextB
-	return true
 }
 
 // Error returns any failure that occurred during iteration, which might have
diff --git a/core/state/snapshot/iterator_test.go b/core/state/snapshot/iterator_test.go
index 752c78216c..c6c6d4928d 100644
--- a/core/state/snapshot/iterator_test.go
+++ b/core/state/snapshot/iterator_test.go
@@ -99,9 +99,6 @@ func TestStorageIteratorBasics(t *testing.T) {
 	for account := range accounts {
 		it, _ := diffLayer.StorageIterator(account, common.Hash{})
 		verifyIterator(t, 100, it, verifyNothing) // Nil is allowed for single layer iterator
-
-		it = diffLayer.newBinaryStorageIterator(account, common.Hash{})
-		verifyIterator(t, 100, it, verifyNothing) // Nil is allowed for single layer iterator
 	}
 
 	diskLayer := diffToDisk(diffLayer)
@@ -569,7 +566,7 @@ func TestAccountIteratorFlattening(t *testing.T) {
 	})
 	t.Run("binary", func(t *testing.T) {
 		testAccountIteratorFlattening(t, func(snaps *Tree, root, seek common.Hash) AccountIterator {
-			return snaps.layers[root].(*diffLayer).newBinaryAccountIterator(common.Hash{})
+			return snaps.layers[root].(*diffLayer).newBinaryAccountIterator(seek)
 
 		})
 	})
@@ -771,7 +768,7 @@ func TestAccountIteratorDeletions(t *testing.T) {
 	})
 	t.Run("binary", func(t *testing.T) {
 		testAccountIteratorDeletions(t, func(snaps *Tree, root, seek common.Hash) AccountIterator {
-			return snaps.layers[root].(*diffLayer).newBinaryAccountIterator(common.Hash{})
+			return snaps.layers[root].(*diffLayer).newBinaryAccountIterator(seek)
 
 		})
 	})
@@ -833,7 +830,7 @@ func TestStorageIteratorDeletions(t *testing.T) {
 	})
 	t.Run("binary", func(t *testing.T) {
 		testStorageIteratorDeletions(t, func(snaps *Tree, root, account, seek common.Hash) StorageIterator {
-			return snaps.layers[root].(*diffLayer).newBinaryStorageIterator(account, common.Hash{})
+			return snaps.layers[root].(*diffLayer).newBinaryStorageIterator(account, seek)
 		})
 	})
 }