[diskann-inmem] Prepare code for quantization and beyond

[Mark Hildebrand (hildebrandmw)]

Prepare diskann-inmem for quantization and beyond.

Note that this is infrastructure work to get the code ready. Quantization is not yet integrated.

Goals

The requirements to support quantization are annoyingly orthogonal:

1. Collections like PQ, scalar, and spherical quantization should be able to run in "quant-only" mode and "quant + full-precision" mode, where the full-precision store is used for reranking. Internally, I would like these to be the same type to cut down on unnecessary monomorphization. This means we need support for at least two collections managed by the same epoch protected Store (one quant, the other full-precision). In this case, the full-precision store can do without the invasive tags used for the primary store, but also needs to be optional and ideally support any of f32, f16, u8, i8 and beyond.

2. For testing purposes, we probably want to retain the ability for PQ to do hybrid pruning (part full-precision, part quantized). This completely breaks the current model used by inmem of managing raw &[u8] slices. While we could technically make it work, if we needed to do something like multi-vector operations, a &[u8] is just not the right approach anyways.

3. We also want to be able to support multi-vectors and other kinds of non-uniform data in an epoch guarded Store, which gets rid of the uniform assumption of the current invasive store.

Supporting all of these required a pretty drastic reorganization of diskann-inmem.

Architecture

This PR is all about moving things up and down. The architecture went from this:

Provider

                   Prune
          +======= ExpandBeam  ===+
          |                       |
   +------|-----------------------|------+
   |      |                       |      |
   |  +-------+      +----------------+  |
   |  | layer |      | invasive-store |  |
   |  +-------+      +----------------+  |
   |                                     |
   +-------------------------------------+

where the provider had a separate Layer and Store and combined the two to make a SearchAccessor and PruneAccessor to this

Provider ------ diskann::graph::glue ----+
   +--- drives --+                       |
                 |     SearchAccessor <--+
                 |     PruneAccessor
                 |             ^
                 V             |
+------------------------------|--------+
| Layer ---- bypses ----+      |        |
|   +--- drives --+     |      |        |
|                 |     +-- ExpandBeam  |
|                 v     |   Prune       |
|  +--------------------|-----+         |
|  | Store - drives -+  |     |         |
|  |                 |  |     |         |
|  |                 v  v     |         |
|  |               +--------+ |         |
|  |               | Plugin | |         |
|  |               +--------+ |         |
|  +--------------------------+         |
+---------------------------------------+

From top to bottom:

• The Provider now contains a Layer instead of a Layer and a Store. Instead, the Store has been moved directly into the Layer. Instead of the Provider being responsible for building a SearchAccessor from pieces exposed by the old Layer, the Provider completely delegates SearchAccessor and PruneAccessor construction to its Layer.

  The rest of the Provider's job is interfacing the diskann::graph::glue API to the simplified Layer API.

• Layer now gains the Store. The Layer family of traits is extended to include logical operations like insert and retire. As mentions above, it is also responsible for building SearchAccessors and PruneAccessors.

• Store receives minor changes - its internal Buffer where it used to manage the invasive data store directly has now been moved to a Plugin trait. In this architecture, the Store is just responsible for driving the Plugin trait's lifecycle API and is completely uninvolved with the mechanics of raw data reading and writing.

• Plugin is a new trait for a slot-based store whose slots are driven by Store. The old invasive store is an example of such a Plugin.

In addition to this hierarchical layering, a Layer is allowed (and indeed, expected) to bypass its immediate Store to the underlying Plugin directly to build the various accessors. For the Invasive store, this works because concurrency tags are embedded directly in the Plugin - readers do not need anything to do with the parent store.

Why This Mostly (Probably) Works

The Plugin trait provides an extension point for managing different types of data. For the quantized case, we reuse Invasive for the quantized data and a simpler Buffer-based one for the full-precision reranking data. Then we can create Plugin consisting of both layers. The short-cut from a Layer to its Plugin means a quantized provider can create it's SearchAccessor/PruneAccessors with knowledge of both. We could even have an Invasive store for both the quantized data and the full-precision data and reuse the existing full-precision infrastructure to support both quantized and full-precision searches over the same `Layer.

Additionally, Plugin makes no requirements on the kind of data store. This allows us to store un-even sized allocations in a Plugin. It's the Layer's job to make sense of everything.

Finally, the Layer knowing the details of its Plugins means we can (with some creativity) still support hybrid pruning.

Suggested Reviewing Order

This is a large PR, but I tried very hard to keep things structured. The reviewing order outlined here is a suggested bottom-up order. Understanding how the lower levels work is important for understanding how the higher ones come together.

• num.rs: A quick warm-up. This PR introduces some strongly typed integers with specific semantics.

• prefetch.rs: Another warm-up. I wasn't satisfied with the safety/flexibility of prefetching in the current in-mem provider. This PR exacerbated the situation, so I introduced a bit more structure on prefetchers.

• store/plugin.rs: This defines the Plugin trait and it's expected lifecycle. I captured the nuances in healthy module level and trait level documentation. This is probably the most nuanced change in this PR.

