From afd38f79ceb5950b393a7ba3fe60a58c22978fb7 Mon Sep 17 00:00:00 2001
From: Diego <diego2737@gmail.com>
Date: Thu, 24 Oct 2024 12:27:40 -0300
Subject: [PATCH] Fix prefixed trie iterator

---
 lib/runtime/storage/trie.go        |  2 +-
 pkg/trie/inmemory/in_memory.go     |  4 +++-
 pkg/trie/inmemory/iterator.go      | 15 ++++++++++++---
 pkg/trie/inmemory/iterator_test.go |  2 ++
 4 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/lib/runtime/storage/trie.go b/lib/runtime/storage/trie.go
index ac2767814a..f4a28e6a1e 100644
--- a/lib/runtime/storage/trie.go
+++ b/lib/runtime/storage/trie.go
@@ -190,7 +190,7 @@ func (t *TrieState) NextKey(key []byte) []byte {
 
 		nextKeyOnState := t.state.PrefixedIter(key).NextKeyFunc(func(nextKey []byte) bool {
 			_, deleted := currentTx.deletes[string(nextKey)]
-			return !deleted
+			return !deleted && !bytes.Equal(nextKey, key)
 		})
 		if nextKeyOnState == nil {
 			return nextKey
diff --git a/pkg/trie/inmemory/in_memory.go b/pkg/trie/inmemory/in_memory.go
index 5326c10704..5f0e3b48e7 100644
--- a/pkg/trie/inmemory/in_memory.go
+++ b/pkg/trie/inmemory/in_memory.go
@@ -55,7 +55,9 @@ func (t *InMemoryTrie) Iter() trie.TrieIterator {
 }
 
 func (t *InMemoryTrie) PrefixedIter(prefix []byte) trie.TrieIterator {
-	return NewInMemoryTrieIterator(WithTrie(t), WithCursorAt(codec.KeyLEToNibbles(prefix)))
+	iter := t.Iter()
+	iter.Seek(codec.KeyLEToNibbles(prefix))
+	return iter
 }
 
 func (t *InMemoryTrie) SetVersion(v trie.TrieLayout) {
diff --git a/pkg/trie/inmemory/iterator.go b/pkg/trie/inmemory/iterator.go
index 284ff2df40..931e1dbe49 100644
--- a/pkg/trie/inmemory/iterator.go
+++ b/pkg/trie/inmemory/iterator.go
@@ -74,9 +74,18 @@ func (t *InMemoryTrieIterator) NextKeyFunc(predicate func(nextKey []byte) bool)
 }
 
 func (t *InMemoryTrieIterator) Seek(targetKey []byte) {
-	t.NextKeyFunc(func(nextKey []byte) bool {
-		return bytes.Compare(nextKey, targetKey) >= 0
-	})
+	var prevEntry *trie.Entry
+	for entry := t.NextEntry(); entry != nil; entry = t.NextEntry() {
+		if bytes.Compare(entry.Key, targetKey) >= 0 {
+			break
+		}
+		prevEntry = entry
+	}
+	if prevEntry != nil {
+		t.cursorAtKey = prevEntry.Key
+		return
+	}
+	t.cursorAtKey = nil
 }
 
 // Entries returns all the key-value pairs in the trie as a map of keys to values
diff --git a/pkg/trie/inmemory/iterator_test.go b/pkg/trie/inmemory/iterator_test.go
index e5a5937a88..4e040a4a9f 100644
--- a/pkg/trie/inmemory/iterator_test.go
+++ b/pkg/trie/inmemory/iterator_test.go
@@ -36,6 +36,7 @@ func TestInMemoryIteratorGetAllKeysWithPrefix(t *testing.T) {
 
 	tt.Put([]byte("services_storage:serviceA:19090"), []byte("0x10"))
 	tt.Put([]byte("services_storage:serviceB:22222"), []byte("0x10"))
+	tt.Put([]byte("account_storage"), []byte("0x10"))
 	tt.Put([]byte("account_storage:ABC:AAA"), []byte("0x10"))
 	tt.Put([]byte("account_storage:ABC:CCC"), []byte("0x10"))
 	tt.Put([]byte("account_storage:ABC:DDD"), []byte("0x10"))
@@ -50,6 +51,7 @@ func TestInMemoryIteratorGetAllKeysWithPrefix(t *testing.T) {
 	}
 
 	expectedKeys := [][]byte{
+		[]byte("account_storage"),
 		[]byte("account_storage:ABC:AAA"),
 		[]byte("account_storage:ABC:CCC"),
 		[]byte("account_storage:ABC:DDD"),