PiPNN 2/6: extract shared RobustPrune

[SeliMeli]

Extracts Vamana RobustPrune into a provider-independent kernel.

Code map

1. diskann/src/graph/prune.rs owns the reusable types: Policy, caller-provided Context, reusable Scratch, and fallible error types.
2. robust_prune first validates alpha and the u16 candidate-position bound, then builds a cache through the caller's lookup closure.
3. The prune loop keeps one State per candidate. last_checked lets later alpha rounds resume occlusion checks instead of recomputing earlier selected-neighbor comparisons.
4. Selected positions are translated back to IDs only after pruning. Optional saturation then walks the original sorted pool, relying on AdjacencyList to reject duplicates.
5. graph/index.rs remains the Vamana adapter: it performs the async provider fill first, then calls the synchronous kernel with provider-backed lookup and distance closures.

Review path

• Compare Policy construction in index.rs with the old Vamana configuration: degree, alpha, prune kind, and saturation conditions must match.
• Trace these branches in robust_prune: excluded/missing candidates, resumed alpha rounds, distance errors, empty input, maximum candidate count, and saturation order.
• Check the async boundary separately: no context mutation occurs before accessor.fill(...).await succeeds, and the extracted kernel introduces no additional await/cancellation point.
• Registered Vamana tests in graph/test/cases/prune.rs characterize externally visible ordering and saturation; pure-kernel tests cover fallibility and capacity boundaries.

Stack 2/6: #1287 → #1290

[Copilot]

Pull request overview

This PR factors DiskANN’s robust-prune logic into a dedicated diskann::graph::prune module, updates the graph index to call the new provider-independent kernel, and adds targeted correctness tests plus a Criterion benchmark to validate and measure pruning behavior.

Changes:

• Moved/rewrote the robust-prune kernel into diskann/src/graph/prune.rs with explicit error handling (RobustPruneError) and supporting scratch/context types.
• Updated DiskANNIndex pruning path to delegate to prune::robust_prune and plumb errors through existing ANNError/ListError machinery.
• Added prune integration test cases and a robust_prune Criterion benchmark; updated mutation-testing exclusions.

Reviewed changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
diskann/src/graph/test/cases/prune.rs	New integration-style prune behavior tests using the test provider.
diskann/src/graph/test/cases/mod.rs	Registers the new prune test module.
diskann/src/graph/prune/tests.rs	New unit tests for the provider-independent robust-prune kernel and error plumbing.
diskann/src/graph/prune.rs	New prune kernel module (policy, scratch/context, robust_prune, list error types).
diskann/src/graph/mod.rs	Exposes the new prune module from graph.
diskann/src/graph/internal/prune.rs	Removes the previous internal prune implementation/types.
diskann/src/graph/internal/mod.rs	Stops exporting the removed internal prune module.
diskann/src/graph/index.rs	Switches occlusion/prune implementation to call the new prune::robust_prune and handles its Result.
diskann/Cargo.toml	Adds Criterion as a dev-dependency and registers a robust_prune benchmark (gated by testing).
diskann/benches/robust_prune.rs	Adds a Criterion benchmark for pruning across candidate sizes, prune kinds, and saturation.
Cargo.lock	Records the new Criterion dependency.
.cargo/mutants.toml	Updates mutant exclusions to include a robust-prune mutation pattern.

---

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

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann/src/graph/prune.rs:33

• Hyphenation/typo in rustdoc: "over-written" should be "overwritten".

/// The actual object passed to the pruning algorithms is [`Context`], which allows
/// sub-fields to be over-written as needed with local state if that is available instead.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 99.52494% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.72%. Comparing base (b046174) to head (60440d8).

Files with missing lines	Patch %	Lines
diskann/src/graph/prune.rs	98.85%	2 Missing ⚠️

Additional details and impacted files

@@                    Coverage Diff                     @@
##           pipnn-stack/01-kernels    #1288      +/-   ##
==========================================================
+ Coverage                   90.66%   90.72%   +0.05%     
==========================================================
  Files                         515      516       +1     
  Lines                       99858   100143     +285     
==========================================================
+ Hits                        90541    90850     +309     
+ Misses                       9317     9293      -24     

Flag	Coverage Δ
miri	90.72% <99.52%> (+0.05%)	⬆️
unittests	90.40% <99.52%> (+0.05%)	⬆️

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	95.79% <100.00%> (+0.04%)	⬆️
diskann/src/graph/test/cases/prune.rs	100.00% <100.00%> (ø)
diskann/src/graph/prune.rs	98.85% <98.85%> (ø)

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

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

diskann/src/graph/mod.rs:23