• store/checked.rs: An implementation of Plugin that aggressively checks that the invariants required for the Plugin API are upheld by Store. Again, there is healthy module-level documentation describing the logic.

• store/invasive.rs: The old invasive data store moved to implement the Plugin trait. This largely preserves what was already in the old Store and is conceptually much simpler than store/checked.rs.

• store/mod.rs: Modifying Store to work against a Plugin generic instead of directly managing the invasive store. Note that there are some changes to initialization. Store::new now takes three distinct arguments:

  • Layout: Description for capacity, number of frozen points, and maximum degree.
  • Config: Configuration state dedicated directly to management of the internal concurrency data structures.
  • PluginConfig: Configuration of the internal plugin. The XConfig traits are used in this PR to perform deferred initialization of large data structure.

• Detour: With the introduction of a modularized Store, changes were made to the following integration-test related files to enable concurrent stress-test of different Store implemntations.

  • src/integration/store/*: Shared boiler plate and implementations for exposing different Store to the integration test framework. This PR exposes wrappers for the Invasive and Checked stores in the invasive and checked modules respectively.
  • integration/store/*: Integration test exposure for the different stores as different jobs in the integration test suite.
    For these, the overall structure of the exposed stores is extremely similar. Most of churn in these files is moving things around so keep the amount of repeated code to a minimum.

• layer/mod.rs: Reworked Layer/Set/Search/Insert traits for the new architecture.

  • Layer: Gains a few life-cycle related items.
  • Set: Is now responsible for also obtaining an internal slot for the inserted element.
  • Search/Insert: Reworked to return SearchAccessor and PruneAccessors directly. This is mainly to allow us to keep the internal Store/Plugin details hidden from the public interface.

  In addition:

  • ExpandBeam: Moves from its old location in provider.rs. Otherwise, is mostly unchanged.
  • Prune: A new trait for pruning. Prune implementation now internally "buffer" items in the prune set. This effectively makes the type of the elements being pruned hidden, allowing for hybrid pruning in the future.

• layer/full.rs: This is where everything comes together. The Config is used to group together full-precision related constructor arguments and Full now gains a Store. Importantly, now that Full knows the full details of its store, we can more aggressively optimize ExpandBeam with fewer bounds and length checks. Note that start point initialization is now managed by Full's constructor.

  The implementation of ExpandBeam is taken pretty much directly from the old provider.rs code. Additionally, the tests have gotten more robust with Miri having more coverage of the ExpandBeam implementation and correctness tests for the various distance specializations.

I highly recommend looking at the public docs

RUSTDOCFLAGS="-D rustdoc::all" cargo doc --locked --no-deps --package diskann-inmem

and the private docs

RUSTDOCFLAGS="-D rustdoc::all" cargo doc --locked --no-deps --package diskann-inmem --all-features --document-private-items

to get a feeling for the public API and internal documentation.

AI Disclosure: An agent was used review changes and implementation details, help brainstorm, and make focused edits to documentation.

[Copilot]

Pull request overview

This PR refactors diskann-inmem’s in-memory provider architecture to prepare for future quantization support by separating slot lifecycle management (Store) from storage mechanics (Plugin), and moving Store ownership into the Layer abstraction.

Changes:

• Introduces a store::plugin module and reworks the in-memory Store into a generic driver over Plugin, enabling multiple storage backends under the same EBR lifecycle.
• Reworks Provider/layers integration so layers construct SearchAccessor/PruneAccessor directly (including new pruning buffering/opaque-key plumbing), and updates call sites (bench + integration).
• Adds/updates concurrency stress tests and CI doc builds to cover the reorganized APIs and internal documentation.

Reviewed changes

Copilot reviewed 34 out of 35 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
diskann-inmem/src/tag.rs	Adds a typed AtomicTag::SIZE constant used by plugins/layout computations.
diskann-inmem/src/store/plugin.rs	Adds the Plugin/Slot traits and lifecycle documentation for store-driven EBR transitions.
diskann-inmem/src/store/mod.rs	New generic Store<P> driving plugin lifecycle + tags/freelist/registry/neighbors; adds tests.
diskann-inmem/src/store/invasive.rs	Ports the prior “invasive store” into a Plugin implementation with reader/slot types.
diskann-inmem/src/store/checked.rs	Adds a Plugin implementation that aggressively validates lifecycle invariants (test/integration).
diskann-inmem/src/store.rs	Removes the old monolithic uniform Store implementation (superseded by store/mod.rs + plugins).
diskann-inmem/src/provider.rs	Reworks provider to delegate storage/accessor construction to layers and updates search/prune paths.
diskann-inmem/src/prefetch.rs	Adds structured, checkable prefetch abstractions and tests.
diskann-inmem/src/num.rs	Adds typed integer wrappers (Capacity, MaxDegree, IdLimit) and new Bytes utilities.
diskann-inmem/src/neighbors.rs	Updates neighbors graph storage to use IdLimit/MaxDegree typed wrappers.
diskann-inmem/src/lib.rs	Exposes the new store module publicly and wires in prefetch.
diskann-inmem/src/layers/mod.rs	Redesigns layer traits around construction + accessors + insert/prune hooks; adds internal prune/expand traits.
diskann-inmem/src/integration/store/mod.rs	Adds shared macro boilerplate for integration-test wrappers over different store plugins.
diskann-inmem/src/integration/store/invasive.rs	Integration-test wrapper for the invasive plugin-backed store.
diskann-inmem/src/integration/store/checked.rs	Integration-test wrapper for the checked plugin-backed store.
diskann-inmem/src/integration/store.rs	Removes the old single-store integration wrapper (replaced with per-plugin wrappers).
diskann-inmem/src/integration/counters.rs	Updates integration counter snapshot docs to match the new public surface.
diskann-inmem/src/ids.rs	Switches ID mapping to typed Capacity and updates bounds logic/tests accordingly.
diskann-inmem/src/freelist.rs	Minor doc tweak while keeping freelist mechanics intact.
diskann-inmem/src/counters.rs	Makes LocalCounters public (still unconstructable externally) to support public APIs.
diskann-inmem/integration/support/datatype.rs	Switches f16 conversions to diskann_wide casting utilities for consistency.
diskann-inmem/integration/store/mod.rs	Replaces the old single store stress benchmark with shared infrastructure + per-plugin benchmarks.
diskann-inmem/integration/store/invasive.rs	Adds invasive-store stress benchmark implementation using the new wrappers.
diskann-inmem/integration/store/checked.rs	Adds checked-store stress benchmark implementation using the new wrappers.
diskann-inmem/integration/store.rs	Removes the old monolithic store stress benchmark implementation.
diskann-inmem/integration/main.rs	Updates integration runner to register the new store benchmarks via store::register.
diskann-inmem/integration/jsons/store-stress.json	Updates benchmark job definitions for separate invasive/checked store stress tests.
diskann-inmem/integration/jsons/store-stress-test.json	Updates smaller test job definitions for separate invasive/checked store stress tests.
diskann-inmem/integration/jsons/integration-baseline.json	Updates baseline counter expectations reflecting the new implementation paths.
diskann-inmem/integration/index/runner.rs	Updates index integration runner to construct providers via new Full::config + typed params.
diskann-inmem/Cargo.toml	Adds hashbrown dependency used by new prune buffering logic.
diskann-benchmark/src/index/inmem2.rs	Updates benchmark provider construction to new Full::config API; adds u8 benchmark.
Cargo.lock	Records the new hashbrown dependency for diskann-inmem.
.github/workflows/ci.yml	Tightens doc builds with RUSTDOCFLAGS=-D rustdoc::all, adds inmem private docs, and includes integration-test feature.

Suppressed comments (1)

diskann-inmem/src/integration/store/mod.rs:20

• Spelling in doc comment: "wraper" should be "wrapper".

        /// A test store wraper.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 94.52852% with 94 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.58%. Comparing base (f399158) to head (e6d8400).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-inmem/src/store/mod.rs	91.91%	33 Missing ⚠️
diskann-inmem/src/layers/full.rs	97.00%	19 Missing ⚠️
diskann-inmem/src/store/invasive.rs	93.60%	17 Missing ⚠️
diskann-inmem/src/layers/mod.rs	50.00%	9 Missing ⚠️
diskann-inmem/src/store/checked.rs	94.69%	7 Missing ⚠️
diskann-inmem/src/prefetch.rs	95.78%	4 Missing ⚠️
diskann-inmem/src/num.rs	80.00%	3 Missing ⚠️
diskann-inmem/src/provider.rs	97.82%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1352      +/-   ##
==========================================
+ Coverage   91.54%   91.58%   +0.04%     
==========================================
  Files         521      523       +2     
  Lines      100347   101158     +811     
==========================================
+ Hits        91863    92650     +787     
- Misses       8484     8508      +24     

Flag	Coverage Δ
miri	91.58% <94.52%> (+0.04%)	⬆️
unittests	91.27% <94.52%> (+0.04%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-inmem/src/counters.rs	100.00% <ø> (ø)
diskann-inmem/src/freelist.rs	98.97% <ø> (+1.53%)	⬆️
diskann-inmem/src/ids.rs	100.00% <100.00%> (ø)
diskann-inmem/src/neighbors.rs	100.00% <100.00%> (ø)
diskann-inmem/src/tag.rs	100.00% <ø> (ø)
diskann-inmem/src/provider.rs	95.52% <97.82%> (+6.59%)	⬆️
diskann-inmem/src/num.rs	98.29% <80.00%> (-1.71%)	⬇️
diskann-inmem/src/prefetch.rs	95.78% <95.78%> (ø)
diskann-inmem/src/store/checked.rs	94.69% <94.69%> (ø)
diskann-inmem/src/layers/mod.rs	52.63% <50.00%> (+2.63%)	⬆️
... and 3 more

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