From 8eb8cdf0f3052af85587b4bf420acfdc558940c3 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sun, 29 Sep 2024 19:44:11 -0700
Subject: [PATCH] Moved insert-specific code to hnswinsert.c

---
 src/hnswinsert.c | 38 +++++++++++++++++++++++++++++++++-
 src/hnswutils.c  | 53 ++++++++----------------------------------------
 2 files changed, 46 insertions(+), 45 deletions(-)

diff --git a/src/hnswinsert.c b/src/hnswinsert.c
index 712eaa1..106428b 100644
--- a/src/hnswinsert.c
+++ b/src/hnswinsert.c
@@ -364,6 +364,32 @@ HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
 	return neighbors;
 }
 
+/*
+ * Load elements for insert
+ */
+static void
+LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation index, FmgrInfo *procinfo, Oid collation)
+{
+	char	   *base = NULL;
+
+	for (int i = 0; i < neighbors->length; i++)
+	{
+		HnswCandidate *hc = &neighbors->items[i];
+		HnswElement element = HnswPtrAccess(base, hc->element);
+		double		distance;
+
+		HnswLoadElement(element, &distance, &q, index, procinfo, collation, true, NULL);
+		hc->distance = distance;
+
+		/* Prune element if being deleted */
+		if (element->heaptidsLength == 0)
+		{
+			*idx = i;
+			break;
+		}
+	}
+}
+
 /*
  * Check if connection already exists
  */
@@ -426,7 +452,17 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
 			 */
 
 			/* Select neighbors */
-			HnswUpdateConnection(NULL, e, hc, neighborNeighbors, lm, &idx, index, procinfo, collation);
+			if (neighbors->length < lm)
+				idx = -2;
+			else
+			{
+				Datum		q = HnswGetValue(base, neighborElement);
+
+				LoadElementsForInsert(neighborNeighbors, q, &idx, index, procinfo, collation);
+
+				if (idx == -1)
+					HnswUpdateConnection(base, e, hc, neighborNeighbors, lm, &idx, index, procinfo, collation);
+			}
 
 			/* New element was not selected as a neighbor */
 			if (idx == -1)
diff --git a/src/hnswutils.c b/src/hnswutils.c
index 019e230..8ac9c08 100644
--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -1074,39 +1074,12 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
 		a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
 }
 
-/*
- * Load elements for insert
- */
-static void
-LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, HnswCandidate * *pruned, Relation index, FmgrInfo *procinfo, Oid collation)
-{
-	char	   *base = NULL;
-
-	for (int i = 0; i < neighbors->length; i++)
-	{
-		HnswCandidate *hc = &neighbors->items[i];
-		HnswElement element = HnswPtrAccess(base, hc->element);
-		double		distance;
-
-		HnswLoadElement(element, &distance, &q, index, procinfo, collation, true, NULL);
-		hc->distance = distance;
-
-		/* Prune element if being deleted */
-		if (element->heaptidsLength == 0)
-		{
-			*pruned = &neighbors->items[i];
-			break;
-		}
-	}
-}
-
 /*
  * Update connections
  */
 void
 HnswUpdateConnection(char *base, HnswElement newElement, HnswCandidate * hc, HnswNeighborArray * neighbors, int lm, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
 {
-	HnswElement hce = HnswPtrAccess(base, hc->element);
 	HnswCandidate newHc;
 
 	HnswPtrStore(base, newHc.element, newElement);
@@ -1123,27 +1096,19 @@ HnswUpdateConnection(char *base, HnswElement newElement, HnswCandidate * hc, Hns
 	else
 	{
 		/* Shrink connections */
+		List	   *c = NIL;
 		HnswCandidate *pruned = NULL;
 
-		/* Load elements on insert */
-		if (index != NULL)
-			LoadElementsForInsert(neighbors, HnswGetValue(base, hce), &pruned, index, procinfo, collation);
+		/* Add candidates */
+		for (int i = 0; i < neighbors->length; i++)
+			c = lappend(c, &neighbors->items[i]);
+		c = lappend(c, &newHc);
 
+		SelectNeighbors(base, c, lm, procinfo, collation, neighbors, &newHc, &pruned, true);
+
+		/* Should not happen */
 		if (pruned == NULL)
-		{
-			List	   *c = NIL;
-
-			/* Add candidates */
-			for (int i = 0; i < neighbors->length; i++)
-				c = lappend(c, &neighbors->items[i]);
-			c = lappend(c, &newHc);
-
-			SelectNeighbors(base, c, lm, procinfo, collation, neighbors, &newHc, &pruned, true);
-
-			/* Should not happen */
-			if (pruned == NULL)
-				return;
-		}
+			return;
 
 		/* Find and replace the pruned element */
 		for (int i = 0; i < neighbors->length; i++)