PiPNN 1/6: extract shared RobustPrune

[weiyaoluo (SeliMeli)]

Purpose

This PR extracts the private Vamana RobustPrune selection loop. PiPNN can reuse this loop in a later PR.

The change adds no public API. It does not change Vamana pruning behavior.

Main changes

• graph::internal::prune::robust_prune now owns the alpha-round selection loop.
• The function accepts one nearest-first cache of source distances and candidate vectors.
• A debug assertion checks cache order with the canonical neighbor distance ordering.
• DiskANNIndex::occlude_list still sorts candidates, loads vectors, maps IDs, and applies saturation.
• State updates, alpha rounds, and unsafe checks retain their original logic.

Preserved behavior

• Candidate order remains unchanged.
• Excluded candidates remain positional None entries.
• Alpha progression and prune-kind behavior remain unchanged.
• Saturation order and duplicate handling remain unchanged.
• Provider errors remain in occlude_list.

Review order

1. Review diskann/src/graph/internal/prune.rs.
2. Compare the extracted loop with the old loop in index.rs.
3. Review the call site in DiskANNIndex::occlude_list.
4. Confirm that preparation, ID mapping, and saturation remain outside the extracted function.

Validation

• All diskann library tests pass.
• All-target Clippy passes with -Dwarnings.
• Existing Vamana prune scenarios pass.

Stack

Stack 1/6. #1287 adds PiPNN numerical kernels.

[Copilot]

Pull request overview

This PR extracts Vamana’s reusable RobustPrune alpha-round occlusion state machine into a new allocation-free kernel (graph/internal/robust_prune.rs) while keeping Vamana-specific scratch space, provider error handling, and adapter logic in graph/internal/vamana_prune.rs. The DiskANNIndex::occlude_list path is updated to prepare “available-only” candidates for the shared kernel, then translate selected positions back to IDs and apply optional saturation.

Changes:

• Introduces internal::robust_prune as a provider-independent RobustPrune kernel over prepared candidates + reusable per-candidate state.
• Moves Vamana-owned scratch/context and ranked provider error types into internal::vamana_prune with co-located tests.
• Refactors graph/index.rs::occlude_list to: validate bounds → prepare available candidates → call shared kernel → write adjacency → saturate from available-only candidates.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
diskann/src/graph/internal/vamana_prune.rs	Adds Vamana-owned scratch/context + provider error ranking and Vamana integration tests.
diskann/src/graph/internal/robust_prune.rs	Adds allocation-free RobustPrune kernel + pure state-machine tests.
diskann/src/graph/internal/prune.rs	Removes the previous combined prune implementation/state.
diskann/src/graph/internal/mod.rs	Rewires internal modules to expose robust_prune and vamana_prune.
diskann/src/graph/index.rs	Switches pruning implementation to prepare candidates and call the shared robust_prune kernel; updates saturation behavior to be available-only.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 98.79518% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 92.55%. Comparing base (16c7451) to head (c3fd818).
⚠️ Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann/src/graph/internal/prune.rs	98.55%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1315      +/-   ##
==========================================
+ Coverage   91.26%   92.55%   +1.28%     
==========================================
  Files         517      522       +5     
  Lines       98511    99556    +1045     
==========================================
+ Hits        89910    92142    +2232     
+ Misses       8601     7414    -1187     

Flag	Coverage Δ
miri	92.55% <98.79%> (+1.28%)	⬆️
unittests	92.51% <98.79%> (+1.57%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann/src/graph/index.rs	96.42% <100.00%> (-0.10%)	⬇️
diskann/src/neighbor/mod.rs	99.04% <100.00%> (+<0.01%)	⬆️
diskann/src/graph/internal/prune.rs	85.32% <98.55%> (+22.82%)	⬆️

... and 77 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Mark Hildebrand (hildebrandmw)]

This PR would have been easy to review, but at nontrivially regresses internal documentation and comments. For such important internal algorithms, please avoid changing comments unnecessarily, particularly when unsafe is involved.

Please restore the deleted or altered comments, making only the necessary updates for the new location. Also, please avoid renaming variables when just moving code around. A simpler diff stands on its own.

[weiyaoluo (SeliMeli)]

Addressed review in c2e6be62:

• kept module and existing types in graph::internal::prune;
• restored original loop comments, variable terminology, and checked-index unsafe justifications;
• changed the helper boundary to &SortedNeighbors so source-distance ordering is enforced by the existing type;
• added local pool/cache/state length checks before unchecked mutation;
• retained Vamana preparation, ID mapping, saturation, errors, and alpha behavior unchanged.

Validation: 311 diskann library tests, all-target Clippy, LSP diagnostics, and git diff --check pass.

[weiyaoluo (SeliMeli)]

Local IAI-Callgrind check against main (16c7451f):

• triangle 750→64: instructions +0.27259%, estimated cycles +0.29260%;
• occluding+saturation 750→64: instructions -0.01482%, estimated cycles +0.01124%.

Measured under WSL2 with the archived local regression fixture; no submitted benchmark target. Full record: experiments/pipnn_local_regression/2026-08-07-wrapper-simplification-callgrind.md in the local analysis workspace.

[weiyaoluo (SeliMeli)]

All review threads now point to superseded lines. The final extraction keeps the helper in graph::internal::prune, accepts SortedNeighbors, restores the original loop/state/unsafe comments and existing debug_assert! sequence, and adds no assertions or pruning behavior changes. I am resolving the outdated threads.

[Wei Wu (wuw92)]

LGTM. Since it is intended to be a focused extraction, I would keep the implementation as close to the original as possible and simplify the PR description.

[Wei Wu (wuw92)]

Could we keep this extraction identical to the old implementation by destructuring (neighbor_distance, neighbor) from cache, as the original loop did?

Once the distance comes from cache, pool has no remaining use inside robust_prune, so it can also be removed from the function signature. This avoids having two sources for the same candidate distance that could disagree.

[Mark Hildebrand (hildebrandmw)]

+1 to this. Drop the pool argument in favor of a documented, sorted cache.

[weiyaoluo (SeliMeli)]

Updated robust_prune to read the source distance from cache and removed the pool parameter. This restores the pre-extraction distance source.

[weiyaoluo (SeliMeli)]

Follow-up: cache is now the only distance source. I retained SortedNeighbors as the sortedness witness required by this private API and added an alignment assertion between the sorted pool and cache. RobustPrune does not read a distance from the pool.

[Wei Wu (wuw92)]

Since pool is only used for a debug-only length check, it does not actually establish that the separately supplied cache has the same ordering. Could we drop pool, document that cache must be sorted by source distance, and add a debug assertion using the canonical distance ordering? This keeps the helper’s contract attached to the data it actually consumes.

[weiyaoluo (SeliMeli)]

Fixed in c3fd818. robust_prune now accepts only the nearest-first cache. The caller retains pool for cache construction and ID mapping. A debug assertion checks adjacent cache distances with the same canonical ordering used by SortedNeighbors.