PiPNN 2/6: extract shared RobustPrune

[SeliMeli]

PiPNN builds candidate edges in batches, but its final graph must obey the same degree and occlusion policy as Vamana. This PR isolates Vamana's reusable RobustPrune state machine while preserving the existing internal layout and behavior.

No public graph-prune API is introduced. Vamana-specific state remains in graph/internal/vamana_prune.rs; the shared allocation-free algorithm lives beside it in graph/internal/robust_prune.rs.

Concepts

RobustPrune starts from available candidates already sorted by distance to one source. It selects at most graph degree R while rejecting candidates occluded by already selected neighbors. Selection starts at current_alpha = 1.0; later rounds follow the existing Vamana alpha progression.

The shared kernel owns only alpha-round selection. Callers own provider lookup, exclusion, allocation, source-distance sorting/capping, ID translation, adjacency mutation, and optional saturation.

This refactor does not add alpha validation. graph::Config remains the sole owner of alpha behavior, matching main.

Code map

1. diskann/src/graph/internal/vamana_prune.rs
   • Retains the state that main already associated with Vamana: Options, reusable Scratch, borrowed Context, FailedVectorRetrieval, and ranked ListError.
   • These types remain visible only inside crate::graph.
2. diskann/src/graph/internal/robust_prune.rs
   • Candidate is one caller-prepared available value.
   • State retains occlusion and selected-prefix cursors across alpha rounds.
   • robust_prune accepts candidate/state slices plus the existing degree, alpha, prune_kind, and distance callback directly.
   • It performs no allocation/provider access and returns selected candidate positions.
   • Structural errors cover u16 position overflow, state/candidate mismatch, and caller distance failure.
3. diskann/src/graph/index.rs
   • Vamana adapter performs async provider fill, exclusion/availability filtering, candidate preparation, allocation, selected-position-to-ID translation, and saturation.
4. Tests are co-located by layer:
   • internal/robust_prune/tests.rs: pure state machine.
   • internal/vamana_prune/tests.rs: provider/Vamana behavior.
5. diskann/benches/benchmarks_iai/robust_prune.rs benchmarks the Vamana adapter through the shared DiskANN IAI target.

End-to-end flow

Vamana caller computes/fetches candidate state → caller applies its source-distance sorting/capping policy → excluded/unavailable candidates are removed during preparation → pure internal::robust_prune::robust_prune selects positions → Vamana maps positions back to IDs → optional saturation appends only prepared available candidates → provider adjacency is written.

PiPNN joins this seam in #1290 with its own contiguous-matrix preparation and adjacency rewrite; it does not use Vamana scratch or provider errors.

Invariants and boundaries

• graph::internal::{vamana_prune,robust_prune} are not externally nameable.
• Vamana scratch/provider errors are separate from the shared algorithm.
• Kernel input contains only caller-prepared available candidates.
• Candidate order remains caller-owned; selected positions index that exact input.
• Candidate/state slices are one-to-one.
• Candidate positions use u16; exactly u16::MAX is accepted, one more is rejected.
• State::last_checked indexes the selected prefix and survives alpha rounds.
• Kernel performs no heap allocation, provider access, adjacency write, saturation, or configuration validation.
• Provider fill completes before synchronous selection.
• Unavailable candidates are omitted from both selection and saturation.
• Empty reused Vamana scratch clears previous adjacency output.

Review path

1. Start with internal/vamana_prune.rs; compare its state/error types with main.
2. Review internal/robust_prune.rs; verify the interface contains only prepared candidates, state, direct Vamana parameters, distance, and selected count.
3. Trace the alpha-round cursor state machine and u16 position representation.
4. Follow graph/index.rs::occlude_list: preparation → pure kernel → ID translation → available-only saturation.
5. Review pure tests, then Vamana tests for unavailable IDs, saturation, max-occlusion capping, equal-distance order, and alpha rounds.
6. Finish with IAI-Callgrind adapter workloads.

Validation

• Six pure-kernel tests cover fallible distance, exact u16::MAX capacity, one-over-capacity rejection, state mismatch, selected-position order, alpha revisits, and empty input.
• Ten Vamana tests cover bounded rows, equal-distance ordering, both prune kinds, unavailable/self candidates, unavailable-plus-saturation, max-occlusion capping, saturation order, and maximum capacity.
• Consolidation scenarios retain transient provider-error coverage.
• PiPNN finalization tests in PiPNN 3/6: add core graph construction #1290 cover exact output and propagated position overflow at the second caller.
• IAI-Callgrind measures 750-to-64 triangle and saturated-occluding adapter workloads; the refactor reduced their instruction counts by 2.8% and 1.0% in local before/after comparison.

Stack relation

Stack 2/6. Depends on #1287. #1290 consumes only internal::robust_prune, while owning separate PiPNN preparation/postprocessing.

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();

[wuw92]

I know you’re working toward integrating PiPNN into the DiskANN workspace. While extracting RobustPrune for reuse, we can reconsider the abstraction boundary.

• Keep the shared kernel limited to the pure pruning algorithm.
• Let upstream callers prepare data and downstream callers handle post-processing.
• Keep allocation and provider-specific optimizations outside the public algorithm API.
• Expose implementation details only when callers genuinely need to control them.

[Copilot]

Pull request overview

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

Suppressed comments (1)

diskann/src/graph/prune.rs:375

• Saturation currently appends IDs from the original pool without checking whether lookup succeeded. If a candidate is unavailable/transient (i.e., lookup returned None and the candidate was excluded during pruning), saturation can reintroduce that ID into neighbors, producing adjacency entries that the provider/view cannot supply vectors for.

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

[Copilot]

Pull request overview

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

Suppressed comments (1)

diskann/src/graph/index.rs:2575

• The occlude_list doc comment’s “Clobbers” bullets reference prune::Context::occlude_factor and prune::Context::last_checked, but those fields no longer exist (and will produce broken intra-doc links). Update the docs to point at the actual scratch fields that are mutated (states and neighbors).

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

[SeliMeli]

Addressed the abstraction feedback in the latest update:

• Main-compatible Vamana state remains under graph/internal/vamana_prune.rs: Options, Scratch, Context, transient retrieval error, and ranked list error.
• The shared allocation-free state machine lives separately in graph/internal/robust_prune.rs. It accepts caller-prepared candidates/state plus the existing degree, alpha, prune_kind, and distance callback, then returns selected positions.
• No new Policy wrapper and no new alpha validation remain; graph::Config behavior is unchanged.
• Vamana owns provider fill, exclusion/availability filtering, allocation, ID translation, and saturation.
• PiPNN owns contiguous-matrix preparation, reusable Rayon-job workspace, and adjacency rewriting.
• Unavailable candidates are excluded from saturation, with a provider-level regression test.
• Tests are co-located as internal/robust_prune/tests.rs and internal/vamana_prune/tests.rs; the scattered graph/test/cases/prune.rs module was removed.

Kernel, Vamana provider, PiPNN finalization, full PiPNN, Clippy, cross-target, Miri, IAI-Callgrind, and VM E2E checks pass.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

Suppressed comments (1)

diskann/src/graph/internal/vamana_prune.rs:21

• Scratch is declared pub(crate), which makes it accessible from anywhere in the crate. The PR description/invariants state that Vamana-specific state should remain visible only inside crate::graph; Options, Context, and the fields already use pub(in crate::graph).

Consider restricting Scratch to pub(in crate::graph) to keep the intended internal boundary consistent.

#[derive(Debug)]
pub(crate) struct Scratch<I>
where