From 9bab3d446e140620b801ad8dff258a6ae5645a0e Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 27 Aug 2025 03:37:46 -0400
Subject: [PATCH 01/15] wip: iterator batching / rust side done

---
 ffi/firewood.h           | 80 ++++++++++++++++++++++++++++++++++++++++
 ffi/src/iterator.rs      | 25 +++++++++----
 ffi/src/lib.rs           | 27 ++++++++++++++
 ffi/src/value.rs         |  2 +-
 ffi/src/value/results.rs | 50 ++++++++++++++++++++++++-
 5 files changed, 174 insertions(+), 10 deletions(-)

diff --git a/ffi/firewood.h b/ffi/firewood.h
index 9c7c71f9b..35b95d0b3 100644
--- a/ffi/firewood.h
+++ b/ffi/firewood.h
@@ -329,6 +329,60 @@ typedef struct KeyValueResult {
   };
 } KeyValueResult;
 
+/**
+ * A Rust-owned vector of bytes that can be passed to C code.
+ *
+ * C callers must free this memory using the respective FFI function for the
+ * concrete type (but not using the `free` function from the C standard library).
+ */
+typedef struct OwnedSlice_OwnedKeyValuePair {
+  struct OwnedKeyValuePair *ptr;
+  size_t len;
+} OwnedSlice_OwnedKeyValuePair;
+
+/**
+ * A result type returned from iterator FFI functions
+ */
+typedef enum KeyValueBatchResult_Tag {
+  /**
+   * The caller provided a null pointer to an iterator handle.
+   */
+  KeyValueBatchResult_NullHandlePointer,
+  /**
+   * The provided root was not found in the database.
+   */
+  KeyValueBatchResult_RevisionNotFound,
+  /**
+   * The next batch of items on iterator are returned.
+   */
+  KeyValueBatchResult_Some,
+  /**
+   * An error occurred and the message is returned as an [`OwnedBytes`]. If
+   * value is guaranteed to contain only valid UTF-8.
+   *
+   * The caller must call [`fwd_free_owned_bytes`] to free the memory
+   * associated with this error.
+   *
+   * [`fwd_free_owned_bytes`]: crate::fwd_free_owned_bytes
+   */
+  KeyValueBatchResult_Err,
+} KeyValueBatchResult_Tag;
+
+typedef struct KeyValueBatchResult {
+  KeyValueBatchResult_Tag tag;
+  union {
+    struct {
+      struct HashKey revision_not_found;
+    };
+    struct {
+      struct OwnedSlice_OwnedKeyValuePair some;
+    };
+    struct {
+      OwnedBytes err;
+    };
+  };
+} KeyValueBatchResult;
+
 /**
  * A result type returned from FFI functions that create an iterator
  */
