PiPNN 1/6: extract shared RobustPrune

[weiyaoluo (SeliMeli)]

DiskANN keeps Vamana prune scratch, provider-error state, and candidate ordering in graph/internal/prune.rs, while the alpha-round selection loop was embedded in DiskANNIndex::occlude_list. PiPNN later needs that same selection loop with its own candidate values. This PR extracts only that loop and leaves the module name, file location, caller flow, comments, and Vamana behavior in place.

No public API or semver surface is introduced.

Code map

1. diskann/src/graph/internal/prune.rs
   • Existing Options, Scratch, Context, State, and ranked provider-error types remain unchanged.
   • robust_prune contains the selection loop moved from index.rs.
   • Its candidate input is SortedNeighbors, the existing type whose constructor sorts and caps by source distance. A raw unsorted candidate slice cannot be passed.
   • The positional cache and state lengths are asserted before any unchecked state mutation.
2. diskann/src/graph/index.rs
   • Retains candidate preparation, provider lookup, cache allocation, ID translation, saturation, assertions, and errors.
   • Calls the extracted loop after constructing the same positional cache from the same SortedNeighbors pool.

End-to-end flow

Vamana fills provider state → occlude_list constructs the positional cache from SortedNeighbors → internal::prune::robust_prune selects positions → occlude_list maps positions through the original pool → the unchanged saturation loop appends candidates → the existing caller writes provider adjacency.

Invariants and preserved behavior

• SortedNeighbors enforces nondecreasing source-distance order at the function boundary.
• Cache and state entries remain one-to-one with sorted candidate positions; O(1) length assertions enforce that alignment before unsafe mutation.
• Excluded and provider-unavailable candidates remain positional None entries and are skipped exactly as before.
• The existing u16::MAX assertion, state initialization, alpha progression, prune-kind inputs, and selected-position representation are unchanged.
• Empty-pool handling, cache allocation, output mapping, saturation order/condition, duplicate handling, and provider-error behavior remain in index.rs.
• NaN/infinite-alpha behavior is intentionally unchanged.
• Loop comments, variable names, and unsafe justifications are retained from the original implementation and relocated beside the code they describe.
• The function remains inside private graph::internal::prune; only graph descendants can call it.

Review path

1. Review the addition in internal/prune.rs: the existing types remain in place and the function body is the loop removed from index.rs.
2. Compare the relocated loop comments, variable names, state updates, and unsafe justifications with main:index.rs::occlude_list.
3. Verify the code before and after the function call in index.rs remains the original preparation, mapping, and saturation flow.
4. Confirm the function accepts SortedNeighbors, not a raw candidate slice, and checks cache/state alignment locally.

Validation

• All 311 diskann library tests pass.
• All-target Clippy and LSP diagnostics pass.
• Existing Vamana scenarios cover bounded rows, equal-distance ordering, alpha rounds, both prune kinds, exclusion/unavailability, saturation, and max-occlusion truncation.

Stack relation

Stack 1/6. #1287 adds numerical kernels next. #1290 later uses the same internal selection function with PiPNN-owned preparation and ID translation.

Stack 1/6 → #1287

[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.

[weiyaoluo (SeliMeli)]

Intentionally unchanged in this refactor. Every states[...] access remains bounds-checked, so malformed/non-finite alpha can panic but cannot write out of bounds. Capping the selection count or changing alpha handling would alter pre-existing Vamana behavior and is outside this behavior-preserving extraction.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

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

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

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1315   +/-   ##
=======================================
  Coverage   91.26%   91.27%           
=======================================
  Files         517      517           
  Lines       98511    98541   +30     
=======================================
+ Hits        89910    89945   +35     
+ Misses       8601     8596    -5     

Flag	Coverage Δ
miri	91.27% <98.80%> (+<0.01%)	⬆️
unittests	90.95% <98.80%> (+<0.01%)	⬆️

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.48% <100.00%> (-0.04%)	⬇️
diskann/src/graph/internal/prune.rs	85.96% <98.64%> (+23.46%)	⬆️

... and 3 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.

[Mark Hildebrand (hildebrandmw)]

This comment is not a valid safety argument as is. While the contract on the function states that candidates and states must have the same length, the function itself is not marked as unsafe and therefore unsafe code cannot rely on that contract being upheld. Nor is this invariant checked in the function preamble (and it need not be with the original justification).

The original comment gave a correct (albeit terse) reason why the indexing was safe: states was already accessed at the given index earlier and would have panicked if index was out of bounds.

Same for the other occurances.

If someone later turns the first destructure into get_unchecked, the cited justification still reads true while soundness silently evaporates.

[weiyaoluo (SeliMeli)]

Fixed in c2e6be6. robust_prune checks pool/cache/state length equality before the loop, and each unsafe comment again relies on the preceding checked states[i] access exactly as the original code did.

[Mark Hildebrand (hildebrandmw)]

Write docs as if they are the current contract ... because they are. If this merged all references to "existing" and "retains" are immediately stale and lose all relevance. This is not the only occurrence of this pattern.

[weiyaoluo (SeliMeli)]

Fixed. Source documentation now states the current contract in present tense. Existing prune types and their documentation remain in graph/internal/prune.rs.

[Mark Hildebrand (hildebrandmw)]

Please restore the deleted comments, relocating the loop-level ones alongside the code they describe. The State doc link needs repointing at the extracted function rather than restoring verbatim.

[weiyaoluo (SeliMeli)]

Fixed. Loop-level algorithm/state comments, inline comments, and variable terminology are restored beside the extracted loop. State now links to robust_prune.

[Mark Hildebrand (hildebrandmw)]

This adds the real requirement that candidates must be sorted by increasing distance, but that is not documented. This is a new function where-as before the invariant was satisfied by construction. Please document the additional requirements at the very least.

[weiyaoluo (SeliMeli)]

Fixed by making order part of the accepted type: robust_prune takes &SortedNeighbors, whose constructor sorts and caps by source distance. A raw unsorted candidate slice cannot be passed through this interface; the function docs state the ordering contract.

[Mark Hildebrand (hildebrandmw)]

Call this module prune. It's one user already does that via aliasing. Also please keep it in the internal module.

[weiyaoluo (SeliMeli)]

Fixed. The helper now remains graph::internal::prune::robust_prune; no new module or public API is introduced.

[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.