From 414021a0683a7992ee8e3a226fbd0ac07b639ee2 Mon Sep 17 00:00:00 2001
From: Csaba Kiraly <csaba.kiraly@gmail.com>
Date: Thu, 10 Apr 2025 09:35:44 +0200
Subject: [PATCH] p2p/enode: add BufferIter to support prefetching

BufferIter wraps an iterator and prefetches up to a given
number of nodes from it.

Signed-off-by: Csaba Kiraly <csaba.kiraly@gmail.com>
---
 p2p/enode/iter.go | 45 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/p2p/enode/iter.go b/p2p/enode/iter.go
index 3735a2ddcf..16f55a7fd9 100644
--- a/p2p/enode/iter.go
+++ b/p2p/enode/iter.go
@@ -232,6 +232,51 @@ func (f *AsyncFilterIter) Close() {
 	})
 }
 
+// BufferIter wraps an iterator and buffers the nodes it returns.
+// The buffer is pre-filled with the given size from the wrapped iterator.
+type BufferIter struct {
+	it     Iterator
+	buffer chan *Node
+	head   *Node
+}
+
+// NewBufferIter creates a new pre-fetch buffer.
+func NewBufferIter(it Iterator, size int) *BufferIter {
+	b := BufferIter{
+		it:     it,
+		buffer: make(chan *Node, size),
+	}
+
+	go func() {
+		defer close(b.buffer)
+		for b.it.Next() {
+			b.buffer <- b.it.Node()
+		}
+	}()
+	return &b
+}
+
+func (b *BufferIter) Next() bool {
+	b.head = <-b.buffer
+	return b.head != nil
+}
+
+func (b *BufferIter) Node() *Node {
+	return b.head
+}
+
+func (b *BufferIter) Close() {
+	b.it.Close()
+	// Wait for the buffer to be consumed.
+	for range b.buffer {
+	}
+	// Close the buffer channel.
+	// close(b.buffer)
+	b.buffer = nil
+	b.head = nil
+	b.it = nil
+}
+
 // FairMix aggregates multiple node iterators. The mixer itself is an iterator which ends
 // only when Close is called. Source iterators added via AddSource are removed from the
 // mix when they end.