From 4fcd1658befcbfe51fe5070e6f1b8cb40b24645d Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sun, 18 Oct 2026 02:21:02 -0700
Subject: [PATCH] Added HnswNextBuffer function [skip ci]

---
 src/hnsw.h      |  7 ++---
 src/hnswutils.c | 70 ++++++++++++++++++++++++++++++++-----------------
 2 files changed, 48 insertions(+), 29 deletions(-)

diff --git a/src/hnsw.h b/src/hnsw.h
index fcc0858..bff4df8 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -417,24 +417,21 @@ typedef union
 	ItemPointerData indextid;
 }			HnswUnvisited;
 
-#ifdef HNSW_USE_READ_STREAM
 typedef struct HnswReadStreamData
 {
 	HnswUnvisited *unvisited;
 	int			unvisitedLength;
 	int			visited;
 }			HnswReadStreamData;
-#endif
 
 typedef struct HnswSearchState
 {
-#ifdef HNSW_USE_READ_STREAM
 	HnswReadStreamData streamData;
+#ifdef HNSW_USE_READ_STREAM
 	ReadStream *stream;
 #elif defined(HNSW_USE_PREFETCH)
 	int			maxPrefetch;
-#else
-	void	   *unused;
+	Relation	index;
 #endif
 }			HnswSearchState;
 
diff --git a/src/hnswutils.c b/src/hnswutils.c
index b9c9fe6..e0ea59d 100644
--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -826,10 +826,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
 			unvisited[(*unvisitedLength)++].indextid = *indextid;
 	}
 
-#ifdef HNSW_USE_READ_STREAM
+#if defined(HNSW_USE_READ_STREAM) || defined(HNSW_USE_PREFETCH)
 	searchState->streamData.unvisited = unvisited;
 	searchState->streamData.unvisitedLength = *unvisitedLength;
 	searchState->streamData.visited = 0;
+#endif
+
+#ifdef HNSW_USE_READ_STREAM
 	read_stream_resume(searchState->stream);
 #elif defined(HNSW_USE_PREFETCH)
 	for (int i = 0, imax = Min(searchState->maxPrefetch, *unvisitedLength); i < imax; i++)
@@ -884,6 +887,8 @@ HnswInitSearchState(HnswSearchState * searchState, Relation index, bool maintena
 			searchState->maxPrefetch = get_tablespace_maintenance_io_concurrency(index->rd_rel->reltablespace);
 		else
 			searchState->maxPrefetch = get_tablespace_io_concurrency(index->rd_rel->reltablespace);
+
+		searchState->index = index;
 	}
 #endif
 }
@@ -900,6 +905,45 @@ HnswEndSearchState(HnswSearchState * searchState)
 #endif
 }
 
+/*
+ * Read next buffer
+ */
+static inline Buffer
+HnswNextBuffer(HnswSearchState * searchState, OffsetNumber *offno)
+{
+	Buffer		buf;
+
+#ifdef HNSW_USE_READ_STREAM
+	void	   *offnoPtr;
+
+	buf = read_stream_next_buffer(searchState->stream, &offnoPtr);
+
+	if (BufferIsValid(buf))
+		*offno = *((OffsetNumber *) offnoPtr);
+
+	return buf;
+#else
+	HnswReadStreamData *streamData = &searchState->streamData;
+	int			i = streamData->visited;
+	HnswUnvisited *uv;
+
+	if (streamData->visited == streamData->unvisitedLength)
+		return InvalidBuffer;
+
+	uv = &streamData->unvisited[streamData->visited++];
+	buf = ReadBuffer(searchState->index, ItemPointerGetBlockNumber(&uv->indextid));
+	*offno = ItemPointerGetOffsetNumber(&uv->indextid);
+
+#ifdef HNSW_USE_PREFETCH
+	/* Prefetch another buffer if possible */
+	if (i + searchState->maxPrefetch < streamData->unvisitedLength)
+		PrefetchBuffer(searchState->index, MAIN_FORKNUM, ItemPointerGetBlockNumber(&streamData->unvisited[i + searchState->maxPrefetch].indextid));
+#endif
+
+	return buf;
+#endif
+}
+
 /*
  * Algorithm 2 from paper
  */
@@ -1009,36 +1053,14 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
 				Buffer		buf;
 				OffsetNumber offno;
 
-#ifdef HNSW_USE_READ_STREAM
-				void	   *offnoPtr;
-
-				buf = read_stream_next_buffer(searchState->stream, &offnoPtr);
-
+				buf = HnswNextBuffer(searchState, &offno);
 				if (!BufferIsValid(buf))
 					break;
 
-				offno = *((OffsetNumber *) offnoPtr);
-#else
-				ItemPointer indextid;
-
-				if (i == unvisitedLength)
-					break;
-
-				indextid = &unvisited[i].indextid;
-				buf = ReadBuffer(index, ItemPointerGetBlockNumber(indextid));
-				offno = ItemPointerGetOffsetNumber(indextid);
-#endif
-
 				/* Avoid any allocations if not adding */
 				eElement = NULL;
 				HnswLoadElementImpl(buf, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
 
-#ifdef HNSW_USE_PREFETCH
-				/* Prefetch another buffer if possible */
-				if (i + searchState->maxPrefetch < unvisitedLength)
-					PrefetchBuffer(index, MAIN_FORKNUM, ItemPointerGetBlockNumber(&unvisited[i + searchState->maxPrefetch].indextid));
-#endif
-
 				if (eElement == NULL)
 					continue;
 			}