fix:hnsw heap polarity (#550)

* fix:hnsw heap polarity

Signed-off-by: cryo <[email protected]>

* fix: go-lint failures

Signed-off-by: cryo <[email protected]>

---------

Signed-off-by: cryo <[email protected]>

Author: cryo-zd

src/semantic-router/pkg/cache/inmemory_cache.go

@@ -771,8 +771,8 @@ func (h *HNSWIndex) searchKNN(queryEmbedding []float32, k, ef int, entries []Cac
 // searchLayer searches for nearest neighbors at a specific layer
 func (h *HNSWIndex) searchLayer(queryEmbedding []float32, entryPoint, ef, layer int, entries []CacheEntry) []int {
 	visited := make(map[int]bool)
-	candidates := newMaxHeap()
-	results := newMinHeap()
+	candidates := newMinHeap() // set of candidates, explore closest candidate first
+	results := newMaxHeap()    // dynamic list of found nearest neighbors, track current frontier, worst distance on top
 
 	// Calculate distance to entry point
 	if entryPoint >= 0 && entryPoint < len(entries) {
@@ -785,11 +785,9 @@ func (h *HNSWIndex) searchLayer(queryEmbedding []float32, entryPoint, ef, layer
 	for candidates.len() > 0 {
 		currentIdx, currentDist := candidates.pop()
 
-		if results.len() > 0 {
-			worstDist := results.peekDist()
-			if currentDist > worstDist {
-				break
-			}
+		// If we have enough results and the current distance is worse than the worst in results, we can stop
+		if results.len() > 0 && currentDist > results.peekDist() {
+			break
 		}
 
 		// Fast O(1) lookup using nodeIndex map
@@ -881,25 +879,10 @@ func (h *minHeap) pop() (int, float32) {
 	return result.index, result.dist
 }
 
-func (h *minHeap) peekDist() float32 {
-	if len(h.data) == 0 {
-		return math.MaxFloat32
-	}
-	return h.data[0].dist
-}
-
 func (h *minHeap) len() int {
 	return len(h.data)
 }
 
-func (h *minHeap) items() []int {
-	result := make([]int, len(h.data))
-	for i, item := range h.data {
-		result[i] = item.index
-	}
-	return result
-}
-
 func (h *minHeap) bubbleUp(i int) {
 	for i > 0 {
 		parent := (i - 1) / 2
@@ -961,6 +944,21 @@ func (h *maxHeap) len() int {
 	return len(h.data)
 }
 
+func (h *maxHeap) peekDist() float32 {
+	if len(h.data) == 0 {
+		return math.MaxFloat32
+	}
+	return h.data[0].dist
+}
+
+func (h *maxHeap) items() []int {
+	result := make([]int, len(h.data))
+	for i, item := range h.data {
+		result[i] = item.index
+	}
+	return result
+}
+
 func (h *maxHeap) bubbleUp(i int) {
 	for i > 0 {
 		parent := (i - 1) / 2

src/semantic-router/pkg/cache/inmemory_hnsw_internal_test.go

@@ -0,0 +1,118 @@
+package cache
+
+import (
+	"slices"
+	"testing"
+)
+
+func TestSearchLayerHeapManagement(t *testing.T) {
+	t.Run("retains the closest neighbor when ef is saturated", func(t *testing.T) {
+		// Regression fixture: with the previous max-heap candidates/min-heap results
+		// mix, trimming to ef would evict the best element instead of the worst.
+		queryEmbedding := []float32{1.0}
+
+		entries := []CacheEntry{
+			{Embedding: []float32{0.1}}, // entry point has low similarity
+			{Embedding: []float32{1.0}}, // neighbor is the true nearest
+		}
+
+		entryNode := &HNSWNode{
+			entryIndex: 0,
+			neighbors: map[int][]int{
+				0: {1},
+			},
+			maxLayer: 0,
+		}
+
+		neighborNode := &HNSWNode{
+			entryIndex: 1,
+			neighbors: map[int][]int{
+				0: {0},
+			},
+			maxLayer: 0,
+		}
+
+		index := &HNSWIndex{
+			nodes: []*HNSWNode{entryNode, neighborNode},
+			nodeIndex: map[int]*HNSWNode{
+				0: entryNode,
+				1: neighborNode,
+			},
+			entryPoint:     0,
+			maxLayer:       0,
+			efConstruction: 2,
+			M:              1,
+			Mmax:           1,
+			Mmax0:          2,
+			ml:             1,
+		}
+
+		results := index.searchLayer(queryEmbedding, index.entryPoint, 1, 0, entries)
+
+		if !slices.Contains(results, 1) {
+			t.Fatalf("expected results to contain best neighbor 1, got %v", results)
+		}
+		if slices.Contains(results, 0) {
+			t.Fatalf("expected results to drop entry point 0 once ef trimmed, got %v", results)
+		}
+	})
+
+	t.Run("continues exploring even when next candidate looks worse", func(t *testing.T) {
+		// Regression fixture: the break condition used the wrong polarity so the
+		// search stopped before expanding the intermediate (worse) vertex, making
+		// the actual best neighbor unreachable.
+		queryEmbedding := []float32{1.0}
+
+		entries := []CacheEntry{
+			{Embedding: []float32{0.2}},  // entry point
+			{Embedding: []float32{0.05}}, // intermediate node with poor similarity
+			{Embedding: []float32{1.0}},  // hidden best match
+		}
+
+		entryNode := &HNSWNode{
+			entryIndex: 0,
+			neighbors: map[int][]int{
+				0: {1},
+			},
+			maxLayer: 0,
+		}
+
+		intermediateNode := &HNSWNode{
+			entryIndex: 1,
+			neighbors: map[int][]int{
+				0: {0, 2},
+			},
+			maxLayer: 0,
+		}
+
+		bestNode := &HNSWNode{
+			entryIndex: 2,
+			neighbors: map[int][]int{
+				0: {1},
+			},
+			maxLayer: 0,
+		}
+
+		index := &HNSWIndex{
+			nodes: []*HNSWNode{entryNode, intermediateNode, bestNode},
+			nodeIndex: map[int]*HNSWNode{
+				0: entryNode,
+				1: intermediateNode,
+				2: bestNode,
+			},
+			entryPoint:     0,
+			maxLayer:       0,
+			efConstruction: 2,
+			M:              1,
+			Mmax:           1,
+			Mmax0:          2,
+			ml:             1,
+		}
+
+		results := index.searchLayer(queryEmbedding, index.entryPoint, 2, 0, entries)
+
+		if !slices.Contains(results, 2) {
+			t.Fatalf("expected results to reach best neighbor 2 via intermediate node, got %v", results)
+		}
+	})
+}