diff --git a/src/CodeIndex.Core/Search/CodeIndexService.cs b/src/CodeIndex.Core/Search/CodeIndexService.cs index b4aa4ec..823b394 100644 --- a/src/CodeIndex.Core/Search/CodeIndexService.cs +++ b/src/CodeIndex.Core/Search/CodeIndexService.cs @@ -390,19 +390,37 @@ public async Task SearchWithStatusAsync( warning = CombineWarnings(warning, RankingDegradedMessage(ex)); } - IReadOnlyList fused = HybridRanker.Fuse(vectorHits, symbolHits, limit); + // Fused deep (branchDepth, not just `limit`) so ResultDiversifier below has more than + // exactly `limit` candidates to redistribute across files — diversifying a list already + // truncated to `limit` would have nothing left to backfill from. See BuildCandidateIndices' + // sibling remarks for the same "don't truncate before the step that needs the depth" idea + // applied to filtering instead of diversification. + IReadOnlyList fused = HybridRanker.Fuse(vectorHits, symbolHits, branchDepth); + + // Every symbol-branch hit is exempt from the per-file cap below — see ResultDiversifier's + // remarks on why a literal identifier match (e.g. a query naming a class whose own const + // fields also matched) must never be sacrificed to make room for an unrelated chunk that + // merely belongs to a still-under-cap file. + HashSet symbolHitIndices = new(symbolHits.Count); + foreach (ScoredIndex hit in symbolHits) + { + symbolHitIndices.Add(hit.Index); + } + + IReadOnlyList diversified = ResultDiversifier.Diversify( + fused, index => snapshot.Chunks[index].FilePath, limit, symbolHitIndices); // Excerpt reads (and the staleness check alongside each one — see IsExcerptPossiblyStaleAsync) // are independent ISourceProvider calls (one file read/stat each), so running them // concurrently instead of one-at-a-time keeps this linear in wall-clock file I/O only for // the slowest read, not the sum of all of them — worth doing now that `limit` (and - // therefore fused.Count) is no longer capped at a small fixed branch depth. + // therefore diversified.Count) is no longer capped at a small fixed branch depth. Dictionary fingerprintByPath = BuildFingerprintLookup(snapshot); - Task[] excerptTasks = new Task[fused.Count]; - Task[] stalenessTasks = new Task[fused.Count]; - for (int i = 0; i < fused.Count; i++) + Task[] excerptTasks = new Task[diversified.Count]; + Task[] stalenessTasks = new Task[diversified.Count]; + for (int i = 0; i < diversified.Count; i++) { - CodeChunk chunk = snapshot.Chunks[fused[i].Index]; + CodeChunk chunk = snapshot.Chunks[diversified[i].Index]; excerptTasks[i] = ReadExcerptAsync(chunk, cancellationToken); stalenessTasks[i] = IsExcerptPossiblyStaleAsync( fingerprintByPath.GetValueOrDefault(chunk.FilePath), chunk.FilePath, cancellationToken); @@ -417,10 +435,10 @@ public async Task SearchWithStatusAsync( string[] excerpts = excerptsTask.Result; bool[] staleFlags = stalenessTask.Result; - List hits = new(fused.Count); - for (int i = 0; i < fused.Count; i++) + List hits = new(diversified.Count); + for (int i = 0; i < diversified.Count; i++) { - ScoredIndex scored = fused[i]; + ScoredIndex scored = diversified[i]; hits.Add(new SearchHit { ChunkId = scored.Index, diff --git a/src/CodeIndex.Core/Search/ResultDiversifier.cs b/src/CodeIndex.Core/Search/ResultDiversifier.cs new file mode 100644 index 0000000..55f8cd3 --- /dev/null +++ b/src/CodeIndex.Core/Search/ResultDiversifier.cs @@ -0,0 +1,206 @@ +namespace CodeIndex.Core.Search; + +/// +/// Caps how many of a single file's chunks may occupy the final result set, so a small limit +/// is not entirely filled by two or three sibling members of the one or two most central files while +/// a genuinely relevant — but slightly lower-ranked — file never gets a slot at all. +/// +/// +/// +/// The problem this targets. fuses purely by rank position and +/// then takes a flat Take(limit) off the +/// front of that fused order. Neither step has any notion of "file" — so when a topic has several +/// sibling declarations in the same file that all score highly on the vector branch alone, they can +/// consume most or all of a small limit before a different file's single, equally relevant +/// chunk is ever reached, even though that chunk ranked only a few places lower. Measured against the +/// real wallet project's index: a natural-language query about the network-unavailable UI flow +/// put three separate members of the same NetworkUnavailableModal.razor.cs class into 3 of 5 +/// result slots at limit=5, leaving no room for XrplSharpClientService.ExecuteIfConnected +/// — a distinct, alternate failure path (the connectivity check that throws +/// NotConnectedException) that ranked 7th overall but 1st among files not already represented. +/// +/// +/// Cap-then-backfill, not a hard exclusion. This walks the already-fused, already-ordered +/// ranking once, keeping every non-exempt candidate (see exemptIndices on +/// ) whose file has not yet reached maxPerFile such selections, and +/// setting aside (not dropping) every one that would exceed it. Once the capped pass is exhausted, +/// the set-aside candidates are appended back in their original rank order until limit is +/// reached. This guarantees the method never returns fewer results than a plain Take(limit) +/// would have. +/// +/// +/// Why symbol-branch hits are exempt from the cap. The first version of this fix capped every +/// candidate uniformly and regressed two existing SearchQualityTests golden queries: "TrustSet" +/// and "LedgerEntry" each legitimately match several sibling declarations in one file (a class via +/// 's exact/prefix band, plus that class's own const fields/properties via +/// its substring band) — exactly the "genuinely several members of one class" case a per-file cap is +/// expected to cost something on. But capping bumped the query's own named class out in favour of a +/// same-branch-depth but otherwise unrelated chunk from a completely different file, which only +/// happened to still be under its own file's cap — a strictly worse result, not a diversified one. +/// The distinguishing signal: checks whether the caller's whole +/// query string is a literal substring of a chunk's symbol/signature/directory, so it essentially only +/// ever matches short, identifier-shaped queries like "TrustSet" — a natural-language sentence (the +/// shape of every query in the real reproduction above) never satisfies it, so the symbol branch +/// contributes nothing to those. Exempting symbol-branch hits from the cap therefore leaves precision +/// identifier lookups exactly as clustered as they earned the right to be, while still catching the +/// vector-only crowding the real defect is made of. +/// +/// +/// Why capping (and not, say, per-directory capping or MMR) was chosen. Per-directory capping +/// was considered — it would also have caught the network-unavailable case above (all three +/// crowding-out hits share both a file and a directory) — but it punishes exactly the "many small +/// files organised under one feature directory" shape this project's own +/// path-match band (see its remarks) already treats as legitimate breadth, not noise. An MMR-style +/// re-rank that penalises embedding similarity to already-selected results was also considered, but +/// it needs the embedding vectors themselves at re-rank time (not just each branch's already-reduced +/// rank position works with), which would mean threading raw vectors +/// through a layer that currently only ever sees . A flat per-file cap needs +/// only two pieces of information — and which candidates the +/// symbol branch already vouched for — that already has in hand for +/// every candidate. +/// +/// +public static class ResultDiversifier +{ + /// + /// The default cap: at most this many non-exempt (see 's + /// exemptIndices) chunks from a single file are taken during the capped pass before that + /// file's further chunks are deferred to the backfill pass. Chosen, not derived: 1 (never more + /// than a single chunk per file) was rejected as too aggressive — a class and its one standout + /// override are routinely both worth showing together, and capping at 1 would separate them + /// purely because they share a file, even when nothing else outranks the second. 2 keeps that + /// common "class + its most relevant member" pairing intact while still stopping a single file + /// from claiming a majority of a limit=5 result set the way the network-unavailable case + /// above did at 3. + /// + public const int DefaultMaxPerFile = 2; + + /// + /// Selects up to entries from — which must + /// already be in final rank order (best first) — preferring breadth across files up to + /// before falling back to 's own order to + /// fill any remaining slots. Relative order is preserved within both the capped and the + /// backfilled portions; only entries that would have exceeded the cap are ever moved later. + /// + /// Fused hits in descending rank order, deep enough to give this method + /// room to work with — see 's branch-depth remarks. Passing a list + /// already truncated to defeats the point: there would be nothing left + /// to backfill from. + /// Resolves a candidate's to the file + /// path used to group candidates. A delegate rather than requiring the caller to pre-join chunk + /// data, so this stays usable directly against , a snapshot + /// lookup, or a test double with no ceremony either way. + /// Maximum number of entries to return. Non-positive returns empty. + /// Indices (see ) that never count + /// against — and are never deferred by — the per-file cap, regardless of how many of their file's + /// chunks have already been selected. Intended for the symbol branch's own hits — see this + /// class's remarks for why a literal identifier match must not be sacrificed to make room for an + /// unrelated chunk elsewhere. (the default) exempts nothing. + /// See . Must be positive; a + /// non-positive value would mean every non-exempt candidate is deferred and the capped pass + /// never selects any of them, silently degrading to backfill-only order for the non-exempt + /// portion — callers that want "no diversification at all" should not call this method rather + /// than pass a value meant to disable it. + public static IReadOnlyList Diversify( + IReadOnlyList ranked, + Func filePathOf, + int limit, + IReadOnlySet? exemptIndices = null, + int maxPerFile = DefaultMaxPerFile) + { + ArgumentNullException.ThrowIfNull(ranked); + ArgumentNullException.ThrowIfNull(filePathOf); + + if (maxPerFile <= 0) + { + throw new ArgumentOutOfRangeException( + nameof(maxPerFile), maxPerFile, "Must be positive."); + } + + if (limit <= 0 || ranked.Count == 0) + { + return []; + } + + if (ranked.Count <= limit) + { + // Nothing to diversify away from: every candidate is already being returned, so a + // capped pass followed by backfill would just reconstruct the same list at extra cost. + return ranked; + } + + // Selection is recorded as flags over `ranked`'s own positions rather than by appending + // to an output list, so the result can be emitted in the input's order at the end. This + // matters because backfill runs after the capped pass: appending a deferred candidate + // directly would place a rank-3 hit *behind* the rank-4 hit that displaced it, and the + // tool's output would stop being ordered by relevance even though its `score` field still + // said otherwise. Diversification is meant to change *which* hits come back, never the + // order they are presented in. + bool[] isSelected = new bool[ranked.Count]; + int selectedCount = 0; + List? deferredOrdinals = null; + Dictionary perFileCount = new(StringComparer.Ordinal); + + for (int ordinal = 0; ordinal < ranked.Count; ordinal++) + { + if (selectedCount == limit) + { + break; + } + + ScoredIndex candidate = ranked[ordinal]; + + if (exemptIndices is not null && exemptIndices.Contains(candidate.Index)) + { + // A literal identifier match earned its place regardless of how many of its + // file's chunks are already selected — and does not itself consume any of the + // file's non-exempt quota, so it can never be the reason a later, genuinely + // vector-only sibling gets deferred. + isSelected[ordinal] = true; + selectedCount++; + continue; + } + + string filePath = filePathOf(candidate.Index); + int countSoFar = perFileCount.GetValueOrDefault(filePath); + + if (countSoFar < maxPerFile) + { + isSelected[ordinal] = true; + selectedCount++; + perFileCount[filePath] = countSoFar + 1; + } + else + { + (deferredOrdinals ??= []).Add(ordinal); + } + } + + // Backfill in rank order, so when the cap leaves slots unfilled the strongest deferred + // candidates are the ones that come back — never an arbitrary subset. + if (selectedCount < limit && deferredOrdinals is not null) + { + foreach (int ordinal in deferredOrdinals) + { + if (selectedCount == limit) + { + break; + } + + isSelected[ordinal] = true; + selectedCount++; + } + } + + List selected = new(selectedCount); + for (int ordinal = 0; ordinal < ranked.Count; ordinal++) + { + if (isSelected[ordinal]) + { + selected.Add(ranked[ordinal]); + } + } + + return selected; + } +} diff --git a/tests/CodeIndex.Core.Tests/Search/CodeIndexServiceTests.cs b/tests/CodeIndex.Core.Tests/Search/CodeIndexServiceTests.cs index 1806174..3732188 100644 --- a/tests/CodeIndex.Core.Tests/Search/CodeIndexServiceTests.cs +++ b/tests/CodeIndex.Core.Tests/Search/CodeIndexServiceTests.cs @@ -557,6 +557,146 @@ public async Task GetChunkAsync_FreshServiceSeedsFromDisk_SoAFailingFirstRefresh Assert.Equal("Acme.A.Widget.DoA", hit!.Chunk.Symbol); } + /// + /// End-to-end reproduction of the diversity defect measured against the real wallet + /// project: a query about app-lifecycle/auto-lock wiring put four hits from two central files + /// (App.xaml.cs, AuthStateProvider.cs) into 4 of 5 result slots at limit=5, + /// leaving room for only one of several genuinely distinct peripheral platform files (iOS + /// background handling, the Windows-only auto-lock wiring in + /// Platforms/Windows/App.xaml.cs, the web app's lifecycle bridge) even though each of + /// them ranked above the dominant file's own 3rd/4th members. This fixture reproduces that shape + /// with standing in for the real embedding model: one + /// dominant file ("A") supplies the four highest-ranked vector hits, three other files ("C", + /// "D", "E") each supply exactly one hit, ranked just below "A"'s. Before + /// existed, SearchWithStatusAsync(limit: 5) returned four + /// hits from "A" and only "C" — "D" and "E" never appeared at all. The fix must cap "A" at two + /// and surface all three peripheral files. + /// + [Fact] + public async Task SearchWithStatusAsync_FileCrowding_DoesNotLetOneFileFillMostOfASmallLimit() + { + InMemorySourceProvider source = new(new Dictionary + { + ["src/A.cs"] = MakeMultiMethodFile("Acme.A", "Dominant", + RankedMarkerEmbeddingClient.MarkerForRank(0), RankedMarkerEmbeddingClient.MarkerForRank(1), + RankedMarkerEmbeddingClient.MarkerForRank(2), RankedMarkerEmbeddingClient.MarkerForRank(3)), + ["src/C.cs"] = MakeMultiMethodFile("Acme.C", "PeripheralC", RankedMarkerEmbeddingClient.MarkerForRank(4)), + ["src/D.cs"] = MakeMultiMethodFile("Acme.D", "PeripheralD", RankedMarkerEmbeddingClient.MarkerForRank(5)), + ["src/E.cs"] = MakeMultiMethodFile("Acme.E", "PeripheralE", RankedMarkerEmbeddingClient.MarkerForRank(6)), + }); + + RankedMarkerEmbeddingClient embedder = new(); + CodeIndexService service = CreateService(source, embedder, out _); + + // Restricted to ChunkKind.Method: each fixture class also produces its own Class-kind + // chunk (whose embed text happens to include its members' names too, so it inherits the + // strongest contained marker's score) — a second, incidental candidate per file that would + // otherwise obscure the one-method-per-file shape this test is deliberately built around. + // BuildCandidateIndices applies this filter before either search branch runs, so it does + // not interact with diversification itself. + SearchResult result = await service.SearchWithStatusAsync( + "a query sentence that matches no symbol literally, so only the vector branch scores anything", + limit: 5, kind: ChunkKind.Method, pathFilter: null, TestContext.Current.CancellationToken); + + Assert.Equal(5, result.Hits.Count); + + Assert.Equal(2, result.Hits.Count(h => h.Chunk.FilePath == "src/A.cs")); + // The two strongest members of the dominant file are still both present... + Assert.Contains(result.Hits, h => h.Chunk.FilePath == "src/A.cs" && + h.Chunk.Symbol.Contains(RankedMarkerEmbeddingClient.MarkerForRank(0), StringComparison.Ordinal)); + Assert.Contains(result.Hits, h => h.Chunk.FilePath == "src/A.cs" && + h.Chunk.Symbol.Contains(RankedMarkerEmbeddingClient.MarkerForRank(1), StringComparison.Ordinal)); + + // ...and now every peripheral file gets its one relevant hit — not just the single one + // ("C") a plain, non-diversified Take(5) would still have reached. + Assert.Contains(result.Hits, h => h.Chunk.FilePath == "src/C.cs"); + Assert.Contains(result.Hits, h => h.Chunk.FilePath == "src/D.cs"); + Assert.Contains(result.Hits, h => h.Chunk.FilePath == "src/E.cs"); + } + + /// A file with several methods, one per marker in — + /// each method's name embeds the marker so can key + /// its vector off it (the chunk's embed text includes the method's own symbol name). + private static string MakeMultiMethodFile(string ns, string className, params string[] methodMarkers) + { + string methods = string.Join('\n', methodMarkers.Select((marker, i) => $$""" + public int Do{{marker}}_{{i}}() + { + return {{i}}; + } + """)); + + return $$""" + namespace {{ns}} + { + public class {{className}} + { + {{methods}} + } + } + """; + } + + /// + /// A fully controllable stand-in for a real embedding backend, giving each of up to markers a distinct, strictly decreasing cosine similarity against every + /// query — unlike (essentially random) or + /// (only three fixed similarities), this is what + /// needs + /// to construct an exact, known rank order spanning more than a couple of files. Every query + /// embeds to the same reference direction; a passage embeds along that same direction at an + /// angle proportional to its marker's rank, so rank 0 is the strongest match and rank + /// - 1 the weakest, with nothing tied. + /// + private sealed class RankedMarkerEmbeddingClient : IEmbeddingClient + { + private const int RankCount = 8; + + /// Small enough that even the weakest rank's cosine similarity stays comfortably + /// positive (well clear of 0), so ordering — not sign — is the only thing under test. + private const double AngleStepRadians = 0.08; + + public static string MarkerForRank(int rank) => $"RANK{rank}MARKER"; + + public int Dimensions => 2; + + public string Model => "ranked-marker-test-model"; + + public Task> EmbedAsync(IReadOnlyList inputs, CancellationToken cancellationToken = default) + { + float[][] vectors = new float[inputs.Count][]; + for (int i = 0; i < inputs.Count; i++) + { + vectors[i] = VectorFor(inputs[i]); + } + + return Task.FromResult>(vectors); + } + + public Task EmbedQueryAsync(string query, CancellationToken cancellationToken = default) => + Task.FromResult(AngleVector(0)); + + private static float[] VectorFor(string text) + { + for (int rank = 0; rank < RankCount; rank++) + { + if (text.Contains(MarkerForRank(rank), StringComparison.Ordinal)) + { + return AngleVector(rank); + } + } + + // No recognised marker: embed far off-axis so it never outranks a marked passage. + return [0f, 1f]; + } + + private static float[] AngleVector(int rank) + { + double angle = rank * AngleStepRadians; + return [(float)Math.Cos(angle), (float)Math.Sin(angle)]; + } + } + [Fact] public async Task SearchWithStatusAsync_RelevanceFloor_ExcludesAWeakVectorMatchButKeepsAStrongOne() { diff --git a/tests/CodeIndex.Core.Tests/Search/ResultDiversifierTests.cs b/tests/CodeIndex.Core.Tests/Search/ResultDiversifierTests.cs new file mode 100644 index 0000000..a4aebcb --- /dev/null +++ b/tests/CodeIndex.Core.Tests/Search/ResultDiversifierTests.cs @@ -0,0 +1,207 @@ +using CodeIndex.Core.Search; +using Xunit; + +namespace CodeIndex.Core.Tests.Search; + +public sealed class ResultDiversifierTests +{ + /// + /// The exact shape of the reproduced defect, reduced to a pure ranking fixture: file "A" (the + /// dominant/central file — standing in for AuthStateProvider.cs/App.xaml.cs in the + /// auto-lock trace, or NetworkUnavailableModal.razor.cs in the network-failure trace) + /// supplies the four highest-ranked candidates, and three other, unrelated-to-each-other files + /// ("C", "D", "E" — standing in for the Windows-only auto-lock wiring in + /// Platforms/Windows/App.xaml.cs, the iOS background-task handling, and similar genuinely + /// distinct peripheral implementations) each supply exactly one relevant candidate, ranked just + /// below "A"'s four. A plain Take(5) over the fused order — the pre-fix behaviour — + /// returns four hits from "A" and only one of the three peripheral files, leaving the other two + /// entirely unrepresented even though each has a hit that outranks "A"'s own 3rd/4th members. With + /// the default cap (2 per file), "A" is limited to its top two, freeing three slots that go to + /// all three peripheral files instead of just one. + /// + [Fact] + public void Diversify_OneFileDominatesTheRanking_SpreadsTheFreedSlotsAcrossPeripheralFilesInsteadOfBackfillingTheSameFile() + { + ScoredIndex[] ranked = + [ + new(0, 0.95f), // A + new(1, 0.90f), // A + new(2, 0.85f), // A — would be 3rd A slot in a plain Take(5) + new(3, 0.80f), // A — would be 4th A slot in a plain Take(5) + new(4, 0.75f), // C — the only peripheral file a plain Take(5) would ever reach + new(5, 0.70f), // D — entirely invisible to a plain Take(5) + new(6, 0.65f), // E — entirely invisible to a plain Take(5) + ]; + + string FilePathOf(int index) => index switch + { + <= 3 => "src/A.cs", + 4 => "src/C.cs", + 5 => "src/D.cs", + _ => "src/E.cs", + }; + + IReadOnlyList preFix = ranked.OrderByDescending(r => r.Score).Take(5).ToArray(); + Assert.Equal(4, preFix.Count(r => FilePathOf(r.Index) == "src/A.cs")); + Assert.DoesNotContain(preFix, r => FilePathOf(r.Index) == "src/D.cs"); + Assert.DoesNotContain(preFix, r => FilePathOf(r.Index) == "src/E.cs"); + + IReadOnlyList diversified = ResultDiversifier.Diversify(ranked, FilePathOf, limit: 5); + + Assert.Equal(5, diversified.Count); + Assert.Equal(2, diversified.Count(r => FilePathOf(r.Index) == "src/A.cs")); + // The two strongest members of the dominant file are still both present... + Assert.Contains(diversified, r => r.Index == 0); + Assert.Contains(diversified, r => r.Index == 1); + // ...and now every peripheral file gets its one relevant hit — not just "C", the one a + // plain Take(5) happened to still reach. + Assert.Contains(diversified, r => r.Index == 4); // C + Assert.Contains(diversified, r => r.Index == 5); // D + Assert.Contains(diversified, r => r.Index == 6); // E + } + + [Fact] + public void Diversify_FewerCandidatesThanLimit_ReturnsAllOfThemUnchanged() + { + ScoredIndex[] ranked = [new(0, 0.9f), new(1, 0.8f)]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, _ => "src/Only.cs", limit: 5); + + Assert.Equal(ranked, diversified); + } + + [Fact] + public void Diversify_MoreCandidatesThanLimitButAllOneFile_StillReturnsLimitEntriesNotFewer() + { + // The "genuinely five members of one class" case: capping must never cause the method to + // return fewer than `limit` results when at least `limit` candidates exist overall — the + // backfill pass has to reach into the same file once every other file is exhausted (here, + // there is no other file at all). + ScoredIndex[] ranked = [new(0, 0.9f), new(1, 0.8f), new(2, 0.7f), new(3, 0.6f), new(4, 0.5f)]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, _ => "src/OneClass.cs", limit: 3); + + Assert.Equal(3, diversified.Count); + // Rank order is preserved: the top three by score, not an arbitrary subset. + Assert.Equal([0, 1, 2], diversified.Select(r => r.Index)); + } + + [Fact] + public void Diversify_BackfilledCandidate_IsReturnedAtItsOwnRankNotAppendedLast() + { + // Regression: backfill runs after the capped pass, so appending a deferred candidate + // directly put it behind every hit selected after it was deferred — index 2 came back + // as [0, 1, 3, 2], a rank-3 hit presented below the rank-4 hit that had displaced it. + // The `score` field still said otherwise, so the output claimed an ordering it did not + // have. Diversification decides *which* hits come back; it must not reorder them. + ScoredIndex[] ranked = [new(0, 0.9f), new(1, 0.8f), new(2, 0.7f), new(3, 0.6f)]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, + index => index == 3 ? "src/Other.cs" : "src/Crowded.cs", + limit: 4, + maxPerFile: 2); + + Assert.Equal([0, 1, 2, 3], diversified.Select(r => r.Index)); + + // Stated as the invariant rather than just the literal expectation above: scores must + // never increase as the caller reads down the list. + float[] scores = [.. diversified.Select(r => r.Score)]; + for (int i = 1; i < scores.Length; i++) + { + Assert.True( + scores[i] <= scores[i - 1], + $"Result {i} scored {scores[i]} against {scores[i - 1]} at position {i - 1} — " + + "diversified output must stay ordered by descending score."); + } + } + + [Fact] + public void Diversify_NoCrowding_PreservesOriginalRankOrder() + { + ScoredIndex[] ranked = [new(0, 0.9f), new(1, 0.8f), new(2, 0.7f), new(3, 0.6f)]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, index => $"src/File{index}.cs", limit: 3); + + Assert.Equal([0, 1, 2], diversified.Select(r => r.Index)); + } + + [Fact] + public void Diversify_LimitZeroOrNegative_ReturnsEmpty() + { + ScoredIndex[] ranked = [new(0, 0.9f)]; + + Assert.Empty(ResultDiversifier.Diversify(ranked, _ => "src/A.cs", limit: 0)); + Assert.Empty(ResultDiversifier.Diversify(ranked, _ => "src/A.cs", limit: -1)); + } + + [Fact] + public void Diversify_EmptyInput_ReturnsEmpty() + { + Assert.Empty(ResultDiversifier.Diversify([], _ => "src/A.cs", limit: 5)); + } + + [Fact] + public void Diversify_NonPositiveMaxPerFile_Throws() + { + ScoredIndex[] ranked = [new(0, 0.9f)]; + + Assert.Throws( + () => ResultDiversifier.Diversify(ranked, _ => "src/A.cs", limit: 5, maxPerFile: 0)); + } + + /// + /// Pins the regression this class's first version caused in SearchQualityTests: a query + /// like "TrustSet" legitimately matches three sibling declarations in one file — the class itself + /// (an exact/prefix symbol-branch hit) plus two of its own const fields (substring symbol-branch + /// hits) — via . An earlier version of this method capped every + /// candidate uniformly, which deferred the class itself (the 3rd hit from that file) once the cap + /// filled with its two fields, and backfilled with an unrelated chunk from a different file that + /// only happened to still be under its own file's cap — trading the query's own named class away + /// for something worse. Marking all three as symbol-branch hits (the exemptIndices + /// argument below) must keep them together rather than repeat that trade. + /// + [Fact] + public void Diversify_SymbolBranchHitsShareAFileBeyondTheCap_AllStayExemptFromCapping() + { + ScoredIndex[] ranked = + [ + new(0, 0.95f), // TrustSetFlags.cs: ClearNoRipple field (symbol-branch substring hit) + new(1, 0.90f), // TrustSetFlags.cs: SetNoRipple field (symbol-branch substring hit) + new(2, 0.85f), // TrustSetFlags.cs: TrustSetFlags class itself (symbol-branch prefix hit) + new(3, 0.50f), // XrplClient.cs: an unrelated chunk that merely isn't over its own cap + ]; + + string FilePathOf(int index) => index <= 2 ? "src/TrustSetFlags.cs" : "src/XrplClient.cs"; + HashSet symbolBranchHits = [0, 1, 2]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, FilePathOf, limit: 3, exemptIndices: symbolBranchHits); + + // All three same-file symbol-branch hits survive, in their original rank order — the + // unrelated 4th-ranked chunk from a different file never gets to displace any of them. + Assert.Equal([0, 1, 2], diversified.Select(r => r.Index)); + } + + [Fact] + public void Diversify_ExemptCandidateDoesNotConsumeItsFilesNonExemptQuota() + { + // A file with one exempt hit (rank 0) and two non-exempt vector-only hits (ranks 1, 2): + // the exempt hit must not count toward the file's cap, so both non-exempt hits still + // compete for the cap exactly as if the exempt one belonged to a different file entirely. + ScoredIndex[] ranked = [new(0, 0.9f), new(1, 0.8f), new(2, 0.7f), new(3, 0.6f)]; + + string FilePathOf(int index) => index <= 2 ? "src/Shared.cs" : "src/Other.cs"; + HashSet exempt = [0]; + + IReadOnlyList diversified = ResultDiversifier.Diversify( + ranked, FilePathOf, limit: 3, exemptIndices: exempt); + + // The exempt hit (0) plus both non-exempt hits (1, 2) fit within the default cap of 2 for + // the non-exempt count, so index 3 (a different file) is never needed to fill the 3 slots. + Assert.Equal([0, 1, 2], diversified.Select(r => r.Index)); + } +}