From eab44c99652d65c68e2cd85fbc1909fc11500fd4 Mon Sep 17 00:00:00 2001
From: Roy Crihfield <roy@manteia.ltd>
Date: Sun, 24 Sep 2023 18:47:14 +0800
Subject: [PATCH] expose bounds of iterator

---
 iterator.go        | 22 +++++++++++++---------
 tracker/tracker.go | 13 ++++++++-----
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/iterator.go b/iterator.go
index 31cde21..fb093fd 100644
--- a/iterator.go
+++ b/iterator.go
@@ -23,16 +23,21 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-// PrefixBoundIterator is a NodeIterator constrained by a lower & upper bound (as hex path prefixes)
-type PrefixBoundIterator struct {
-	trie.NodeIterator
-	EndPath []byte
-}
-
 // IteratorConstructor is a constructor returning a NodeIterator, which is used to decouple this
 // code from the trie implementation.
 type IteratorConstructor = func(startKey []byte) trie.NodeIterator
 
+// PrefixBoundIterator is a NodeIterator constrained by a lower & upper bound (as hex path prefixes)
+type PrefixBoundIterator struct {
+	trie.NodeIterator
+	StartPath, EndPath []byte
+}
+
+// NewPrefixBoundIterator returns an iterator with an upper bound value (hex path prefix)
+func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) *PrefixBoundIterator {
+	return &PrefixBoundIterator{NodeIterator: it, StartPath: it.Path(), EndPath: to}
+}
+
 func (it *PrefixBoundIterator) Next(descend bool) bool {
 	if it.EndPath == nil {
 		return it.NodeIterator.Next(descend)
@@ -49,9 +54,8 @@ func (it *PrefixBoundIterator) Next(descend bool) bool {
 	return bytes.Compare(it.Path(), it.EndPath) <= 0
 }
 
-// NewPrefixBoundIterator returns an iterator with an upper bound value (hex path prefix)
-func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) *PrefixBoundIterator {
-	return &PrefixBoundIterator{NodeIterator: it, EndPath: to}
+func (it *PrefixBoundIterator) Bounds() ([]byte, []byte) {
+	return it.StartPath, it.EndPath
 }
 
 // generates nibble slice prefixes at uniform intervals
diff --git a/tracker/tracker.go b/tracker/tracker.go
index b0996da..cb7d4f6 100644
--- a/tracker/tracker.go
+++ b/tracker/tracker.go
@@ -69,11 +69,7 @@ func (tr *Tracker) dump() error {
 	log.Debug("Dumping recovery state", "to", tr.recoveryFile)
 	var rows [][]string
 	for it := range tr.started {
-		var endPath []byte
-		if impl, ok := it.NodeIterator.(*iter.PrefixBoundIterator); ok {
-			endPath = impl.EndPath
-		}
-
+		_, endPath := it.Bounds()
 		rows = append(rows, []string{
 			fmt.Sprintf("%x", it.Path()),
 			fmt.Sprintf("%x", endPath),
@@ -184,6 +180,13 @@ func (it *Iterator) Next(descend bool) bool {
 	return ret
 }
 
+func (it *Iterator) Bounds() ([]byte, []byte) {
+	if impl, ok := it.NodeIterator.(*iter.PrefixBoundIterator); ok {
+		return impl.Bounds()
+	}
+	return nil, nil
+}
+
 // Rewinds to the path of the previous (pre-order) node:
 // If the last byte of the path is zero, pops it. Otherwise, decrements it
 // and pads with 0xF to 64 bytes (e.g. [1] => [0 f f f ...]).