From 02ba3dd9b3462061020a6e9a8c5263a4d201e486 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Wed, 25 Mar 2020 09:59:22 +0100
Subject: [PATCH] core/state/snapshot: start fixing disk iterator seek

---
 common/bytes.go                       | 11 ++++++
 common/bytes_test.go                  | 19 +++++++++
 core/state/snapshot/disklayer_test.go | 56 +++++++++++++++++++++++++++
 core/state/snapshot/iterator.go       |  4 +-
 4 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/common/bytes.go b/common/bytes.go
index fa457b92cf..634041804d 100644
--- a/common/bytes.go
+++ b/common/bytes.go
@@ -145,3 +145,14 @@ func TrimLeftZeroes(s []byte) []byte {
 	}
 	return s[idx:]
 }
+
+// TrimRightZeroes returns a subslice of s without trailing zeroes
+func TrimRightZeroes(s []byte) []byte {
+	idx := len(s)
+	for ; idx > 0; idx-- {
+		if s[idx-1] != 0 {
+			break
+		}
+	}
+	return s[:idx]
+}
diff --git a/common/bytes_test.go b/common/bytes_test.go
index 7cf6553377..0e3ec974ee 100644
--- a/common/bytes_test.go
+++ b/common/bytes_test.go
@@ -105,3 +105,22 @@ func TestNoPrefixShortHexOddLength(t *testing.T) {
 		t.Errorf("Expected %x got %x", expected, result)
 	}
 }
+
+func TestTrimRightZeroes(t *testing.T) {
+	tests := []struct {
+		arr []byte
+		exp []byte
+	}{
+		{FromHex("0x00ffff00ff0000"), FromHex("0x00ffff00ff")},
+		{FromHex("0x00000000000000"), []byte{}},
+		{FromHex("0xff"), FromHex("0xff")},
+		{[]byte{}, []byte{}},
+		{FromHex("0x00ffffffffffff"), FromHex("0x00ffffffffffff")},
+	}
+	for i, test := range tests {
+		got := TrimRightZeroes(test.arr)
+		if !bytes.Equal(got, test.exp) {
+			t.Errorf("test %d, got %x exp %x", i, got, test.exp)
+		}
+	}
+}
diff --git a/core/state/snapshot/disklayer_test.go b/core/state/snapshot/disklayer_test.go
index aae2aa6b56..26fe8a5339 100644
--- a/core/state/snapshot/disklayer_test.go
+++ b/core/state/snapshot/disklayer_test.go
@@ -432,4 +432,60 @@ func TestDiskPartialMerge(t *testing.T) {
 // This test case is a tiny specialized case of TestDiskPartialMerge, which tests
 // some very specific cornercases that random tests won't ever trigger.
 func TestDiskMidAccountPartialMerge(t *testing.T) {
+	// TODO(@karalabe) ?
+}
+
+// TestDiskSeek tests that seek-operations work on the disk layer
+func TestDiskSeek(t *testing.T) {
+	// Create some accounts in the disk layer
+	db := memorydb.New()
+	// Fill even keys [0,2,4...]
+	for i := 0; i < 0xff; i += 2 {
+		acc := common.Hash{byte(i)}
+		rawdb.WriteAccountSnapshot(db, acc, acc[:])
+	}
+	baseRoot := randomHash()
+	rawdb.WriteSnapshotRoot(db, baseRoot)
+
+	snaps := &Tree{
+		layers: map[common.Hash]snapshot{
+			baseRoot: &diskLayer{
+				diskdb: db,
+				cache:  fastcache.New(500 * 1024),
+				root:   baseRoot,
+			},
+		},
+	}
+	// Test some different seek positions
+	type testcase struct {
+		pos    byte
+		expkey byte
+	}
+	var cases = []testcase{
+		{0xff, 0x55}, // this should exit immediately without checking key
+		{0x01, 0x02},
+		{0xfe, 0xfe},
+		{0xfd, 0xfe},
+		{0x00, 0x00},
+	}
+	for i, tc := range cases {
+		it, err := snaps.AccountIterator(baseRoot, common.Hash{tc.pos})
+		if err != nil {
+			t.Fatalf("case %d, error: %v", i, err)
+		}
+		count := 0
+		for it.Next() {
+			count++
+			k, v, err := it.Hash()[0], it.Account()[0], it.Error()
+			if err != nil {
+				t.Fatalf("test %d, error: %v", i, err)
+			}
+			if k != tc.expkey {
+				t.Fatalf("test %d, got %v exp %v", i, k, tc.expkey)
+			}
+			if v != k {
+				t.Fatalf("test %d, value wrong, got %v exp %v", i, v, k)
+			}
+		}
+	}
 }
diff --git a/core/state/snapshot/iterator.go b/core/state/snapshot/iterator.go
index e062298fa4..e6d08c8106 100644
--- a/core/state/snapshot/iterator.go
+++ b/core/state/snapshot/iterator.go
@@ -148,10 +148,10 @@ type diskAccountIterator struct {
 
 // AccountIterator creates an account iterator over a disk layer.
 func (dl *diskLayer) AccountIterator(seek common.Hash) AccountIterator {
-	// TODO: Fix seek position, or remove seek parameter
+	pos := append(rawdb.SnapshotAccountPrefix, common.TrimRightZeroes(seek[:])...)
 	return &diskAccountIterator{
 		layer: dl,
-		it:    dl.diskdb.NewIteratorWithPrefix(rawdb.SnapshotAccountPrefix),
+		it:    dl.diskdb.NewIteratorWithPrefix(pos),
 	}
 }