feat(ffi-iterator): Implementation of Iterator in Go (2/4)

ffi/firewood.go

@@ -214,6 +214,17 @@ func (db *Database) Root() ([]byte, error) {
 	return bytes, err
 }
 
+func (db *Database) LatestRevision() (*Revision, error) {
+	root, err := db.Root()
+	if err != nil {
+		return nil, err
+	}
+	if bytes.Equal(root, EmptyRoot) {
+		return nil, errRevisionNotFound
+	}
+	return db.Revision(root)
+}
+
 // Revision returns a historical revision of the database.
 func (db *Database) Revision(root []byte) (*Revision, error) {
 	if root == nil || len(root) != RootLength {

ffi/firewood_test.go

@@ -1136,3 +1136,144 @@ func TestGetFromRootParallel(t *testing.T) {
 		r.NoError(err, "Parallel operation failed")
 	}
 }
+
+func assertIteratorYields(r *require.Assertions, it *Iterator, keys [][]byte, vals [][]byte) {
+	i := 0
+	for ; it.Next(); i += 1 {
+		r.Equal(keys[i], it.Key())
+		r.Equal(vals[i], it.Value())
+	}
+	r.NoError(it.Err())
+	r.Equal(len(keys), i)
+}
+
+// Tests that basic iterator functionality works
+func TestIter(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForTest(10)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
+
+	rev, err := db.LatestRevision()
+	r.NoError(err)
+	it, err := rev.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(it.Drop())
+		r.NoError(rev.Drop())
+	})
+
+	assertIteratorYields(r, it, keys, vals)
+}
+
+// Tests that iterators on different roots work fine
+func TestIterOnRoot(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	// Commit 10 key-value pairs.
+	keys, vals := kvForTest(20)
+	keys = keys[:10]
+	vals1, vals2 := vals[:10], vals[10:]
+
+	firstRoot, err := db.Update(keys, vals1)
+	r.NoError(err)
+
+	// we use the same keys, but update the values
+	secondRoot, err := db.Update(keys, vals2)
+	r.NoError(err)
+
+	r1, err := db.Revision(firstRoot)
+	r.NoError(err)
+	h1, err := r1.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(h1.Drop())
+		r.NoError(r1.Drop())
+	})
+
+	r2, err := db.Revision(secondRoot)
+	r.NoError(err)
+	h2, err := r2.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(h2.Drop())
+		r.NoError(r2.Drop())
+	})
+
+	assertIteratorYields(r, h1, keys, vals1)
+	assertIteratorYields(r, h2, keys, vals2)
+}
+
+// Tests that basic iterator functionality works for proposal
+func TestIterOnProposal(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForTest(10)
+	p, err := db.Propose(keys, vals)
+	r.NoError(err)
+
+	it, err := p.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(it.Drop())
+	})
+
+	assertIteratorYields(r, it, keys, vals)
+}
+
+// Tests that the iterator still works after proposal is committed
+func TestIterAfterProposalCommit(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForTest(10)
+	p, err := db.Propose(keys, vals)
+	r.NoError(err)
+
+	it, err := p.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(it.Drop())
+	})
+
+	err = p.Commit()
+	r.NoError(err)
+
+	// iterate after commit
+	// because iterator hangs on the nodestore reference of proposal
+	// the nodestore won't be dropped until we drop the iterator
+	assertIteratorYields(r, it, keys, vals)
+}
+
+// Tests that the iterator on latest revision works properly after a proposal commit
+func TestIterUpdate(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForTest(10)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
+
+	// get an iterator on latest revision
+	rev, err := db.LatestRevision()
+	r.NoError(err)
+	it, err := rev.Iter(nil)
+	r.NoError(err)
+	t.Cleanup(func() {
+		r.NoError(it.Drop())
+		r.NoError(rev.Drop())
+	})
+
+	// update the database
+	keys2, vals2 := kvForTest(10)
+	_, err = db.Update(keys2, vals2)
+	r.NoError(err)
+
+	// iterate after commit
+	// because iterator is fixed on the revision hash, it should return the initial values
+	assertIteratorYields(r, it, keys, vals)
+}

ffi/iterator.go

@@ -0,0 +1,92 @@
+// Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE.md for licensing terms.
+
+package ffi
+
+// #include <stdlib.h>
+// #include "firewood.h"
+import "C"
+
+import (
+	"fmt"
+	"unsafe"
+)
+
+type Iterator struct {
+	// handle is an opaque pointer to the iterator within Firewood. It should be
+	// passed to the C FFI functions that operate on iterators
+	//
+	// It is not safe to call these methods with a nil handle.
+	handle *C.IteratorHandle
+
+	// currentKey is the current key retrieved from the iterator
+	currentKey []byte
+	// currentVal is the current value retrieved from the iterator
+	currentVal []byte
+	// err is the error from the iterator, if any
+	err error
+}
+
+// Next proceeds to the next item on the iterator, and returns true
+// if succeeded and there is a pair available.
+// The new pair could be retrieved with Key and Value methods.
+func (it *Iterator) Next() bool {
+	kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
+	it.err = e
+	if kv == nil || e != nil {
+		return false
+	}
+	k, v, e := kv.Consume()
+	it.currentKey = k
+	it.currentVal = v
+	it.err = e
+	return e == nil
+}
+
+// Key returns the key of the current pair
+func (it *Iterator) Key() []byte {
+	if (it.currentKey == nil && it.currentVal == nil) || it.err != nil {
+		return nil
+	}
+	return it.currentKey
+}
+
+// Value returns the value of the current pair
+func (it *Iterator) Value() []byte {
+	if (it.currentKey == nil && it.currentVal == nil) || it.err != nil {
+		return nil
+	}
+	return it.currentVal
+}
+
+// Err returns the error if Next failed
+func (it *Iterator) Err() error {
+	return it.err
+}
+
+// Drop drops the iterator and releases the resources
+func (it *Iterator) Drop() error {
+	if it.handle != nil {
+		return getErrorFromVoidResult(C.fwd_free_iterator(it.handle))
+	}
+	return nil
+}
+
+// getIteratorFromIteratorResult converts a C.IteratorResult to an Iterator or error.
+func getIteratorFromIteratorResult(result C.IteratorResult) (*Iterator, error) {
+	switch result.tag {
+	case C.IteratorResult_NullHandlePointer:
+		return nil, errDBClosed
+	case C.IteratorResult_Ok:
+		body := (*C.IteratorResult_Ok_Body)(unsafe.Pointer(&result.anon0))
+		proposal := &Iterator{
+			handle: body.handle,
+		}
+		return proposal, nil
+	case C.IteratorResult_Err:
+		err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError()
+		return nil, err
+	default:
+		return nil, fmt.Errorf("unknown C.IteratorResult tag: %d", result.tag)
+	}
+}

