diff --git a/core/lsh/index.go b/core/lsh/index.go index 619c401..e5b5471 100644 --- a/core/lsh/index.go +++ b/core/lsh/index.go @@ -53,22 +53,33 @@ func (idx *LSHIndex) BuildIndex() error { return nil } // FindCandidates retrieves candidate fragment IDs that share at least one band bucket. func (idx *LSHIndex) FindCandidates(signature *MinHashSignature) []string { + return idx.FindCandidatesLimit(signature, 0) +} + +// FindCandidatesLimit retrieves candidate fragment IDs that share at least one +// band bucket, stopping as soon as maxCandidates distinct IDs have been +// collected so dense buckets never materialize in full. maxCandidates <= 0 +// disables the cap. Traversal order is deterministic — band order, then bucket +// insertion order — so capped queries keep the earliest-encountered +// candidates rather than an arbitrary subset. +func (idx *LSHIndex) FindCandidatesLimit(signature *MinHashSignature, maxCandidates int) []string { if signature == nil || len(signature.signatures) == 0 { return []string{} } - ids := make(map[string]struct{}) - bands := idx.computeBandKeys(signature) - for _, key := range bands { - if bucket, ok := idx.buckets[key]; ok { - for _, id := range bucket { - ids[id] = struct{}{} + seen := make(map[string]struct{}) + out := []string{} + for _, key := range idx.computeBandKeys(signature) { + for _, id := range idx.buckets[key] { + if maxCandidates > 0 && len(out) >= maxCandidates { + return out } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + out = append(out, id) } } - out := make([]string, 0, len(ids)) - for id := range ids { - out = append(out, id) - } return out } diff --git a/core/lsh/index_test.go b/core/lsh/index_test.go index c15eda1..c39c9b5 100644 --- a/core/lsh/index_test.go +++ b/core/lsh/index_test.go @@ -1,6 +1,7 @@ package lsh import ( + "fmt" "testing" ) @@ -216,3 +217,75 @@ func BenchmarkLSH_AddAndFind(b *testing.B) { idx.FindCandidates(sigs[0]) } } + +// bandSig builds a raw signature for band-collision tests (bands*rows values). +func bandSig(values ...uint64) *MinHashSignature { + return &MinHashSignature{signatures: values, numHashes: len(values)} +} + +func TestFindCandidatesLimit_BandOrderSelection(t *testing.T) { + idx := NewLSHIndex(2, 2) + query := bandSig(1, 2, 3, 4) + // "1" collides with the query only in band 1, "10" only in band 0. + if err := idx.AddFragment("1", bandSig(8, 8, 3, 4)); err != nil { + t.Fatalf("add 1: %v", err) + } + if err := idx.AddFragment("10", bandSig(1, 2, 9, 9)); err != nil { + t.Fatalf("add 10: %v", err) + } + + got := idx.FindCandidatesLimit(query, 0) + if len(got) != 2 || got[0] != "10" || got[1] != "1" { + t.Fatalf("uncapped candidates = %v, want [10 1] in band order", got) + } + + // A cap of 1 must keep the band-0 candidate, not an ID-sorted subset. + got = idx.FindCandidatesLimit(query, 1) + if len(got) != 1 || got[0] != "10" { + t.Fatalf("capped candidates = %v, want [10]", got) + } +} + +func TestFindCandidatesLimit_StopsAtCapInDenseBucket(t *testing.T) { + mh := NewMinHasher(128) + sig := mh.ComputeSignature([]string{"same", "feature", "set"}) + + idx := NewLSHIndex(32, 4) + want := []string{} + for i := 0; i < 100; i++ { + id := fmt.Sprintf("f%03d", i) + if err := idx.AddFragment(id, sig); err != nil { + t.Fatalf("add %s: %v", id, err) + } + if i < 8 { + want = append(want, id) + } + } + + got := idx.FindCandidatesLimit(sig, 8) + if len(got) != len(want) { + t.Fatalf("capped candidate count = %d, want %d", len(got), len(want)) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("capped candidates = %v, want first-inserted %v", got, want) + } + } +} + +func TestFindCandidatesLimit_DedupAcrossBands(t *testing.T) { + idx := NewLSHIndex(2, 2) + query := bandSig(1, 2, 3, 4) + // "both" collides in both bands but must occupy a single cap slot. + if err := idx.AddFragment("both", bandSig(1, 2, 3, 4)); err != nil { + t.Fatalf("add both: %v", err) + } + if err := idx.AddFragment("late", bandSig(8, 8, 3, 4)); err != nil { + t.Fatalf("add late: %v", err) + } + + got := idx.FindCandidatesLimit(query, 2) + if len(got) != 2 || got[0] != "both" || got[1] != "late" { + t.Fatalf("candidates = %v, want [both late]", got) + } +} diff --git a/jscan/internal/analyzer/lsh_index.go b/jscan/internal/analyzer/lsh_index.go index 9d34941..b41f609 100644 --- a/jscan/internal/analyzer/lsh_index.go +++ b/jscan/internal/analyzer/lsh_index.go @@ -34,7 +34,7 @@ func (idx *lshCandidateIndex) AddFragment(id int, signature *corelsh.MinHashSign } func (idx *lshCandidateIndex) FindCandidates(signature *corelsh.MinHashSignature) []int { - candidates := idx.index.FindCandidates(signature) + candidates := idx.index.FindCandidatesLimit(signature, idx.maxCandidates) ids := make([]int, 0, len(candidates)) for _, candidate := range candidates { id, err := strconv.Atoi(candidate) @@ -43,8 +43,5 @@ func (idx *lshCandidateIndex) FindCandidates(signature *corelsh.MinHashSignature } } sort.Ints(ids) - if len(ids) > idx.maxCandidates { - ids = ids[:idx.maxCandidates] - } return ids } diff --git a/jscan/internal/analyzer/lsh_index_test.go b/jscan/internal/analyzer/lsh_index_test.go index ff07a43..8651166 100644 --- a/jscan/internal/analyzer/lsh_index_test.go +++ b/jscan/internal/analyzer/lsh_index_test.go @@ -18,7 +18,9 @@ func TestLSHCandidateIndexConvertsSortsAndCapsIDs(t *testing.T) { } } - want := []int{1, 2, 3} + // The cap keeps the first candidates in traversal order ({4, 2, 3}), + // not the smallest IDs; the capped set is then sorted ascending. + want := []int{2, 3, 4} for i := 0; i < 10; i++ { if got := lsh.FindCandidates(sig); !reflect.DeepEqual(got, want) { t.Fatalf("candidate mismatch: got %v want %v", got, want) @@ -36,8 +38,10 @@ func TestLSHCandidateIndexUsesDefaultCap(t *testing.T) { } } + // Descending insertion: the default cap keeps the first-inserted + // defaultLSHMaxCandidates candidates (IDs 1024..1), sorted ascending. got := lsh.FindCandidates(sig) - if len(got) != defaultLSHMaxCandidates || got[0] != 0 || got[len(got)-1] != defaultLSHMaxCandidates-1 { + if len(got) != defaultLSHMaxCandidates || got[0] != 1 || got[len(got)-1] != defaultLSHMaxCandidates { t.Fatalf("default cap or order mismatch: len=%d first=%d last=%d", len(got), got[0], got[len(got)-1]) } }