From eff7cfbb03b74816513d415c0ecfe93ae83f4096 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Wed, 4 Mar 2020 13:38:55 +0100
Subject: [PATCH] core/state/snapshot: handle deleted accounts in fast iterator

---
 core/state/snapshot/iterator_fast.go | 34 +++++++++++++-----
 core/state/snapshot/iterator_test.go | 53 +++++++++++++++++++++++++++-
 2 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/core/state/snapshot/iterator_fast.go b/core/state/snapshot/iterator_fast.go
index ef0212ac2..99734ec91 100644
--- a/core/state/snapshot/iterator_fast.go
+++ b/core/state/snapshot/iterator_fast.go
@@ -164,17 +164,35 @@ func (fi *fastAccountIterator) Next() bool {
 		fi.curAccount = fi.iterators[0].it.Account()
 		if innerErr := fi.iterators[0].it.Error(); innerErr != nil {
 			fi.fail = innerErr
+			return false
 		}
-		return fi.Error() == nil
+		if fi.curAccount != nil {
+			return true
+		}
+		// Implicit else: we've hit a nil-account, and need to fall through to the
+		// loop below to land on something non-nil
 	}
-	if !fi.next(0) {
-		return false
+	// If an account is deleted in one of the layers, the key will still be there,
+	// but the actual value will be nil. However, the iterator should not
+	// export nil-values (but instead simply omit the key), so we need to loop
+	// here until we either
+	//  - get a non-nil value,
+	//  - hit an error,
+	//  - or exhaust the iterator
+	for {
+		if !fi.next(0) {
+			return false // exhausted
+		}
+		fi.curAccount = fi.iterators[0].it.Account()
+		if innerErr := fi.iterators[0].it.Error(); innerErr != nil {
+			fi.fail = innerErr
+			return false // error
+		}
+		if fi.curAccount != nil {
+			break // non-nil value found
+		}
 	}
-	fi.curAccount = fi.iterators[0].it.Account()
-	if innerErr := fi.iterators[0].it.Error(); innerErr != nil {
-		fi.fail = innerErr
-	}
-	return fi.Error() == nil
+	return true
 }
 
 // next handles the next operation internally and should be invoked when we know
diff --git a/core/state/snapshot/iterator_test.go b/core/state/snapshot/iterator_test.go
index 832be10a4..935fafc2f 100644
--- a/core/state/snapshot/iterator_test.go
+++ b/core/state/snapshot/iterator_test.go
@@ -130,9 +130,13 @@ func verifyIterator(t *testing.T, expCount int, it AccountIterator) {
 		last  = common.Hash{}
 	)
 	for it.Next() {
-		if hash := it.Hash(); bytes.Compare(last[:], hash[:]) >= 0 {
+		hash := it.Hash()
+		if bytes.Compare(last[:], hash[:]) >= 0 {
 			t.Errorf("wrong order: %x >= %x", last, hash)
 		}
+		if it.Account() == nil {
+			t.Errorf("iterator returned nil-value for hash %x", hash)
+		}
 		count++
 	}
 	if count != expCount {
@@ -377,6 +381,53 @@ func TestAccountIteratorSeek(t *testing.T) {
 	verifyIterator(t, 0, it) // expected: nothing
 }
 
+// TestIteratorDeletions tests that the iterator behaves correct when there are
+// deleted accounts (where the Account() value is nil). The iterator
+// should not output any accounts or nil-values for those cases.
+func TestIteratorDeletions(t *testing.T) {
+	// Create an empty base layer and a snapshot tree out of it
+	base := &diskLayer{
+		diskdb: rawdb.NewMemoryDatabase(),
+		root:   common.HexToHash("0x01"),
+		cache:  fastcache.New(1024 * 500),
+	}
+	snaps := &Tree{
+		layers: map[common.Hash]snapshot{
+			base.root: base,
+		},
+	}
+	// Stack three diff layers on top with various overlaps
+	snaps.Update(common.HexToHash("0x02"), common.HexToHash("0x01"),
+		randomAccountSet("0x11", "0x22", "0x33"), nil)
+
+	set := randomAccountSet("0x11", "0x22", "0x33")
+	deleted := common.HexToHash("0x22")
+	set[deleted] = nil
+	snaps.Update(common.HexToHash("0x03"), common.HexToHash("0x02"), set, nil)
+
+	snaps.Update(common.HexToHash("0x04"), common.HexToHash("0x03"),
+		randomAccountSet("0x33", "0x44", "0x55"), nil)
+
+	// The output should be 11,33,44,55
+	it, _ := snaps.AccountIterator(common.HexToHash("0x04"), common.Hash{})
+	// Do a quick check
+	verifyIterator(t, 4, it)
+	it.Release()
+
+	// And a more detailed verification that we indeed do not see '0x22'
+	it, _ = snaps.AccountIterator(common.HexToHash("0x04"), common.Hash{})
+	defer it.Release()
+	for it.Next() {
+		hash := it.Hash()
+		if it.Account() == nil {
+			t.Errorf("iterator returned nil-value for hash %x", hash)
+		}
+		if hash == deleted {
+			t.Errorf("expected deleted elem %x to not be returned by iterator", deleted)
+		}
+	}
+}
+
 // BenchmarkAccountIteratorTraversal is a bit a bit notorious -- all layers contain the
 // exact same 200 accounts. That means that we need to process 2000 items, but
 // only spit out 200 values eventually.