Fix prefixed trie iterator

ChainSafe · Oct 24, 2024 · afd38f7 · afd38f7
1 parent 2577544
commit afd38f7
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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"),