• graph::prune is now a public module (pub mod prune;), and it contains several pub items (e.g., Scratch, Context, Policy, robust_prune). That’s a new externally visible API surface for the diskann crate and is hard to retract later; if this kernel is intended to be internal-only for now, it should stay crate-private to avoid accidental downstream coupling (and potential semver implications).

pub mod index;
pub use index::DiskANNIndex;

pub mod prune;

mod start_point;
pub use start_point::{SampleableForStart, StartPointStrategy};

diskann/src/graph/index.rs:2578

• The occlude_list doc comment immediately above still links to prune::Context::occlude_factor and prune::Context::last_checked, but those fields no longer exist on prune::Context after the refactor (they’re on prune::State). This creates broken intra-doc links and makes the comment misleading.

    fn occlude_list<M, C, F>(
        &self,
        computer: &C,
        context: &mut prune::Context<'_, DP::InternalId>,
        map: M,
        exclude: F,
        options: prune::Options,
    ) -> Result<(), prune::RobustPruneError>

diskann/src/graph/test/cases/prune.rs:309

• maximum_u16_candidate_pool_is_supported constructs 65k vectors (and a transient set of ~65k IDs) via the test provider. That’s an unusually heavy fixture for a correctness test and is likely to slow CI or cause memory pressure. The exact u16::MAX boundary is already covered in diskann/src/graph/prune/tests.rs, so this integration test can be scaled down while still exercising the Vamana/provider seam.

