From c0ca60b3f4c1d5266a116bc503dad629065589bf Mon Sep 17 00:00:00 2001
From: Roy Crihfield <roy@manteia.ltd>
Date: Thu, 3 Aug 2023 13:57:32 +0800
Subject: [PATCH] yield value nodes pre-order

---
 trie/iterator.go      | 42 ++++++++++++++++++++++++++++++++++--------
 trie/iterator_test.go | 14 +++++++-------
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/trie/iterator.go b/trie/iterator.go
index 214db1c292..5338aed8e7 100644
--- a/trie/iterator.go
+++ b/trie/iterator.go
@@ -449,7 +449,7 @@ func (it *nodeIterator) findChild(n *fullNode, index int, ancestor common.Hash)
 		state     *nodeIteratorState
 		childPath []byte
 	)
-	for ; index < len(n.Children); index++ {
+	for ; index < len(n.Children); index = nextChildIndex(index) {
 		if n.Children[index] != nil {
 			child = n.Children[index]
 			hash, _ := child.cache()
@@ -471,8 +471,8 @@ func (it *nodeIterator) nextChild(parent *nodeIteratorState, ancestor common.Has
 	switch node := parent.node.(type) {
 	case *fullNode:
 		// Full node, move to the first non-nil child.
-		if child, state, path, index := it.findChild(node, parent.index+1, ancestor); child != nil {
-			parent.index = index - 1
+		if child, state, path, index := it.findChild(node, nextChildIndex(parent.index), ancestor); child != nil {
+			parent.index = prevChildIndex(index)
 			return state, path, true
 		}
 	case *shortNode:
@@ -498,23 +498,23 @@ func (it *nodeIterator) nextChildAt(parent *nodeIteratorState, ancestor common.H
 	switch n := parent.node.(type) {
 	case *fullNode:
 		// Full node, move to the first non-nil child before the desired key position
-		child, state, path, index := it.findChild(n, parent.index+1, ancestor)
+		child, state, path, index := it.findChild(n, nextChildIndex(parent.index), ancestor)
 		if child == nil {
 			// No more children in this fullnode
 			return parent, it.path, false
 		}
 		// If the child we found is already past the seek position, just return it.
 		if bytes.Compare(path, key) >= 0 {
-			parent.index = index - 1
+			parent.index = prevChildIndex(index)
 			return state, path, true
 		}
 		// The child is before the seek position. Try advancing
 		for {
-			nextChild, nextState, nextPath, nextIndex := it.findChild(n, index+1, ancestor)
+			nextChild, nextState, nextPath, nextIndex := it.findChild(n, nextChildIndex(index), ancestor)
 			// If we run out of children, or skipped past the target, return the
 			// previous one
 			if nextChild == nil || bytes.Compare(nextPath, key) >= 0 {
-				parent.index = index - 1
+				parent.index = prevChildIndex(index)
 				return state, path, true
 			}
 			// We found a better child closer to the target
@@ -541,7 +541,7 @@ func (it *nodeIterator) push(state *nodeIteratorState, parentIndex *int, path []
 	it.path = path
 	it.stack = append(it.stack, state)
 	if parentIndex != nil {
-		*parentIndex++
+		*parentIndex = nextChildIndex(*parentIndex)
 	}
 }
 
@@ -563,6 +563,32 @@ func reachedPath(path, target []byte) bool {
 	return bytes.Compare(path, target) >= 0
 }
 
+func prevChildIndex(index int) int {
+	switch index {
+	case 0:
+		return 16
+	case 16:
+		return -1
+	case 17:
+		return 15
+	default:
+		return index - 1
+	}
+}
+
+func nextChildIndex(index int) int {
+	switch index {
+	case -1:
+		return 16
+	case 15:
+		return 17
+	case 16:
+		return 0
+	default:
+		return index + 1
+	}
+}
+
 func compareNodes(a, b NodeIterator) int {
 	if cmp := bytes.Compare(a.Path(), b.Path()); cmp != 0 {
 		return cmp
diff --git a/trie/iterator_test.go b/trie/iterator_test.go
index 5540c907cd..47dc18c9db 100644
--- a/trie/iterator_test.go
+++ b/trie/iterator_test.go
@@ -182,14 +182,14 @@ func testNodeIteratorCoverage(t *testing.T, scheme string) {
 type kvs struct{ k, v string }
 
 var testdata1 = []kvs{
+	{"bar", "b"},
 	{"barb", "ba"},
 	{"bard", "bc"},
 	{"bars", "bb"},
-	{"bar", "b"},
 	{"fab", "z"},
+	{"foo", "a"},
 	{"food", "ab"},
 	{"foos", "aa"},
-	{"foo", "a"},
 }
 
 var testdata2 = []kvs{
@@ -218,7 +218,7 @@ func TestIteratorSeek(t *testing.T) {
 
 	// Seek to a non-existent key.
 	it = NewIterator(trie.MustNodeIterator([]byte("barc")))
-	if err := checkIteratorOrder(testdata1[1:], it); err != nil {
+	if err := checkIteratorOrder(testdata1[2:], it); err != nil {
 		t.Fatal(err)
 	}
 
@@ -230,7 +230,7 @@ func TestIteratorSeek(t *testing.T) {
 
 	// Seek to a key for which a prefixing key exists.
 	it = NewIterator(trie.MustNodeIterator([]byte("food")))
-	if err := checkIteratorOrder(testdata1[5:], it); err != nil {
+	if err := checkIteratorOrder(testdata1[6:], it); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -317,16 +317,16 @@ func TestUnionIterator(t *testing.T) {
 
 	all := []struct{ k, v string }{
 		{"aardvark", "c"},
+		{"bar", "b"},
 		{"barb", "ba"},
 		{"barb", "bd"},
 		{"bard", "bc"},
 		{"bars", "bb"},
 		{"bars", "be"},
-		{"bar", "b"},
 		{"fab", "z"},
+		{"foo", "a"},
 		{"food", "ab"},
 		{"foos", "aa"},
-		{"foo", "a"},
 		{"jars", "d"},
 	}
 
@@ -517,7 +517,7 @@ func testIteratorContinueAfterSeekError(t *testing.T, memonly bool, scheme strin
 		rawdb.WriteTrieNode(diskdb, common.Hash{}, barNodePath, barNodeHash, barNodeBlob, triedb.Scheme())
 	}
 	// Check that iteration produces the right set of values.
-	if err := checkIteratorOrder(testdata1[2:], NewIterator(it)); err != nil {
+	if err := checkIteratorOrder(testdata1[3:], NewIterator(it)); err != nil {
 		t.Fatal(err)
 	}
 }