From b40114b7396a5b924d96d9153154b2edc5421855 Mon Sep 17 00:00:00 2001
From: prathamesh0 <42446521+prathamesh0@users.noreply.github.com>
Date: Wed, 3 Aug 2022 16:12:45 +0530
Subject: [PATCH] Store start path in prefix bound iterator object (#17)

---
 iterator.go      | 11 ++++++-----
 iterator_test.go |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/iterator.go b/iterator.go
index ccd160a..b5b752a 100644
--- a/iterator.go
+++ b/iterator.go
@@ -27,7 +27,8 @@ import (
 // PrefixBoundIterator is a NodeIterator constrained by a lower & upper bound (as hex path prefixes)
 type PrefixBoundIterator struct {
 	trie.NodeIterator
-	EndPath []byte
+	StartPath []byte
+	EndPath   []byte
 }
 
 func (it *PrefixBoundIterator) Next(descend bool) bool {
@@ -43,8 +44,8 @@ func (it *PrefixBoundIterator) Next(descend bool) bool {
 }
 
 // Iterator with an upper bound value (hex path prefix)
-func NewPrefixBoundIterator(it trie.NodeIterator, to []byte) *PrefixBoundIterator {
-	return &PrefixBoundIterator{NodeIterator: it, EndPath: to}
+func NewPrefixBoundIterator(it trie.NodeIterator, from []byte, to []byte) *PrefixBoundIterator {
+	return &PrefixBoundIterator{NodeIterator: it, StartPath: from, EndPath: to}
 }
 
 // generates nibble slice prefixes at uniform intervals
@@ -133,7 +134,7 @@ func SubtrieIterators(tree state.Trie, nbins uint) []trie.NodeIterator {
 	var iters []trie.NodeIterator
 	eachPrefixRange(nil, nbins, func(from []byte, to []byte) {
 		it := tree.NodeIterator(HexToKeyBytes(from))
-		iters = append(iters, NewPrefixBoundIterator(it, to))
+		iters = append(iters, NewPrefixBoundIterator(it, from, to))
 	})
 	return iters
 }
@@ -148,7 +149,7 @@ func (fac *SubtrieIteratorFactory) Length() int { return len(fac.startPaths) }
 
 func (fac *SubtrieIteratorFactory) IteratorAt(bin uint) *PrefixBoundIterator {
 	it := fac.tree.NodeIterator(HexToKeyBytes(fac.startPaths[bin]))
-	return NewPrefixBoundIterator(it, fac.endPaths[bin])
+	return NewPrefixBoundIterator(it, fac.startPaths[bin], fac.endPaths[bin])
 }
 
 // Cut a trie by path prefix, returning `nbins` iterators covering its subtries
diff --git a/iterator_test.go b/iterator_test.go
index ba6e34b..50660ec 100644
--- a/iterator_test.go
+++ b/iterator_test.go
@@ -59,7 +59,7 @@ func TestIterator(t *testing.T) {
 		}
 
 		runCase := func(t *testing.T, tc testCase) {
-			it := iter.NewPrefixBoundIterator(tree.NodeIterator(iter.HexToKeyBytes(tc.lower)), tc.upper)
+			it := iter.NewPrefixBoundIterator(tree.NodeIterator(iter.HexToKeyBytes(tc.lower)), tc.lower, tc.upper)
 			for it.Next(true) {
 				if bytes.Compare(it.Path(), tc.lower) < 0 {
 					t.Fatalf("iterator outside lower bound: %v", it.Path())