#[tokio::test(flavor = "current_thread")]
async fn maximum_u16_candidate_pool_is_supported() {
    let num_candidates = u16::MAX as usize;
    let vectors = (0..=num_candidates)
        .map(|position| vec![position as f32])

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

diskann/src/graph/index.rs:2601

• occlude_list allocates a fresh Vec cache (let mut cache = Vec::new()) on every prune invocation, which defeats the surrounding intent to minimize allocations and will add per-call heap churn in hot paths like prune_range / multi-insert. This cache should be reusable across calls (e.g., thread a cache buffer through the call stack or attach a reusable buffer to the existing prune scratch) so capacity can be retained between prunes.

        let policy = prune::Policy::new(
            self.config.pruned_degree().get(),
            self.config.alpha(),
            self.config.prune_kind(),
            options.force_saturate
                || (self.config.saturate_after_prune() && self.config.alpha() > 1.0),
        );
        let mut cache = Vec::new();
        prune::robust_prune(
            context,
            policy,
            &mut cache,
            |id| map.get(id),
            |neighbor, selected| {
                Ok(computer.evaluate_similarity((*neighbor).reborrow(), selected.reborrow()))
            },
            exclude,
        )

diskann/src/graph/prune.rs:300

• robust_prune returns RobustPruneError::Allocation for some workspace reserves, but building the output list is still potentially panicking: AdjacencyList::resize uses Vec::resize (panics on allocation failure/capacity overflow) and saturation later uses neighbors.push (also may allocate/panic). That makes the function not fully fallible despite exposing an allocation error variant.

    let mut guard = neighbors.resize(found);
    std::iter::zip(guard.iter_mut(), states.iter()).for_each(|(destination, state)| {
        *destination = *pool[state.neighbor.into_usize()].id();
    });
    guard.finish(found);

diskann/src/graph/test/cases/prune.rs:178

• This test asserts a specific neighbor order for equal-distance candidates, but the candidate sorting pipeline uses SortedNeighbors::new which ultimately sorts with an unstable comparator over distance-only ties. For equal distances, the relative order is not a defined contract and can change across Rust versions/platforms, making this test potentially flaky unless tie-breaking is made explicit (e.g., distance then id) or the assertion is relaxed to avoid depending on tie order.

async fn equal_distances_keep_current_sorted_neighbor_order() {
    let case = PruneCase::new(
        vec![
            vec![0.0, 0.0, 0.0],
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
        ],
        [3, 1, 2],
        PruneConfig {
            metric: Metric::L2,
            source: 0,
            degree: 2,
            alpha: 1.2,
            prune_kind: PruneKind::TriangleInequality,
            saturate: false,
            max_occlusion_size: 10,
        },
    );

    assert_eq!(&*case.run(&test_provider::Strategy::new()).await, &[2, 1]);

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann/src/graph/prune.rs:60

• Scratch::as_context docs say it only truncates the pool, but SortedNeighbors::new also sorts the retained candidates by distance (and can reorder self.pool). Since this is a public API surface, callers need this behavior documented to avoid assuming original insertion order is preserved.

    /// Convert `self` into a `Context`, truncating the internal `pool` list to a length of
    /// `max_candidates`.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (3)

diskann/src/graph/index.rs:2578

• The rustdoc above occlude_list still mentions prune::Context::occlude_factor / prune::Context::last_checked, but the extracted prune::Context no longer has those fields (it has pool, states, neighbors). This makes the adapter docs misleading.

Update the "Clobbers" bullets to refer to prune::Context::states (which holds the per-candidate State::{occlude_factor,last_checked} tracking).

        map: M,
        exclude: F,
        options: prune::Options,
    ) -> Result<(), prune::RobustPruneError>

diskann/src/graph/index.rs:2592

• occlude_list allocates a fresh Vec for the lookup cache on every call (let mut cache = Vec::new();). This negates the surrounding intent to minimize allocations in this hot path, especially since prune is called per-node during graph construction.

Consider moving this cache into prune::Scratch (or threading a &mut Vec<_> through the call chain) so the allocation is amortized across calls.

            options.force_saturate
                || (self.config.saturate_after_prune() && self.config.alpha() > 1.0),
        );
        let mut cache = Vec::new();
        prune::robust_prune(

diskann/src/graph/test/cases/prune.rs:288

• This test constructs a fixture with u16::MAX + 1 separate Vec<f32> allocations (one per point), plus a 65k-sized adjacency list. That is likely to add noticeable runtime and allocator pressure to the default unit test suite.

Given the kernel-level unit tests already cover the u16 boundary, consider marking this integration test as ignored by default (or gating it behind a feature) so CI doesn't pay this cost on every run.

#[tokio::test(flavor = "current_thread")]
async fn maximum_u16_candidate_pool_is_supported() {
    let num_candidates = u16::MAX as usize;
    let vectors = (0..=num_candidates)
        .map(|position| vec![position as f32])
        .collect();

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann/src/graph/prune.rs:369

• Saturation currently ignores whether a candidate was actually available (i.e., lookup returned None). Since View::get is explicitly fallible, this can re-introduce unavailable IDs into the final adjacency list when policy.saturate is enabled, even though they were excluded during pruning.

    if policy.saturate {
        for neighbor in pool.iter() {
            if neighbors.len() >= policy.degree {
                break;
            }

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann/src/graph/prune.rs:369

• Saturation currently only re-checks exclude(...) and will append candidate IDs even when lookup(...) returned None earlier (unavailable vectors). That can reintroduce candidates the prune loop intentionally skipped, producing adjacency entries that were not actually usable during this prune pass.

            if !exclude(*neighbor.id()) {
                neighbors.push(*neighbor.id());
            }
        }
    }

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann/src/graph/test/cases/prune.rs:309

• The maximum_u16_candidate_pool_is_supported test constructs u16::MAX + 1 separate Vec<f32> vectors (one per point) to build a full DiskANNIndex fixture. This creates ~65k heap allocations and a large provider state, which is likely to slow CI and make the test suite more brittle/time-sensitive. The u16 candidate-pool boundary is already exercised at the kernel level in graph::prune unit tests, so this case can be covered without building a full provider.

async fn maximum_u16_candidate_pool_is_supported() {
    let num_candidates = u16::MAX as usize;
    let vectors = (0..=num_candidates)
        .map(|position| vec![position as f32])
        .collect();

[Copilot]

Pull request overview

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

Suppressed comments (2)

diskann/src/graph/index.rs:2575

• The doc comment for occlude_list still claims it clobbers prune::Context::occlude_factor / last_checked, but those fields no longer exist on Context after moving occlusion bookkeeping into prune::State. This is now misleading for maintainers reading the function contract.

    fn occlude_list<M, C, F>(
        &self,
        computer: &C,
        context: &mut prune::Context<'_, DP::InternalId>,
        map: M,

diskann/src/graph/index.rs:2592

• occlude_list creates a fresh Vec cache each call (let mut cache = Vec::new();), but prune::robust_prune is designed to accept a caller-owned cache specifically to reuse capacity across repeated prunes. As written, this reintroduces per-call allocations for large candidate pools and defeats the reuse hook.

            options.force_saturate
                || (self.config.saturate_after_prune() && self.config.alpha() > 1.0),
        );
        let mut cache = Vec::new();
        prune::robust_prune(

[Copilot]

Pull request overview

Copilot reviewed 10 out of 11 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann/src/graph/prune.rs:364

• When policy.saturate is enabled, saturation currently pushes every non-excluded ID from pool even if that candidate was unavailable (lookup returned None). This can reintroduce missing/transient candidates into the output adjacency list, contradicting the earlier “unavailable candidates are excluded” behavior in the main prune loop.

    if policy.saturate {
        for neighbor in pool.iter() {
            if neighbors.len() >= policy.degree {
                break;
            }