@@ -815,6 +869,32 @@ struct ValueResult fwd_get_latest(const struct DatabaseHandle *db, BorrowedBytes
  */
 struct KeyValueResult fwd_iter_next(struct IteratorHandle *handle);
 
+/**
+ * Retrieves the next batch of items from the iterator
+ *
+ * # Arguments
+ *
+ * * `handle` - The iterator handle returned by [`fwd_iter_on_root`] or
+ *   [`fwd_iter_on_proposal`].
+ *
+ * # Returns
+ *
+ * - [`KeyValueResult::NullHandlePointer`] if the provided iterator handle is null.
+ * - [`KeyValueResult::None`] if the iterator doesn't have any remaining values/exhausted.
+ * - [`KeyValueResult::Some`] if the next item on iterator was retrieved, with the associated
+ *   key value pair.
+ * - [`KeyValueResult::Err`] if an error occurred while retrieving the next item on iterator.
+ *
+ * # Safety
+ *
+ * The caller must:
+ * * ensure that `handle` is a valid pointer to a [`IteratorHandle`].
+ * * call [`fwd_free_owned_bytes`] on [`OwnedKeyValuePair::key`] and [`OwnedKeyValuePair::value`]
+ *   to free the memory associated with the returned error or value.
+ *
+ */
+struct KeyValueBatchResult fwd_iter_next_n(struct IteratorHandle *handle, size_t n);
+
 /**
  * Return an iterator on proposal optionally starting from a key
  *
diff --git a/ffi/src/iterator.rs b/ffi/src/iterator.rs
index ff5d1c64d..0ea9dc105 100644
--- a/ffi/src/iterator.rs
+++ b/ffi/src/iterator.rs
@@ -2,18 +2,14 @@
 // See the file LICENSE.md for licensing terms.
 
 use firewood::merkle;
-use firewood::v2::api::{self};
 use std::fmt::{Debug, Formatter};
+use firewood::v2::api;
 
-// This wrapper doesn't do anything special now, but we need to figure out how to handle
-// proposal commits, drops, revision cleaning
-// TODO(amin): figure this out
-
-type KeyValueItem = Result<(merkle::Key, merkle::Value), api::Error>;
+type KeyValueItem = (merkle::Key, merkle::Value);
 
 /// An opaque wrapper around an Iterator.
 pub struct IteratorHandle<'db> {
-    pub iterator: Box<dyn Iterator<Item = KeyValueItem> + 'db>,
+    pub iterator: Box<dyn Iterator<Item = Result<KeyValueItem, api::Error>> + 'db>,
 }
 
 impl Debug for IteratorHandle<'_> {
@@ -23,9 +19,22 @@ impl Debug for IteratorHandle<'_> {
 }
 
 impl IteratorHandle<'_> {
-    pub fn iter_next(&mut self) -> Option<KeyValueItem> {
+    pub fn iter_next(&mut self) -> Option<Result<KeyValueItem, api::Error>> {
         self.iterator.next()
     }
+
+    pub fn iter_next_n(&mut self, n: usize) -> Result<Vec<KeyValueItem>, api::Error> {
+        let mut items = Vec::<KeyValueItem>::new();
+        for _ in 0..n {
+            let item = self.iter_next();
+            if let Some(item) = item {
+                items.push(item?);
+            } else {
+                break;
+            }
+        }
+        Ok(items)
+    }
 }
 
 #[derive(Debug)]
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 54bea9cb7..9e65a73a1 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -198,6 +198,33 @@ pub unsafe extern "C" fn fwd_iter_next(handle: Option<&mut IteratorHandle<'_>>)
     invoke_with_handle(handle, IteratorHandle::iter_next)
 }
 
+/// Retrieves the next batch of items from the iterator
+///
+/// # Arguments
+///
+/// * `handle` - The iterator handle returned by [`fwd_iter_on_root`] or
+///   [`fwd_iter_on_proposal`].
+///
+/// # Returns
+///
+/// - [`KeyValueResult::NullHandlePointer`] if the provided iterator handle is null.
+/// - [`KeyValueResult::None`] if the iterator doesn't have any remaining values/exhausted.
+/// - [`KeyValueResult::Some`] if the next item on iterator was retrieved, with the associated
+///   key value pair.
+/// - [`KeyValueResult::Err`] if an error occurred while retrieving the next item on iterator.
+///
+/// # Safety
+///
+/// The caller must:
+/// * ensure that `handle` is a valid pointer to a [`IteratorHandle`].
+/// * call [`fwd_free_owned_bytes`] on [`OwnedKeyValuePair::key`] and [`OwnedKeyValuePair::value`]
+///   to free the memory associated with the returned error or value.
+///
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn fwd_iter_next_n(handle: Option<&mut IteratorHandle<'_>>, n: usize) -> KeyValueBatchResult {
+    invoke_with_handle(handle, |it| it.iter_next_n(n))
+}
+
 /// Consumes the [`IteratorHandle`], destroys the iterator, and frees the memory.
 ///
 /// # Arguments
diff --git a/ffi/src/value.rs b/ffi/src/value.rs
index e32eadf09..1a56f4da3 100644
--- a/ffi/src/value.rs
+++ b/ffi/src/value.rs
@@ -15,6 +15,6 @@ pub use self::kvp::{KeyValuePair, OwnedKeyValuePair};
 pub use self::owned::{OwnedBytes, OwnedSlice};
 pub(crate) use self::results::{CResult, NullHandleResult};
 pub use self::results::{
-    HandleResult, HashResult, IteratorResult, KeyValueResult, ProposalResult, ValueResult,
+    HandleResult, HashResult, IteratorResult, KeyValueResult, KeyValueBatchResult, ProposalResult, ValueResult,
     VoidResult,
 };
diff --git a/ffi/src/value/results.rs b/ffi/src/value/results.rs
index ccd99d377..3dcb321d2 100644
--- a/ffi/src/value/results.rs
+++ b/ffi/src/value/results.rs
@@ -7,7 +7,7 @@ use std::fmt;
 
 use crate::iterator::{CreateIteratorResult, IteratorHandle};
 use crate::value::kvp::OwnedKeyValuePair;
-use crate::{CreateProposalResult, HashKey, OwnedBytes, ProposalHandle};
+use crate::{CreateProposalResult, HashKey, OwnedBytes, OwnedSlice, ProposalHandle};
 
 /// The result type returned from an FFI function that returns no value but may
 /// return an error.
@@ -261,6 +261,42 @@ impl From<Option<Result<(merkle::Key, merkle::Value), api::Error>>> for KeyValue
     }
 }
 
+
+/// A result type returned from iterator FFI functions
+#[derive(Debug)]
+#[repr(C)]
+pub enum KeyValueBatchResult {
+    /// The caller provided a null pointer to an iterator handle.
+    NullHandlePointer,
+    /// The provided root was not found in the database.
+    RevisionNotFound(HashKey),
+    /// The next batch of items on iterator are returned.
+    Some(OwnedSlice<OwnedKeyValuePair>),
+    /// An error occurred and the message is returned as an [`OwnedBytes`]. If
+    /// value is guaranteed to contain only valid UTF-8.
+    ///
+    /// The caller must call [`fwd_free_owned_bytes`] to free the memory
+    /// associated with this error.
+    ///
+    /// [`fwd_free_owned_bytes`]: crate::fwd_free_owned_bytes
+    Err(OwnedBytes),
+}
+
+impl From<Result<Vec<(merkle::Key, merkle::Value)>, api::Error>> for KeyValueBatchResult {
+    fn from(value: Result<Vec<(merkle::Key, merkle::Value)>, api::Error>) -> Self {
+        match value {
+                Ok(pairs) => {
+                    let values: Vec<_> = pairs.into_iter().map(|(k, v)| OwnedKeyValuePair {key:k.into(), value:v.into()}).collect();
+                    KeyValueBatchResult::Some(values.into())
+                },
+                Err(api::Error::RevisionNotFound { provided }) => KeyValueBatchResult::RevisionNotFound(
+                    HashKey::from(provided.unwrap_or_else(api::HashKey::empty)),
+                ),
+                Err(err) => KeyValueBatchResult::Err(err.to_string().into_bytes().into()),
+        }
+    }
+}
+
 impl<'db, E: fmt::Display> From<Result<CreateIteratorResult<'db>, E>> for IteratorResult<'db> {
     fn from(value: Result<CreateIteratorResult<'db>, E>) -> Self {
         match value {
@@ -395,6 +431,18 @@ impl CResult for KeyValueResult {
     }
 }
 
+impl NullHandleResult for KeyValueBatchResult {
+    fn null_handle_pointer_error() -> Self {
+        Self::NullHandlePointer
+    }
+}
+
+impl CResult for KeyValueBatchResult {
+    fn from_err(err: impl ToString) -> Self {
+        Self::Err(err.to_string().into_bytes().into())
+    }
+}
+
 enum Panic {
     Static(&'static str),
     Formatted(String),

From a5ff25314dfc768740e81dc8c5b7dc21fffcb885 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 27 Aug 2025 16:11:21 -0400
Subject: [PATCH 02/15] wip: iterator batching / go side done with tests

---
 ffi/firewood.h           |  7 +++-
 ffi/firewood_test.go     | 61 +++++++++++++++++++++++++++++++++++
 ffi/iterator.go          | 63 +++++++++++++++++++++++++++++-------
 ffi/memory.go            | 69 +++++++++++++++++++++++++++++++++++++++-
 ffi/src/value.rs         |  2 +-
 ffi/src/value/kvp.rs     |  5 ++-
 ffi/src/value/results.rs |  6 ++--
 7 files changed, 195 insertions(+), 18 deletions(-)

diff --git a/ffi/firewood.h b/ffi/firewood.h
index 35b95d0b3..09fd85c95 100644
--- a/ffi/firewood.h
+++ b/ffi/firewood.h
@@ -340,6 +340,11 @@ typedef struct OwnedSlice_OwnedKeyValuePair {
   size_t len;
 } OwnedSlice_OwnedKeyValuePair;
 
+/**
+ * A type alias for a rust-owned byte slice.
+ */
+typedef struct OwnedSlice_OwnedKeyValuePair OwnedKeyValueBatch;
+
 /**
  * A result type returned from iterator FFI functions
  */
@@ -375,7 +380,7 @@ typedef struct KeyValueBatchResult {
       struct HashKey revision_not_found;
     };
     struct {
-      struct OwnedSlice_OwnedKeyValuePair some;
+      OwnedKeyValueBatch some;
     };
     struct {
       OwnedBytes err;
diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index f88bb07de..71da23257 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -7,10 +7,12 @@ import (
 	"bytes"
 	"crypto/rand"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
+	"sort"
 	"strconv"
 	"strings"
 	"sync"
@@ -254,6 +256,37 @@ func randomBytes(n int) []byte {
 	return b
 }
 
+// sortKV sorts keys lexicographically and keeps vals paired.
+func sortKV(keys, vals [][]byte) error {
+	if len(keys) != len(vals) {
+		return errors.New("keys/vals length mismatch")
+	}
+	n := len(keys)
+	if n <= 1 {
+		return nil
+	}
+	ord := make([]int, n)
+	for i := range ord {
+		ord[i] = i
+	}
+	sort.Slice(ord, func(i, j int) bool {
+		return bytes.Compare(keys[ord[i]], keys[ord[j]]) < 0
+	})
+	perm := make([]int, n)
+	for dest, orig := range ord {
+		perm[orig] = dest
+	}
+	for i := 0; i < n; i++ {
+		for perm[i] != i {
+			j := perm[i]
+			keys[i], keys[j] = keys[j], keys[i]
+			vals[i], vals[j] = vals[j], vals[i]
+			perm[i], perm[j] = perm[j], j
+		}
+	}
+	return nil
+}
+
 func kvForBench(num int) ([][]byte, [][]byte) {
 	keys := make([][]byte, num)
 	vals := make([][]byte, num)
@@ -262,6 +295,7 @@ func kvForBench(num int) ([][]byte, [][]byte) {
 		keys[i] = randomBytes(32)
 		vals[i] = randomBytes(128)
 	}
+	_ = sortKV(keys, vals)
 	return keys, vals
 }
 
@@ -1138,3 +1172,30 @@ func TestIterOnProposal(t *testing.T) {
 	}
 	r.NoError(it.Err())
 }
+
+// Tests that batched iterator functionality works
+func TestIterBatched(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForBench(1000)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
+
+	it, err := db.IterLatest(nil)
+	r.NoError(err)
+	it.SetBatchSize(100)
+	it2, err := db.IterLatest(nil)
+	r.NoError(err)
+
+	i := 0
+	for ; it.Next() && it2.Next(); i += 1 {
+		r.Equal(it.Key(), it2.Key())
+		r.Equal(it.Value(), it2.Value())
+		r.Equal(keys[i], it.Key())
+		r.Equal(vals[i], it.Value())
+	}
+	r.NoError(it.Err())
+	r.NoError(it2.Err())
+	r.Equal(i, 1000)
+}
diff --git a/ffi/iterator.go b/ffi/iterator.go
index 04c76aab7..ba9f6b784 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -19,39 +19,80 @@ type Iterator struct {
 	// It is not safe to call these methods with a nil handle.
 	handle *C.IteratorHandle
 
-	// currentKey is the current key retrieved from the iterator
+	// batchSize is the number of items that are loaded at once from ffi
+	// to reduce ffi call overheads
+	batchSize int
+	// loadedPairs is the latest loaded key value pairs retrieved
+	// from the iterator, not yet consumed by user
+	loadedPairs []*ownedKeyValue
+	// currentPair is the current pair retrieved from the iterator
+	currentPair *ownedKeyValue
+	// currentKey is the current pair retrieved from the iterator
 	currentKey []byte
-	// currentVal is the current value retrieved from the iterator
-	currentVal []byte
+	// currentValue is the current pair retrieved from the iterator
+	currentValue []byte
 	// err is the error from the iterator, if any
 	err error
 }
 
+func (it *Iterator) nextInternal() error {
+	if len(it.loadedPairs) == 0 {
+		if it.batchSize < 1 {
+			kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
+			if e != nil {
+				return e
+			}
+			if kv != nil {
+				// kv is nil when done
+				it.loadedPairs = append(it.loadedPairs, kv)
+			}
+		} else {
+			batch, e := getKeyValueBatchFromKeyValueBatchResult(C.fwd_iter_next_n(it.handle, C.size_t(it.batchSize)))
+			if e != nil {
+				return e
+			}
+			it.loadedPairs = batch.Copied()
+			if e = batch.Free(); e != nil {
+				return e
+			}
+		}
+	}
+	if len(it.loadedPairs) > 0 {
+		it.currentPair, it.loadedPairs = it.loadedPairs[0], it.loadedPairs[1:]
+	} else {
+		it.currentPair = nil
+	}
+	return nil
+}
+
+func (it *Iterator) SetBatchSize(batchSize int) {
+	it.batchSize = batchSize
+}
+
 func (it *Iterator) Next() bool {
-	kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
-	it.err = e
-	if kv == nil || e != nil {
+	it.err = it.nextInternal()
+	if it.currentPair == nil || it.err != nil {
 		return false
 	}
-	k, v, e := kv.Consume()
+	k, v, e := it.currentPair.Consume()
 	it.currentKey = k
-	it.currentVal = v
+	it.currentValue = v
 	it.err = e
 	return e == nil
 }
 
 func (it *Iterator) Key() []byte {
-	if (it.currentKey == nil && it.currentVal == nil) || it.err != nil {
+	if it.currentPair == nil || it.err != nil {
 		return nil
 	}
 	return it.currentKey
 }
 
 func (it *Iterator) Value() []byte {
-	if (it.currentKey == nil && it.currentVal == nil) || it.err != nil {
+	if it.currentPair == nil || it.err != nil {
 		return nil
 	}
-	return it.currentVal
+	return it.currentValue
 }
 
 func (it *Iterator) Err() error {
diff --git a/ffi/memory.go b/ffi/memory.go
index 63e5f106a..eb96494e7 100644
--- a/ffi/memory.go
+++ b/ffi/memory.go
@@ -350,6 +350,49 @@ func getValueFromValueResult(result C.ValueResult) ([]byte, error) {
 	}
 }
 
+type ownedKeyValueBatch struct {
+	owned C.OwnedKeyValueBatch
+}
+
+func (b *ownedKeyValueBatch) Copied() []*ownedKeyValue {
+	if b.owned.ptr == nil {
+		return nil
+	}
+	borrowed := b.Borrow()
+	copied := make([]*ownedKeyValue, len(borrowed))
+	for i, borrow := range borrowed {
+		copied[i] = newOwnedKeyValue(borrow)
+	}
+	return copied
+}
+
+func (b *ownedKeyValueBatch) Borrow() []C.OwnedKeyValuePair {
+	if b.owned.ptr == nil {
+		return nil
+	}
+
+	return unsafe.Slice((*C.OwnedKeyValuePair)(unsafe.Pointer(b.owned.ptr)), b.owned.len)
+}
+
+func (b *ownedKeyValueBatch) Free() error {
+	if b.owned.ptr == nil {
+		return nil
+	}
+
+	// TODO
+	return nil
+}
+
+// newOwnedKeyValueBatch creates a ownedKeyValueBatch from a C.OwnedKeyValueBatch.
+//
+// The caller is responsible for calling Free() on the returned ownedKeyValue
+// when it is no longer needed otherwise memory will leak.
+func newOwnedKeyValueBatch(owned C.OwnedKeyValueBatch) *ownedKeyValueBatch {
+	return &ownedKeyValueBatch{
+		owned: owned,
+	}
+}
+
 type ownedKeyValue struct {
 	key   *ownedBytes
 	value *ownedBytes
@@ -396,7 +439,7 @@ func getKeyValueFromKeyValueResult(result C.KeyValueResult) (*ownedKeyValue, err
 	case C.KeyValueResult_Some:
 		ownedKvp := newOwnedKeyValue(*(*C.OwnedKeyValuePair)(unsafe.Pointer(&result.anon0)))
 		return ownedKvp, nil
-	case C.ValueResult_Err:
+	case C.KeyValueResult_Err:
 		err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError()
 		return nil, err
 	default:
@@ -404,6 +447,30 @@ func getKeyValueFromKeyValueResult(result C.KeyValueResult) (*ownedKeyValue, err
 	}
 }
 
+// getKeyValueBatchFromKeyValueBatchResult converts a C.KeyValueResult to a key value pair or error.
+//
+// It returns nil, nil if the result is None.
+// It returns a *ownedKeyValueBatch, nil if the result is Some.
+// It returns an error if the result is an error.
+func getKeyValueBatchFromKeyValueBatchResult(result C.KeyValueBatchResult) (*ownedKeyValueBatch, error) {
+	switch result.tag {
+	case C.KeyValueBatchResult_NullHandlePointer:
+		return nil, errDBClosed
+	case C.KeyValueBatchResult_RevisionNotFound:
+		// NOTE: the result value contains the provided root hash, we could use
+		// it in the error message if needed.
+		return nil, errRevisionNotFound
+	case C.KeyValueBatchResult_Some:
+		ownedBatch := newOwnedKeyValueBatch(*(*C.OwnedKeyValueBatch)(unsafe.Pointer(&result.anon0)))
+		return ownedBatch, nil
+	case C.KeyValueBatchResult_Err:
+		err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError()
+		return nil, err
+	default:
+		return nil, fmt.Errorf("unknown C.KeyValueBatchResult 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.
diff --git a/ffi/src/value.rs b/ffi/src/value.rs
index 1a56f4da3..146d6081e 100644
--- a/ffi/src/value.rs
+++ b/ffi/src/value.rs
@@ -11,7 +11,7 @@ mod results;
 pub use self::borrowed::{BorrowedBytes, BorrowedKeyValuePairs, BorrowedSlice};
 use self::display_hex::DisplayHex;
 pub use self::hash_key::HashKey;
-pub use self::kvp::{KeyValuePair, OwnedKeyValuePair};
+pub use self::kvp::{KeyValuePair, OwnedKeyValuePair, OwnedKeyValueBatch};
 pub use self::owned::{OwnedBytes, OwnedSlice};
 pub(crate) use self::results::{CResult, NullHandleResult};
 pub use self::results::{
diff --git a/ffi/src/value/kvp.rs b/ffi/src/value/kvp.rs
index 24b59e47b..c24da609c 100644
--- a/ffi/src/value/kvp.rs
+++ b/ffi/src/value/kvp.rs
@@ -3,10 +3,13 @@
 
 use std::fmt;
 
-use crate::OwnedBytes;
+use crate::{OwnedBytes, OwnedSlice};
 use crate::value::BorrowedBytes;
 use firewood::v2::api;
 
+/// A type alias for a rust-owned byte slice.
+pub type OwnedKeyValueBatch = OwnedSlice<OwnedKeyValuePair>;
+
 /// A `KeyValue` represents a key-value pair, passed to the FFI.
 #[repr(C)]
 #[derive(Debug, Clone, Copy)]
diff --git a/ffi/src/value/results.rs b/ffi/src/value/results.rs
index 3dcb321d2..5d7fdc5cc 100644
--- a/ffi/src/value/results.rs
+++ b/ffi/src/value/results.rs
@@ -6,8 +6,8 @@ use firewood::v2::api;
 use std::fmt;
 
 use crate::iterator::{CreateIteratorResult, IteratorHandle};
-use crate::value::kvp::OwnedKeyValuePair;
-use crate::{CreateProposalResult, HashKey, OwnedBytes, OwnedSlice, ProposalHandle};
+use crate::value::kvp::{OwnedKeyValueBatch, OwnedKeyValuePair};
+use crate::{CreateProposalResult, HashKey, OwnedBytes, ProposalHandle};
 
 /// The result type returned from an FFI function that returns no value but may
 /// return an error.
@@ -271,7 +271,7 @@ pub enum KeyValueBatchResult {
     /// The provided root was not found in the database.
     RevisionNotFound(HashKey),
     /// The next batch of items on iterator are returned.
-    Some(OwnedSlice<OwnedKeyValuePair>),
+    Some(OwnedKeyValueBatch),
     /// An error occurred and the message is returned as an [`OwnedBytes`]. If
     /// value is guaranteed to contain only valid UTF-8.
     ///

From 351e2f43055833a42012749788ac9833d4bc3257 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 27 Aug 2025 16:24:38 -0400
Subject: [PATCH 03/15] wip: iterator batching / better docs

---
 ffi/iterator.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/ffi/iterator.go b/ffi/iterator.go
index ba9f6b784..bc418df8e 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -19,18 +19,23 @@ type Iterator struct {
 	// It is not safe to call these methods with a nil handle.
 	handle *C.IteratorHandle
 
-	// batchSize is the number of items that are loaded at once from ffi
+	// batchSize is the number of items that are loaded at once
 	// to reduce ffi call overheads
 	batchSize int
+
 	// loadedPairs is the latest loaded key value pairs retrieved
 	// from the iterator, not yet consumed by user
 	loadedPairs []*ownedKeyValue
+
 	// currentPair is the current pair retrieved from the iterator
 	currentPair *ownedKeyValue
+
 	// currentKey is the current pair retrieved from the iterator
 	currentKey []byte
+
 	// currentValue is the current pair retrieved from the iterator
 	currentValue []byte
+
 	// err is the error from the iterator, if any
 	err error
 }
@@ -65,10 +70,14 @@ func (it *Iterator) nextInternal() error {
 	return nil
 }
 
+// SetBatchSize sets the max number of pairs to be retrieved in one ffi call.
 func (it *Iterator) SetBatchSize(batchSize int) {
 	it.batchSize = batchSize
 }
 
+// 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 {
 	it.err = it.nextInternal()
 	if it.currentPair == nil || it.err != nil {
@@ -81,6 +90,7 @@ func (it *Iterator) Next() bool {
 	return e == nil
 }
 
+// Key returns the key of the current pair
 func (it *Iterator) Key() []byte {
 	if it.currentPair == nil || it.err != nil {
 		return nil
@@ -88,6 +98,7 @@ func (it *Iterator) Key() []byte {
 	return it.currentKey
 }
 
+// Value returns the value of the current pair
 func (it *Iterator) Value() []byte {
 	if it.currentPair == nil || it.err != nil {
 		return nil
@@ -95,6 +106,7 @@ func (it *Iterator) Value() []byte {
 	return it.currentValue
 }
 
+// Err returns the error if Next failed
 func (it *Iterator) Err() error {
 	return it.err
 }

From 695db4c064da5a1127205521b4ccba04357427b9 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Tue, 2 Sep 2025 20:18:44 -0400
Subject: [PATCH 04/15] feat: borrowed next

---
 ffi/firewood_test.go | 27 +++++++++++++++++++++++++++
 ffi/iterator.go      | 20 ++++++++++++++++++++
 ffi/memory.go        | 20 +++++++++++++++-----
 3 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index 71da23257..9fff1f781 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1199,3 +1199,30 @@ func TestIterBatched(t *testing.T) {
 	r.NoError(it2.Err())
 	r.Equal(i, 1000)
 }
+
+// Tests that batched iterator functionality works
+func TestIterBatchedBorrowed(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+
+	keys, vals := kvForBench(1000)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
+
+	it, err := db.IterLatest(nil)
+	r.NoError(err)
+	it.SetBatchSize(100)
+	it2, err := db.IterLatest(nil)
+	r.NoError(err)
+
+	i := 0
+	for ; it.NextBorrowed() && it2.NextBorrowed(); i += 1 {
+		r.Equal(it.Key(), it2.Key())
+		r.Equal(it.Value(), it2.Value())
+		r.Equal(keys[i], it.Key())
+		r.Equal(vals[i], it.Value())
+	}
+	r.NoError(it.Err())
+	r.NoError(it2.Err())
+	r.Equal(i, 1000)
+}
diff --git a/ffi/iterator.go b/ffi/iterator.go
index bc418df8e..9b03239ce 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -41,6 +41,12 @@ type Iterator struct {
 }
 
 func (it *Iterator) nextInternal() error {
+	if it.currentPair != nil {
+		err := it.currentPair.Free()
+		if err != nil {
+			return err
+		}
+	}
 	if len(it.loadedPairs) == 0 {
 		if it.batchSize < 1 {
 			kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
@@ -90,6 +96,20 @@ func (it *Iterator) Next() bool {
 	return e == nil
 }
 
+// NextBorrowed retrieves the next item on the iterator similar to Next
+// the difference is that returned bytes in Key and Value are not copied
+// and will be freed on next call to Next or NextBorrowed
+func (it *Iterator) NextBorrowed() bool {
+	it.err = it.nextInternal()
+	if it.currentPair == nil || it.err != nil {
+		return false
+	}
+	it.currentKey = it.currentPair.key.BorrowedBytes()
+	it.currentValue = it.currentPair.value.BorrowedBytes()
+	it.err = nil
+	return true
+}
+
 // Key returns the key of the current pair
 func (it *Iterator) Key() []byte {
 	if it.currentPair == nil || it.err != nil {
diff --git a/ffi/memory.go b/ffi/memory.go
index eb96494e7..9493c3665 100644
--- a/ffi/memory.go
+++ b/ffi/memory.go
@@ -400,16 +400,26 @@ type ownedKeyValue struct {
 
 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)
+	e := kv.Free()
+	if e != nil {
+		return nil, nil, e
 	}
 	return key, value, nil
 }
 
+func (kv *ownedKeyValue) Free() error {
+	err := kv.key.Free()
+	if err != nil {
+		return fmt.Errorf("%w: %w", errFreeingValue, err)
+	}
+	err = kv.value.Free()
+	if err != nil {
+		return fmt.Errorf("%w: %w", errFreeingValue, err)
+	}
+	return nil
+}
+
 // newOwnedKeyValue creates a ownedKeyValue from a C.OwnedKeyValuePair.
 //
 // The caller is responsible for calling Free() on the returned ownedKeyValue

From f76202c3c9b754b2ed0b619b9bc0a3047eefb3b7 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 01:46:51 -0400
Subject: [PATCH 05/15] feat: better tests

---
 ffi/firewood_test.go | 179 ++++++++++++++++++++-----------------------
 ffi/iterator.go      |   2 +-
 2 files changed, 84 insertions(+), 97 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index 9fff1f781..edfa9746e 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1106,123 +1106,110 @@ func TestIterEmptyDb(t *testing.T) {
 	r.Error(err)
 }
 
-// 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)
+type kvIter interface {
+	Next() bool
+	Key() []byte
+	Value() []byte
+	Err() error
+}
+type borrowIter struct{ it *Iterator }
 
-	it, err := db.IterLatest(nil)
-	r.NoError(err)
+func (b borrowIter) Next() bool    { return b.it.NextBorrowed() }
+func (b borrowIter) Key() []byte   { return b.it.Key() }
+func (b borrowIter) Value() []byte { return b.it.Value() }
+func (b borrowIter) Err() error    { return b.it.Err() }
 
-	for i := 0; it.Next(); i += 1 {
+func assertIteratorYields(r *require.Assertions, it kvIter, 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 iterators on different roots work fine
-func TestIterOnRoot(t *testing.T) {
+// Tests that basic iterator functionality works
+func TestIter(t *testing.T) {
 	r := require.New(t)
 	db := newTestDatabase(t)
 
-	// Commit 10 key-value pairs.
-	keys, vals := kvForTest(20)
-	firstRoot, err := db.Update(keys[:10], vals[:10])
-	r.NoError(err)
-
-	secondRoot, err := db.Update(keys[:10], vals[10:])
-	r.NoError(err)
-
-	h1, err := db.IterOnRoot(firstRoot, nil)
-	r.NoError(err)
-
-	h2, err := db.IterOnRoot(secondRoot, nil)
-	r.NoError(err)
-
-	for i := 0; h1.Next() && h2.Next(); i += 1 {
-		r.Equal(keys[i], h1.Key())
-		r.Equal(keys[i], h2.Key())
-		r.Equal(vals[i], h1.Value())
-		r.Equal(vals[i+10], h2.Value())
+	dataModes := []struct {
+		name     string
+		configFn func(it *Iterator) kvIter
+	}{
+		{"Owned", func(it *Iterator) kvIter { return it }},
+		{"Borrowed", func(it *Iterator) kvIter { return borrowIter{it: it} }},
 	}
-	r.NoError(h1.Err())
-	r.NoError(h2.Err())
-}
-
-// 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)
-
-	for i := 0; it.Next(); i += 1 {
-		r.Equal(keys[i], it.Key())
-		r.Equal(vals[i], it.Value())
+	batchModes := []struct {
+		name     string
+		configFn func(it *Iterator)
+	}{
+		{"Single", func(it *Iterator) {
+			it.SetBatchSize(1)
+		}},
+		{"Batched", func(it *Iterator) {
+			it.SetBatchSize(100)
+		}},
 	}
-	r.NoError(it.Err())
-}
 
-// Tests that batched iterator functionality works
-func TestIterBatched(t *testing.T) {
-	r := require.New(t)
-	db := newTestDatabase(t)
-
-	keys, vals := kvForBench(1000)
-	_, err := db.Update(keys, vals)
+	keys, vals := kvForBench(240)
+	firstRoot, err := db.Update(keys[:80], vals[:80])
 	r.NoError(err)
-
-	it, err := db.IterLatest(nil)
+	secondRoot, err := db.Update(keys[80:160], vals[80:160])
 	r.NoError(err)
-	it.SetBatchSize(100)
-	it2, err := db.IterLatest(nil)
+	thirdRoot, err := db.Update(keys[160:], vals[160:])
 	r.NoError(err)
 
-	i := 0
-	for ; it.Next() && it2.Next(); i += 1 {
-		r.Equal(it.Key(), it2.Key())
-		r.Equal(it.Value(), it2.Value())
-		r.Equal(keys[i], it.Key())
-		r.Equal(vals[i], it.Value())
-	}
-	r.NoError(it.Err())
-	r.NoError(it2.Err())
-	r.Equal(i, 1000)
-}
-
-// Tests that batched iterator functionality works
-func TestIterBatchedBorrowed(t *testing.T) {
-	r := require.New(t)
-	db := newTestDatabase(t)
+	for _, dataMode := range dataModes {
+		for _, batchMode := range batchModes {
+			t.Run(fmt.Sprintf("Latest/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
+				r := require.New(t)
+				it, err := db.IterLatest(nil)
+				r.NoError(err)
 
-	keys, vals := kvForBench(1000)
-	_, err := db.Update(keys, vals)
-	r.NoError(err)
+				batchMode.configFn(it)
+				assertIteratorYields(r, dataMode.configFn(it), keys, vals)
+			})
 
-	it, err := db.IterLatest(nil)
-	r.NoError(err)
-	it.SetBatchSize(100)
-	it2, err := db.IterLatest(nil)
-	r.NoError(err)
+			t.Run(fmt.Sprintf("OnRoot/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
+				r := require.New(t)
+				h1, err := db.IterOnRoot(firstRoot, nil)
+				r.NoError(err)
+				h2, err := db.IterOnRoot(secondRoot, nil)
+				r.NoError(err)
+				h3, err := db.IterOnRoot(thirdRoot, nil)
+				r.NoError(err)
+				batchMode.configFn(h1)
+				batchMode.configFn(h2)
+				batchMode.configFn(h3)
+				assertIteratorYields(r, dataMode.configFn(h1), keys[:80], vals[:80])
+				assertIteratorYields(r, dataMode.configFn(h2), keys[:160], vals[:160])
+				assertIteratorYields(r, dataMode.configFn(h3), keys, vals)
+			})
+
+			t.Run(fmt.Sprintf("OnProposal/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
+				r := require.New(t)
+				updatedValues := make([][]byte, len(vals))
+				copy(updatedValues, vals)
+
+				changedKeys := make([][]byte, 0)
+				changedVals := make([][]byte, 0)
+				for i := 0; i < len(vals); i += 4 {
+					changedKeys = append(changedKeys, keys[i])
+					newVal := []byte{byte(i)}
+					changedVals = append(changedVals, newVal)
+					updatedValues[i] = newVal
+				}
+				p, err := db.Propose(changedKeys, changedVals)
+				r.NoError(err)
+				it, err := p.Iter(nil)
+				r.NoError(err)
 
-	i := 0
-	for ; it.NextBorrowed() && it2.NextBorrowed(); i += 1 {
-		r.Equal(it.Key(), it2.Key())
-		r.Equal(it.Value(), it2.Value())
-		r.Equal(keys[i], it.Key())
-		r.Equal(vals[i], it.Value())
+				batchMode.configFn(it)
+				assertIteratorYields(r, dataMode.configFn(it), keys, updatedValues)
+			})
+		}
 	}
-	r.NoError(it.Err())
-	r.NoError(it2.Err())
-	r.Equal(i, 1000)
 }
diff --git a/ffi/iterator.go b/ffi/iterator.go
index 9b03239ce..d065bac1f 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -48,7 +48,7 @@ func (it *Iterator) nextInternal() error {
 		}
 	}
 	if len(it.loadedPairs) == 0 {
-		if it.batchSize < 1 {
+		if it.batchSize <= 1 {
 			kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
 			if e != nil {
 				return e

From c54582c290d0daf5e74d272c4be4202451655919 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 01:56:34 -0400
Subject: [PATCH 06/15] feat: better test syntax

---
 ffi/firewood_test.go | 96 +++++++++++++++++++++++---------------------
 1 file changed, 50 insertions(+), 46 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index edfa9746e..a2fcc3d08 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1162,54 +1162,58 @@ func TestIter(t *testing.T) {
 	thirdRoot, err := db.Update(keys[160:], vals[160:])
 	r.NoError(err)
 
-	for _, dataMode := range dataModes {
-		for _, batchMode := range batchModes {
-			t.Run(fmt.Sprintf("Latest/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
-				r := require.New(t)
-				it, err := db.IterLatest(nil)
-				r.NoError(err)
+	runForAllModes := func(parentT *testing.T, name string, fn func(*testing.T, func(it *Iterator) kvIter)) {
+		for _, dataMode := range dataModes {
+			for _, batchMode := range batchModes {
+				parentT.Run(fmt.Sprintf("%s/%s/%s", name, dataMode.name, batchMode.name), func(t *testing.T) {
+					fn(t, func(it *Iterator) kvIter {
+						batchMode.configFn(it)
+						return dataMode.configFn(it)
+					})
+				})
+			}
+		}
+	}
+	runForAllModes(t, "Latest", func(t *testing.T, fn func(it *Iterator) kvIter) {
+		r := require.New(t)
+		it, err := db.IterLatest(nil)
+		r.NoError(err)
 
-				batchMode.configFn(it)
-				assertIteratorYields(r, dataMode.configFn(it), keys, vals)
-			})
+		assertIteratorYields(r, fn(it), keys, vals)
+	})
 
-			t.Run(fmt.Sprintf("OnRoot/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
-				r := require.New(t)
-				h1, err := db.IterOnRoot(firstRoot, nil)
-				r.NoError(err)
-				h2, err := db.IterOnRoot(secondRoot, nil)
-				r.NoError(err)
-				h3, err := db.IterOnRoot(thirdRoot, nil)
-				r.NoError(err)
-				batchMode.configFn(h1)
-				batchMode.configFn(h2)
-				batchMode.configFn(h3)
-				assertIteratorYields(r, dataMode.configFn(h1), keys[:80], vals[:80])
-				assertIteratorYields(r, dataMode.configFn(h2), keys[:160], vals[:160])
-				assertIteratorYields(r, dataMode.configFn(h3), keys, vals)
-			})
-
-			t.Run(fmt.Sprintf("OnProposal/%s/%s", dataMode.name, batchMode.name), func(t *testing.T) {
-				r := require.New(t)
-				updatedValues := make([][]byte, len(vals))
-				copy(updatedValues, vals)
-
-				changedKeys := make([][]byte, 0)
-				changedVals := make([][]byte, 0)
-				for i := 0; i < len(vals); i += 4 {
-					changedKeys = append(changedKeys, keys[i])
-					newVal := []byte{byte(i)}
-					changedVals = append(changedVals, newVal)
-					updatedValues[i] = newVal
-				}
-				p, err := db.Propose(changedKeys, changedVals)
-				r.NoError(err)
-				it, err := p.Iter(nil)
-				r.NoError(err)
+	runForAllModes(t, "OnRoot", func(t *testing.T, fn func(it *Iterator) kvIter) {
+		r := require.New(t)
+		h1, err := db.IterOnRoot(firstRoot, nil)
+		r.NoError(err)
+		h2, err := db.IterOnRoot(secondRoot, nil)
+		r.NoError(err)
+		h3, err := db.IterOnRoot(thirdRoot, nil)
+		r.NoError(err)
+
+		assertIteratorYields(r, fn(h1), keys[:80], vals[:80])
+		assertIteratorYields(r, fn(h2), keys[:160], vals[:160])
+		assertIteratorYields(r, fn(h3), keys, vals)
+	})
 
-				batchMode.configFn(it)
-				assertIteratorYields(r, dataMode.configFn(it), keys, updatedValues)
-			})
+	runForAllModes(t, "OnProposal", func(t *testing.T, fn func(it *Iterator) kvIter) {
+		r := require.New(t)
+		updatedValues := make([][]byte, len(vals))
+		copy(updatedValues, vals)
+
+		changedKeys := make([][]byte, 0)
+		changedVals := make([][]byte, 0)
+		for i := 0; i < len(vals); i += 4 {
+			changedKeys = append(changedKeys, keys[i])
+			newVal := []byte{byte(i)}
+			changedVals = append(changedVals, newVal)
+			updatedValues[i] = newVal
 		}
-	}
+		p, err := db.Propose(changedKeys, changedVals)
+		r.NoError(err)
+		it, err := p.Iter(nil)
+		r.NoError(err)
+
+		assertIteratorYields(r, fn(it), keys, updatedValues)
+	})
 }

From ad00f2083049ecf0229f50ac3a01b7ff9138def7 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 01:57:39 -0400
Subject: [PATCH 07/15] feat: better test syntax

---
 ffi/firewood_test.go | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index a2fcc3d08..89a965a9f 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1174,15 +1174,15 @@ func TestIter(t *testing.T) {
 			}
 		}
 	}
-	runForAllModes(t, "Latest", func(t *testing.T, fn func(it *Iterator) kvIter) {
+	runForAllModes(t, "Latest", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
 		r := require.New(t)
 		it, err := db.IterLatest(nil)
 		r.NoError(err)
 
-		assertIteratorYields(r, fn(it), keys, vals)
+		assertIteratorYields(r, configureIterator(it), keys, vals)
 	})
 
-	runForAllModes(t, "OnRoot", func(t *testing.T, fn func(it *Iterator) kvIter) {
+	runForAllModes(t, "OnRoot", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
 		r := require.New(t)
 		h1, err := db.IterOnRoot(firstRoot, nil)
 		r.NoError(err)
@@ -1191,12 +1191,12 @@ func TestIter(t *testing.T) {
 		h3, err := db.IterOnRoot(thirdRoot, nil)
 		r.NoError(err)
 
-		assertIteratorYields(r, fn(h1), keys[:80], vals[:80])
-		assertIteratorYields(r, fn(h2), keys[:160], vals[:160])
-		assertIteratorYields(r, fn(h3), keys, vals)
+		assertIteratorYields(r, configureIterator(h1), keys[:80], vals[:80])
+		assertIteratorYields(r, configureIterator(h2), keys[:160], vals[:160])
+		assertIteratorYields(r, configureIterator(h3), keys, vals)
 	})
 
-	runForAllModes(t, "OnProposal", func(t *testing.T, fn func(it *Iterator) kvIter) {
+	runForAllModes(t, "OnProposal", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
 		r := require.New(t)
 		updatedValues := make([][]byte, len(vals))
 		copy(updatedValues, vals)
@@ -1214,6 +1214,6 @@ func TestIter(t *testing.T) {
 		it, err := p.Iter(nil)
 		r.NoError(err)
 
-		assertIteratorYields(r, fn(it), keys, updatedValues)
+		assertIteratorYields(r, configureIterator(it), keys, updatedValues)
 	})
 }

From 273c64e14af4f6a2364e8e17d85d2d56ad04847c Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 11:41:33 -0400
Subject: [PATCH 08/15] feat: free resources

---
 ffi/firewood.h  | 69 ++++++++++++++++++++++++++++++++-----------------
 ffi/iterator.go | 30 ++++++++++++---------
 ffi/memory.go   | 15 ++++++-----
 ffi/src/lib.rs  | 23 +++++++++++++++++
 4 files changed, 94 insertions(+), 43 deletions(-)

diff --git a/ffi/firewood.h b/ffi/firewood.h
index 09fd85c95..815a2cf2e 100644
--- a/ffi/firewood.h
+++ b/ffi/firewood.h
@@ -222,6 +222,30 @@ typedef struct VoidResult {
   };
 } VoidResult;
 
+/**
+ * Owned version of `KeyValuePair`, returned to the FFI.
+ */
+typedef struct OwnedKeyValuePair {
+  OwnedBytes key;
+  OwnedBytes value;
+} OwnedKeyValuePair;
+
+/**
+ * A Rust-owned vector of bytes that can be passed to C code.
+ *
+ * C callers must free this memory using the respective FFI function for the
+ * concrete type (but not using the `free` function from the C standard library).
+ */
+typedef struct OwnedSlice_OwnedKeyValuePair {
+  struct OwnedKeyValuePair *ptr;
+  size_t len;
+} OwnedSlice_OwnedKeyValuePair;
+
+/**
+ * A type alias for a rust-owned byte slice.
+ */
+typedef struct OwnedSlice_OwnedKeyValuePair OwnedKeyValueBatch;
+
 /**
  * A result type returned from FFI functions that retrieve a single value.
  */
@@ -274,14 +298,6 @@ typedef struct ValueResult {
   };
 } ValueResult;
 
-/**
- * Owned version of `KeyValuePair`, returned to the FFI.
- */
-typedef struct OwnedKeyValuePair {
-  OwnedBytes key;
-  OwnedBytes value;
-} OwnedKeyValuePair;
-
 /**
  * A result type returned from iterator FFI functions
  */
@@ -329,22 +345,6 @@ typedef struct KeyValueResult {
   };
 } KeyValueResult;
 
-/**
- * A Rust-owned vector of bytes that can be passed to C code.
- *
- * C callers must free this memory using the respective FFI function for the
- * concrete type (but not using the `free` function from the C standard library).
- */
-typedef struct OwnedSlice_OwnedKeyValuePair {
-  struct OwnedKeyValuePair *ptr;
-  size_t len;
-} OwnedSlice_OwnedKeyValuePair;
-
-/**
- * A type alias for a rust-owned byte slice.
- */
-typedef struct OwnedSlice_OwnedKeyValuePair OwnedKeyValueBatch;
-
 /**
  * A result type returned from iterator FFI functions
  */
@@ -718,6 +718,27 @@ struct VoidResult fwd_free_iterator(struct IteratorHandle *iterator);
  */
 struct VoidResult fwd_free_owned_bytes(OwnedBytes bytes);
 
+/**
+ * Consumes the [`OwnedKeyValueBatch`] and frees the memory associated with it.
+ *
+ * # Arguments
+ *
+ * * `batch` - The [`OwnedKeyValueBatch`] struct to free, previously returned from any
+ *   function from this library.
+ *
+ * # Returns
+ *
+ * - [`VoidResult::Ok`] if the memory was successfully freed.
+ * - [`VoidResult::Err`] if the process panics while freeing the memory.
+ *
+ * # Safety
+ *
+ * The caller must ensure that the `batch` struct is valid and that the memory
+ * it points to is uniquely owned by this object. However, if `batch.ptr` is null,
+ * this function does nothing.
+ */
+struct VoidResult fwd_free_owned_key_value_batch(OwnedKeyValueBatch batch);
+
 /**
  * Consumes the [`ProposalHandle`], cancels the proposal, and frees the memory.
  *
diff --git a/ffi/iterator.go b/ffi/iterator.go
index d065bac1f..ea85f9d84 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -38,16 +38,23 @@ type Iterator struct {
 
 	// err is the error from the iterator, if any
 	err error
+
+	// currentResource is a reference to a freeable resource to clean up
+	currentResource interface{ Free() error }
 }
 
-func (it *Iterator) nextInternal() error {
-	if it.currentPair != nil {
-		err := it.currentPair.Free()
-		if err != nil {
-			return err
-		}
+func (it *Iterator) Release() error {
+	if it.currentResource == nil {
+		return nil
 	}
+	return it.currentResource.Free()
+}
+
+func (it *Iterator) nextInternal() error {
 	if len(it.loadedPairs) == 0 {
+		if e := it.Release(); e != nil {
+			return e
+		}
 		if it.batchSize <= 1 {
 			kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
 			if e != nil {
@@ -57,15 +64,14 @@ func (it *Iterator) nextInternal() error {
 				// kv is nil when done
 				it.loadedPairs = append(it.loadedPairs, kv)
 			}
+			it.currentResource = kv
 		} else {
 			batch, e := getKeyValueBatchFromKeyValueBatchResult(C.fwd_iter_next_n(it.handle, C.size_t(it.batchSize)))
 			if e != nil {
 				return e
 			}
 			it.loadedPairs = batch.Copied()
-			if e = batch.Free(); e != nil {
-				return e
-			}
+			it.currentResource = batch
 		}
 	}
 	if len(it.loadedPairs) > 0 {
@@ -89,11 +95,11 @@ func (it *Iterator) Next() bool {
 	if it.currentPair == nil || it.err != nil {
 		return false
 	}
-	k, v, e := it.currentPair.Consume()
+	k, v := it.currentPair.Copy()
 	it.currentKey = k
 	it.currentValue = v
-	it.err = e
-	return e == nil
+	it.err = nil
+	return true
 }
 
 // NextBorrowed retrieves the next item on the iterator similar to Next
diff --git a/ffi/memory.go b/ffi/memory.go
index 9493c3665..07cf7477b 100644
--- a/ffi/memory.go
+++ b/ffi/memory.go
@@ -379,7 +379,12 @@ func (b *ownedKeyValueBatch) Free() error {
 		return nil
 	}
 
-	// TODO
+	if err := getErrorFromVoidResult(C.fwd_free_owned_key_value_batch(b.owned)); err != nil {
+		return fmt.Errorf("%w: %w", errFreeingValue, err)
+	}
+
+	b.owned = C.OwnedKeyValueBatch{}
+
 	return nil
 }
 
@@ -398,14 +403,10 @@ type ownedKeyValue struct {
 	value *ownedBytes
 }
 
-func (kv *ownedKeyValue) Consume() ([]byte, []byte, error) {
+func (kv *ownedKeyValue) Copy() ([]byte, []byte) {
 	key := kv.key.CopiedBytes()
 	value := kv.value.CopiedBytes()
-	e := kv.Free()
-	if e != nil {
-		return nil, nil, e
-	}
-	return key, value, nil
+	return key, value
 }
 
 func (kv *ownedKeyValue) Free() error {
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 9e65a73a1..64083bf30 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -639,3 +639,26 @@ pub unsafe extern "C" fn fwd_close_db(db: Option<Box<DatabaseHandle>>) -> VoidRe
 pub unsafe extern "C" fn fwd_free_owned_bytes(bytes: OwnedBytes) -> VoidResult {
     invoke(move || drop(bytes))
 }
+
+
+/// Consumes the [`OwnedKeyValueBatch`] and frees the memory associated with it.
+///
+/// # Arguments
+///
+/// * `batch` - The [`OwnedKeyValueBatch`] struct to free, previously returned from any
+///   function from this library.
+///
+/// # Returns
+///
+/// - [`VoidResult::Ok`] if the memory was successfully freed.
+/// - [`VoidResult::Err`] if the process panics while freeing the memory.
+///
+/// # Safety
+///
+/// The caller must ensure that the `batch` struct is valid and that the memory
+/// it points to is uniquely owned by this object. However, if `batch.ptr` is null,
+/// this function does nothing.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn fwd_free_owned_key_value_batch(batch: OwnedKeyValueBatch) -> VoidResult {
+    invoke(move || drop(batch))
+}
\ No newline at end of file

From d1fc2564793e3a46ee8632669005d3db3eb0306a Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 12:45:57 -0400
Subject: [PATCH 09/15] fmt

---
 ffi/src/iterator.rs      |  3 ++-
 ffi/src/lib.rs           |  8 +++++---
 ffi/src/value.rs         |  6 +++---
 ffi/src/value/kvp.rs     |  2 +-
 ffi/src/value/results.rs | 25 ++++++++++++++++---------
 5 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/ffi/src/iterator.rs b/ffi/src/iterator.rs
index 0ea9dc105..798146999 100644
--- a/ffi/src/iterator.rs
+++ b/ffi/src/iterator.rs
@@ -2,8 +2,8 @@
 // See the file LICENSE.md for licensing terms.
 
 use firewood::merkle;
-use std::fmt::{Debug, Formatter};
 use firewood::v2::api;
+use std::fmt::{Debug, Formatter};
 
 type KeyValueItem = (merkle::Key, merkle::Value);
 
@@ -18,6 +18,7 @@ impl Debug for IteratorHandle<'_> {
     }
 }
 
+#[expect(clippy::missing_errors_doc)]
 impl IteratorHandle<'_> {
     pub fn iter_next(&mut self) -> Option<Result<KeyValueItem, api::Error>> {
         self.iterator.next()
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index a12603663..f095696d0 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -221,7 +221,10 @@ pub unsafe extern "C" fn fwd_iter_next(handle: Option<&mut IteratorHandle<'_>>)
 ///   to free the memory associated with the returned error or value.
 ///
 #[unsafe(no_mangle)]
-pub unsafe extern "C" fn fwd_iter_next_n(handle: Option<&mut IteratorHandle<'_>>, n: usize) -> KeyValueBatchResult {
+pub unsafe extern "C" fn fwd_iter_next_n(
+    handle: Option<&mut IteratorHandle<'_>>,
+    n: usize,
+) -> KeyValueBatchResult {
     invoke_with_handle(handle, |it| it.iter_next_n(n))
 }
 
@@ -640,7 +643,6 @@ pub unsafe extern "C" fn fwd_free_owned_bytes(bytes: OwnedBytes) -> VoidResult {
     invoke(move || drop(bytes))
 }
 
-
 /// Consumes the [`OwnedKeyValueBatch`] and frees the memory associated with it.
 ///
 /// # Arguments
@@ -661,4 +663,4 @@ pub unsafe extern "C" fn fwd_free_owned_bytes(bytes: OwnedBytes) -> VoidResult {
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn fwd_free_owned_key_value_batch(batch: OwnedKeyValueBatch) -> VoidResult {
     invoke(move || drop(batch))
-}
\ No newline at end of file
+}
diff --git a/ffi/src/value.rs b/ffi/src/value.rs
index 146d6081e..bc0ae3f8a 100644
--- a/ffi/src/value.rs
+++ b/ffi/src/value.rs
@@ -11,10 +11,10 @@ mod results;
 pub use self::borrowed::{BorrowedBytes, BorrowedKeyValuePairs, BorrowedSlice};
 use self::display_hex::DisplayHex;
 pub use self::hash_key::HashKey;
-pub use self::kvp::{KeyValuePair, OwnedKeyValuePair, OwnedKeyValueBatch};
+pub use self::kvp::{KeyValuePair, OwnedKeyValueBatch, OwnedKeyValuePair};
 pub use self::owned::{OwnedBytes, OwnedSlice};
 pub(crate) use self::results::{CResult, NullHandleResult};
 pub use self::results::{
-    HandleResult, HashResult, IteratorResult, KeyValueResult, KeyValueBatchResult, ProposalResult, ValueResult,
-    VoidResult,
+    HandleResult, HashResult, IteratorResult, KeyValueBatchResult, KeyValueResult, ProposalResult,
+    ValueResult, VoidResult,
 };
diff --git a/ffi/src/value/kvp.rs b/ffi/src/value/kvp.rs
index c24da609c..967d182f4 100644
--- a/ffi/src/value/kvp.rs
+++ b/ffi/src/value/kvp.rs
@@ -3,8 +3,8 @@
 
 use std::fmt;
 
-use crate::{OwnedBytes, OwnedSlice};
 use crate::value::BorrowedBytes;
+use crate::{OwnedBytes, OwnedSlice};
 use firewood::v2::api;
 
 /// A type alias for a rust-owned byte slice.
diff --git a/ffi/src/value/results.rs b/ffi/src/value/results.rs
index 5d7fdc5cc..012950a5d 100644
--- a/ffi/src/value/results.rs
+++ b/ffi/src/value/results.rs
@@ -261,7 +261,6 @@ impl From<Option<Result<(merkle::Key, merkle::Value), api::Error>>> for KeyValue
     }
 }
 
-
 /// A result type returned from iterator FFI functions
 #[derive(Debug)]
 #[repr(C)]
@@ -285,14 +284,22 @@ pub enum KeyValueBatchResult {
 impl From<Result<Vec<(merkle::Key, merkle::Value)>, api::Error>> for KeyValueBatchResult {
     fn from(value: Result<Vec<(merkle::Key, merkle::Value)>, api::Error>) -> Self {
         match value {
-                Ok(pairs) => {
-                    let values: Vec<_> = pairs.into_iter().map(|(k, v)| OwnedKeyValuePair {key:k.into(), value:v.into()}).collect();
-                    KeyValueBatchResult::Some(values.into())
-                },
-                Err(api::Error::RevisionNotFound { provided }) => KeyValueBatchResult::RevisionNotFound(
-                    HashKey::from(provided.unwrap_or_else(api::HashKey::empty)),
-                ),
-                Err(err) => KeyValueBatchResult::Err(err.to_string().into_bytes().into()),
+            Ok(pairs) => {
+                let values: Vec<_> = pairs
+                    .into_iter()
+                    .map(|(k, v)| OwnedKeyValuePair {
+                        key: k.into(),
+                        value: v.into(),
+                    })
+                    .collect();
+                KeyValueBatchResult::Some(values.into())
+            }
+            Err(api::Error::RevisionNotFound { provided }) => {
+                KeyValueBatchResult::RevisionNotFound(HashKey::from(
+                    provided.unwrap_or_else(api::HashKey::empty),
+                ))
+            }
+            Err(err) => KeyValueBatchResult::Err(err.to_string().into_bytes().into()),
         }
     }
 }

From 73d5c577fa16b4531bb60b34a644094f19284bc2 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Wed, 3 Sep 2025 18:53:54 -0400
Subject: [PATCH 10/15] fix: golangci-lint issues

---
 ffi/firewood_test.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index a35103aba..6f4a959a9 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -5,13 +5,14 @@ package ffi
 
 import (
 	"bytes"
+	"crypto/rand"
 	"encoding/hex"
 	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
-	"sort"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -268,8 +269,8 @@ func sortKV(keys, vals [][]byte) error {
 	for i := range ord {
 		ord[i] = i
 	}
-	sort.Slice(ord, func(i, j int) bool {
-		return bytes.Compare(keys[ord[i]], keys[ord[j]]) < 0
+	slices.SortFunc(ord, func(i, j int) int {
+		return bytes.Compare(keys[i], keys[j])
 	})
 	perm := make([]int, n)
 	for dest, orig := range ord {

From 7a76f3dea8892423679d3ba505a61a6119294a69 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Mon, 8 Sep 2025 17:16:38 -0400
Subject: [PATCH 11/15] feat: update tests, update drop logic

---
 ffi/firewood_test.go | 122 +++++++++++++++++++++++++------------------
 ffi/iterator.go      |  13 +++--
 2 files changed, 80 insertions(+), 55 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index b66e0abff..5a9fd1145 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1107,6 +1107,7 @@ func TestIterEmptyDb(t *testing.T) {
 }
 
 type kvIter interface {
+	SetBatchSize(int)
 	Next() bool
 	Key() []byte
 	Value() []byte
@@ -1115,11 +1116,12 @@ type kvIter interface {
 }
 type borrowIter struct{ it *Iterator }
 
-func (b borrowIter) Next() bool    { return b.it.NextBorrowed() }
-func (b borrowIter) Key() []byte   { return b.it.Key() }
-func (b borrowIter) Value() []byte { return b.it.Value() }
-func (b borrowIter) Err() error    { return b.it.Err() }
-func (b borrowIter) Drop() error   { return b.it.Drop() }
+func (b borrowIter) SetBatchSize(int) { b.it.SetBatchSize(1) }
+func (b borrowIter) Next() bool       { return b.it.NextBorrowed() }
+func (b borrowIter) Key() []byte      { return b.it.Key() }
+func (b borrowIter) Value() []byte    { return b.it.Value() }
+func (b borrowIter) Err() error       { return b.it.Err() }
+func (b borrowIter) Drop() error      { return b.it.Drop() }
 
 func assertIteratorYields(r *require.Assertions, it kvIter, keys [][]byte, vals [][]byte) {
 	i := 0
@@ -1132,31 +1134,63 @@ func assertIteratorYields(r *require.Assertions, it kvIter, keys [][]byte, vals
 	r.NoError(it.Drop())
 }
 
-// Tests that basic iterator functionality works
+type iteratorConfigFn = func(it kvIter) kvIter
+
+var iterConfigs = map[string]iteratorConfigFn{
+	"Owned":    func(it kvIter) kvIter { return it },
+	"Borrowed": func(it kvIter) kvIter { return borrowIter{it: it.(*Iterator)} },
+	"Single": func(it kvIter) kvIter {
+		it.SetBatchSize(1)
+		return it
+	},
+	"Batched": func(it kvIter) kvIter {
+		it.SetBatchSize(100)
+		return it
+	},
+}
+
+func runIteratorTestForModes(parentT *testing.T, fn func(*testing.T, iteratorConfigFn), modes ...string) {
+	r := require.New(parentT)
+	testName := strings.Join(modes, "/")
+	parentT.Run(testName, func(t *testing.T) {
+		fn(t, func(it kvIter) kvIter {
+			for _, m := range modes {
+				config, ok := iterConfigs[m]
+				r.Truef(ok, "specified config mode %s does not exist", m)
+				it = config(it)
+			}
+			return it
+		})
+	})
+}
+
+func runIteratorTestForAllModes(parentT *testing.T, fn func(*testing.T, iteratorConfigFn)) {
+	for _, dataMode := range []string{"Owned", "Borrowed"} {
+		for _, batchMode := range []string{"Single", "Batched"} {
+			runIteratorTestForModes(parentT, fn, batchMode, dataMode)
+		}
+	}
+}
+
 func TestIter(t *testing.T) {
 	r := require.New(t)
 	db := newTestDatabase(t)
+	keys, vals := kvForBench(100)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
 
-	dataModes := []struct {
-		name     string
-		configFn func(it *Iterator) kvIter
-	}{
-		{"Owned", func(it *Iterator) kvIter { return it }},
-		{"Borrowed", func(it *Iterator) kvIter { return borrowIter{it: it} }},
-	}
+	runIteratorTestForAllModes(t, func(t *testing.T, cfn iteratorConfigFn) {
+		r := require.New(t)
+		it, err := db.Iter(nil)
+		r.NoError(err)
 
-	batchModes := []struct {
-		name     string
-		configFn func(it *Iterator)
-	}{
-		{"Single", func(it *Iterator) {
-			it.SetBatchSize(1)
-		}},
-		{"Batched", func(it *Iterator) {
-			it.SetBatchSize(100)
-		}},
-	}
+		assertIteratorYields(r, cfn(it), keys, vals)
+	})
+}
 
+func TestIterOnRoot(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
 	keys, vals := kvForBench(240)
 	firstRoot, err := db.Update(keys[:80], vals[:80])
 	r.NoError(err)
@@ -1165,27 +1199,7 @@ func TestIter(t *testing.T) {
 	thirdRoot, err := db.Update(keys[160:], vals[160:])
 	r.NoError(err)
 
-	runForAllModes := func(parentT *testing.T, name string, fn func(*testing.T, func(it *Iterator) kvIter)) {
-		for _, dataMode := range dataModes {
-			for _, batchMode := range batchModes {
-				parentT.Run(fmt.Sprintf("%s/%s/%s", name, dataMode.name, batchMode.name), func(t *testing.T) {
-					fn(t, func(it *Iterator) kvIter {
-						batchMode.configFn(it)
-						return dataMode.configFn(it)
-					})
-				})
-			}
-		}
-	}
-	runForAllModes(t, "Latest", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
-		r := require.New(t)
-		it, err := db.Iter(nil)
-		r.NoError(err)
-
-		assertIteratorYields(r, configureIterator(it), keys, vals)
-	})
-
-	runForAllModes(t, "OnRoot", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
+	runIteratorTestForAllModes(t, func(t *testing.T, cfn iteratorConfigFn) {
 		r := require.New(t)
 		h1, err := db.IterOnRoot(firstRoot, nil)
 		r.NoError(err)
@@ -1194,12 +1208,20 @@ func TestIter(t *testing.T) {
 		h3, err := db.IterOnRoot(thirdRoot, nil)
 		r.NoError(err)
 
-		assertIteratorYields(r, configureIterator(h1), keys[:80], vals[:80])
-		assertIteratorYields(r, configureIterator(h2), keys[:160], vals[:160])
-		assertIteratorYields(r, configureIterator(h3), keys, vals)
+		assertIteratorYields(r, cfn(h1), keys[:80], vals[:80])
+		assertIteratorYields(r, cfn(h2), keys[:160], vals[:160])
+		assertIteratorYields(r, cfn(h3), keys, vals)
 	})
+}
+
+func TestIterOnProposal(t *testing.T) {
+	r := require.New(t)
+	db := newTestDatabase(t)
+	keys, vals := kvForBench(240)
+	_, err := db.Update(keys, vals)
+	r.NoError(err)
 
-	runForAllModes(t, "OnProposal", func(t *testing.T, configureIterator func(it *Iterator) kvIter) {
+	runIteratorTestForAllModes(t, func(t *testing.T, cfn iteratorConfigFn) {
 		r := require.New(t)
 		updatedValues := make([][]byte, len(vals))
 		copy(updatedValues, vals)
@@ -1217,7 +1239,7 @@ func TestIter(t *testing.T) {
 		it, err := p.Iter(nil)
 		r.NoError(err)
 
-		assertIteratorYields(r, configureIterator(it), keys, updatedValues)
+		assertIteratorYields(r, cfn(it), keys, updatedValues)
 	})
 }
 
diff --git a/ffi/iterator.go b/ffi/iterator.go
index 7b24cc055..2cea68d5e 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -6,6 +6,7 @@ package ffi
 // #include <stdlib.h>
 // #include "firewood.h"
 import "C"
+import "errors"
 
 type Iterator struct {
 	// The database this iterator is associated with. We hold onto this to ensure
@@ -43,7 +44,7 @@ type Iterator struct {
 	currentResource interface{ Free() error }
 }
 
-func (it *Iterator) Release() error {
+func (it *Iterator) freeCurrentAllocation() error {
 	if it.currentResource == nil {
 		return nil
 	}
@@ -52,7 +53,7 @@ func (it *Iterator) Release() error {
 
 func (it *Iterator) nextInternal() error {
 	if len(it.loadedPairs) == 0 {
-		if e := it.Release(); e != nil {
+		if e := it.freeCurrentAllocation(); e != nil {
 			return e
 		}
 		if it.batchSize <= 1 {
@@ -98,7 +99,6 @@ func (it *Iterator) Next() bool {
 	k, v := it.currentPair.Copy()
 	it.currentKey = k
 	it.currentValue = v
-	it.err = nil
 	return true
 }
 
@@ -139,8 +139,11 @@ func (it *Iterator) Err() error {
 
 // Drop drops the iterator and releases the resources
 func (it *Iterator) Drop() error {
+	e1 := it.freeCurrentAllocation()
 	if it.handle != nil {
-		return getErrorFromVoidResult(C.fwd_free_iterator(it.handle))
+		return errors.Join(
+			e1,
+			getErrorFromVoidResult(C.fwd_free_iterator(it.handle)))
 	}
-	return nil
+	return e1
 }

From 05a4877e0f4f2966efdf0d2be4b747bb428df7a7 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Mon, 8 Sep 2025 20:19:47 -0400
Subject: [PATCH 12/15] feat: apply austin's comments

---
 ffi/iterator.go | 110 +++++++++++++++++++++++++++++++-----------------
 ffi/memory.go   |  51 +++++++---------------
 2 files changed, 87 insertions(+), 74 deletions(-)

diff --git a/ffi/iterator.go b/ffi/iterator.go
index 2cea68d5e..d22cd9959 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -6,7 +6,12 @@ package ffi
 // #include <stdlib.h>
 // #include "firewood.h"
 import "C"
-import "errors"
+
+import (
+	"errors"
+	"fmt"
+	"unsafe"
+)
 
 type Iterator struct {
 	// The database this iterator is associated with. We hold onto this to ensure
@@ -28,58 +33,59 @@ type Iterator struct {
 	// from the iterator, not yet consumed by user
 	loadedPairs []*ownedKeyValue
 
-	// currentPair is the current pair retrieved from the iterator
-	currentPair *ownedKeyValue
-
-	// currentKey is the current pair retrieved from the iterator
-	currentKey []byte
-
-	// currentValue is the current pair retrieved from the iterator
+	// current* fields correspond to the current cursor state
+	// nil/empty if not started or exhausted; refreshed on each Next().
+	currentPair  *ownedKeyValue
+	currentKey   []byte
 	currentValue []byte
+	// FFI resource for current pair or batch to free on advance or drop
+	currentResource interface{ free() error }
 
 	// err is the error from the iterator, if any
 	err error
-
-	// currentResource is a reference to a freeable resource to clean up
-	currentResource interface{ Free() error }
 }
 
 func (it *Iterator) freeCurrentAllocation() error {
 	if it.currentResource == nil {
 		return nil
 	}
-	return it.currentResource.Free()
+	e := it.currentResource.free()
+	it.currentResource = nil
+	return e
 }
 
 func (it *Iterator) nextInternal() error {
-	if len(it.loadedPairs) == 0 {
-		if e := it.freeCurrentAllocation(); e != nil {
+	if len(it.loadedPairs) > 0 {
+		it.currentPair, it.loadedPairs = it.loadedPairs[0], it.loadedPairs[1:]
+		return nil
+	}
+
+	// current resources should **only** be freed, on the next call to the FFI
+	// this is to make sure we don't invalidate a batch in between iteration
+	if e := it.freeCurrentAllocation(); e != nil {
+		return e
+	}
+	if it.batchSize <= 1 {
+		kv, e := getKeyValueFromResult(C.fwd_iter_next(it.handle))
+		if e != nil {
+			return e
+		}
+		it.currentPair = kv
+		it.currentResource = kv
+	} else {
+		batch, e := getKeyValueBatchFromResult(C.fwd_iter_next_n(it.handle, C.size_t(it.batchSize)))
+		if e != nil {
 			return e
 		}
-		if it.batchSize <= 1 {
-			kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle))
-			if e != nil {
-				return e
-			}
-			if kv != nil {
-				// kv is nil when done
-				it.loadedPairs = append(it.loadedPairs, kv)
-			}
-			it.currentResource = kv
+		pairs := batch.copy()
+		if len(pairs) > 0 {
+			it.currentPair, it.loadedPairs = pairs[0], pairs[1:]
 		} else {
-			batch, e := getKeyValueBatchFromKeyValueBatchResult(C.fwd_iter_next_n(it.handle, C.size_t(it.batchSize)))
-			if e != nil {
-				return e
-			}
-			it.loadedPairs = batch.Copied()
-			it.currentResource = batch
+			it.currentPair = nil
 		}
+		it.currentResource = batch
 	}
-	if len(it.loadedPairs) > 0 {
-		it.currentPair, it.loadedPairs = it.loadedPairs[0], it.loadedPairs[1:]
-	} else {
-		it.currentPair = nil
-	}
+
 	return nil
 }
 
@@ -96,15 +102,21 @@ func (it *Iterator) Next() bool {
 	if it.currentPair == nil || it.err != nil {
 		return false
 	}
-	k, v := it.currentPair.Copy()
+	k, v := it.currentPair.copy()
 	it.currentKey = k
 	it.currentValue = v
 	return true
 }
 
-// NextBorrowed retrieves the next item on the iterator similar to Next
-// the difference is that returned bytes in Key and Value are not copied
-// and will be freed on next call to Next or NextBorrowed
+// NextBorrowed is like Next, but Key and Value **borrow** rust-owned buffers.
+//
+// ⚠️ Lifetime: the returned slices are valid **only until** the next call to
+// Next, NextBorrowed, Close, or any operation that advances/invalidates the iterator.
+// They alias FFI-owned memory that will be **freed or reused** on the next advance.
+//
+// Do **not** retain, store, or modify these slices.
+// **Copy** or use Next if you need to keep them.
+// Misuse can read freed memory and cause corruption or crashes.
 func (it *Iterator) NextBorrowed() bool {
 	it.err = it.nextInternal()
 	if it.currentPair == nil || it.err != nil {
@@ -147,3 +159,23 @@ func (it *Iterator) Drop() error {
 	}
 	return e1
 }
+
+// getIteratorFromIteratorResult converts a C.IteratorResult to an Iterator or error.
+func getIteratorFromIteratorResult(result C.IteratorResult, db *Database) (*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{
+			db:     db,
+			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)
+	}
+}
diff --git a/ffi/memory.go b/ffi/memory.go
index e2703183b..0184616fe 100644
--- a/ffi/memory.go
+++ b/ffi/memory.go
@@ -313,11 +313,11 @@ type ownedKeyValueBatch struct {
 	owned C.OwnedKeyValueBatch
 }
 
-func (b *ownedKeyValueBatch) Copied() []*ownedKeyValue {
+func (b *ownedKeyValueBatch) copy() []*ownedKeyValue {
 	if b.owned.ptr == nil {
 		return nil
 	}
-	borrowed := b.Borrow()
+	borrowed := b.borrow()
 	copied := make([]*ownedKeyValue, len(borrowed))
 	for i, borrow := range borrowed {
 		copied[i] = newOwnedKeyValue(borrow)
@@ -325,7 +325,7 @@ func (b *ownedKeyValueBatch) Copied() []*ownedKeyValue {
 	return copied
 }
 
-func (b *ownedKeyValueBatch) Borrow() []C.OwnedKeyValuePair {
+func (b *ownedKeyValueBatch) borrow() []C.OwnedKeyValuePair {
 	if b.owned.ptr == nil {
 		return nil
 	}
@@ -333,8 +333,9 @@ func (b *ownedKeyValueBatch) Borrow() []C.OwnedKeyValuePair {
 	return unsafe.Slice((*C.OwnedKeyValuePair)(unsafe.Pointer(b.owned.ptr)), b.owned.len)
 }
 
-func (b *ownedKeyValueBatch) Free() error {
-	if b.owned.ptr == nil {
+func (b *ownedKeyValueBatch) free() error {
+	if b == nil || b.owned.ptr == nil {
+		// we want ownedKeyValueBatch to be typed-nil safe
 		return nil
 	}
 
@@ -362,18 +363,18 @@ type ownedKeyValue struct {
 	value *ownedBytes
 }
 
-func (kv *ownedKeyValue) Copy() ([]byte, []byte) {
+func (kv *ownedKeyValue) copy() ([]byte, []byte) {
 	key := kv.key.CopiedBytes()
 	value := kv.value.CopiedBytes()
 	return key, value
 }
 
-func (kv *ownedKeyValue) Free() error {
-	err := kv.key.Free()
-	if err != nil {
-		return fmt.Errorf("%w: %w", errFreeingValue, err)
+func (kv *ownedKeyValue) free() error {
+	if kv == nil {
+		// we want ownedKeyValue to be typed-nil safe
+		return nil
 	}
-	err = kv.value.Free()
+	err := errors.Join(kv.key.Free(), kv.value.Free())
 	if err != nil {
 		return fmt.Errorf("%w: %w", errFreeingValue, err)
 	}
@@ -391,12 +392,12 @@ func newOwnedKeyValue(owned C.OwnedKeyValuePair) *ownedKeyValue {
 	}
 }
 
-// getKeyValueFromKeyValueResult converts a C.KeyValueResult to a key value pair or error.
+// getKeyValueFromResult 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) {
+func getKeyValueFromResult(result C.KeyValueResult) (*ownedKeyValue, error) {
 	switch result.tag {
 	case C.KeyValueResult_NullHandlePointer:
 		return nil, errDBClosed
@@ -417,12 +418,12 @@ func getKeyValueFromKeyValueResult(result C.KeyValueResult) (*ownedKeyValue, err
 	}
 }
 
-// getKeyValueBatchFromKeyValueBatchResult converts a C.KeyValueResult to a key value pair or error.
+// getKeyValueBatchFromResult converts a C.KeyValueBatchResult to a key value batch or error.
 //
 // It returns nil, nil if the result is None.
 // It returns a *ownedKeyValueBatch, nil if the result is Some.
 // It returns an error if the result is an error.
-func getKeyValueBatchFromKeyValueBatchResult(result C.KeyValueBatchResult) (*ownedKeyValueBatch, error) {
+func getKeyValueBatchFromResult(result C.KeyValueBatchResult) (*ownedKeyValueBatch, error) {
 	switch result.tag {
 	case C.KeyValueBatchResult_NullHandlePointer:
 		return nil, errDBClosed
@@ -487,23 +488,3 @@ func getChangeProofFromChangeProofResult(result C.ChangeProofResult) (*ChangePro
 		return nil, fmt.Errorf("unknown C.ChangeProofResult tag: %d", result.tag)
 	}
 }
-
-// getIteratorFromIteratorResult converts a C.IteratorResult to an Iterator or error.
-func getIteratorFromIteratorResult(result C.IteratorResult, db *Database) (*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{
-			db:     db,
-			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)
-	}
-}

From 7009ea7c73d683b91e81f880e30e1739ae8bbaf7 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Tue, 9 Sep 2025 14:51:56 -0400
Subject: [PATCH 13/15] fix: apply austin's suggestions

---
 ffi/firewood_test.go | 46 ++++++++++++++------------------------------
 ffi/iterator.go      | 12 +++++++-----
 2 files changed, 21 insertions(+), 37 deletions(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index 8f4022398..98e2519ae 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -5,7 +5,6 @@ package ffi
 
 import (
 	"bytes"
-	"crypto/rand"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -247,15 +246,10 @@ func kvForTest(num int) ([][]byte, [][]byte) {
 		keys[i] = keyForTest(i)
 		vals[i] = valForTest(i)
 	}
+	_ = sortKV(keys, vals)
 	return keys, vals
 }
 
-func randomBytes(n int) []byte {
-	b := make([]byte, n)
-	_, _ = rand.Read(b)
-	return b
-}
-
 // sortKV sorts keys lexicographically and keeps vals paired.
 func sortKV(keys, vals [][]byte) error {
 	if len(keys) != len(vals) {
@@ -287,18 +281,6 @@ func sortKV(keys, vals [][]byte) error {
 	return nil
 }
 
-func kvForBench(num int) ([][]byte, [][]byte) {
-	keys := make([][]byte, num)
-	vals := make([][]byte, num)
-
-	for i := range keys {
-		keys[i] = randomBytes(32)
-		vals[i] = randomBytes(128)
-	}
-	_ = sortKV(keys, vals)
-	return keys, vals
-}
-
 // Tests that 100 key-value pairs can be inserted and retrieved.
 // This happens in three ways:
 // 1. By calling [Database.Propose] and then [Proposal.Commit].
@@ -1120,12 +1102,12 @@ type kvIter interface {
 }
 type borrowIter struct{ it *Iterator }
 
-func (b borrowIter) SetBatchSize(int) { b.it.SetBatchSize(1) }
-func (b borrowIter) Next() bool       { return b.it.NextBorrowed() }
-func (b borrowIter) Key() []byte      { return b.it.Key() }
-func (b borrowIter) Value() []byte    { return b.it.Value() }
-func (b borrowIter) Err() error       { return b.it.Err() }
-func (b borrowIter) Drop() error      { return b.it.Drop() }
+func (b *borrowIter) SetBatchSize(batchSize int) { b.it.SetBatchSize(batchSize) }
+func (b *borrowIter) Next() bool                 { return b.it.NextBorrowed() }
+func (b *borrowIter) Key() []byte                { return b.it.Key() }
+func (b *borrowIter) Value() []byte              { return b.it.Value() }
+func (b *borrowIter) Err() error                 { return b.it.Err() }
+func (b *borrowIter) Drop() error                { return b.it.Drop() }
 
 func assertIteratorYields(r *require.Assertions, it kvIter, keys [][]byte, vals [][]byte) {
 	i := 0
@@ -1141,7 +1123,7 @@ type iteratorConfigFn = func(it kvIter) kvIter
 
 var iterConfigs = map[string]iteratorConfigFn{
 	"Owned":    func(it kvIter) kvIter { return it },
-	"Borrowed": func(it kvIter) kvIter { return borrowIter{it: it.(*Iterator)} },
+	"Borrowed": func(it kvIter) kvIter { return &borrowIter{it: it.(*Iterator)} },
 	"Single": func(it kvIter) kvIter {
 		it.SetBatchSize(1)
 		return it
@@ -1152,10 +1134,10 @@ var iterConfigs = map[string]iteratorConfigFn{
 	},
 }
 
-func runIteratorTestForModes(parentT *testing.T, fn func(*testing.T, iteratorConfigFn), modes ...string) {
-	r := require.New(parentT)
+func runIteratorTestForModes(t *testing.T, fn func(*testing.T, iteratorConfigFn), modes ...string) {
+	r := require.New(t)
 	testName := strings.Join(modes, "/")
-	parentT.Run(testName, func(t *testing.T) {
+	t.Run(testName, func(t *testing.T) {
 		fn(t, func(it kvIter) kvIter {
 			for _, m := range modes {
 				config, ok := iterConfigs[m]
@@ -1178,7 +1160,7 @@ func runIteratorTestForAllModes(parentT *testing.T, fn func(*testing.T, iterator
 func TestIter(t *testing.T) {
 	r := require.New(t)
 	db := newTestDatabase(t)
-	keys, vals := kvForBench(100)
+	keys, vals := kvForTest(100)
 	_, err := db.Update(keys, vals)
 	r.NoError(err)
 
@@ -1197,7 +1179,7 @@ func TestIter(t *testing.T) {
 func TestIterOnRoot(t *testing.T) {
 	r := require.New(t)
 	db := newTestDatabase(t)
-	keys, vals := kvForBench(240)
+	keys, vals := kvForTest(240)
 	firstRoot, err := db.Update(keys[:80], vals[:80])
 	r.NoError(err)
 	secondRoot, err := db.Update(keys[80:160], vals[80:160])
@@ -1234,7 +1216,7 @@ func TestIterOnRoot(t *testing.T) {
 func TestIterOnProposal(t *testing.T) {
 	r := require.New(t)
 	db := newTestDatabase(t)
-	keys, vals := kvForBench(240)
+	keys, vals := kvForTest(240)
 	_, err := db.Update(keys, vals)
 	r.NoError(err)
 
diff --git a/ffi/iterator.go b/ffi/iterator.go
index d22cd9959..d222900f5 100644
--- a/ffi/iterator.go
+++ b/ffi/iterator.go
@@ -110,13 +110,13 @@ func (it *Iterator) Next() bool {
 
 // NextBorrowed is like Next, but Key and Value **borrow** rust-owned buffers.
 //
-// ⚠️ Lifetime: the returned slices are valid **only until** the next call to
+// Lifetime: the returned slices are valid **only until** the next call to
 // Next, NextBorrowed, Close, or any operation that advances/invalidates the iterator.
 // They alias FFI-owned memory that will be **freed or reused** on the next advance.
 //
 // Do **not** retain, store, or modify these slices.
 // **Copy** or use Next if you need to keep them.
-// Misuse can read freed memory and cause corruption or crashes.
+// Misuse can read freed memory and cause undefined behavior.
 func (it *Iterator) NextBorrowed() bool {
 	it.err = it.nextInternal()
 	if it.currentPair == nil || it.err != nil {
@@ -151,13 +151,15 @@ func (it *Iterator) Err() error {
 
 // Drop drops the iterator and releases the resources
 func (it *Iterator) Drop() error {
-	e1 := it.freeCurrentAllocation()
+	err := it.freeCurrentAllocation()
 	if it.handle != nil {
+		// Always free the iterator even if releasing the current KV/batch failed.
+		// The iterator holds a NodeStore ref that must be dropped.
 		return errors.Join(
-			e1,
+			err,
 			getErrorFromVoidResult(C.fwd_free_iterator(it.handle)))
 	}
-	return e1
+	return err
 }
 
 // getIteratorFromIteratorResult converts a C.IteratorResult to an Iterator or error.

From cfaf9a427532ca728f872c5b029cf4aef2b90dfc Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Tue, 9 Sep 2025 15:11:18 -0400
Subject: [PATCH 14/15] fix: require

---
 ffi/firewood_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go
index 98e2519ae..c1b1ed8c5 100644
--- a/ffi/firewood_test.go
+++ b/ffi/firewood_test.go
@@ -1135,9 +1135,9 @@ var iterConfigs = map[string]iteratorConfigFn{
 }
 
 func runIteratorTestForModes(t *testing.T, fn func(*testing.T, iteratorConfigFn), modes ...string) {
-	r := require.New(t)
 	testName := strings.Join(modes, "/")
 	t.Run(testName, func(t *testing.T) {
+		r := require.New(t)
 		fn(t, func(it kvIter) kvIter {
 			for _, m := range modes {
 				config, ok := iterConfigs[m]

From a817c92bd937c5fae1eb8cb1d94bbf59cc91f9e3 Mon Sep 17 00:00:00 2001
From: Amin Rezaei <amin.rezaei@avalabs.org>
Date: Thu, 2 Oct 2025 05:33:14 -0400
Subject: [PATCH 15/15] fmt: ...

---
 ffi/firewood.h   | 2 +-
 ffi/src/lib.rs   | 2 +-
 ffi/src/value.rs | 5 +++--
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/ffi/firewood.h b/ffi/firewood.h
index d47575a81..b5fe634af 100644
--- a/ffi/firewood.h
+++ b/ffi/firewood.h
@@ -1663,7 +1663,7 @@ struct KeyValueResult fwd_iter_next(struct IteratorHandle *handle);
  *
  * # Arguments
  *
- * * `handle` - The iterator handle returned by [`fwd_iter_on_root`] or
+ * * `handle` - The iterator handle returned by [`fwd_iter_on_revision`] or
  *   [`fwd_iter_on_proposal`].
  *
  * # Returns
diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs
index 3a5099384..f6279b4a9 100644
--- a/ffi/src/lib.rs
+++ b/ffi/src/lib.rs
@@ -199,7 +199,7 @@ pub unsafe extern "C" fn fwd_iter_next(handle: Option<&mut IteratorHandle<'_>>)
 ///
 /// # Arguments
 ///
-/// * `handle` - The iterator handle returned by [`fwd_iter_on_root`] or
+/// * `handle` - The iterator handle returned by [`fwd_iter_on_revision`] or
 ///   [`fwd_iter_on_proposal`].
 ///
 /// # Returns
diff --git a/ffi/src/value.rs b/ffi/src/value.rs
index 031ca7640..c80f2a610 100644
--- a/ffi/src/value.rs
+++ b/ffi/src/value.rs
@@ -15,8 +15,9 @@ pub use self::kvp::{KeyValuePair, OwnedKeyValueBatch, OwnedKeyValuePair};
 pub use self::owned::{OwnedBytes, OwnedSlice};
 pub(crate) use self::results::{CResult, NullHandleResult};
 pub use self::results::{
-    ChangeProofResult, HandleResult, HashResult, IteratorResult, KeyValueBatchResult, KeyValueResult,
-    NextKeyRangeResult, ProposalResult, RangeProofResult, RevisionResult, ValueResult, VoidResult,
+    ChangeProofResult, HandleResult, HashResult, IteratorResult, KeyValueBatchResult,
+    KeyValueResult, NextKeyRangeResult, ProposalResult, RangeProofResult, RevisionResult,
+    ValueResult, VoidResult,
 };
 
 /// Maybe is a C-compatible optional type using a tagged union pattern.