Maven/internal/store/memory_search_test.go

package store

import (
	"context"
	"fmt"
	"math/rand"
	"sort"
	"testing"

	"github.com/kami/maven/internal/memory"
)

// naiveSearch is the implementation Search replaced: score every row into a
// slice, sort the whole slice, truncate. It stays in the test file as the
// reference the bounded-heap version is judged against, because "recall must
// not change" is a claim about output, not about the code that produces it.
func naiveSearch(t *testing.T, m *MemoryStore, vec []float32, topK int) []memory.Result {
	t.Helper()
	rows, err := m.db.QueryContext(context.Background(),
		`SELECT id, vec FROM memory_vectors WHERE id NOT LIKE ? ESCAPE '\'`,
		escapeLike(memory.NonRecallPrefix)+"%")
	if err != nil {
		t.Fatalf("naive scan: %v", err)
	}
	defer rows.Close()
	var out []memory.Result
	for rows.Next() {
		var id string
		var blob []byte
		if err := rows.Scan(&id, &blob); err != nil {
			t.Fatalf("naive row: %v", err)
		}
		out = append(out, memory.Result{ID: id, Score: dot(vec, decodeVec(blob))})
	}
	if err := rows.Err(); err != nil {
		t.Fatalf("naive rows: %v", err)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Score > out[j].Score })
	if topK < len(out) {
		out = out[:topK]
	}
	return out
}

// TestMemoryStoreSearchMatchesNaive is the constraint on V-643: the bounded
// heap must return exactly what a full scan and sort returned. Distinct random
// vectors, so no two scores tie and the ranking is total — a mismatch here is
// arithmetic or heap logic, not a tie-break difference.
func TestMemoryStoreSearchMatchesNaive(t *testing.T) {
	ctx := context.Background()
	m := newMemTestStore(t).VectorMemory()

	rng := rand.New(rand.NewSource(7))
	const rows, dim = 500, 64
	for i := 0; i < rows; i++ {
		if err := m.Insert(ctx, fmt.Sprintf("n%d", i), randUnitVec(rng, dim), map[string]string{
			"text": fmt.Sprintf("note %d", i),
		}); err != nil {
			t.Fatalf("Insert %d: %v", i, err)
		}
	}

	for _, topK := range []int{1, 3, 10, 50, rows, rows + 100} {
		q := randUnitVec(rng, dim)
		got, err := m.Search(ctx, q, topK)
		if err != nil {
			t.Fatalf("Search topK=%d: %v", topK, err)
		}
		want := naiveSearch(t, m, q, topK)
		if len(got) != len(want) {
			t.Fatalf("topK=%d: got %d results, naive returned %d", topK, len(got), len(want))
		}
		for i := range want {
			if got[i].ID != want[i].ID {
				t.Errorf("topK=%d rank %d: got %q, naive says %q", topK, i, got[i].ID, want[i].ID)
			}
			if got[i].Score != want[i].Score {
				t.Errorf("topK=%d rank %d (%s): score %v, naive says %v",
					topK, i, got[i].ID, got[i].Score, want[i].Score)
			}
		}
		if len(got) > 0 && got[0].Meta["text"] == "" {
			t.Errorf("topK=%d: survivor %s has no meta — it was never unmarshalled", topK, got[0].ID)
		}
	}
}

// TestDotBlobMatchesDot pins the claim in dotBlob's doc comment: reading the
// vector out of its stored bytes is bit-identical to decoding it first. Scores
// feed a gate with a 0.008 margin, so "close enough" is not the bar.
func TestDotBlobMatchesDot(t *testing.T) {
	rng := rand.New(rand.NewSource(11))
	for i := 0; i < 200; i++ {
		a := randUnitVec(rng, 384)
		b := randUnitVec(rng, 384)
		if got, want := dotBlob(a, encodeVec(b)), dot(a, b); got != want {
			t.Fatalf("dotBlob = %v, dot = %v", got, want)
		}
	}
	// Length mismatch is 0 in both, and so is an empty vector.
	if got := dotBlob([]float32{1, 0}, encodeVec([]float32{1, 0, 0})); got != 0 {
		t.Errorf("mismatched lengths scored %v, want 0", got)
	}
	if got := dotBlob(nil, nil); got != 0 {
		t.Errorf("empty scored %v, want 0", got)
	}
}