ffi/memory.go

@@ -309,6 +309,56 @@ func getValueFromValueResult(result C.ValueResult) ([]byte, error) {
 	}
 }
 
+type ownedKeyValue struct {
+	key   *ownedBytes
+	value *ownedBytes
+}
+
+func (kv *ownedKeyValue) Consume() ([]byte, []byte, error) {
+	key := kv.key.CopiedBytes()
+	if err := kv.key.Free(); err != nil {
+		return nil, nil, fmt.Errorf("%w: %w", errFreeingValue, err)
+	}
+	value := kv.value.CopiedBytes()
+	if err := kv.value.Free(); err != nil {
+		return nil, nil, fmt.Errorf("%w: %w", errFreeingValue, err)
+	}
+	return key, value, nil
+}
+
+// newOwnedKeyValue creates a ownedKeyValue from a C.OwnedKeyValuePair.
+//
+// The caller is responsible for calling Free() on the returned ownedKeyValue
+// when it is no longer needed otherwise memory will leak.
+func newOwnedKeyValue(owned C.OwnedKeyValuePair) *ownedKeyValue {
+	return &ownedKeyValue{
+		key:   newOwnedBytes(owned.key),
+		value: newOwnedBytes(owned.value),
+	}
+}
+
+// getKeyValueFromKeyValueResult converts a C.KeyValueResult to a key value pair or error.
+//
+// It returns nil, nil if the result is None.
+// It returns a *ownedKeyValue, nil if the result is Some.
+// It returns an error if the result is an error.
+func getKeyValueFromKeyValueResult(result C.KeyValueResult) (*ownedKeyValue, error) {
+	switch result.tag {
+	case C.KeyValueResult_NullHandlePointer:
+		return nil, errDBClosed
+	case C.KeyValueResult_None:
+		return nil, nil
+	case C.KeyValueResult_Some:
+		ownedKvp := newOwnedKeyValue(*(*C.OwnedKeyValuePair)(unsafe.Pointer(&result.anon0)))
+		return ownedKvp, nil
+	case C.ValueResult_Err:
+		err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError()
+		return nil, err
+	default:
+		return nil, fmt.Errorf("unknown C.KeyValueResult tag: %d", result.tag)
+	}
+}
+
 // getDatabaseFromHandleResult converts a C.HandleResult to a Database or error.
 //
 // If the C.HandleResult is an error, it returns an error instead of a Database.

ffi/proposal.go

@@ -58,6 +58,21 @@ func (p *Proposal) Get(key []byte) ([]byte, error) {
 	return getValueFromValueResult(C.fwd_get_from_proposal(p.handle, newBorrowedBytes(key, &pinner)))
 }
 
+// Iter creates and iterator starting from the provided key on proposal.
+// pass empty slice to start from beginning
+func (p *Proposal) Iter(key []byte) (*Iterator, error) {
+	if p.handle == nil {
+		return nil, errDBClosed
+	}
+
+	var pinner runtime.Pinner
+	defer pinner.Unpin()
+
+	itResult := C.fwd_iter_on_proposal(p.handle, newBorrowedBytes(key, &pinner))
+
+	return getIteratorFromIteratorResult(itResult)
+}
+
 // Propose creates a new proposal with the given keys and values.
 // The proposal is not committed until Commit is called.
 func (p *Proposal) Propose(keys, vals [][]byte) (*Proposal, error) {

ffi/revision.go

@@ -53,6 +53,21 @@ func (r *Revision) Get(key []byte) ([]byte, error) {
 	))
 }
 
+// Iter creates an iterator starting from the provided key on revision.
+// pass empty slice to start from beginning
+func (r *Revision) Iter(key []byte) (*Iterator, error) {
+	if r.handle == nil {
+		return nil, errDroppedRevision
+	}
+
+	var pinner runtime.Pinner
+	defer pinner.Unpin()
+
+	itResult := C.fwd_iter_on_revision(r.handle, newBorrowedBytes(key, &pinner))
+
+	return getIteratorFromIteratorResult(itResult)
+}
+
 // Drop releases the resources backed by the revision handle.
 //
 // It is safe to call Drop multiple times; subsequent calls after the first are no-ops.