From b5e92f3d1bce0dc6531e4f89e26ebad16b8bd8bb Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Wed, 29 Jul 2026 15:09:19 -0700
Subject: [PATCH 01/34] Prepare to swap in `Invasive`.

---
 diskann-inmem/src/store/invasive.rs          | 255 +++++++++++++++++++
 diskann-inmem/src/{store.rs => store/mod.rs} |   6 +-
 diskann-inmem/src/store/plugin.rs            |  66 +++++
 diskann-inmem/src/tag.rs                     |   5 +
 4 files changed, 329 insertions(+), 3 deletions(-)
 create mode 100644 diskann-inmem/src/store/invasive.rs
 rename diskann-inmem/src/{store.rs => store/mod.rs} (99%)
 create mode 100644 diskann-inmem/src/store/plugin.rs

diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
new file mode 100644
index 0000000000..37125065c1
--- /dev/null
+++ b/diskann-inmem/src/store/invasive.rs
@@ -0,0 +1,255 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{num::NonZeroUsize, sync::atomic::Ordering};
+
+use diskann::utils::IntoUsize;
+
+use crate::{
+    buffer::{Buffer, RawSlice},
+    num::{Bytes, Align},
+    tag::{Tag, AtomicTag},
+    epoch,
+};
+
+#[derive(Debug)]
+pub(crate) struct Invasive {
+    buffer: Buffer,
+    unpadded: Bytes,
+}
+
+const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
+
+impl Invasive {
+    pub(crate) fn new(entries: usize, bytes: Bytes) -> Self {
+        let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
+        let padded_bytes = unpadded
+            .checked_next_multiple_of(Bytes::CACHELINE.div(TWO)).unwrap();
+
+        Self {
+            buffer: Buffer::new(entries, padded_bytes, Align::_128).unwrap(),
+            unpadded,
+        }
+    }
+
+    pub(crate) unsafe fn reader<'a>(&'a self, guard: epoch::Guard<'a>) -> Reader<'a> {
+        Reader {
+            buffer: &self.buffer,
+            unpadded: self.unpadded,
+            _guard: guard,
+        }
+    }
+
+    /// Return the data at position `i` without bound-checking.
+    ///
+    /// # Safety
+    ///
+    /// The index `i` must be less then `self.buffer.len()`.
+    unsafe fn data_unchecked(&self, i: usize) -> (&AtomicTag, RawSlice<'_>) {
+        // SAFETY: inherited from caller.
+        let (data, mirror) = unsafe { self.buffer.get_unchecked(i) }
+            .truncate(self.unpadded)
+            .split(self.unpadded.unchecked_sub(AtomicTag::SIZE));
+        (
+            // SAFETY: We're careful in this module to ensure the inline tags are only
+            // ever accessed atomically.
+            unsafe { AtomicTag::from_ptr(mirror.as_mut_ptr().cast()) },
+            data,
+        )
+    }
+
+    fn data(&self, i: usize) -> Option<(&AtomicTag, RawSlice<'_>)> {
+        if i >= self.buffer.len() {
+            None
+        } else {
+            Some(unsafe { self.data_unchecked(i) })
+        }
+    }
+}
+
+impl super::plugin::Plugin for Invasive {
+    type Slot<'a> = Slot<'a>;
+
+    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
+        let Some((tag, data)) = self.data(i.into_usize()) else {
+            panic!("index {i} is out-of-bounds");
+        };
+
+        // This is a pessimistic check to ensure that the caller is correctly using the
+        // `plugin` API.
+        assert_eq!(
+            tag.load(Ordering::Relaxed),
+            Tag::AVAILABLE,
+            "concurrency violation",
+        );
+
+        // While we can leave this tag as `Tag::AVAILABLE` since it's just a mirror, setting
+        // it to `Tag::OWNED` lets us more precisely detect misuse from the caller.
+        tag.store(Tag::OWNED, Ordering::Relaxed);
+        Slot { tag, data }
+    }
+
+    fn reclaim(&self, i: u32) {
+        let Some((tag, _)) = self.data(i.into_usize()) else {
+            panic!("index {i} is out-of-bounds");
+        };
+
+        tag.store(Tag::AVAILABLE, Ordering::Release);
+    }
+
+    fn retire(&self, i: u32) {
+        let Some((tag, _)) = self.data(i.into_usize()) else {
+            panic!("index {i} is out-of-bounds");
+        };
+
+        tag.store(Tag::RETIRING, Ordering::Relaxed);
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Reader<'a> {
+    buffer: &'a Buffer,
+    unpadded: Bytes,
+    _guard: epoch::Guard<'a>,
+}
+
+impl<'a> Reader<'a> {
+    /// Attempt to read the value at index `i`. This can fail for any of the
+    /// following reasons:
+    ///
+    /// 1. Index `i` is out-of-bounds.
+    /// 2. The read cannot be guaranteed to be race-free.
+    #[inline]
+    pub(crate) fn read(&self, i: usize) -> Option<&[u8]> {
+        if self.is_in_bounds(i) {
+            // SAFETY: `i` is in-bounds.
+            unsafe { self.read_in_bounds(i) }
+        } else {
+            None
+        }
+    }
+
+    /// Return `true` if the index `i` is in-bounds.
+    #[inline]
+    #[must_use = "this function has no side-effects"]
+    pub(crate) fn is_in_bounds(&self, i: usize) -> bool {
+        i < self.buffer.len()
+    }
+
+    /// Return `true` if it is safe to read the data at position `i`.
+    ///
+    /// This guarantee only holds while `self` is alive. Construction of a new [`Reader`]
+    /// requires a separate check.
+    #[cfg_attr(
+        not(test),
+        expect(
+            dead_code,
+            reason = "this is non-trivial method that likely be used in the future"
+        )
+    )]
+    pub(crate) fn can_read(&self, i: usize) -> Option<bool> {
+        if !self.is_in_bounds(i) {
+            return None;
+        }
+
+        // SAFETY: We've checked that `i` is in-bounds.
+        //
+        // Further, we guarantee that `self.unpadded >= AtomicTag:::SIZE`, so the pointer
+        // arithmetic is in-bounds.
+        let tag_ptr = unsafe {
+            self.buffer
+                .get_unchecked(i)
+                .as_mut_ptr()
+                .add(self.unpadded.unchecked_sub(AtomicTag::SIZE).value())
+        };
+
+        // SAFETY: We only access tag pointers atomically.
+        let can_read = unsafe { AtomicTag::from_ptr(tag_ptr.cast()) }
+            .load(Ordering::Acquire)
+            .can_read();
+
+        Some(can_read)
+    }
+
+    /// Read the data as position `i` if it is guaranteed to be race-free without bounds
+    /// checking.
+    ///
+    /// # Safety
+    ///
+    /// The index `i` must satisfy [`Self::is_in_bounds`].
+    #[inline]
+    pub(crate) unsafe fn read_in_bounds(&self, i: usize) -> Option<&[u8]> {
+        debug_assert!(self.is_in_bounds(i));
+
+        // SAFETY:
+        //
+        // * The caller asserts `i` is in-bounds.
+        // * We maintain an internal invariant that `self.buffer.stride() <= self.unpadded`.
+        // * Further, we maintain that `self.unpadded >= AtomicTag::SIZE`.
+        let (data, tag_ptr) = unsafe {
+            self.buffer
+                .get_unchecked(i)
+                .truncate_unchecked(self.unpadded)
+                .split_unchecked(self.unpadded.unchecked_sub(AtomicTag::SIZE))
+        };
+
+        // NOTE: Must be `Acquire` to correctly synchronize with writes.
+        //
+        // SAFETY: We are careful in this module to ensure that inline tags are only accessed
+        // atomically.
+        let can_read = unsafe { AtomicTag::from_ptr(tag_ptr.as_mut_ptr().cast()) }
+            .load(Ordering::Acquire)
+            .can_read();
+
+        if can_read {
+            // SAFETY: We've passed the `can_read` check - `_guard` will ensure the read
+            // slice is valid and race-free.
+            Some(unsafe { data.as_slice() })
+        } else {
+            None
+        }
+    }
+
+    /// Return the raw data slice for index `i` without any race guarantees.
+    ///
+    /// # Safety
+    ///
+    /// The index `i` must be satisfy [`Self::is_in_bounds`].
+    #[inline]
+    pub(crate) unsafe fn read_raw_unchecked(&self, i: usize) -> RawSlice<'_> {
+        // SAFETY: Inherited from caller: `i` is inbounds.
+        unsafe { self.buffer.get_unchecked(i) }.truncate(self.unpadded)
+    }
+
+    /// Return the number of bytes for each entry.
+    pub(crate) fn bytes(&self) -> Bytes {
+        self.unpadded
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Slot<'a> {
+    tag: &'a AtomicTag,
+    data: RawSlice<'a>,
+}
+
+impl<'a> Slot<'a> {
+    pub(crate) unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe { self.data.as_mut_slice() }
+    }
+}
+
+impl super::plugin::Slot for Slot<'_> {
+    fn publish(self) {
+        self.tag.store(Tag::PUBLISHED, Ordering::Release);
+    }
+    fn freeze(self) {
+        self.tag.store(Tag::FROZEN, Ordering::Release);
+    }
+    fn abort(self) {
+        self.tag.store(Tag::AVAILABLE, Ordering::Release);
+    }
+}
+
diff --git a/diskann-inmem/src/store.rs b/diskann-inmem/src/store/mod.rs
similarity index 99%
rename from diskann-inmem/src/store.rs
rename to diskann-inmem/src/store/mod.rs
index cd845230b4..0213fe8d6f 100644
--- a/diskann-inmem/src/store.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -70,6 +70,9 @@ use crate::{
     tag::{AtomicTag, Tag},
 };
 
+pub(crate) mod plugin;
+pub(crate) mod invasive;
+
 /// Configuration for the concurrenct store.
 #[derive(Debug)]
 pub(crate) struct Config {
@@ -170,9 +173,6 @@ impl Store {
     /// within the store and must be compatible the the number of bytes in `config`.
     pub(crate) fn new(
         config: Config,
-        // entries: usize,
-        // bytes: Bytes,
-        // max_neighbors: usize,
         init: MatrixView<'_, u8>,
     ) -> Result<Self, StoreError> {
         let Config {
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
new file mode 100644
index 0000000000..73ed5667c5
--- /dev/null
+++ b/diskann-inmem/src/store/plugin.rs
@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{mem::ManuallyDrop, fmt::Debug, sync::atomic::Ordering};
+
+use crate::{
+    buffer::{Buffer, RawSlice},
+    epoch,
+    num::Bytes,
+    tag::AtomicTag,
+};
+
+pub(crate) trait Plugin: 'static {
+    type Slot<'a>: Slot;
+
+    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_>;
+    fn reclaim(&self, i: u32);
+    fn retire(&self, i: u32);
+}
+
+pub(crate) trait Slot: Debug {
+    fn publish(self);
+    fn freeze(self);
+    fn abort(self);
+}
+
+#[derive(Debug)]
+pub(crate) struct ManagedSlot<T>
+where
+    T: Slot,
+{
+    slot: ManuallyDrop<T>
+}
+
+impl<T> ManagedSlot<T>
+where
+    T: Slot,
+{
+    fn new(slot: T) -> Self {
+        Self {
+            slot: ManuallyDrop::new(slot),
+        }
+    }
+
+    unsafe fn publish(self) {
+        let mut me = ManuallyDrop::new(self);
+        unsafe { ManuallyDrop::take(&mut me.slot).publish() }
+    }
+
+    unsafe fn freeze(self) {
+        let mut me = ManuallyDrop::new(self);
+        unsafe { ManuallyDrop::take(&mut me.slot).freeze() }
+    }
+}
+
+impl<T> Drop for ManagedSlot<T>
+where
+    T: Slot,
+{
+    fn drop(&mut self) {
+        unsafe { ManuallyDrop::take(&mut self.slot).abort() }
+    }
+}
+
diff --git a/diskann-inmem/src/tag.rs b/diskann-inmem/src/tag.rs
index 4840583fe0..27240bbe34 100644
--- a/diskann-inmem/src/tag.rs
+++ b/diskann-inmem/src/tag.rs
@@ -16,6 +16,8 @@
 
 use std::sync::atomic::{AtomicU8, Ordering};
 
+use crate::num::Bytes;
+
 /// A tag for controlling concurrent access to data.
 ///
 /// Tag updates and reads should use [`AtomicTag`].
@@ -169,6 +171,9 @@ impl std::fmt::Display for Tag {
 pub(crate) struct AtomicTag(AtomicU8);
 
 impl AtomicTag {
+    /// The size of an [`AtomicTag`].
+    pub(crate) const SIZE: Bytes = Bytes::size_of::<Self>();
+
     /// Construct a new [`AtomicTag`] initialized to `tag`.
     pub(crate) const fn new(tag: Tag) -> Self {
         Self(AtomicU8::new(tag.value()))

From 2c4625579da32352383c36e1af2b049c30a07bae Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 30 Jul 2026 09:27:56 -0700
Subject: [PATCH 02/34] Move store backend into a plugin system.

---
 diskann-inmem/src/integration/store.rs |   2 +-
 diskann-inmem/src/provider.rs          |  18 +-
 diskann-inmem/src/store/invasive.rs    |  17 +-
 diskann-inmem/src/store/mod.rs         | 301 ++++++-------------------
 diskann-inmem/src/store/plugin.rs      |  48 +---
 diskann-inmem/src/store/stacked.rs     |  74 ++++++
 6 files changed, 173 insertions(+), 287 deletions(-)
 create mode 100644 diskann-inmem/src/store/stacked.rs

diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 352e36572b..9a8b9865f7 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -107,7 +107,7 @@ impl<'a> Reader<'a> {
     }
 
     pub fn read(&self, i: usize) -> Option<&[u8]> {
-        self.reader.read(i)
+        self.reader.inner().read(i)
     }
 }
 
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 6167276186..14de88fce5 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -419,7 +419,7 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     {
         let work = move || {
             for p in self.start_points.clone() {
-                match self.reader.read(p.into_usize()) {
+                match self.reader.inner().read(p.into_usize()) {
                     Some(point) => {
                         // Counters are no-ops without `integration-test`.
                         self.counters.get_vector(1);
@@ -636,16 +636,16 @@ where
     T: layers::QueryDistance,
 {
     debug_assert!(
-        BYTES + store::TAG_SIZE.value() <= reader.bytes().value(),
+        BYTES + store::TAG_SIZE.value() <= reader.inner().bytes().value(),
         "we really rely on this: {}, bytes = {}",
         BYTES + store::TAG_SIZE.value(),
-        reader.bytes()
+        reader.inner().bytes()
     );
 
     debug_assert!(buffer.len() >= list.len());
 
     let bytes = if BYTES == 0 {
-        reader.bytes().value()
+        reader.inner().bytes().value()
     } else {
         BYTES + store::TAG_SIZE.value()
     };
@@ -659,6 +659,7 @@ where
         unsafe {
             prefetch(
                 reader
+                    .inner()
                     .read_raw_unchecked(list.get_unchecked(j).into_usize())
                     .as_ptr()
                     .cast(),
@@ -677,6 +678,7 @@ where
             unsafe {
                 prefetch(
                     reader
+                        .inner()
                         .read_raw_unchecked(list.get_unchecked(j).into_usize())
                         .as_ptr()
                         .cast(),
@@ -687,7 +689,7 @@ where
         }
 
         // SAFETY: Caller asserts that `i` is in-bounds.
-        if let Some(data) = unsafe { reader.read_in_bounds(i.into_usize()) } {
+        if let Some(data) = unsafe { reader.inner().read_in_bounds(i.into_usize()) } {
             // SAFETY: Inherited from caller.
             *unsafe { buffer.get_unchecked_mut(processed) } = (i, distance.evaluate(data)?);
             processed += 1;
@@ -837,7 +839,7 @@ impl workingset::View<u32> for &PruneAccessor<'_> {
     where
         Self: 'a;
     fn get(&self, id: u32) -> Option<&[u8]> {
-        match self.reader.read(id.into_usize()) {
+        match self.reader.inner().read(id.into_usize()) {
             Some(data) => {
                 self.counters.get_vector_ref(1);
                 Some(data)
@@ -873,7 +875,7 @@ where
             &provider.layer,
             query,
             ExpandBeamVisitor {
-                bytes: provider.store.bytes(),
+                bytes: provider.store.plugin().bytes(),
                 prefetch_lookahead: provider.config.prefetch_lookahead.map_or(0, |x| x.get()),
             },
         )?;
@@ -1030,7 +1032,7 @@ where
     {
         let work = move || {
             let reader = provider.store.reader()?;
-            let data = match reader.read(id.into_usize()) {
+            let data = match reader.inner().read(id.into_usize()) {
                 Some(data) => data,
                 None => {
                     return Err(ANNError::message(
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index 37125065c1..de92811cd4 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -9,14 +9,19 @@ use diskann::utils::IntoUsize;
 
 use crate::{
     buffer::{Buffer, RawSlice},
-    num::{Bytes, Align},
-    tag::{Tag, AtomicTag},
     epoch,
+    num::{Align, Bytes},
+    tag::{AtomicTag, Tag},
 };
 
+/// The invasive store where concurrency tags are stored inline with the data.
 #[derive(Debug)]
 pub(crate) struct Invasive {
+    // The inline tags are `AtomicTag`s stored after the data.
     buffer: Buffer,
+
+    // The unpadded size of each row in `buffer`. This includes both the data **and** the
+    // 1-byte tag. Tags are located at byte `unpadded - 1`.
     unpadded: Bytes,
 }
 
@@ -26,7 +31,8 @@ impl Invasive {
     pub(crate) fn new(entries: usize, bytes: Bytes) -> Self {
         let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
         let padded_bytes = unpadded
-            .checked_next_multiple_of(Bytes::CACHELINE.div(TWO)).unwrap();
+            .checked_next_multiple_of(Bytes::CACHELINE.div(TWO))
+            .unwrap();
 
         Self {
             buffer: Buffer::new(entries, padded_bytes, Align::_128).unwrap(),
@@ -34,6 +40,10 @@ impl Invasive {
         }
     }
 
+    pub(crate) fn bytes(&self) -> Bytes {
+        self.unpadded
+    }
+
     pub(crate) unsafe fn reader<'a>(&'a self, guard: epoch::Guard<'a>) -> Reader<'a> {
         Reader {
             buffer: &self.buffer,
@@ -252,4 +262,3 @@ impl super::plugin::Slot for Slot<'_> {
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
-
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 0213fe8d6f..7f08cf9b3c 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -53,6 +53,7 @@
 
 use std::{
     iter::repeat_n,
+    mem::ManuallyDrop,
     num::{NonZeroU32, NonZeroUsize},
     sync::atomic::Ordering,
 };
@@ -62,17 +63,20 @@ use diskann_utils::views::MatrixView;
 use thiserror::Error;
 
 use crate::{
-    buffer::{Buffer, BufferError, RawSlice},
+    buffer::{BufferError},
     epoch::{self, Registry},
     freelist::{self, Freelist},
     neighbors::{Neighbors, NeighborsError},
-    num::{Align, Bytes},
+    num::{Bytes},
     tag::{AtomicTag, Tag},
 };
 
 pub(crate) mod plugin;
+pub(crate) mod stacked;
 pub(crate) mod invasive;
 
+pub(crate) const TAG_SIZE: Bytes = AtomicTag::SIZE;
+
 /// Configuration for the concurrenct store.
 #[derive(Debug)]
 pub(crate) struct Config {
@@ -134,17 +138,8 @@ impl Config {
 /// A concurrent data and graph store.
 #[derive(Debug)]
 pub(crate) struct Store {
-    // The invasive store where concurrency tags are stored inline with the data.
-    //
-    // These tags are mirrored from `tags` - with the latter being used for secondary scans
-    // offering slightly better locality.
-    //
-    // The inline tags are stored after the data.
-    buffer: Buffer,
-
-    // The unpadded size of each row in `buffer`. This includes both the data **and** the
-    // 1-byte tag. Tags are located at byte `unpadded - 1`.
-    unpadded: Bytes,
+    // This is a temporary concrete type until [`Store`] is properly parameterized by its plugin.
+    plugin: invasive::Invasive,
 
     // The number of unfrozen points. This is guaranteed to be less than `buffer`.
     unfrozen: usize,
@@ -160,21 +155,13 @@ pub(crate) struct Store {
     neighbors: Neighbors,
 }
 
-/// The number of bytes occupied by the in-line concurrency tag.
-pub(crate) const TAG_SIZE: Bytes = Bytes::size_of::<AtomicTag>();
-
-const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
-
 // TODO: This is a guess and probably needs tuning.
 const RETRY_LIMIT: usize = 20;
 
 impl Store {
     /// Create a new [`Store`]. The entries within `init` will be used as frozen points
     /// within the store and must be compatible the the number of bytes in `config`.
-    pub(crate) fn new(
-        config: Config,
-        init: MatrixView<'_, u8>,
-    ) -> Result<Self, StoreError> {
+    pub(crate) fn new(config: Config, init: MatrixView<'_, u8>) -> Result<Self, StoreError> {
         let Config {
             entries,
             bytes,
@@ -191,25 +178,6 @@ impl Store {
             return Err(StoreError::need_frozen_point());
         }
 
-        #[expect(
-            clippy::expect_used,
-            reason = "we expect `init` to have at least one row, so this should never happen"
-        )]
-        let unpadded = bytes
-            .checked_add(TAG_SIZE)
-            .expect("unreachable because `init` cannot exceed `isize::MAX` bytes");
-
-        // Pad to half a cache line. When data occupies just part of a cache line, this
-        // results in the same total number of cache lines being fetched while potentially
-        // enabling more compact memory.
-        #[expect(
-            clippy::expect_used,
-            reason = "we expect `init` to have at least one row, so this should never happen"
-        )]
-        let padded_bytes = unpadded
-            .checked_next_multiple_of(Bytes::CACHELINE.div(TWO))
-            .expect("unreachable because `init` cannot exceed `isize::MAX` bytes");
-
         let too_many_entries = || StoreError::too_many_entries(entries, init.nrows());
 
         // We have a hard upper-bound of `u32::MAX` total slots.
@@ -226,8 +194,7 @@ impl Store {
             .map_err(|_| StoreError::too_many_neighbors(max_neighbors))?;
 
         let me = Self {
-            buffer: Buffer::new(total.into_usize(), padded_bytes, Align::_128)?,
-            unpadded,
+            plugin: invasive::Invasive::new(total.into_usize(), bytes),
             unfrozen: entries.into_usize(),
             tags: repeat_n(Tag::AVAILABLE, total.into_usize())
                 .map(AtomicTag::new)
@@ -256,14 +223,13 @@ impl Store {
         Ok(me)
     }
 
-    /// Return the range of slots containing frozen items in `self`.
-    pub(crate) fn frozen(&self) -> std::ops::Range<u32> {
-        (self.unfrozen as u32)..(self.buffer.len() as u32)
+    pub(crate) fn plugin(&self) -> &invasive::Invasive {
+        &self.plugin
     }
 
-    /// Return the number of bytes occupied by each entry.
-    pub(crate) fn bytes(&self) -> Bytes {
-        self.unpadded
+    /// Return the range of slots containing frozen items in `self`.
+    pub(crate) fn frozen(&self) -> std::ops::Range<u32> {
+        (self.unfrozen as u32)..self.neighbors.entries()
     }
 
     /// Return the maximum degree that can be stored in the graph.
@@ -275,7 +241,21 @@ impl Store {
     ///
     /// If successful, returns the number of slots reclaimed.
     pub(crate) fn try_drain(&self) -> Option<usize> {
-        fn release(tag: &AtomicTag, kind: &'static str) {
+        let drain = self.registry.try_advance()?;
+        let items = drain.len();
+        for i in drain {
+            let Some(tag) = self.tags.get(i.into_usize()) else {
+                panic!(
+                    "received an invalid ID ({}) while reclaiming slots - max allowed is {}",
+                    i,
+                    self.neighbors.entries(),
+                );
+            };
+
+            // We release the plugin before the main tag. The other direction would
+            // prematurely advertise availability.
+            plugin::Plugin::reclaim(self.plugin(), i);
+
             // Use `Release` ordering to ensure that the store to the mirror cannot get moved
             // after the store to the authoritative list.
             //
@@ -287,30 +267,10 @@ impl Store {
             assert_eq!(
                 tag.load(Ordering::Relaxed),
                 Tag::RETIRING,
-                "CONCURRENCY VIOLATION: {}",
-                kind,
+                "CONCURRENCY VIOLATION",
             );
 
             tag.store(Tag::AVAILABLE, Ordering::Release);
-        }
-
-        let drain = self.registry.try_advance()?;
-        let items = drain.len();
-        for i in drain {
-            assert!(
-                i.into_usize() < self.buffer.len(),
-                "received an invalid ID ({}) while reclaiming slots - max allowed is {}",
-                i,
-                self.buffer.len(),
-            );
-
-            // We release the mirror before the main tag. The other direction would
-            // prematurely advertise availability.
-            //
-            // SAFETY: We've verified that `i` is in-bounds.
-            let (mirror, _) = unsafe { self.data_unchecked(i.into_usize()) };
-            release(mirror, "mirror");
-            release(&self.tags[i.into_usize()], "tag");
             self.freelist.push(i);
         }
         Some(items)
@@ -323,10 +283,8 @@ impl Store {
     /// Returns [`epoch::Unavailable`] if there are too many active readers.
     pub(crate) fn reader(&self) -> Result<Reader<'_>, epoch::Unavailable> {
         Ok(Reader {
-            buffer: &self.buffer,
-            unpadded: self.unpadded,
+            inner: unsafe { self.plugin().reader(self.registry.guard()?) },
             neighbors: &self.neighbors,
-            _guard: self.registry.guard()?,
         })
     }
 
@@ -390,10 +348,7 @@ impl Store {
         match tag.compare_exchange(current, retiring, Ordering::Relaxed, Ordering::Relaxed) {
             Ok(_) => {
                 // Set the metadata in the mirror as well.
-                //
-                // SAFETY: We've checked that `i` is in-bounds.
-                let (mirror, _) = unsafe { self.data_unchecked(i) };
-                mirror.store(retiring, Ordering::Relaxed);
+                plugin::Plugin::retire(self.plugin(), i.try_into().unwrap());
                 guard.retire(i as u32);
                 Ok(())
             }
@@ -488,12 +443,10 @@ impl Store {
             Ordering::Relaxed,
         ) {
             Ok(_) => {
-                // SAFETY: Inherited from caller - `slot` is in-bounds.
-                let (mirror, data) = unsafe { self.data_unchecked(slot.into_usize()) };
+                let data = unsafe { plugin::Plugin::acquire(self.plugin(), slot) };
                 Some(Slot {
                     tag,
-                    mirror,
-                    data,
+                    data: ManuallyDrop::new(data),
                     slot,
                 })
             }
@@ -501,24 +454,6 @@ impl Store {
         }
     }
 
-    /// Return the data at position `i` without bound-checking.
-    ///
-    /// # Safety
-    ///
-    /// The index `i` must be less then `self.buffer.len()`.
-    unsafe fn data_unchecked(&self, i: usize) -> (&AtomicTag, RawSlice<'_>) {
-        // SAFETY: inherited from caller.
-        let (data, mirror) = unsafe { self.buffer.get_unchecked(i) }
-            .truncate(self.unpadded)
-            .split(self.unpadded.unchecked_sub(TAG_SIZE));
-        (
-            // SAFETY: We're careful in this module to ensure the inline tags are only
-            // ever accessed atomically.
-            unsafe { AtomicTag::from_ptr(mirror.as_mut_ptr().cast()) },
-            data,
-        )
-    }
-
     /// Return whether or not it is probably okay to read from the slot `i`.
     ///
     /// This check is approximate and non-synchronizing. To fully check, [`Reader::can_read`]
@@ -618,124 +553,21 @@ pub(crate) enum RetireError {
 /// Created via [`Store::reader`].
 #[derive(Debug)]
 pub(crate) struct Reader<'a> {
-    buffer: &'a Buffer,
-    unpadded: Bytes,
+    inner: invasive::Reader<'a>,
     neighbors: &'a Neighbors,
-    // It's important that we hold onto this, even if we don't use it.
-    _guard: epoch::Guard<'a>,
 }
 
 impl<'a> Reader<'a> {
-    /// Attempt to read the value at index `i`. This can fail for any of the
-    /// following reasons:
-    ///
-    /// 1. Index `i` is out-of-bounds.
-    /// 2. The read cannot be guaranteed to be race-free.
-    #[inline]
-    pub(crate) fn read(&self, i: usize) -> Option<&[u8]> {
-        if self.is_in_bounds(i) {
-            // SAFETY: `i` is in-bounds.
-            unsafe { self.read_in_bounds(i) }
-        } else {
-            None
-        }
-    }
-
     /// Return `true` if the index `i` is in-bounds.
     #[inline]
     #[must_use = "this function has no side-effects"]
     pub(crate) fn is_in_bounds(&self, i: usize) -> bool {
-        i < self.buffer.len()
-    }
-
-    /// Return `true` if it is safe to read the data at position `i`.
-    ///
-    /// This guarantee only holds while `self` is alive. Construction of a new [`Reader`]
-    /// requires a separate check.
-    #[cfg_attr(
-        not(test),
-        expect(
-            dead_code,
-            reason = "this is non-trivial method that likely be used in the future"
-        )
-    )]
-    pub(crate) fn can_read(&self, i: usize) -> Option<bool> {
-        if !self.is_in_bounds(i) {
-            return None;
-        }
-
-        // SAFETY: We've checked that `i` is in-bounds.
-        //
-        // Further, we guarantee that `self.unpadded >= TAG_SIZE`, so the pointer arithmetic
-        // is in-bounds.
-        let tag_ptr = unsafe {
-            self.buffer
-                .get_unchecked(i)
-                .as_mut_ptr()
-                .add(self.unpadded.unchecked_sub(TAG_SIZE).value())
-        };
-
-        // SAFETY: We only access tag pointers atomically.
-        let can_read = unsafe { AtomicTag::from_ptr(tag_ptr.cast()) }
-            .load(Ordering::Acquire)
-            .can_read();
-
-        Some(can_read)
-    }
-
-    /// Read the data as position `i` if it is guaranteed to be race-free without bounds
-    /// checking.
-    ///
-    /// # Safety
-    ///
-    /// The index `i` must satisfy [`Self::is_in_bounds`].
-    #[inline]
-    pub(crate) unsafe fn read_in_bounds(&self, i: usize) -> Option<&[u8]> {
-        debug_assert!(self.is_in_bounds(i));
-
-        // SAFETY:
-        //
-        // * The caller asserts `i` is in-bounds.
-        // * We maintain an internal invariant that `self.buffer.stride() <= self.unpadded`.
-        // * Further, we maintain that `self.unpadded >= TAG_SIZE`.
-        let (data, tag_ptr) = unsafe {
-            self.buffer
-                .get_unchecked(i)
-                .truncate_unchecked(self.unpadded)
-                .split_unchecked(self.unpadded.unchecked_sub(TAG_SIZE))
-        };
-
-        // NOTE: Must be `Acquire` to correctly synchronize with writes.
-        //
-        // SAFETY: We are careful in this module to ensure that inline tags are only accessed
-        // atomically.
-        let can_read = unsafe { AtomicTag::from_ptr(tag_ptr.as_mut_ptr().cast()) }
-            .load(Ordering::Acquire)
-            .can_read();
-
-        if can_read {
-            // SAFETY: We've passed the `can_read` check - `_guard` will ensure the read
-            // slice is valid and race-free.
-            Some(unsafe { data.as_slice() })
-        } else {
-            None
-        }
+        i < self.neighbors.entries().into_usize()
     }
 
-    /// Return the raw data slice for index `i` without any race guarantees.
-    ///
-    /// # Safety
-    ///
-    /// The index `i` must be satisfy [`Self::is_in_bounds`].
     #[inline]
-    pub(crate) unsafe fn read_raw_unchecked(&self, i: usize) -> RawSlice<'_> {
-        // SAFETY: Inherited from caller: `i` is inbounds.
-        unsafe { self.buffer.get_unchecked(i) }.truncate(self.unpadded)
-    }
-
-    /// Return the number of bytes for each entry.
-    pub(crate) fn bytes(&self) -> Bytes {
-        self.unpadded
+    pub(crate) fn inner(&self) -> &invasive::Reader<'_> {
+        &self.inner
     }
 
     /// Return [`Neighbors`].
@@ -748,8 +580,7 @@ impl<'a> Reader<'a> {
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
     tag: &'a AtomicTag,
-    mirror: &'a AtomicTag,
-    data: RawSlice<'a>,
+    data: ManuallyDrop<invasive::Slot<'a>>,
     slot: u32,
 }
 
@@ -766,8 +597,13 @@ impl<'a> Slot<'a> {
     }
 
     fn freeze(self) {
-        let me = std::mem::ManuallyDrop::new(self);
-        me.mirror.store(Tag::FROZEN, Ordering::Release);
+        // Suppress normal `Drop`.
+        let mut me = ManuallyDrop::new(self);
+
+        // Freeze the inner slot.
+        plugin::Slot::freeze(unsafe { ManuallyDrop::take(&mut me.data) });
+
+        // Update the authoritative store.
         me.tag.store(Tag::FROZEN, Ordering::Release);
     }
 
@@ -776,8 +612,14 @@ impl<'a> Slot<'a> {
     /// Return the internal slot ID.
     pub(crate) fn publish(self) -> u32 {
         let id = self.slot();
-        let me = std::mem::ManuallyDrop::new(self);
-        me.mirror.store(Tag::PUBLISHED, Ordering::Release);
+
+        // Suppress normal `Drop`.
+        let mut me = ManuallyDrop::new(self);
+
+        // Publish the inner slot.
+        plugin::Slot::publish(unsafe { ManuallyDrop::take(&mut me.data) });
+
+        // Update the authoritative store.
         me.tag.store(Tag::PUBLISHED, Ordering::Release);
         id
     }
@@ -785,7 +627,7 @@ impl<'a> Slot<'a> {
 
 impl Drop for Slot<'_> {
     fn drop(&mut self) {
-        self.mirror.store(Tag::AVAILABLE, Ordering::Release);
+        plugin::Slot::abort(unsafe { ManuallyDrop::take(&mut self.data) });
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
@@ -878,21 +720,21 @@ mod tests {
         let reader = s.reader().unwrap();
         for i in 0..4 {
             assert!(!s.can_read_approximate(i).unwrap());
-            assert!(!reader.can_read(i).unwrap());
-            assert!(reader.read(i).is_none());
+            assert!(!reader.inner().can_read(i).unwrap());
+            assert!(reader.inner().read(i).is_none());
         }
 
         assert!(s.can_read_approximate(4).unwrap());
-        assert!(reader.can_read(4).unwrap());
-        assert_eq!(reader.read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
+        assert!(reader.inner().can_read(4).unwrap());
+        assert_eq!(reader.inner().read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
 
         assert!(s.can_read_approximate(5).unwrap());
-        assert!(reader.can_read(5).unwrap());
-        assert_eq!(reader.read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
+        assert!(reader.inner().can_read(5).unwrap());
+        assert_eq!(reader.inner().read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
 
         assert!(s.can_read_approximate(6).is_none());
-        assert!(reader.can_read(6).is_none());
-        assert!(reader.read(6).is_none());
+        assert!(reader.inner().can_read(6).is_none());
+        assert!(reader.inner().read(6).is_none());
     }
 
     ///////////////
@@ -912,13 +754,16 @@ mod tests {
                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
 
             // Before the slot is dropped - we should not be able to read it.
-            assert!(reader.read(idx).is_none());
+            assert!(reader.inner().read(idx).is_none());
             assert!(!s.can_read_approximate(idx).unwrap());
             slot.publish();
             idx
         };
 
-        assert_eq!(reader.read(idx), Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice()));
+        assert_eq!(
+            reader.inner().read(idx),
+            Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice())
+        );
         assert!(s.can_read_approximate(idx).unwrap());
     }
 
@@ -935,14 +780,14 @@ mod tests {
                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
 
             // Before the slot is dropped - we should not be able to read it.
-            assert!(reader.read(idx).is_none());
+            assert!(reader.inner().read(idx).is_none());
             assert!(!s.can_read_approximate(idx).unwrap());
 
             // NOTE: We do not explicitly publish the slot.
             idx
         };
 
-        assert!(reader.read(idx).is_none());
+        assert!(reader.inner().read(idx).is_none());
         assert!(!s.can_read_approximate(idx).unwrap());
     }
 
@@ -1003,8 +848,8 @@ mod tests {
 
         // A reader opened after retirement must not observe the retired slot.
         let reader = s.reader().unwrap();
-        assert_eq!(reader.read(idx), None);
-        assert_eq!(reader.can_read(idx), Some(false));
+        assert_eq!(reader.inner().read(idx), None);
+        assert_eq!(reader.inner().can_read(idx), Some(false));
 
         // The slot can also not be retired again.
         assert!(matches!(
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index 73ed5667c5..3a314ff3af 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -3,16 +3,9 @@
  * Licensed under the MIT license.
  */
 
-use std::{mem::ManuallyDrop, fmt::Debug, sync::atomic::Ordering};
+use std::{fmt::Debug};
 
-use crate::{
-    buffer::{Buffer, RawSlice},
-    epoch,
-    num::Bytes,
-    tag::AtomicTag,
-};
-
-pub(crate) trait Plugin: 'static {
+pub(crate) trait Plugin: Debug + 'static {
     type Slot<'a>: Slot;
 
     unsafe fn acquire(&self, i: u32) -> Self::Slot<'_>;
@@ -26,41 +19,4 @@ pub(crate) trait Slot: Debug {
     fn abort(self);
 }
 
-#[derive(Debug)]
-pub(crate) struct ManagedSlot<T>
-where
-    T: Slot,
-{
-    slot: ManuallyDrop<T>
-}
-
-impl<T> ManagedSlot<T>
-where
-    T: Slot,
-{
-    fn new(slot: T) -> Self {
-        Self {
-            slot: ManuallyDrop::new(slot),
-        }
-    }
-
-    unsafe fn publish(self) {
-        let mut me = ManuallyDrop::new(self);
-        unsafe { ManuallyDrop::take(&mut me.slot).publish() }
-    }
-
-    unsafe fn freeze(self) {
-        let mut me = ManuallyDrop::new(self);
-        unsafe { ManuallyDrop::take(&mut me.slot).freeze() }
-    }
-}
-
-impl<T> Drop for ManagedSlot<T>
-where
-    T: Slot,
-{
-    fn drop(&mut self) {
-        unsafe { ManuallyDrop::take(&mut self.slot).abort() }
-    }
-}
 
diff --git a/diskann-inmem/src/store/stacked.rs b/diskann-inmem/src/store/stacked.rs
new file mode 100644
index 0000000000..eb36ca2f8a
--- /dev/null
+++ b/diskann-inmem/src/store/stacked.rs
@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use super::plugin::{self, Plugin};
+
+/// A [`super::Plugin`] for cascading multiple plugins together.
+#[derive(Debug)]
+pub(crate) struct Stacked<T, U> {
+    first: T,
+    rest: U,
+}
+
+impl<T, U> Plugin for Stacked<T, U>
+where
+    T: Plugin,
+    U: Plugin,
+{
+    type Slot<'a> = Slot<T::Slot<'a>, U::Slot<'a>>;
+
+    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
+        Slot::new(unsafe { self.first.acquire(i) }, unsafe { self.rest.acquire(i) })
+    }
+
+    fn reclaim(&self, i: u32) {
+        self.first.reclaim(i);
+        self.rest.reclaim(i);
+    }
+
+    fn retire(&self, i: u32) {
+        self.first.retire(i);
+        self.rest.retire(i);
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Slot<T, U> {
+    first: T,
+    rest: U,
+}
+
+impl<T, U> Slot<T, U> {
+    fn new(first: T, rest: U) -> Self {
+        Self { first, rest }
+    }
+
+    pub(crate) fn first(&self) -> &T {
+        &self.first
+    }
+
+    pub(crate) fn rest(&self) -> &U {
+        &self.rest
+    }
+}
+
+impl<T, U> plugin::Slot for Slot<T, U>
+where
+    T: plugin::Slot,
+    U: plugin::Slot,
+{
+    fn publish(self) {
+        self.first.publish();
+        self.rest.publish();
+    }
+    fn freeze(self) {
+        self.first.freeze();
+        self.rest.freeze();
+    }
+    fn abort(self) {
+        self.first.abort();
+        self.rest.abort();
+    }
+}

From a4a3943a93137f83680d18058dabf7fceeda8a60 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 30 Jul 2026 10:21:49 -0700
Subject: [PATCH 03/34] Prepare for merge.

---
 diskann-inmem/src/freelist.rs | 1 +
 diskann-inmem/src/provider.rs | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/diskann-inmem/src/freelist.rs b/diskann-inmem/src/freelist.rs
index 33e5521b6a..3f196473b8 100644
--- a/diskann-inmem/src/freelist.rs
+++ b/diskann-inmem/src/freelist.rs
@@ -17,6 +17,7 @@
 //! bounded to conserve memory.
 //!
 //! ## Minted
+//!
 //! If no slots live in the recycled queue, new slots can be "minted" up to the configured
 //! maximum. This simply tracks the maximum slot ID that has been yielded so far and returns
 //! the next one.
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 14de88fce5..e062acea8e 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -62,6 +62,10 @@ pub trait Id: Send + Sync + Hash + Eq + Clone + 'static {}
 
 impl<T> Id for T where T: Send + Sync + Hash + Eq + Clone + 'static {}
 
+// pub trait Index: Send + Sync + 'static {
+//     type Query<'a>: Copy + Send + Sync + 'static;
+// }
+
 /// An in-memory data-provider for DiskANN's graph indexing algorithms.
 ///
 /// The first type parameter `L` is a [`layers::Layer`] for describing the kind of data

From 5dbcbe14f391f47df7bdc53ece73b109267d44ba Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 30 Jul 2026 15:35:50 -0700
Subject: [PATCH 04/34] Making a mess.

---
 diskann-inmem/src/arch.rs          |  46 ++
 diskann-inmem/src/layers/full.rs   | 649 ++++++++++++++++-------------
 diskann-inmem/src/layers/mod.rs    |  94 +++--
 diskann-inmem/src/lib.rs           |   1 +
 diskann-inmem/src/provider.rs      | 464 ++++++++++-----------
 diskann-inmem/src/store/mod.rs     |  15 +-
 diskann-inmem/src/store/plugin.rs  |   4 +-
 diskann-inmem/src/store/stacked.rs |   4 +-
 8 files changed, 722 insertions(+), 555 deletions(-)
 create mode 100644 diskann-inmem/src/arch.rs

diff --git a/diskann-inmem/src/arch.rs b/diskann-inmem/src/arch.rs
new file mode 100644
index 0000000000..a40f9c0f00
--- /dev/null
+++ b/diskann-inmem/src/arch.rs
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use crate::num::Bytes;
+
+/// Prefetch `len` bytes beginning at `ptr`.
+///
+/// The last cache line prefetched first, followed by the rest in ascending order.
+///
+/// # Safety
+///
+/// The memory range `[ptr, ptr.add(len))` must be valid.
+#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
+#[inline(always)]
+pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
+    use std::arch::x86_64::*;
+
+    // Fetch the last cache line (the one with the tag) first.
+    let stride = Bytes::CACHELINE.value();
+    let ptr = ptr.cast::<i8>();
+    let lines = len.div_ceil(stride);
+    if lines == 0 {
+        return;
+    }
+
+    // SAFETY: Inherited from caller.
+    unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
+    for i in 0..(lines - 1) {
+        // SAFETY: Inherited from caller.
+        unsafe {
+            _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
+        }
+    }
+}
+
+/// Prefetch `len` bytes beginning at `ptr`.
+///
+/// The last cache line prefetched first, followed by the rest in ascending order.
+///
+/// # Safety
+///
+/// The memory range `[ptr, ptr.add(len))` must be valid.
+#[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
+pub(crate) unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 9cdd94c529..116be8c7f6 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -5,7 +5,7 @@
 
 use std::{fmt::Debug, marker::PhantomData};
 
-use diskann::{ANNError, ANNResult};
+use diskann::{ANNError, ANNResult, utils::IntoUsize};
 use diskann_vector::{
     UnalignedSlice,
     conversion::SliceCast,
@@ -21,7 +21,11 @@ use diskann_wide::{
 use half::f16;
 use thiserror::Error;
 
-use crate::{Hidden, layers, num::Bytes};
+use crate::{
+    Hidden, layers,
+    num::Bytes,
+    store::{self, Store},
+};
 
 /// A useful trait bound for types compatible with [`Full`].
 ///
@@ -32,14 +36,12 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<Self>;
 
     #[doc(hidden)]
-    fn __query_distance<'a, V>(
+    fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         query: &'a [Self],
-        visitor: V,
-    ) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>;
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>>;
 }
 
 /// Full-precision data layer.
@@ -162,11 +164,13 @@ where
 {
     type Query<'a> = &'a [T];
 
-    fn query_distance<'a, V>(&'a self, query: &'a [T], visitor: V) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>,
-    {
-        T::__query_distance(Hidden::new(), self, query, visitor)
+    fn __search_expand_beam<'a>(
+        &'a self,
+        query: Self::Query<'a>,
+        store: &'a Store,
+        _: Hidden,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
+        T::__expand_beam(Hidden::new(), self, query, store)
     }
 }
 
@@ -283,7 +287,10 @@ impl<T> std::ops::Deref for Calf<'_, T> {
 /// would otherwise be needed.
 #[derive(Debug)]
 struct QueryDistance<'a, T, U, D> {
+    // The original query.
     query: Calf<'a, T>,
+    // A reader into a layer's store.
+    reader: store::invasive::Reader<'a>,
     // The type of the data in the original dataset.
     _data: PhantomData<U>,
     // The type of the `PureDistanceFunction` used for the implementation.
@@ -291,9 +298,10 @@ struct QueryDistance<'a, T, U, D> {
 }
 
 impl<'a, T, U, D> QueryDistance<'a, T, U, D> {
-    fn new(query: Calf<'a, T>) -> Self {
+    fn new(query: Calf<'a, T>, reader: store::invasive::Reader<'a>) -> Self {
         Self {
             query,
+            reader,
             _data: PhantomData,
             _distance: PhantomData,
         }
@@ -312,9 +320,28 @@ impl<'a, T, U, D> QueryDistance<'a, T, U, D> {
 
         Err(ANNError::new(error))
     }
+
+    // TODO: Since we control the reader - we can avoid the length check.
+    #[inline(always)]
+    fn run(&self, x: &[u8]) -> ANNResult<f32>
+    where
+        D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>,
+    {
+        if x.len() != self.bytes() {
+            self.error(x.len())
+        } else {
+            // SAFETY: We've validated that `x` has the correct length.
+            let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
+            Ok(D::run(ARCH, (*self.query).into(), x))
+        }
+    }
 }
 
-impl<T, U, D> layers::QueryDistance for QueryDistance<'_, T, U, D>
+// TEMPORARY DEFINITIONS
+const LOOKAHEAD: usize = 8;
+const BYTES: usize = 0;
+
+impl<T, U, D> layers::__ExpandBeam for QueryDistance<'_, T, U, D>
 where
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
@@ -323,18 +350,83 @@ where
         + Sync
         + Debug,
 {
-    #[inline(always)]
-    fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
-        if x.len() != self.bytes() {
-            self.error(x.len())
+    fn __evaluate(&self, i: u32, _: Hidden) -> ANNResult<Option<f32>> {
+        if !self.reader.is_in_bounds(i.into_usize()) {
+            return Err(ANNError::new(OutOfBounds(i)));
         } else {
-            // SAFETY: We've validated that `x` has the correct length.
-            let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
-            Ok(D::run(ARCH, (*self.query).into(), x))
+            match unsafe { self.reader.read_in_bounds(i.into_usize()) } {
+                Some(data) => Ok(Some(self.run(data)?)),
+                None => Ok(None),
+            }
         }
     }
+
+    unsafe fn __expand_beam(
+        &self,
+        list: &[u32],
+        buffer: &mut [(u32, f32)],
+        _: Hidden,
+    ) -> ANNResult<usize> {
+        let len = list.len();
+        let lookahead = LOOKAHEAD.min(len);
+
+        let bytes = if BYTES == 0 {
+            self.reader.bytes().value()
+        } else {
+            BYTES + store::TAG_SIZE.value()
+        };
+
+        for j in 0..lookahead {
+            // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
+            // as the validity of the prefetch bounds.
+            unsafe {
+                crate::arch::prefetch(
+                    self.reader
+                        .read_raw_unchecked(list.get_unchecked(j).into_usize())
+                        .as_ptr()
+                        .cast(),
+                    bytes,
+                )
+            }
+        }
+
+        // Disable prefetching if the lookahead is 0.
+        let mut j = if lookahead == 0 { len } else { lookahead };
+        let mut processed = 0;
+        for &i in list.iter() {
+            if j != len {
+                // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
+                // well as the validity of the prefetch bounds.
+                unsafe {
+                    crate::arch::prefetch(
+                        self.reader
+                            .read_raw_unchecked(list.get_unchecked(j).into_usize())
+                            .as_ptr()
+                            .cast(),
+                        bytes,
+                    )
+                }
+                j += 1;
+            }
+
+            // SAFETY: Caller asserts that `i` is in-bounds.
+            if let Some(data) = unsafe { self.reader.read_in_bounds(i.into_usize()) } {
+                // SAFETY: Inherited from caller.
+                *unsafe { buffer.get_unchecked_mut(processed) } = (i, self.run(data)?);
+                processed += 1;
+            }
+        }
+
+        Ok(processed)
+    }
 }
 
+#[derive(Debug, Error)]
+#[error("index {} is out-of-bounds", self.0)]
+struct OutOfBounds(u32);
+
+diskann::convert_error!(OutOfBounds);
+
 #[derive(Debug, Error)]
 #[error(
     "expected slice of length {} - instead got {}",
@@ -349,19 +441,18 @@ struct QueryDistanceError {
 diskann::convert_error!(QueryDistanceError);
 
 macro_rules! mint {
-    ($query:ident, $visitor:ident, $T:ty => { $N:literal, $f:ident }) => {{
-        mint!($query, $visitor, { $T, $T } => { $N, $f })
-    }};
-    ($query:ident, $visitor:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
-        let inner = QueryDistance::<$T, $U, Specialize<$N, $f>>::new($query);
-        $visitor.visit_sized::<{ $N * std::mem::size_of::<$U>() }, _>(inner)
-    }};
-    ($query:ident, $visitor:ident, $T:ty => $f:ident) => {{
-        mint!($query, $visitor, { $T, $T } => $f)
+    // ($query:ident, $reader:ident, $T:ty => { $N:literal, $f:ident }) => {{
+    //     mint!($query, $visitor, { $T, $T } => { $N, $f })
+    // }};
+    // ($query:ident, $reader:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
+    //     let inner = Box::new(QueryDistance::<$T, $U, Specialize<$N, $f>>::new($query));
+    //     $visitor.visit_sized::<{ $N * std::mem::size_of::<$U>() }, _>(inner)
+    // }};
+    ($query:ident, $reader:ident, $T:ty => $f:ident) => {{
+        mint!($query, $reader, { $T, $T } => $f)
     }};
-    ($query:ident, $visitor:ident, { $T:ty, $U:ty } => $f:ident) => {{
-        let inner = QueryDistance::<$T, $U, $f>::new($query);
-        $visitor.visit(inner)
+    ($query:ident, $reader:ident, { $T:ty, $U:ty } => $f:ident) => {{
+        Box::new(QueryDistance::<$T, $U, $f>::new($query, $reader))
     }};
 }
 
@@ -370,32 +461,30 @@ impl FullPrecision for f32 {
         Full::from_distance_provider(dim, metric)
     }
 
-    fn __query_distance<'a, V>(
+    fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f32>,
         query: &'a [f32],
-        visitor: V,
-    ) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>,
-    {
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
+        let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output = match full.metric {
+        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 100 {
-                    mint!(query, visitor, f32 => { 100, SquaredL2 })
-                } else {
-                    mint!(query, visitor, f32 => SquaredL2)
-                }
+                // if full.dim() == 100 {
+                //     mint!(query, visitor, f32 => { 100, SquaredL2 })
+                // } else {
+                mint!(query, reader, f32 => SquaredL2)
+                // }
             }
             Metric::InnerProduct => {
-                mint!(query, visitor, f32 => InnerProduct)
+                mint!(query, reader, f32 => InnerProduct)
             }
-            Metric::Cosine => mint!(query, visitor, f32 => Cosine),
-            Metric::CosineNormalized => mint!(query, visitor, f32 => CosineNormalized),
+            Metric::Cosine => mint!(query, reader, f32 => Cosine),
+            Metric::CosineNormalized => mint!(query, reader, f32 => CosineNormalized),
         };
 
         Ok(output)
@@ -407,32 +496,30 @@ impl FullPrecision for f16 {
         Full::from_distance_provider(dim, metric)
     }
 
-    fn __query_distance<'a, V>(
+    fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f16>,
         query: &'a [f16],
-        visitor: V,
-    ) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>,
-    {
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
+        let reader = store.temp_inner_reader()?;
 
         let mut as_f32: Box<[f32]> = std::iter::repeat_n(0.0, full.dim()).collect();
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
         let query = Calf::Owned(as_f32);
 
-        let output = match full.metric {
+        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 100 {
-                    mint!(query, visitor, { f32, f16 } => { 100, SquaredL2 })
-                } else {
-                    mint!(query, visitor, { f32, f16 } => SquaredL2)
-                }
+                // if full.dim() == 100 {
+                //     mint!(query, visitor, { f32, f16 } => { 100, SquaredL2 })
+                // } else {
+                mint!(query, reader, { f32, f16 } => SquaredL2)
+                // }
             }
-            Metric::InnerProduct => mint!(query, visitor, { f32, f16 } => InnerProduct),
-            Metric::Cosine => mint!(query, visitor, { f32, f16 } => Cosine),
-            Metric::CosineNormalized => mint!(query, visitor, { f32, f16 } => CosineNormalized),
+            Metric::InnerProduct => mint!(query, reader, { f32, f16 } => InnerProduct),
+            Metric::Cosine => mint!(query, reader, { f32, f16 } => Cosine),
+            Metric::CosineNormalized => mint!(query, reader, { f32, f16 } => CosineNormalized),
         };
 
         Ok(output)
@@ -444,30 +531,28 @@ impl FullPrecision for u8 {
         Full::from_distance_provider(dim, metric)
     }
 
-    fn __query_distance<'a, V>(
+    fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<u8>,
         query: &'a [u8],
-        visitor: V,
-    ) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>,
-    {
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
+        let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output = match full.metric {
+        let output: Box::<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 128 {
-                    mint!(query, visitor, u8 => { 128, SquaredL2 })
-                } else {
-                    mint!(query, visitor, u8 => SquaredL2)
-                }
+                // if full.dim() == 128 {
+                //     mint!(query, visitor, u8 => { 128, SquaredL2 })
+                // } else {
+                mint!(query, reader, u8 => SquaredL2)
+                // }
             }
-            Metric::InnerProduct => mint!(query, visitor, u8 => InnerProduct),
-            Metric::Cosine => mint!(query, visitor, u8 => Cosine),
-            Metric::CosineNormalized => mint!(query, visitor, u8 => Cosine),
+            Metric::InnerProduct => mint!(query, reader, u8 => InnerProduct),
+            Metric::Cosine => mint!(query, reader, u8 => Cosine),
+            Metric::CosineNormalized => mint!(query, reader, u8 => Cosine),
         };
 
         Ok(output)
@@ -479,24 +564,22 @@ impl FullPrecision for i8 {
         Full::from_distance_provider(dim, metric)
     }
 
-    fn __query_distance<'a, V>(
+    fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<i8>,
         query: &'a [i8],
-        visitor: V,
-    ) -> ANNResult<V::Output>
-    where
-        V: layers::QueryVisitor<'a>,
-    {
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
+        let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output = match full.metric {
-            Metric::L2 => mint!(query, visitor, i8 => SquaredL2),
-            Metric::InnerProduct => mint!(query, visitor, i8 => InnerProduct),
-            Metric::Cosine => mint!(query, visitor, i8 => Cosine),
-            Metric::CosineNormalized => mint!(query, visitor, i8 => Cosine),
+        let output: Box<dyn layers::__ExpandBeam + 'a> = match full.metric {
+            Metric::L2 => mint!(query, reader, i8 => SquaredL2),
+            Metric::InnerProduct => mint!(query, reader, i8 => InnerProduct),
+            Metric::Cosine => mint!(query, reader, i8 => Cosine),
+            Metric::CosineNormalized => mint!(query, reader, i8 => Cosine),
         };
 
         Ok(output)
@@ -507,197 +590,197 @@ impl FullPrecision for i8 {
 // Tests //
 ///////////
 
-#[cfg(test)]
-#[cfg(not(miri))]
-mod tests {
-    use std::fmt::Display;
-
-    use rand::{Rng, SeedableRng, rngs::StdRng};
-
-    use super::*;
-    // Bring the inherent-call traits into method scope. The `Distance` / `QueryDistance`
-    // traits are not imported: their methods are reached through `&dyn _` trait objects,
-    // which does not require the trait to be in scope.
-    use crate::layers::{AsDistance as _, QueryVisitor, Search as _, Set as _};
-
-    /// Generate random elements of a layer's data type from a seeded RNG.
-    trait Sample: bytemuck::Pod {
-        fn sample<R: Rng>(rng: &mut R) -> Self;
-    }
-
-    impl Sample for f32 {
-        fn sample<R: Rng>(rng: &mut R) -> Self {
-            rng.random_range(-1.0f32..1.0f32)
-        }
-    }
-
-    impl Sample for f16 {
-        fn sample<R: Rng>(rng: &mut R) -> Self {
-            f16::from_f32(rng.random_range(-1.0f32..1.0f32))
-        }
-    }
-
-    impl Sample for u8 {
-        fn sample<R: Rng>(rng: &mut R) -> Self {
-            rng.random()
-        }
-    }
-
-    impl Sample for i8 {
-        fn sample<R: Rng>(rng: &mut R) -> Self {
-            rng.random()
-        }
-    }
-
-    fn gen_vec<T: Sample, R: Rng>(rng: &mut R, dim: usize) -> Vec<T> {
-        (0..dim).map(|_| T::sample(rng)).collect()
-    }
-
-    /// A [`QueryVisitor`] that simply boxes the minted kernel so the test can probe it
-    /// directly. Exercises both `visit` (dynamic) and `visit_sized` (specialized) paths.
-    struct Collect;
-
-    impl<'a> QueryVisitor<'a> for Collect {
-        type Output = Box<dyn layers::QueryDistance + 'a>;
-
-        fn visit<Q>(self, distance: Q) -> Self::Output
-        where
-            Q: layers::QueryDistance + 'a,
-        {
-            Box::new(distance)
-        }
-    }
-
-    /// Compare two distances allowing for floating-point reassociation between the
-    /// specialized / converted kernels and the dynamic reference.
-    fn approx_eq(got: f32, want: f32) -> bool {
-        (got - want).abs() <= 1e-3 + 1e-4 * want.abs()
-    }
-
-    /// Exercise every `Full<T>` API across dimensions `1..=max_dim`.
-    ///
-    /// For each dimension we check that `bytes`/`set` agree, that `distance` and
-    /// `query_distance` are consistent with `DistanceProvider`, and that all of these
-    /// reject byte slices that are too long or too short.
-    fn test_impl<T>(max_dim: usize, ctx: &dyn Display)
-    where
-        T: FullPrecision + Sample + DistanceProvider<T>,
-    {
-        let mut rng = StdRng::seed_from_u64(0x0D15_0ACE ^ max_dim as u64);
-        let metrics = [
-            Metric::L2,
-            Metric::InnerProduct,
-            Metric::Cosine,
-            Metric::CosineNormalized,
-        ];
-
-        for dim in 1..=max_dim {
-            let a = gen_vec::<T, _>(&mut rng, dim);
-            let b = gen_vec::<T, _>(&mut rng, dim);
-
-            // `bytes` and `set` agree: the encoded buffer equals the raw cast bytes.
-            let layer = Full::<T>::new(dim, Metric::L2);
-            assert_eq!(
-                layer.bytes().value(),
-                dim * std::mem::size_of::<T>(),
-                "{ctx}: dim {dim}: unexpected byte length",
-            );
-
-            let mut a_bytes = vec![0u8; layer.bytes().value()];
-            layer.set(&a, &mut a_bytes).unwrap();
-            assert_eq!(
-                a_bytes.as_slice(),
-                bytemuck::cast_slice::<T, u8>(&a),
-                "{ctx}: dim {dim}: set mismatch",
-            );
-
-            let mut b_bytes = vec![0u8; layer.bytes().value()];
-            layer.set(&b, &mut b_bytes).unwrap();
-
-            for metric in metrics {
-                let full = Full::<T>::new(dim, metric);
-
-                // Reference value straight from `DistanceProvider`.
-                let reference =
-                    <T as DistanceProvider<T>>::distance_comparer(metric, Some(dim)).call(&a, &b);
-
-                // `distance` is built from the same comparer, so it must match exactly.
-                let distance = full.as_distance();
-                let via_distance = distance.evaluate(&a_bytes, &b_bytes).unwrap();
-                assert_eq!(
-                    via_distance, reference,
-                    "{ctx}: dim {dim}, metric {metric:?}: distance != DistanceProvider",
-                );
-
-                // `query_distance` computes the same geometry. Specialized and f16-converted
-                // kernels may reassociate the summation, so compare approximately.
-                let query = full.query_distance(a.as_slice(), Collect).unwrap();
-                let via_query = query.evaluate(&b_bytes).unwrap();
-                assert!(
-                    approx_eq(via_query, via_distance),
-                    "{ctx}: dim {dim}, metric {metric:?}: query {via_query} != distance {via_distance}",
-                );
-
-                // Every distance API rejects byte slices that are too long or too short.
-                let short = &a_bytes[..a_bytes.len() - 1];
-                let mut long = a_bytes.clone();
-                long.push(0);
-
-                assert!(distance.evaluate(short, &b_bytes).is_err());
-                assert!(distance.evaluate(&long, &b_bytes).is_err());
-                assert!(distance.evaluate(&a_bytes, short).is_err());
-                assert!(distance.evaluate(&a_bytes, &long).is_err());
-
-                assert!(query.evaluate(short).is_err());
-                assert!(query.evaluate(&long).is_err());
-            }
-
-            // `set` rejects mis-sized element and buffer slices.
-            let mut buf = vec![0u8; layer.bytes().value()];
-            let too_many = gen_vec::<T, _>(&mut rng, dim + 1);
-            assert!(
-                layer.set(&too_many, &mut buf).is_err(),
-                "{ctx}: dim {dim}: set accepted an over-long element slice",
-            );
-
-            assert!(
-                layer.query_distance(&too_many, Collect).is_err(),
-                "{ctx}: dim {dim}: incorrect query lengths should be rejected"
-            );
-
-            let mut short_buf = vec![0u8; layer.bytes().value().saturating_sub(1)];
-            assert!(
-                layer.set(&a, &mut short_buf).is_err(),
-                "{ctx}: dim {dim}: set accepted an under-sized buffer",
-            );
-
-            let too_few = gen_vec::<T, _>(&mut rng, dim - 1);
-            assert!(
-                layer.query_distance(&too_few, Collect).is_err(),
-                "{ctx}: dim {dim}: incorrect query lengths should be rejected"
-            );
-        }
-    }
-
-    // `max_dim` must exceed the largest specialized dimension for each type so the
-    // const-generic (`visit_sized`) paths are covered alongside the dynamic ones.
-    #[test]
-    fn full_f32() {
-        test_impl::<f32>(256, &"f32");
-    }
-
-    #[test]
-    fn full_f16() {
-        test_impl::<f16>(256, &"f16");
-    }
-
-    #[test]
-    fn full_u8() {
-        test_impl::<u8>(160, &"u8");
-    }
-
-    #[test]
-    fn full_i8() {
-        test_impl::<i8>(160, &"i8");
-    }
-}
+// #[cfg(test)]
+// #[cfg(not(miri))]
+// mod tests {
+//     use std::fmt::Display;
+//
+//     use rand::{Rng, SeedableRng, rngs::StdRng};
+//
+//     use super::*;
+//     // Bring the inherent-call traits into method scope. The `Distance` / `QueryDistance`
+//     // traits are not imported: their methods are reached through `&dyn _` trait objects,
+//     // which does not require the trait to be in scope.
+//     use crate::layers::{AsDistance as _, QueryVisitor, Search as _, Set as _};
+//
+//     /// Generate random elements of a layer's data type from a seeded RNG.
+//     trait Sample: bytemuck::Pod {
+//         fn sample<R: Rng>(rng: &mut R) -> Self;
+//     }
+//
+//     impl Sample for f32 {
+//         fn sample<R: Rng>(rng: &mut R) -> Self {
+//             rng.random_range(-1.0f32..1.0f32)
+//         }
+//     }
+//
+//     impl Sample for f16 {
+//         fn sample<R: Rng>(rng: &mut R) -> Self {
+//             f16::from_f32(rng.random_range(-1.0f32..1.0f32))
+//         }
+//     }
+//
+//     impl Sample for u8 {
+//         fn sample<R: Rng>(rng: &mut R) -> Self {
+//             rng.random()
+//         }
+//     }
+//
+//     impl Sample for i8 {
+//         fn sample<R: Rng>(rng: &mut R) -> Self {
+//             rng.random()
+//         }
+//     }
+//
+//     fn gen_vec<T: Sample, R: Rng>(rng: &mut R, dim: usize) -> Vec<T> {
+//         (0..dim).map(|_| T::sample(rng)).collect()
+//     }
+//
+//     /// A [`QueryVisitor`] that simply boxes the minted kernel so the test can probe it
+//     /// directly. Exercises both `visit` (dynamic) and `visit_sized` (specialized) paths.
+//     struct Collect;
+//
+//     impl<'a> QueryVisitor<'a> for Collect {
+//         type Output = Box<dyn layers::QueryDistance + 'a>;
+//
+//         fn visit<Q>(self, distance: Q) -> Self::Output
+//         where
+//             Q: layers::QueryDistance + 'a,
+//         {
+//             Box::new(distance)
+//         }
+//     }
+//
+//     /// Compare two distances allowing for floating-point reassociation between the
+//     /// specialized / converted kernels and the dynamic reference.
+//     fn approx_eq(got: f32, want: f32) -> bool {
+//         (got - want).abs() <= 1e-3 + 1e-4 * want.abs()
+//     }
+//
+//     /// Exercise every `Full<T>` API across dimensions `1..=max_dim`.
+//     ///
+//     /// For each dimension we check that `bytes`/`set` agree, that `distance` and
+//     /// `query_distance` are consistent with `DistanceProvider`, and that all of these
+//     /// reject byte slices that are too long or too short.
+//     fn test_impl<T>(max_dim: usize, ctx: &dyn Display)
+//     where
+//         T: FullPrecision + Sample + DistanceProvider<T>,
+//     {
+//         let mut rng = StdRng::seed_from_u64(0x0D15_0ACE ^ max_dim as u64);
+//         let metrics = [
+//             Metric::L2,
+//             Metric::InnerProduct,
+//             Metric::Cosine,
+//             Metric::CosineNormalized,
+//         ];
+//
+//         for dim in 1..=max_dim {
+//             let a = gen_vec::<T, _>(&mut rng, dim);
+//             let b = gen_vec::<T, _>(&mut rng, dim);
+//
+//             // `bytes` and `set` agree: the encoded buffer equals the raw cast bytes.
+//             let layer = Full::<T>::new(dim, Metric::L2);
+//             assert_eq!(
+//                 layer.bytes().value(),
+//                 dim * std::mem::size_of::<T>(),
+//                 "{ctx}: dim {dim}: unexpected byte length",
+//             );
+//
+//             let mut a_bytes = vec![0u8; layer.bytes().value()];
+//             layer.set(&a, &mut a_bytes).unwrap();
+//             assert_eq!(
+//                 a_bytes.as_slice(),
+//                 bytemuck::cast_slice::<T, u8>(&a),
+//                 "{ctx}: dim {dim}: set mismatch",
+//             );
+//
+//             let mut b_bytes = vec![0u8; layer.bytes().value()];
+//             layer.set(&b, &mut b_bytes).unwrap();
+//
+//             for metric in metrics {
+//                 let full = Full::<T>::new(dim, metric);
+//
+//                 // Reference value straight from `DistanceProvider`.
+//                 let reference =
+//                     <T as DistanceProvider<T>>::distance_comparer(metric, Some(dim)).call(&a, &b);
+//
+//                 // `distance` is built from the same comparer, so it must match exactly.
+//                 let distance = full.as_distance();
+//                 let via_distance = distance.evaluate(&a_bytes, &b_bytes).unwrap();
+//                 assert_eq!(
+//                     via_distance, reference,
+//                     "{ctx}: dim {dim}, metric {metric:?}: distance != DistanceProvider",
+//                 );
+//
+//                 // `query_distance` computes the same geometry. Specialized and f16-converted
+//                 // kernels may reassociate the summation, so compare approximately.
+//                 let query = full.query_distance(a.as_slice(), Collect).unwrap();
+//                 let via_query = query.evaluate(&b_bytes).unwrap();
+//                 assert!(
+//                     approx_eq(via_query, via_distance),
+//                     "{ctx}: dim {dim}, metric {metric:?}: query {via_query} != distance {via_distance}",
+//                 );
+//
+//                 // Every distance API rejects byte slices that are too long or too short.
+//                 let short = &a_bytes[..a_bytes.len() - 1];
+//                 let mut long = a_bytes.clone();
+//                 long.push(0);
+//
+//                 assert!(distance.evaluate(short, &b_bytes).is_err());
+//                 assert!(distance.evaluate(&long, &b_bytes).is_err());
+//                 assert!(distance.evaluate(&a_bytes, short).is_err());
+//                 assert!(distance.evaluate(&a_bytes, &long).is_err());
+//
+//                 assert!(query.evaluate(short).is_err());
+//                 assert!(query.evaluate(&long).is_err());
+//             }
+//
+//             // `set` rejects mis-sized element and buffer slices.
+//             let mut buf = vec![0u8; layer.bytes().value()];
+//             let too_many = gen_vec::<T, _>(&mut rng, dim + 1);
+//             assert!(
+//                 layer.set(&too_many, &mut buf).is_err(),
+//                 "{ctx}: dim {dim}: set accepted an over-long element slice",
+//             );
+//
+//             assert!(
+//                 layer.query_distance(&too_many, Collect).is_err(),
+//                 "{ctx}: dim {dim}: incorrect query lengths should be rejected"
+//             );
+//
+//             let mut short_buf = vec![0u8; layer.bytes().value().saturating_sub(1)];
+//             assert!(
+//                 layer.set(&a, &mut short_buf).is_err(),
+//                 "{ctx}: dim {dim}: set accepted an under-sized buffer",
+//             );
+//
+//             let too_few = gen_vec::<T, _>(&mut rng, dim - 1);
+//             assert!(
+//                 layer.query_distance(&too_few, Collect).is_err(),
+//                 "{ctx}: dim {dim}: incorrect query lengths should be rejected"
+//             );
+//         }
+//     }
+//
+//     // `max_dim` must exceed the largest specialized dimension for each type so the
+//     // const-generic (`visit_sized`) paths are covered alongside the dynamic ones.
+//     #[test]
+//     fn full_f32() {
+//         test_impl::<f32>(256, &"f32");
+//     }
+//
+//     #[test]
+//     fn full_f16() {
+//         test_impl::<f16>(256, &"f16");
+//     }
+//
+//     #[test]
+//     fn full_u8() {
+//         test_impl::<u8>(160, &"u8");
+//     }
+//
+//     #[test]
+//     fn full_i8() {
+//         test_impl::<i8>(160, &"i8");
+//     }
+// }
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 13179d9bc5..35591f96dd 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -28,7 +28,7 @@
 
 use diskann::ANNResult;
 
-use crate::num::Bytes;
+use crate::{Hidden, num::Bytes, store::Store};
 
 mod full;
 pub use full::{Full, FullPrecision};
@@ -76,6 +76,25 @@ pub trait QueryDistance: Send + Sync + std::fmt::Debug {
     fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
 }
 
+#[doc(hidden)]
+pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
+    /// Evaluate a raw distance against index `i`.
+    fn __evaluate(&self, i: u32, _: Hidden) -> ANNResult<Option<f32>>;
+
+    /// Compute the distance between the query and each neighbor in `list`.
+    ///
+    /// # Safety
+    ///
+    /// * All items in `list` must in-bounds with respect to `reader`.
+    /// * `buffer.len() >= list.len()`.
+    unsafe fn __expand_beam(
+        &self,
+        list: &[u32],
+        buffer: &mut [(u32, f32)],
+        _: Hidden,
+    ) -> ANNResult<usize>;
+}
+
 /// Enable search over vectors defined by a [`Layer`].
 pub trait Search: Send + Sync + 'static {
     /// The type of the query. This should be equivalent to the generic parameter in
@@ -83,45 +102,56 @@ pub trait Search: Send + Sync + 'static {
     /// design.
     type Query<'a>;
 
-    /// Create a distance computer specialized for `query` and provide it to `visitor`.
-    fn query_distance<'a, V>(&'a self, query: Self::Query<'a>, visitor: V) -> ANNResult<V::Output>
-    where
-        V: QueryVisitor<'a>;
-}
+    #[doc(hidden)]
+    fn __search_expand_beam<'a>(
+        &'a self,
+        query: Self::Query<'a>,
+        store: &'a Store,
+        _: Hidden,
+    ) -> ANNResult<Box<dyn __ExpandBeam + 'a>>;
 
-/// Specialize a kernel around a [`QueryDistance`] implementation.
-pub trait QueryVisitor<'a>: Sized {
-    /// The type of the type-erased output.
-    type Output;
-
-    /// Specialize [`Self::Output`] for `distance`.
-    fn visit<T>(self, distance: T) -> Self::Output
-    where
-        T: QueryDistance + 'a;
-
-    /// Specialize [`Self::Output`] for `distance` accepting a hint that `distance` has been
-    /// specialized to work on data elements of exactly `BYTES` bytes long.
-    ///
-    /// This can be used to tailor surrounding code (e.g. software prefetches) for exactly
-    /// the length of the data being processed.
-    fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
-    where
-        T: QueryDistance + 'a,
-    {
-        self.visit(distance)
-    }
+    // /// Create a distance computer specialized for `query` and provide it to `visitor`.
+    // fn query_distance<'a, V>(&'a self, query: Self::Query<'a>, visitor: V) -> ANNResult<V::Output>
+    // where
+    //     V: QueryVisitor<'a>;
 }
 
+// /// Specialize a kernel around a [`QueryDistance`] implementation.
+// pub trait QueryVisitor<'a>: Sized {
+//     /// The type of the type-erased output.
+//     type Output;
+//
+//     /// Specialize [`Self::Output`] for `distance`.
+//     fn visit<T>(self, distance: T) -> Self::Output
+//     where
+//         T: QueryDistance + 'a;
+//
+//     /// Specialize [`Self::Output`] for `distance` accepting a hint that `distance` has been
+//     /// specialized to work on data elements of exactly `BYTES` bytes long.
+//     ///
+//     /// This can be used to tailor surrounding code (e.g. software prefetches) for exactly
+//     /// the length of the data being processed.
+//     fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
+//     where
+//         T: QueryDistance + 'a,
+//     {
+//         self.visit(distance)
+//     }
+// }
+
 /// A insert-specific specialization of [`Search`].
 ///
 /// Note that the bounds for this trait are unnecessarily complicated, but rely on changes
 /// to `diskann` to full resolve.
 pub trait Insert: Search + for<'a> Set<Self::Query<'a>> + AsDistance {
     /// A specialization of [`Search::query_distance`] targeting vector insert specifically.
-    fn insert_distance<'a, V>(&'a self, query: Self::Query<'a>, visitor: V) -> ANNResult<V::Output>
-    where
-        V: QueryVisitor<'a>,
-    {
-        self.query_distance(query, visitor)
+    #[doc(hidden)]
+    fn __insert_expand_beam<'a, V>(
+        &'a self,
+        query: Self::Query<'a>,
+        store: &'a Store,
+        _: Hidden,
+    ) -> ANNResult<Box<dyn __ExpandBeam + 'a>> {
+        self.__search_expand_beam(query, store, Hidden::new())
     }
 }
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 8f0570dc30..b179fc45ff 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -9,6 +9,7 @@
 
 pub mod num;
 
+mod arch;
 mod buffer;
 mod counters;
 mod epoch;
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index b59fa4e0e4..ca4f00ec30 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -50,9 +50,11 @@ use diskann_utils::views::Matrix;
 use thiserror::Error;
 
 use crate::{
+    Hidden,
     counters::{Counters, LocalCounters},
     ids::IdMap,
     layers::{self, QueryDistance},
+    neighbors::Neighbors,
     num::Bytes,
     store::{self, Store},
 };
@@ -62,10 +64,6 @@ pub trait Id: Send + Sync + Hash + Eq + Clone + 'static {}
 
 impl<T> Id for T where T: Send + Sync + Hash + Eq + Clone + 'static {}
 
-// pub trait Index: Send + Sync + 'static {
-//     type Query<'a>: Copy + Send + Sync + 'static;
-// }
-
 /// An in-memory data-provider for DiskANN's graph indexing algorithms.
 ///
 /// The first type parameter `L` is a [`layers::Layer`] for describing the kind of data
@@ -387,9 +385,9 @@ where
 /// kernel once and reuse it to balance compile times and performance.
 #[derive(Debug)]
 pub struct SearchAccessor<'a> {
-    reader: store::Reader<'a>,
+    neighbors: &'a Neighbors,
     ids: AdjacencyList<u32>,
-    expand_beam: Box<dyn ExpandBeam + 'a>,
+    expand_beam: Box<dyn layers::__ExpandBeam + 'a>,
     buffer: Vec<(u32, f32)>,
 
     // The parent provider for the accessor.
@@ -418,13 +416,12 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     {
         let work = move || {
             for p in self.start_points.clone() {
-                match self.reader.inner().read(p.into_usize()) {
-                    Some(point) => {
+                match self.expand_beam.__evaluate(p, Hidden::new())? {
+                    Some(distance) => {
                         // Counters are no-ops without `integration-test`.
                         self.counters.get_vector(1);
                         self.counters.query_distance(1);
-
-                        f(p, self.expand_beam.evaluate(point)?);
+                        f(p, distance);
                     }
                     None => {
                         return Err(ANNError::message("could not retrieve start point"));
@@ -450,22 +447,25 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     {
         let work = move || -> ANNResult<()> {
             for i in ids {
-                self.reader.neighbors().get(i, &mut self.ids)?;
+                self.neighbors.get(i, &mut self.ids)?;
                 self.counters.get_neighbors(1);
 
                 // Filter out unvisited IDs and ensure that all the IDs we are about
                 self.ids
-                    .retain(|i| pred.eval_mut(i) && self.reader.is_in_bounds(i.into_usize()));
+                    .retain(|i| pred.eval_mut(i) && *i < self.neighbors.entries());
 
                 // This should always hold, but let's double check.
                 assert!(self.buffer.len() >= self.ids.len());
 
                 // SAFETY: We've verified that each entry in `self.ids` is in-bounds and the
                 // `self.buffer` is long enough to hold all the IDs.
-                let processed = unsafe {
-                    self.expand_beam
-                        .expand_beam(&self.ids, &self.reader, &mut self.buffer)
-                }?;
+                let processed =
+                    unsafe { self.expand_beam.__expand_beam(&self.ids, &mut self.buffer, Hidden::new()) }?;
+
+                // let processed = unsafe {
+                //     self.expand_beam
+                //         .expand_beam(&self.ids, &self.reader, &mut self.buffer)
+                // }?;
 
                 self.counters.get_vector(processed as u64);
                 self.counters.query_distance(processed as u64);
@@ -483,217 +483,217 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     }
 }
 
-trait ExpandBeam: Send + Sync + std::fmt::Debug {
-    /// Evaluate a raw distance function.
-    fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
-
-    /// Compute the distance between the query and each neighbor in `list`.
-    ///
-    /// # Safety
-    ///
-    /// * All items in `list` must in-bounds with respect to `reader`.
-    /// * `buffer.len() >= list.len()`.
-    unsafe fn expand_beam(
-        &self,
-        list: &[u32],
-        reader: &store::Reader<'_>,
-        buffer: &mut [(u32, f32)],
-    ) -> ANNResult<usize>;
-}
-
-#[derive(Debug)]
-struct ExpandBeamImpl<T, const BYTES: usize> {
-    inner: T,
-    prefetch_lookahead: usize,
-}
-
-impl<T, const BYTES: usize> ExpandBeamImpl<T, BYTES> {
-    fn new(inner: T, prefetch_lookahead: usize) -> Self {
-        Self {
-            inner,
-            prefetch_lookahead,
-        }
-    }
-}
-
-impl<T, const BYTES: usize> ExpandBeam for ExpandBeamImpl<T, BYTES>
-where
-    T: layers::QueryDistance,
-{
-    fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
-        self.inner.evaluate(x)
-    }
-
-    unsafe fn expand_beam(
-        &self,
-        list: &[u32],
-        reader: &store::Reader<'_>,
-        buffer: &mut [(u32, f32)],
-    ) -> ANNResult<usize> {
-        // SAFETY: Inherited from caller.
-        unsafe {
-            expand_beam_inner::<T, BYTES>(
-                &self.inner,
-                list,
-                self.prefetch_lookahead,
-                reader,
-                buffer,
-            )
-        }
-    }
-}
-
-#[derive(Debug)]
-struct ExpandBeamVisitor {
-    bytes: Bytes,
-    prefetch_lookahead: usize,
-}
-
-impl<'a> layers::QueryVisitor<'a> for ExpandBeamVisitor {
-    type Output = Box<dyn ExpandBeam + 'a>;
-
-    fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
-    where
-        T: QueryDistance + 'a,
-    {
-        // This is critical to ensure we emit the correct number of prefetches.
-        assert!(Bytes::new(BYTES + store::TAG_SIZE.value()) <= self.bytes);
-        Box::new(ExpandBeamImpl::<_, BYTES>::new(
-            distance,
-            self.prefetch_lookahead,
-        ))
-    }
-
-    fn visit<T>(self, distance: T) -> Self::Output
-    where
-        T: QueryDistance + 'a,
-    {
-        Box::new(ExpandBeamImpl::<_, 0>::new(
-            distance,
-            self.prefetch_lookahead,
-        ))
-    }
-}
-
-/// Prefetch `len` bytes beginning at `ptr`.
-///
-/// The last cache line prefetched first, followed by the rest in ascending order.
-///
-/// # Safety
-///
-/// The memory range `[ptr, ptr.add(len))` must be valid.
-#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
-#[inline(always)]
-unsafe fn prefetch(ptr: *const u8, len: usize) {
-    use std::arch::x86_64::*;
-
-    // Fetch the last cache line (the one with the tag) first.
-    let stride = Bytes::CACHELINE.value();
-    let ptr = ptr.cast::<i8>();
-    let lines = len.div_ceil(stride);
-    if lines == 0 {
-        return;
-    }
-
-    // SAFETY: Inherited from caller.
-    unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
-    for i in 0..(lines - 1) {
-        // SAFETY: Inherited from caller.
-        unsafe {
-            _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
-        }
-    }
-}
-
-/// Prefetch `len` bytes beginning at `ptr`.
-///
-/// The last cache line prefetched first, followed by the rest in ascending order.
-///
-/// # Safety
-///
-/// The memory range `[ptr, ptr.add(len))` must be valid.
-#[cfg(not(any(target_arch = "x86_64", target_feature = "avx2")))]
-unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
-
-/// # Safety
-///
-/// * All items in `list` must in-bounds with respect to `reader`.
-/// * The number of bytes associated with `N` cache lines must "make sense".
-/// * `buffer.len() >= list.len()`.
-#[inline]
-unsafe fn expand_beam_inner<T, const BYTES: usize>(
-    distance: &T,
-    list: &[u32],
-    lookahead: usize,
-    reader: &store::Reader<'_>,
-    buffer: &mut [(u32, f32)],
-) -> ANNResult<usize>
-where
-    T: layers::QueryDistance,
-{
-    debug_assert!(
-        BYTES + store::TAG_SIZE.value() <= reader.inner().bytes().value(),
-        "we really rely on this: {}, bytes = {}",
-        BYTES + store::TAG_SIZE.value(),
-        reader.inner().bytes()
-    );
-
-    debug_assert!(buffer.len() >= list.len());
-
-    let bytes = if BYTES == 0 {
-        reader.inner().bytes().value()
-    } else {
-        BYTES + store::TAG_SIZE.value()
-    };
-
-    let len = list.len();
-    let lookahead = lookahead.min(len);
-
-    for j in 0..lookahead {
-        // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
-        // as the validity of the prefetch bounds.
-        unsafe {
-            prefetch(
-                reader
-                    .inner()
-                    .read_raw_unchecked(list.get_unchecked(j).into_usize())
-                    .as_ptr()
-                    .cast(),
-                bytes,
-            )
-        }
-    }
-
-    // Disable prefetching if the lookahead is 0.
-    let mut j = if lookahead == 0 { len } else { lookahead };
-    let mut processed = 0;
-    for &i in list.iter() {
-        if j != len {
-            // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
-            // well as the validity of the prefetch bounds.
-            unsafe {
-                prefetch(
-                    reader
-                        .inner()
-                        .read_raw_unchecked(list.get_unchecked(j).into_usize())
-                        .as_ptr()
-                        .cast(),
-                    bytes,
-                )
-            }
-            j += 1;
-        }
-
-        // SAFETY: Caller asserts that `i` is in-bounds.
-        if let Some(data) = unsafe { reader.inner().read_in_bounds(i.into_usize()) } {
-            // SAFETY: Inherited from caller.
-            *unsafe { buffer.get_unchecked_mut(processed) } = (i, distance.evaluate(data)?);
-            processed += 1;
-        }
-    }
-
-    Ok(processed)
-}
+// trait ExpandBeam: Send + Sync + std::fmt::Debug {
+//     /// Evaluate a raw distance function.
+//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
+//
+//     /// Compute the distance between the query and each neighbor in `list`.
+//     ///
+//     /// # Safety
+//     ///
+//     /// * All items in `list` must in-bounds with respect to `reader`.
+//     /// * `buffer.len() >= list.len()`.
+//     unsafe fn expand_beam(
+//         &self,
+//         list: &[u32],
+//         reader: &store::Reader<'_>,
+//         buffer: &mut [(u32, f32)],
+//     ) -> ANNResult<usize>;
+// }
+//
+// #[derive(Debug)]
+// struct ExpandBeamImpl<T, const BYTES: usize> {
+//     inner: T,
+//     prefetch_lookahead: usize,
+// }
+//
+// impl<T, const BYTES: usize> ExpandBeamImpl<T, BYTES> {
+//     fn new(inner: T, prefetch_lookahead: usize) -> Self {
+//         Self {
+//             inner,
+//             prefetch_lookahead,
+//         }
+//     }
+// }
+//
+// impl<T, const BYTES: usize> ExpandBeam for ExpandBeamImpl<T, BYTES>
+// where
+//     T: layers::QueryDistance,
+// {
+//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
+//         self.inner.evaluate(x)
+//     }
+//
+//     unsafe fn expand_beam(
+//         &self,
+//         list: &[u32],
+//         reader: &store::Reader<'_>,
+//         buffer: &mut [(u32, f32)],
+//     ) -> ANNResult<usize> {
+//         // SAFETY: Inherited from caller.
+//         unsafe {
+//             expand_beam_inner::<T, BYTES>(
+//                 &self.inner,
+//                 list,
+//                 self.prefetch_lookahead,
+//                 reader,
+//                 buffer,
+//             )
+//         }
+//     }
+// }
+//
+// #[derive(Debug)]
+// struct ExpandBeamVisitor {
+//     bytes: Bytes,
+//     prefetch_lookahead: usize,
+// }
+//
+// impl<'a> layers::QueryVisitor<'a> for ExpandBeamVisitor {
+//     type Output = Box<dyn ExpandBeam + 'a>;
+//
+//     fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
+//     where
+//         T: QueryDistance + 'a,
+//     {
+//         // This is critical to ensure we emit the correct number of prefetches.
+//         assert!(Bytes::new(BYTES + store::TAG_SIZE.value()) <= self.bytes);
+//         Box::new(ExpandBeamImpl::<_, BYTES>::new(
+//             distance,
+//             self.prefetch_lookahead,
+//         ))
+//     }
+//
+//     fn visit<T>(self, distance: T) -> Self::Output
+//     where
+//         T: QueryDistance + 'a,
+//     {
+//         Box::new(ExpandBeamImpl::<_, 0>::new(
+//             distance,
+//             self.prefetch_lookahead,
+//         ))
+//     }
+// }
+//
+// /// Prefetch `len` bytes beginning at `ptr`.
+// ///
+// /// The last cache line prefetched first, followed by the rest in ascending order.
+// ///
+// /// # Safety
+// ///
+// /// The memory range `[ptr, ptr.add(len))` must be valid.
+// #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
+// #[inline(always)]
+// unsafe fn prefetch(ptr: *const u8, len: usize) {
+//     use std::arch::x86_64::*;
+//
+//     // Fetch the last cache line (the one with the tag) first.
+//     let stride = Bytes::CACHELINE.value();
+//     let ptr = ptr.cast::<i8>();
+//     let lines = len.div_ceil(stride);
+//     if lines == 0 {
+//         return;
+//     }
+//
+//     // SAFETY: Inherited from caller.
+//     unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
+//     for i in 0..(lines - 1) {
+//         // SAFETY: Inherited from caller.
+//         unsafe {
+//             _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
+//         }
+//     }
+// }
+//
+// /// Prefetch `len` bytes beginning at `ptr`.
+// ///
+// /// The last cache line prefetched first, followed by the rest in ascending order.
+// ///
+// /// # Safety
+// ///
+// /// The memory range `[ptr, ptr.add(len))` must be valid.
+// #[cfg(not(any(target_arch = "x86_64", target_feature = "avx2")))]
+// unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
+//
+// /// # Safety
+// ///
+// /// * All items in `list` must in-bounds with respect to `reader`.
+// /// * The number of bytes associated with `N` cache lines must "make sense".
+// /// * `buffer.len() >= list.len()`.
+// #[inline]
+// unsafe fn expand_beam_inner<T, const BYTES: usize>(
+//     distance: &T,
+//     list: &[u32],
+//     lookahead: usize,
+//     reader: &store::Reader<'_>,
+//     buffer: &mut [(u32, f32)],
+// ) -> ANNResult<usize>
+// where
+//     T: layers::QueryDistance,
+// {
+//     debug_assert!(
+//         BYTES + store::TAG_SIZE.value() <= reader.inner().bytes().value(),
+//         "we really rely on this: {}, bytes = {}",
+//         BYTES + store::TAG_SIZE.value(),
+//         reader.inner().bytes()
+//     );
+//
+//     debug_assert!(buffer.len() >= list.len());
+//
+//     let bytes = if BYTES == 0 {
+//         reader.inner().bytes().value()
+//     } else {
+//         BYTES + store::TAG_SIZE.value()
+//     };
+//
+//     let len = list.len();
+//     let lookahead = lookahead.min(len);
+//
+//     for j in 0..lookahead {
+//         // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
+//         // as the validity of the prefetch bounds.
+//         unsafe {
+//             prefetch(
+//                 reader
+//                     .inner()
+//                     .read_raw_unchecked(list.get_unchecked(j).into_usize())
+//                     .as_ptr()
+//                     .cast(),
+//                 bytes,
+//             )
+//         }
+//     }
+//
+//     // Disable prefetching if the lookahead is 0.
+//     let mut j = if lookahead == 0 { len } else { lookahead };
+//     let mut processed = 0;
+//     for &i in list.iter() {
+//         if j != len {
+//             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
+//             // well as the validity of the prefetch bounds.
+//             unsafe {
+//                 prefetch(
+//                     reader
+//                         .inner()
+//                         .read_raw_unchecked(list.get_unchecked(j).into_usize())
+//                         .as_ptr()
+//                         .cast(),
+//                     bytes,
+//                 )
+//             }
+//             j += 1;
+//         }
+//
+//         // SAFETY: Caller asserts that `i` is in-bounds.
+//         if let Some(data) = unsafe { reader.inner().read_in_bounds(i.into_usize()) } {
+//             // SAFETY: Inherited from caller.
+//             *unsafe { buffer.get_unchecked_mut(processed) } = (i, distance.evaluate(data)?);
+//             processed += 1;
+//         }
+//     }
+//
+//     Ok(processed)
+// }
 
 ////////////
 // Insert //
@@ -867,17 +867,15 @@ where
         query: L::Query<'a>,
     ) -> ANNResult<SearchAccessor<'a>> {
         let reader = provider.store.reader()?;
-        let expand_beam = <L as layers::Search>::query_distance(
+        let expand_beam = <L as layers::Search>::__search_expand_beam(
             &provider.layer,
             query,
-            ExpandBeamVisitor {
-                bytes: provider.store.plugin().bytes(),
-                prefetch_lookahead: provider.config.prefetch_lookahead.map_or(0, |x| x.get()),
-            },
+            &provider.store,
+            Hidden::new(),
         )?;
 
         let accessor = SearchAccessor {
-            reader,
+            neighbors: provider.store.temp_neighbors(),
             ids: AdjacencyList::new(),
             expand_beam,
             buffer: vec![(0, 0.0); provider.max_degree()],
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 7f08cf9b3c..35eb35b2bb 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -63,17 +63,17 @@ use diskann_utils::views::MatrixView;
 use thiserror::Error;
 
 use crate::{
-    buffer::{BufferError},
+    buffer::BufferError,
     epoch::{self, Registry},
     freelist::{self, Freelist},
     neighbors::{Neighbors, NeighborsError},
-    num::{Bytes},
+    num::Bytes,
     tag::{AtomicTag, Tag},
 };
 
+pub(crate) mod invasive;
 pub(crate) mod plugin;
 pub(crate) mod stacked;
-pub(crate) mod invasive;
 
 pub(crate) const TAG_SIZE: Bytes = AtomicTag::SIZE;
 
@@ -237,6 +237,10 @@ impl Store {
         self.neighbors.max_length()
     }
 
+    pub(crate) fn temp_neighbors(&self) -> &Neighbors {
+        &self.neighbors
+    }
+
     /// Attempt to reclaim retired slots.
     ///
     /// If successful, returns the number of slots reclaimed.
@@ -288,6 +292,11 @@ impl Store {
         })
     }
 
+    // TODO: Rework neighbor storage.
+    pub(crate) fn temp_inner_reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
+        Ok(unsafe { self.plugin().reader(self.registry.guard()?) })
+    }
+
     /// Attempt to acquire a new [`Slot`] for writing.
     ///
     /// This method first consults the freelist and falls back to scanning the tags list
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index 3a314ff3af..4e243cc817 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-use std::{fmt::Debug};
+use std::fmt::Debug;
 
 pub(crate) trait Plugin: Debug + 'static {
     type Slot<'a>: Slot;
@@ -18,5 +18,3 @@ pub(crate) trait Slot: Debug {
     fn freeze(self);
     fn abort(self);
 }
-
-
diff --git a/diskann-inmem/src/store/stacked.rs b/diskann-inmem/src/store/stacked.rs
index eb36ca2f8a..65a3a8b56c 100644
--- a/diskann-inmem/src/store/stacked.rs
+++ b/diskann-inmem/src/store/stacked.rs
@@ -20,7 +20,9 @@ where
     type Slot<'a> = Slot<T::Slot<'a>, U::Slot<'a>>;
 
     unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
-        Slot::new(unsafe { self.first.acquire(i) }, unsafe { self.rest.acquire(i) })
+        Slot::new(unsafe { self.first.acquire(i) }, unsafe {
+            self.rest.acquire(i)
+        })
     }
 
     fn reclaim(&self, i: u32) {

From 5b661a349b0373612dfb029acb13aacbdfc60545 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Fri, 31 Jul 2026 12:49:51 -0700
Subject: [PATCH 05/34] Checkpoint.

---
 diskann-benchmark/src/index/inmem2.rs         |  2 +
 diskann-inmem/integration/support/datatype.rs | 17 +++---
 diskann-inmem/src/layers/full.rs              | 59 ++++++++++---------
 diskann-inmem/src/layers/mod.rs               | 24 +++++---
 diskann-inmem/src/provider.rs                 | 10 ++--
 diskann/src/error/ann_error.rs                |  1 +
 6 files changed, 65 insertions(+), 48 deletions(-)

diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs
index 93c100c902..8622afe243 100644
--- a/diskann-benchmark/src/index/inmem2.rs
+++ b/diskann-benchmark/src/index/inmem2.rs
@@ -47,6 +47,8 @@ use crate::{
 
 pub(crate) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> {
     registry.register("inmem2-f32", Build::<f32>::new())?;
+    registry.register("inmem2-u8", Build::<u8>::new())?;
+
     // registry.register("inmem2-f16", Build::<f16>::new())?;
     registry.register("inmem2-f32-stream", StreamingBenchmark::<f32>::new())?;
     Ok(())
diff --git a/diskann-inmem/integration/support/datatype.rs b/diskann-inmem/integration/support/datatype.rs
index fe61de5398..f0e63debe4 100644
--- a/diskann-inmem/integration/support/datatype.rs
+++ b/diskann-inmem/integration/support/datatype.rs
@@ -7,6 +7,7 @@ use diskann_utils::{
     sampling::medoid::ComputeMedoid,
     views::{Matrix, MatrixView, MutMatrixView},
 };
+use diskann_wide::{cast_f16_to_f32, cast_f32_to_f16};
 use half::f16;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
@@ -146,8 +147,8 @@ where
 }
 
 fn f32_to_f16(x: f32) -> Result<f16, ()> {
-    let y = f16::from_f32(x);
-    let z = f32::from(y);
+    let y = cast_f32_to_f16(x);
+    let z = cast_f16_to_f32(y);
     if z != x { Err(()) } else { Ok(y) }
 }
 
@@ -164,11 +165,11 @@ fn f32_to_i8(x: f32) -> Result<i8, ()> {
 }
 
 fn f16_to_u8(x: f16) -> Result<u8, ()> {
-    f32_to_u8(x.into())
+    f32_to_u8(cast_f16_to_f32(x))
 }
 
 fn f16_to_i8(x: f16) -> Result<i8, ()> {
-    f32_to_i8(x.into())
+    f32_to_i8(cast_f16_to_f32(x))
 }
 
 impl<'a> SliceMut<'a> {
@@ -192,14 +193,14 @@ impl<'a> SliceMut<'a> {
 
         match (self, rhs) {
             (SliceMut::F32(dst), Slice::F32(src)) => dst.copy_from_slice(src),
-            (SliceMut::F32(dst), Slice::F16(src)) => map(dst, src, |x| x.into()),
+            (SliceMut::F32(dst), Slice::F16(src)) => map(dst, src, |x| cast_f16_to_f32(x)),
             (SliceMut::F32(dst), Slice::U8(src)) => map(dst, src, |x| x.into()),
             (SliceMut::F32(dst), Slice::I8(src)) => map(dst, src, |x| x.into()),
 
             (SliceMut::F16(dst), Slice::F32(src)) => try_map(dst, src, f32_to_f16)?,
             (SliceMut::F16(dst), Slice::F16(src)) => dst.copy_from_slice(src),
-            (SliceMut::F16(dst), Slice::U8(src)) => map(dst, src, |x| x.into()),
-            (SliceMut::F16(dst), Slice::I8(src)) => map(dst, src, |x| x.into()),
+            (SliceMut::F16(dst), Slice::U8(src)) => map(dst, src, |x| cast_f32_to_f16(x.into())),
+            (SliceMut::F16(dst), Slice::I8(src)) => map(dst, src, |x| cast_f32_to_f16(x.into())),
 
             (SliceMut::U8(dst), Slice::F32(src)) => try_map(dst, src, f32_to_u8)?,
             (SliceMut::U8(dst), Slice::F16(src)) => try_map(dst, src, f16_to_u8)?,
@@ -544,7 +545,7 @@ mod tests {
         SliceMut::from(dst.as_mut_slice())
             .convert_lossless(Slice::from(src))
             .unwrap();
-        assert_eq!(dst, [f16::from_f32(-5.0), f16::from_f32(7.0)]);
+        assert_eq!(dst, [cast_f32_to_f16(-5.0), cast_f32_to_f16(7.0)]);
     }
 
     #[test]
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 116be8c7f6..eb6dfa9d4b 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -25,6 +25,7 @@ use crate::{
     Hidden, layers,
     num::Bytes,
     store::{self, Store},
+    tag::AtomicTag,
 };
 
 /// A useful trait bound for types compatible with [`Full`].
@@ -286,7 +287,7 @@ impl<T> std::ops::Deref for Calf<'_, T> {
 /// allow `f16` queries to be pre-converted to `f32`, saving on-the-fly conversion that
 /// would otherwise be needed.
 #[derive(Debug)]
-struct QueryDistance<'a, T, U, D> {
+struct QueryDistance<'a, const PREFETCH: usize, T, U, D> {
     // The original query.
     query: Calf<'a, T>,
     // A reader into a layer's store.
@@ -297,8 +298,9 @@ struct QueryDistance<'a, T, U, D> {
     _distance: PhantomData<D>,
 }
 
-impl<'a, T, U, D> QueryDistance<'a, T, U, D> {
+impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
     fn new(query: Calf<'a, T>, reader: store::invasive::Reader<'a>) -> Self {
+        // TODO: Check PREFETCH and `query` with the reader's size.
         Self {
             query,
             reader,
@@ -341,7 +343,7 @@ impl<'a, T, U, D> QueryDistance<'a, T, U, D> {
 const LOOKAHEAD: usize = 8;
 const BYTES: usize = 0;
 
-impl<T, U, D> layers::__ExpandBeam for QueryDistance<'_, T, U, D>
+impl<const PREFETCH: usize, T, U, D> layers::__ExpandBeam for QueryDistance<'_, PREFETCH, T, U, D>
 where
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
@@ -370,10 +372,10 @@ where
         let len = list.len();
         let lookahead = LOOKAHEAD.min(len);
 
-        let bytes = if BYTES == 0 {
+        let bytes = if PREFETCH == 0 {
             self.reader.bytes().value()
         } else {
-            BYTES + store::TAG_SIZE.value()
+            PREFETCH * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
         };
 
         for j in 0..lookahead {
@@ -441,18 +443,17 @@ struct QueryDistanceError {
 diskann::convert_error!(QueryDistanceError);
 
 macro_rules! mint {
-    // ($query:ident, $reader:ident, $T:ty => { $N:literal, $f:ident }) => {{
-    //     mint!($query, $visitor, { $T, $T } => { $N, $f })
-    // }};
-    // ($query:ident, $reader:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
-    //     let inner = Box::new(QueryDistance::<$T, $U, Specialize<$N, $f>>::new($query));
-    //     $visitor.visit_sized::<{ $N * std::mem::size_of::<$U>() }, _>(inner)
-    // }};
+    ($query:ident, $reader:ident, $T:ty => { $N:literal, $f:ident }) => {{
+        mint!($query, $reader, { $T, $T } => { $N, $f })
+    }};
+    ($query:ident, $reader:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
+        Box::new(QueryDistance::<$N, $T, $U, Specialize<$N, $f>>::new($query, $reader))
+    }};
     ($query:ident, $reader:ident, $T:ty => $f:ident) => {{
         mint!($query, $reader, { $T, $T } => $f)
     }};
     ($query:ident, $reader:ident, { $T:ty, $U:ty } => $f:ident) => {{
-        Box::new(QueryDistance::<$T, $U, $f>::new($query, $reader))
+        Box::new(QueryDistance::<0, $T, $U, $f>::new($query, $reader))
     }};
 }
 
@@ -474,11 +475,11 @@ impl FullPrecision for f32 {
 
         let output: Box<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 100 {
-                //     mint!(query, visitor, f32 => { 100, SquaredL2 })
-                // } else {
-                mint!(query, reader, f32 => SquaredL2)
-                // }
+                if full.dim() == 100 {
+                    mint!(query, reader, f32 => { 100, SquaredL2 })
+                } else {
+                    mint!(query, reader, f32 => SquaredL2)
+                }
             }
             Metric::InnerProduct => {
                 mint!(query, reader, f32 => InnerProduct)
@@ -511,11 +512,11 @@ impl FullPrecision for f16 {
 
         let output: Box<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 100 {
-                //     mint!(query, visitor, { f32, f16 } => { 100, SquaredL2 })
-                // } else {
-                mint!(query, reader, { f32, f16 } => SquaredL2)
-                // }
+                if full.dim() == 100 {
+                    mint!(query, reader, { f32, f16 } => { 100, SquaredL2 })
+                } else {
+                    mint!(query, reader, { f32, f16 } => SquaredL2)
+                }
             }
             Metric::InnerProduct => mint!(query, reader, { f32, f16 } => InnerProduct),
             Metric::Cosine => mint!(query, reader, { f32, f16 } => Cosine),
@@ -542,13 +543,13 @@ impl FullPrecision for u8 {
 
         let query = Calf::Borrowed(query);
 
-        let output: Box::<dyn layers::__ExpandBeam> = match full.metric {
+        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 128 {
-                //     mint!(query, visitor, u8 => { 128, SquaredL2 })
-                // } else {
-                mint!(query, reader, u8 => SquaredL2)
-                // }
+                if full.dim() == 128 {
+                    mint!(query, reader, u8 => { 128, SquaredL2 })
+                } else {
+                    mint!(query, reader, u8 => SquaredL2)
+                }
             }
             Metric::InnerProduct => mint!(query, reader, u8 => InnerProduct),
             Metric::Cosine => mint!(query, reader, u8 => Cosine),
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 35591f96dd..198b1d97d8 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -33,19 +33,29 @@ use crate::{Hidden, num::Bytes, store::Store};
 mod full;
 pub use full::{Full, FullPrecision};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Status {
+    Available,
+    Published,
+    Retiring,
+    Frozen,
+}
+
 /// Base layer for data representations.
 pub trait Layer: Send + Sync + 'static {
-    /// Return the number of bytes needed by this layer representation.
-    ///
-    /// To be well-behaved, this function must be idempotent.
+    // /// Return the number of entries present in this layer.
+    // fn entries(&self) -> usize;
+
+    // /// Attempt to freeze entry `i`.
+    // unsafe fn freeze(&self, i: u32) -> ANNResult<()>;
+
+    // /// Attempt to delete entry `i`.
+    // unsafe fn delete(&self, i: u32) -> ANNResult<bool>;
+
     fn bytes(&self) -> Bytes;
 }
 
-/// Store an element of type `T` into a raw byte buffer.
-///
-/// Implementations may assume that `bytes.len()` is equal to [`Layer::bytes`].
 pub trait Set<T>: Layer {
-    /// Write into the stored representation.
     fn set(&self, element: T, bytes: &mut [u8]) -> ANNResult<()>;
 }
 
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index ca4f00ec30..658ca21190 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -459,8 +459,10 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
 
                 // SAFETY: We've verified that each entry in `self.ids` is in-bounds and the
                 // `self.buffer` is long enough to hold all the IDs.
-                let processed =
-                    unsafe { self.expand_beam.__expand_beam(&self.ids, &mut self.buffer, Hidden::new()) }?;
+                let processed = unsafe {
+                    self.expand_beam
+                        .__expand_beam(&self.ids, &mut self.buffer, Hidden::new())
+                }?;
 
                 // let processed = unsafe {
                 //     self.expand_beam
@@ -890,10 +892,10 @@ where
 // This is a utility for helping inspect the generated code for `ExpandBeam`.
 //
 pub fn test_function<'a>(
-    x: &'a Provider<layers::Full<f32>>,
+    x: &'a Provider<layers::Full<u8>>,
     strategy: &'a Strategy,
     context: &'a Context,
-    query: &'a [f32],
+    query: &'a [u8],
 ) -> ANNResult<SearchAccessor<'a>> {
     glue::SearchStrategy::search_accessor(strategy, x, context, query)
 }
diff --git a/diskann/src/error/ann_error.rs b/diskann/src/error/ann_error.rs
index bfd211281d..be12669ef9 100644
--- a/diskann/src/error/ann_error.rs
+++ b/diskann/src/error/ann_error.rs
@@ -242,6 +242,7 @@ macro_rules! convert_error {
     ($T:ty) => {
         impl From<$T> for $crate::ANNError {
             #[track_caller]
+            #[inline]
             fn from(e: $T) -> $crate::ANNError {
                 $crate::ANNError::new(e)
             }

From 36f801ca36a5ef1abc497a3fef60ed26d6a70c7e Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 10:45:06 -0700
Subject: [PATCH 06/34] Checkpoint.

---
 Cargo.lock                      |  1 +
 diskann-inmem/Cargo.toml        |  1 +
 diskann-inmem/src/iter.rs       | 55 +++++++++++++++++++++++++++++++++
 diskann-inmem/src/layers/mod.rs | 29 +++++++++++------
 diskann-inmem/src/lib.rs        |  1 +
 diskann-inmem/src/provider.rs   |  2 +-
 6 files changed, 78 insertions(+), 11 deletions(-)
 create mode 100644 diskann-inmem/src/iter.rs

diff --git a/Cargo.lock b/Cargo.lock
index 50ed17b118..b74bdbede1 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -634,6 +634,7 @@ dependencies = [
  "diskann-vector",
  "diskann-wide",
  "half",
+ "hashbrown 0.16.1",
  "parking_lot",
  "rand",
  "serde",
diff --git a/diskann-inmem/Cargo.toml b/diskann-inmem/Cargo.toml
index 00e02dbf26..1277ff0e0c 100644
--- a/diskann-inmem/Cargo.toml
+++ b/diskann-inmem/Cargo.toml
@@ -27,6 +27,7 @@ anyhow = { workspace = true, optional = true }
 rand = { workspace = true, optional = true }
 diskann-benchmark-core = { workspace = true, optional = true }
 tokio = { workspace = true, optional = true }
+hashbrown.workspace = true
 
 [lints.clippy]
 undocumented_unsafe_blocks = "warn"
diff --git a/diskann-inmem/src/iter.rs b/diskann-inmem/src/iter.rs
new file mode 100644
index 0000000000..57197d808b
--- /dev/null
+++ b/diskann-inmem/src/iter.rs
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::mem::MaybeUninit;
+
+#[derive(Debug)]
+pub(crate) struct StackBuffer<T: Copy, const N: usize>([MaybeUninit<T>; N]);
+
+impl<T, const N: usize> StackBuffer<T, N>
+where
+    T: Copy,
+{
+    pub(crate) fn new() -> Self {
+        Self(core::array::from_fn(|_| MaybeUninit::uninit()))
+    }
+
+    pub(crate) fn as_mut_slice(&mut self) -> StackSlice<'_, T> {
+        StackSlice(&self.0)
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct StackSlice<'a, T: Copy>(&'a mut [MaybeUninit<T>]);
+
+pub(crate) trait Chunked<T>: std::fmt::Debug
+where
+    T: Copy,
+{
+    fn next<'a>(&'a mut self, buffer: StackSlice<'a, T>) -> &'a [T];
+}
+
+#[derive(Debug)]
+pub(crate) struct Iter<I>(pub(crate) I);
+
+impl<I> Chunked<I::Item> for Iter<I>
+where
+    I: Iterator + std::fmt::Debug,
+    I::Item: Copy,
+{
+    fn next<'a>(&'a mut self, buffer: StackSlice<'a, I::Item>) -> &'a [I::Item] {
+        let raw = buffer.0;
+
+        let count = std::iter::zip(
+            raw.iter_mut(),
+            self.0.by_ref(),
+        ).map(|(dst, src)| {
+            dst.write(src);
+        }).count();
+
+        unsafe { raw[..count].assume_init_ref() }
+    }
+}
+
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 198b1d97d8..aa8d5a4c2c 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -75,17 +75,18 @@ pub trait AsDistance: Send + Sync + std::fmt::Debug {
     fn as_distance(&self) -> &dyn Distance;
 }
 
-/// A unary query distance on raw byte slices.
-///
-/// When paired with [`Layer`] via helpers like [`Search`], implementations may assume
-/// that `x` has length [`Layer::bytes`].
-///
-/// No alignment guarantees are made for `x`, though in practice it is likely to be
-/// aligned to 32 or 64 bytes.
-pub trait QueryDistance: Send + Sync + std::fmt::Debug {
-    fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
-}
+// /// A unary query distance on raw byte slices.
+// ///
+// /// When paired with [`Layer`] via helpers like [`Search`], implementations may assume
+// /// that `x` has length [`Layer::bytes`].
+// ///
+// /// No alignment guarantees are made for `x`, though in practice it is likely to be
+// /// aligned to 32 or 64 bytes.
+// pub trait QueryDistance: Send + Sync + std::fmt::Debug {
+//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
+// }
 
+// TODO: Try to hide?
 #[doc(hidden)]
 pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Evaluate a raw distance against index `i`.
@@ -105,6 +106,14 @@ pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
     ) -> ANNResult<usize>;
 }
 
+// TODO: Try to hide?
+#[doc(hidden)]
+pub trait __Prune: Send + Sync + std::fmt::Debug {
+    // fn prepare(
+    //     &mut self,
+    //     hashbrown::hash_map::IterMut<'_, Option
+}
+
 /// Enable search over vectors defined by a [`Layer`].
 pub trait Search: Send + Sync + 'static {
     /// The type of the query. This should be equivalent to the generic parameter in
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index b179fc45ff..d222489f61 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -13,6 +13,7 @@ mod arch;
 mod buffer;
 mod counters;
 mod epoch;
+mod iter;
 mod freelist;
 mod ids;
 mod neighbors;
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 658ca21190..9cde2dcf56 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -53,7 +53,7 @@ use crate::{
     Hidden,
     counters::{Counters, LocalCounters},
     ids::IdMap,
-    layers::{self, QueryDistance},
+    layers,
     neighbors::Neighbors,
     num::Bytes,
     store::{self, Store},

From 76d9832ee124ab3322aa2e5ad5f22aeab9875e33 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 12:51:03 -0700
Subject: [PATCH 07/34] Checkpoint.

---
 diskann-inmem/src/layers/full.rs | 139 +++++++++++++++++++++++++++----
 diskann-vector/src/unaligned.rs  |   3 +
 2 files changed, 124 insertions(+), 18 deletions(-)

diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index ee81466ffd..12f7dcc5b4 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -185,7 +185,18 @@ where
     }
 }
 
-impl<T> layers::Insert for Full<T> where T: FullPrecision {}
+impl<T> layers::Insert for Full<T>
+where
+    T: FullPrecision,
+{
+    fn __prune<'a>(
+        &'a self,
+        store: &'a Store,
+        _: Hidden,
+    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+        T::__prune(Hidden::new(), self, store)
+    }
+}
 
 ///////////
 // Prune //
@@ -197,44 +208,68 @@ struct Prune<'a, T, D> {
     buffer: Vec<Option<UnalignedSlice<'a, T>>>,
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
-    // The expected dim.
-    dim: usize,
-    // Tye type of the data in the original dataset.
-    _data: PhantomData<T>,
     // Type type of the `PureDistanceFunction` used for the implementation.
     _distance: PhantomData<D>,
 }
 
+impl<'a, T, D> Prune<'a, T, D> {
+    fn new(reader: store::invasive::Reader<'a>) -> Self {
+        assert!(
+            reader
+                .bytes()
+                .value()
+                .is_multiple_of(std::mem::size_of::<T>()),
+            "internal inveriant violated",
+        );
+
+        Self {
+            buffer: Vec::new(),
+            reader,
+            _distance: PhantomData,
+        }
+    }
+}
+
 impl<T, D> layers::__Prune for Prune<'_, T, D>
 where
-    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, T>>,
+    T: Send + Sync + 'static + Debug,
+    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, T>>
+        + Send
+        + Sync
+        + Debug,
 {
     fn __prepare(&mut self, iter: &mut dyn crate::iter::Chunked<u32>) -> ANNResult<()> {
         let mut stack = crate::iter::StackBuffer::<u32, CHUNK_SIZE>::new();
         self.buffer.clear();
         loop {
-            let next = iter.next();
+            let next = iter.next(stack.as_mut_slice());
             if next.is_empty() {
                 break;
             }
 
-            self.buffer
-                .extend(next.iter().map(|id| match self.reader.read(i) {
-                    Some(v) => unsafe {
-                        Some(UnalignedSlice::new(
-                            v.as_ptr().cast::<T>(),
-                            reader.bytes().value() / std::mem::size_of::<T>(),
-                        ))
-                    },
-                    None => None,
-                }));
+            self.buffer.extend(
+                next.iter()
+                    .map(|id| match self.reader.read(id.into_usize()) {
+                        Some(v) => unsafe {
+                            Some(UnalignedSlice::new(
+                                v.as_ptr().cast::<T>(),
+                                self.reader.bytes().value() / std::mem::size_of::<T>(),
+                            ))
+                        },
+                        None => None,
+                    }),
+            );
         }
 
         Ok(())
     }
 
     fn __evaluate(&self, a: u32, b: u32) -> f32 {
-        D::run(ARCH, self.buffer[a], self.buffer[b])
+        D::run(
+            ARCH,
+            self.buffer[a.into_usize()].unwrap(),
+            self.buffer[b.into_usize()].unwrap(),
+        )
     }
 }
 
@@ -551,6 +586,23 @@ impl FullPrecision for f32 {
 
         Ok(output)
     }
+
+    fn __prune<'a>(
+        _: Hidden,
+        full: &'a Full<Self>,
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+        let reader = store.temp_inner_reader()?;
+
+        let output: Box<dyn layers::__Prune> = match full.metric {
+            Metric::L2 => Box::new(Prune::<f32, SquaredL2>::new(reader)),
+            Metric::InnerProduct => Box::new(Prune::<f32, InnerProduct>::new(reader)),
+            Metric::Cosine => Box::new(Prune::<f32, Cosine>::new(reader)),
+            Metric::CosineNormalized => Box::new(Prune::<f32, CosineNormalized>::new(reader)),
+        };
+
+        Ok(output)
+    }
 }
 
 impl FullPrecision for f16 {
@@ -586,6 +638,23 @@ impl FullPrecision for f16 {
 
         Ok(output)
     }
+
+    fn __prune<'a>(
+        _: Hidden,
+        full: &'a Full<Self>,
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+        let reader = store.temp_inner_reader()?;
+
+        let output: Box<dyn layers::__Prune> = match full.metric {
+            Metric::L2 => Box::new(Prune::<f16, SquaredL2>::new(reader)),
+            Metric::InnerProduct => Box::new(Prune::<f16, InnerProduct>::new(reader)),
+            Metric::Cosine => Box::new(Prune::<f16, Cosine>::new(reader)),
+            Metric::CosineNormalized => Box::new(Prune::<f16, CosineNormalized>::new(reader)),
+        };
+
+        Ok(output)
+    }
 }
 
 impl FullPrecision for u8 {
@@ -619,6 +688,23 @@ impl FullPrecision for u8 {
 
         Ok(output)
     }
+
+    fn __prune<'a>(
+        _: Hidden,
+        full: &'a Full<Self>,
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+        let reader = store.temp_inner_reader()?;
+
+        let output: Box<dyn layers::__Prune> = match full.metric {
+            Metric::L2 => Box::new(Prune::<u8, SquaredL2>::new(reader)),
+            Metric::InnerProduct => Box::new(Prune::<u8, InnerProduct>::new(reader)),
+            Metric::Cosine => Box::new(Prune::<u8, Cosine>::new(reader)),
+            Metric::CosineNormalized => Box::new(Prune::<u8, CosineNormalized>::new(reader)),
+        };
+
+        Ok(output)
+    }
 }
 
 impl FullPrecision for i8 {
@@ -646,6 +732,23 @@ impl FullPrecision for i8 {
 
         Ok(output)
     }
+
+    fn __prune<'a>(
+        _: Hidden,
+        full: &'a Full<Self>,
+        store: &'a Store,
+    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+        let reader = store.temp_inner_reader()?;
+
+        let output: Box<dyn layers::__Prune> = match full.metric {
+            Metric::L2 => Box::new(Prune::<i8, SquaredL2>::new(reader)),
+            Metric::InnerProduct => Box::new(Prune::<i8, InnerProduct>::new(reader)),
+            Metric::Cosine => Box::new(Prune::<i8, Cosine>::new(reader)),
+            Metric::CosineNormalized => Box::new(Prune::<i8, CosineNormalized>::new(reader)),
+        };
+
+        Ok(output)
+    }
 }
 
 ///////////
diff --git a/diskann-vector/src/unaligned.rs b/diskann-vector/src/unaligned.rs
index 26dec240a3..4871821984 100644
--- a/diskann-vector/src/unaligned.rs
+++ b/diskann-vector/src/unaligned.rs
@@ -81,6 +81,9 @@ impl<'a, T, const N: usize> From<&'a [T; N]> for UnalignedSlice<'a, T> {
     }
 }
 
+unsafe impl<T> Send for UnalignedSlice<'_, T> where T: Sync {}
+unsafe impl<T> Sync for UnalignedSlice<'_, T> where T: Sync {}
+
 /// View `self` as an [`UnalignedSlice`].
 pub trait AsUnaligned {
     /// The element type of the slice.

From 4e82f465f573bc97512ce69ab7e375a041ac6744 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 16:52:08 -0700
Subject: [PATCH 08/34] Closer to new architecture.

---
 .../jsons/integration-baseline.json           |   8 +-
 diskann-inmem/src/layers/full.rs              | 183 +++---------------
 diskann-inmem/src/layers/mod.rs               | 106 +++++-----
 diskann-inmem/src/provider.rs                 |  56 +++---
 diskann-inmem/src/store/invasive.rs           |   2 +-
 5 files changed, 114 insertions(+), 241 deletions(-)

diff --git a/diskann-inmem/integration/jsons/integration-baseline.json b/diskann-inmem/integration/jsons/integration-baseline.json
index 39bd0485d5..941a5cb752 100644
--- a/diskann-inmem/integration/jsons/integration-baseline.json
+++ b/diskann-inmem/integration/jsons/integration-baseline.json
@@ -48,7 +48,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -169,7 +169,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -290,7 +290,7 @@
         "append_neighbors": 96949,
         "distance": 2867876,
         "get_neighbors": 352139,
-        "get_vector": 3067092,
+        "get_vector": 2781513,
         "query_distance": 2240744,
         "set_neighbors": 23599,
         "set_vector": 10000
@@ -414,7 +414,7 @@
         "append_neighbors": 97055,
         "distance": 2867292,
         "get_neighbors": 352087,
-        "get_vector": 3064106,
+        "get_vector": 2778779,
         "query_distance": 2238420,
         "set_neighbors": 23587,
         "set_vector": 10000
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 12f7dcc5b4..0515aad49b 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -36,9 +36,6 @@ const CHUNK_SIZE: usize = 16;
 /// This encompasses *everything* required for `Full: layers::Insert` and can be used as
 /// a single bound.
 pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
-    #[doc(hidden)]
-    fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<Self>;
-
     #[doc(hidden)]
     fn __expand_beam<'a>(
         _: Hidden,
@@ -61,8 +58,9 @@ pub struct Full<T>
 where
     T: 'static,
 {
-    distance: Distance<T>,
+    dim: usize,
     metric: Metric,
+    _type: PhantomData<T>,
 }
 
 impl<T> Full<T>
@@ -74,24 +72,16 @@ where
     where
         T: FullPrecision,
     {
-        T::__new(Hidden::new(), dim, metric)
-    }
-
-    fn from_distance_provider(dim: usize, metric: Metric) -> Self
-    where
-        T: DistanceProvider<T>,
-    {
-        let distance = Distance {
-            f: T::distance_comparer(metric, Some(dim)),
+        Self {
             dim,
-        };
-
-        Self { distance, metric }
+            metric,
+            _type: PhantomData,
+        }
     }
 
     /// Return the logical dimension of the data handled by this [`layers::Layer`].
     pub fn dim(&self) -> usize {
-        self.distance.dim
+        self.dim
     }
 
     /// Return the number of bytes of the data handles by this [`layers::Layer`].
@@ -160,15 +150,6 @@ enum SetError {
 
 diskann::convert_error!(SetError);
 
-impl<T> layers::AsDistance for Full<T>
-where
-    T: FullPrecision,
-{
-    fn as_distance(&self) -> &dyn layers::Distance {
-        &self.distance
-    }
-}
-
 impl<T> layers::Search for Full<T>
 where
     T: FullPrecision,
@@ -205,7 +186,7 @@ where
 #[derive(Debug)]
 struct Prune<'a, T, D> {
     // Buffered data to prune over.
-    buffer: Vec<Option<UnalignedSlice<'a, T>>>,
+    buffer: Vec<UnalignedSlice<'a, T>>,
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
     // Type type of the `PureDistanceFunction` used for the implementation.
@@ -238,123 +219,39 @@ where
         + Sync
         + Debug,
 {
-    fn __prepare(&mut self, iter: &mut dyn crate::iter::Chunked<u32>) -> ANNResult<()> {
-        let mut stack = crate::iter::StackBuffer::<u32, CHUNK_SIZE>::new();
+    fn __prepare(
+        &mut self,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
+    ) -> ANNResult<layers::PruneKey> {
+        let mut counter = layers::PruneKey::counter();
         self.buffer.clear();
-        loop {
-            let next = iter.next(stack.as_mut_slice());
-            if next.is_empty() {
-                break;
+        self.buffer.reserve(items.len());
+
+        for (id, key) in items {
+            if let Some(v) = self.reader.read(id.into_usize()) {
+                self.buffer.push(unsafe { UnalignedSlice::new(
+                    v.as_ptr().cast::<T>(),
+                    self.reader.bytes().value() / std::mem::size_of::<T>(),
+                )});
+
+                *key = Some(counter);
+
+                // Potential overflow issue - but it's exceedingly unlikely that
+                // someone will provide a prune list exceeding `u16::MAX`.
+                //
+                // In addition, `diskann` limits this bound as well.
+                counter = counter.inc()?;
             }
-
-            self.buffer.extend(
-                next.iter()
-                    .map(|id| match self.reader.read(id.into_usize()) {
-                        Some(v) => unsafe {
-                            Some(UnalignedSlice::new(
-                                v.as_ptr().cast::<T>(),
-                                self.reader.bytes().value() / std::mem::size_of::<T>(),
-                            ))
-                        },
-                        None => None,
-                    }),
-            );
         }
 
-        Ok(())
+        Ok(counter)
     }
 
-    fn __evaluate(&self, a: u32, b: u32) -> f32 {
-        D::run(
-            ARCH,
-            self.buffer[a.into_usize()].unwrap(),
-            self.buffer[b.into_usize()].unwrap(),
-        )
+    fn __evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
+        D::run(ARCH, self.buffer[a.index()], self.buffer[b.index()])
     }
 }
 
-//////////////
-// Distance //
-//////////////
-
-#[derive(Debug)]
-#[doc(hidden)]
-pub struct Distance<T, U = T>
-where
-    T: 'static,
-    U: 'static,
-{
-    f: distance::Distance<T, U>,
-    dim: usize,
-}
-
-impl<T, U> Clone for Distance<T, U> {
-    fn clone(&self) -> Self {
-        *self
-    }
-}
-
-impl<T, U> Copy for Distance<T, U> {}
-
-impl<T, U> Distance<T, U>
-where
-    T: 'static,
-    U: 'static,
-{
-    #[cold]
-    #[inline(never)]
-    fn error(&self, x: &[u8], y: &[u8]) -> ANNResult<f32> {
-        let error = DistanceError {
-            expected: self.bytes(),
-            xlen: x.len(),
-            ylen: y.len(),
-        };
-
-        Err(ANNError::new(error))
-    }
-
-    fn dim(&self) -> usize {
-        self.dim
-    }
-
-    fn bytes(&self) -> usize {
-        self.dim() * std::mem::size_of::<U>()
-    }
-}
-
-impl<T> layers::Distance for Distance<T>
-where
-    T: Debug + 'static,
-{
-    fn evaluate(&self, x: &[u8], y: &[u8]) -> ANNResult<f32> {
-        let bytes = self.bytes();
-        if x.len() != bytes || y.len() != bytes {
-            self.error(x, y)
-        } else {
-            // SAFETY: We've checked that both `x` and `y` are valid for
-            // `size_of::<T>() * self.dim` bytes.
-            let ux = unsafe { UnalignedSlice::new(x.as_ptr().cast::<T>(), self.dim) };
-
-            // SAFETY: Same as above
-            let uy = unsafe { UnalignedSlice::new(y.as_ptr().cast::<T>(), self.dim) };
-            Ok(self.f.call_unaligned(ux, uy))
-        }
-    }
-}
-
-#[derive(Debug, Error)]
-#[error(
-    "expected slices of length {} - instead got {} and {}",
-    self.expected,
-    self.xlen,
-    self.ylen
-)]
-struct DistanceError {
-    expected: usize,
-    xlen: usize,
-    ylen: usize,
-}
-
 ///////////////////
 // QueryDistance //
 ///////////////////
@@ -554,10 +451,6 @@ macro_rules! mint {
 }
 
 impl FullPrecision for f32 {
-    fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<f32> {
-        Full::from_distance_provider(dim, metric)
-    }
-
     fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f32>,
@@ -606,10 +499,6 @@ impl FullPrecision for f32 {
 }
 
 impl FullPrecision for f16 {
-    fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<f16> {
-        Full::from_distance_provider(dim, metric)
-    }
-
     fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f16>,
@@ -658,10 +547,6 @@ impl FullPrecision for f16 {
 }
 
 impl FullPrecision for u8 {
-    fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<u8> {
-        Full::from_distance_provider(dim, metric)
-    }
-
     fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<u8>,
@@ -708,10 +593,6 @@ impl FullPrecision for u8 {
 }
 
 impl FullPrecision for i8 {
-    fn __new(_: Hidden, dim: usize, metric: Metric) -> Full<i8> {
-        Full::from_distance_provider(dim, metric)
-    }
-
     fn __expand_beam<'a>(
         _: Hidden,
         full: &'a Full<i8>,
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index d13f5b2d5f..b1c924fed0 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -26,7 +26,10 @@
 //!
 //! 2. Keep the number of specializations bounded for compile time reasons.
 
+use std::num::NonZeroU16;
+
 use diskann::ANNResult;
+use thiserror::Error;
 
 use crate::{Hidden, num::Bytes, store::Store};
 
@@ -59,33 +62,6 @@ pub trait Set<T>: Layer {
     fn set(&self, element: T, bytes: &mut [u8]) -> ANNResult<()>;
 }
 
-/// A distance computation on raw byte slices.
-///
-/// When paired with [`Layer`] via helpers like [`AsDistance`], implementations may assume
-/// that `x` and `y` have length [`Layer::bytes`].
-///
-/// No alignment guarantees are made for `x` and `y`, though in practice they are likely
-/// to be aligned to 32 or 64 bytes.
-pub trait Distance: Send + Sync + std::fmt::Debug {
-    fn evaluate(&self, x: &[u8], y: &[u8]) -> ANNResult<f32>;
-}
-
-/// Return a [`Distance`] function for a [`Layer`].
-pub trait AsDistance: Send + Sync + std::fmt::Debug {
-    fn as_distance(&self) -> &dyn Distance;
-}
-
-// /// A unary query distance on raw byte slices.
-// ///
-// /// When paired with [`Layer`] via helpers like [`Search`], implementations may assume
-// /// that `x` has length [`Layer::bytes`].
-// ///
-// /// No alignment guarantees are made for `x`, though in practice it is likely to be
-// /// aligned to 32 or 64 bytes.
-// pub trait QueryDistance: Send + Sync + std::fmt::Debug {
-//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
-// }
-
 // TODO: Try to hide?
 #[doc(hidden)]
 pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
@@ -106,12 +82,54 @@ pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
     ) -> ANNResult<usize>;
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct PruneKey(NonZeroU16);
+
+impl PruneKey {
+    const ONE: Self = Self(NonZeroU16::new(1).unwrap());
+
+    pub(crate) fn counter() -> Self {
+        Self::ONE
+    }
+
+    pub(crate) fn inc(self) -> Result<Self, Overflow> {
+        match self.0.checked_add(1) {
+            Some(v) => Ok(Self(v)),
+            None => Err(Overflow),
+        }
+    }
+
+    pub(crate) fn as_u64(self) -> u64 {
+        u64::from(self.0.get()) - 1
+    }
+
+    pub(crate) fn index(self) -> usize {
+        usize::from(self.0.get()) - 1
+    }
+}
+
+impl<'a> diskann_utils::Reborrow<'a> for PruneKey {
+    type Target = PruneKey;
+    fn reborrow(&'a self) -> Self::Target {
+        *self
+    }
+}
+
+#[derive(Debug, Error)]
+#[error("prune list exceeded u16::MAX")]
+struct Overflow;
+
+diskann::convert_error!(Overflow);
+
 // TODO: Try to hide?
 #[doc(hidden)]
 pub(crate) trait __Prune: Send + Sync + std::fmt::Debug {
-    fn __prepare(&mut self, items: &mut dyn crate::iter::Chunked<u32>) -> ANNResult<()>;
+    fn __prepare(
+        &mut self,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
+    ) -> ANNResult<PruneKey>;
 
-    fn __evaluate(&self, a: u32, b: u32) -> f32;
+    fn __evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
 }
 
 /// Enable search over vectors defined by a [`Layer`].
@@ -128,41 +146,13 @@ pub trait Search: Send + Sync + 'static {
         store: &'a Store,
         _: Hidden,
     ) -> ANNResult<Box<dyn __ExpandBeam + 'a>>;
-
-    // /// Create a distance computer specialized for `query` and provide it to `visitor`.
-    // fn query_distance<'a, V>(&'a self, query: Self::Query<'a>, visitor: V) -> ANNResult<V::Output>
-    // where
-    //     V: QueryVisitor<'a>;
 }
 
-// /// Specialize a kernel around a [`QueryDistance`] implementation.
-// pub trait QueryVisitor<'a>: Sized {
-//     /// The type of the type-erased output.
-//     type Output;
-//
-//     /// Specialize [`Self::Output`] for `distance`.
-//     fn visit<T>(self, distance: T) -> Self::Output
-//     where
-//         T: QueryDistance + 'a;
-//
-//     /// Specialize [`Self::Output`] for `distance` accepting a hint that `distance` has been
-//     /// specialized to work on data elements of exactly `BYTES` bytes long.
-//     ///
-//     /// This can be used to tailor surrounding code (e.g. software prefetches) for exactly
-//     /// the length of the data being processed.
-//     fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
-//     where
-//         T: QueryDistance + 'a,
-//     {
-//         self.visit(distance)
-//     }
-// }
-
 /// A insert-specific specialization of [`Search`].
 ///
 /// Note that the bounds for this trait are unnecessarily complicated, but rely on changes
 /// to `diskann` to full resolve.
-pub trait Insert: Search + for<'a> Set<Self::Query<'a>> + AsDistance {
+pub trait Insert: Search + for<'a> Set<Self::Query<'a>> {
     /// A specialization of [`Search::query_distance`] targeting vector insert specifically.
     #[doc(hidden)]
     fn __insert_expand_beam<'a, V>(
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 8fee891264..3c5aff07b1 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -706,21 +706,22 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
 /// This type implements zero-copy access to the data within its parent provider during prunes.
 #[derive(Debug)]
 pub struct PruneAccessor<'a> {
-    reader: store::Reader<'a>,
-    distance: &'a dyn layers::Distance,
+    prune: Box<dyn layers::__Prune + 'a>,
+    keys: hashbrown::HashMap<u32, Option<layers::PruneKey>>,
+    neighbors: &'a Neighbors,
     counters: LocalCounters<'a>,
 }
 
 /// The distance computer for [`PruneAccessor`].
 #[derive(Debug)]
 pub struct Distance<'a> {
-    distance: &'a dyn layers::Distance,
+    prune: &'a dyn layers::__Prune,
     counters: LocalCounters<'a>,
 }
 
 impl<'a> Distance<'a> {
-    fn new(distance: &'a dyn layers::Distance, counters: LocalCounters<'a>) -> Self {
-        Self { distance, counters }
+    fn new(prune: &'a dyn layers::__Prune, counters: LocalCounters<'a>) -> Self {
+        Self { prune, counters }
     }
 }
 
@@ -728,11 +729,11 @@ impl<'a> Distance<'a> {
     clippy::unwrap_used,
     reason = "prune does not allow fallible distance functions yet"
 )]
-impl diskann_vector::DistanceFunction<&[u8], &[u8], f32> for Distance<'_> {
+impl diskann_vector::DistanceFunction<layers::PruneKey, layers::PruneKey, f32> for Distance<'_> {
     #[inline]
-    fn evaluate_similarity(&self, x: &[u8], y: &[u8]) -> f32 {
+    fn evaluate_similarity(&self, x: layers::PruneKey, y: layers::PruneKey) -> f32 {
         self.counters.distance_ref(1);
-        self.distance.evaluate(x, y).unwrap()
+        self.prune.__evaluate(x, y)
     }
 }
 
@@ -746,7 +747,7 @@ impl glue::PruneAccessor for PruneAccessor<'_> {
     where
         Self: 'a;
 
-    type ElementRef<'a> = &'a [u8];
+    type ElementRef<'a> = layers::PruneKey;
 
     type View<'a>
         = &'a Self
@@ -764,12 +765,17 @@ impl glue::PruneAccessor for PruneAccessor<'_> {
 
     async fn fill<'a, Itr>(
         &'a mut self,
-        _itr: Itr,
+        itr: Itr,
     ) -> ANNResult<(Self::View<'a>, Self::Distance<'a>)>
     where
         Itr: ExactSizeIterator<Item = Self::Id> + Clone + Send + Sync,
     {
-        Ok((self, Distance::new(self.distance, self.counters.fork())))
+        self.keys.clear();
+        self.keys.extend(itr.map(|i| (i, None)));
+        let count = self.prune.__prepare(self.keys.iter_mut())?;
+        self.counters.get_vector(count.as_u64());
+
+        Ok((self, Distance::new(&*self.prune, self.counters.fork())))
     }
 }
 
@@ -781,7 +787,7 @@ impl provider::NeighborAccessor for PruneAccessor<'_> {
     ) -> impl std::future::Future<Output = ANNResult<()>> + Send {
         let work = move || {
             self.counters.get_neighbors(1);
-            Ok(self.reader.neighbors().get(id, neighbors)?)
+            Ok(self.neighbors.get(id, neighbors)?)
         };
         ready(work)
     }
@@ -795,7 +801,7 @@ impl provider::NeighborAccessorMut for PruneAccessor<'_> {
     ) -> impl std::future::Future<Output = ANNResult<()>> + Send {
         let work = move || {
             self.counters.set_neighbors(1);
-            Ok(self.reader.neighbors().set(id, neighbors)?)
+            Ok(self.neighbors.set(id, neighbors)?)
         };
         ready(work)
     }
@@ -807,7 +813,7 @@ impl provider::NeighborAccessorMut for PruneAccessor<'_> {
     ) -> impl std::future::Future<Output = ANNResult<()>> + Send {
         let work = move || -> ANNResult<()> {
             self.counters.append_vector(1);
-            let lock = self.reader.neighbors().lock(id)?;
+            let lock = self.neighbors.lock(id)?;
 
             // Due to race conditions between calls to `get_neighbors` and `append_vector`
             // in `diskann` - it's possible that the state of the adjacency list has changed
@@ -831,19 +837,14 @@ impl provider::NeighborAccessorMut for PruneAccessor<'_> {
 }
 
 impl workingset::View<u32> for &PruneAccessor<'_> {
-    type ElementRef<'a> = &'a [u8];
+    type ElementRef<'a> = layers::PruneKey;
     type Element<'a>
-        = &'a [u8]
+        = layers::PruneKey
     where
         Self: 'a;
-    fn get(&self, id: u32) -> Option<&[u8]> {
-        match self.reader.inner().read(id.into_usize()) {
-            Some(data) => {
-                self.counters.get_vector_ref(1);
-                Some(data)
-            }
-            None => None,
-        }
+
+    fn get(&self, id: u32) -> Option<layers::PruneKey> {
+        *self.keys.get(&id)?
     }
 }
 
@@ -965,7 +966,7 @@ where
 
 impl<L, M> glue::PruneStrategy<Provider<L, M>> for Strategy
 where
-    L: layers::Layer + layers::AsDistance,
+    L: layers::Insert,
     M: Id,
 {
     type PruneAccessor<'a> = PruneAccessor<'a>;
@@ -978,8 +979,9 @@ where
         _capacity: usize,
     ) -> ANNResult<PruneAccessor<'a>> {
         Ok(PruneAccessor {
-            reader: provider.store.reader()?,
-            distance: <L as layers::AsDistance>::as_distance(&provider.layer),
+            prune: <L as layers::Insert>::__prune(&provider.layer, &provider.store, Hidden::new())?,
+            keys: hashbrown::HashMap::new(),
+            neighbors: provider.store.temp_neighbors(),
             counters: provider.local_counters(),
         })
     }
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index de92811cd4..5015aaac0b 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -235,7 +235,7 @@ impl<'a> Reader<'a> {
 
     /// Return the number of bytes for each entry.
     pub(crate) fn bytes(&self) -> Bytes {
-        self.unpadded
+        self.unpadded.unchecked_sub(AtomicTag::SIZE)
     }
 }
 

From ea7433e226b20614edc162834b67799d1c006ee0 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 17:52:45 -0700
Subject: [PATCH 09/34] Inching closer.

---
 diskann-inmem/src/layers/full.rs | 115 +++++++++++++++++++++++--------
 diskann-inmem/src/layers/mod.rs  |  59 ++++++++--------
 diskann-inmem/src/provider.rs    |  68 +++++++++++-------
 3 files changed, 156 insertions(+), 86 deletions(-)

diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 0515aad49b..c9f0630208 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -22,7 +22,9 @@ use half::f16;
 use thiserror::Error;
 
 use crate::{
-    Hidden, layers,
+    Hidden,
+    counters::LocalCounters,
+    layers,
     num::Bytes,
     store::{self, Store},
     tag::AtomicTag,
@@ -35,9 +37,8 @@ const CHUNK_SIZE: usize = 16;
 ///
 /// This encompasses *everything* required for `Full: layers::Insert` and can be used as
 /// a single bound.
-pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
-    #[doc(hidden)]
-    fn __expand_beam<'a>(
+pub(crate) trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
+    fn make_expand_beam<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         query: &'a [Self],
@@ -45,13 +46,66 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>>;
 
     #[doc(hidden)]
-    fn __prune<'a>(
+    fn make_prune<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
     ) -> ANNResult<Box<dyn layers::__Prune + 'a>>;
 }
 
+pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
+    #[doc(hidden)]
+    fn __search_accessor<'a>(
+        layer: &'a Full<Self>,
+        query: &'a [Self],
+        store: &'a Store,
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
+
+    #[doc(hidden)]
+    fn __prune_accessor<'a>(
+        layer: &'a Full<Self>,
+        store: &'a Store,
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
+}
+
+impl<T> FullPrecision for T
+where
+    T: FullPrecisionImpl,
+{
+    fn __search_accessor<'a>(
+        layer: &'a Full<Self>,
+        query: &'a [Self],
+        store: &'a Store,
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
+        let expand_beam = T::make_expand_beam(Hidden::new(), layer, query, store)?;
+        Ok(crate::provider::SearchAccessor::new(
+            store.temp_neighbors(),
+            expand_beam,
+            provider,
+            store.frozen(),
+            counters,
+        ))
+    }
+
+    fn __prune_accessor<'a>(
+        layer: &'a Full<Self>,
+        store: &'a Store,
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
+        let prune = T::make_prune(Hidden::new(), layer, store)?;
+        Ok(crate::provider::PruneAccessor::new(
+            prune,
+            store.temp_neighbors(),
+            counters,
+        ))
+    }
+}
+
 /// Full-precision data layer.
 #[derive(Debug)]
 pub struct Full<T>
@@ -156,13 +210,14 @@ where
 {
     type Query<'a> = &'a [T];
 
-    fn __search_expand_beam<'a>(
+    fn search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
         store: &'a Store,
-        _: Hidden,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
-        T::__expand_beam(Hidden::new(), self, query, store)
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
+        T::__search_accessor(self, query, store, provider, counters)
     }
 }
 
@@ -170,12 +225,12 @@ impl<T> layers::Insert for Full<T>
 where
     T: FullPrecision,
 {
-    fn __prune<'a>(
+    fn prune_accessor<'a>(
         &'a self,
         store: &'a Store,
-        _: Hidden,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
-        T::__prune(Hidden::new(), self, store)
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
+        T::__prune_accessor(self, store, counters)
     }
 }
 
@@ -229,10 +284,12 @@ where
 
         for (id, key) in items {
             if let Some(v) = self.reader.read(id.into_usize()) {
-                self.buffer.push(unsafe { UnalignedSlice::new(
-                    v.as_ptr().cast::<T>(),
-                    self.reader.bytes().value() / std::mem::size_of::<T>(),
-                )});
+                self.buffer.push(unsafe {
+                    UnalignedSlice::new(
+                        v.as_ptr().cast::<T>(),
+                        self.reader.bytes().value() / std::mem::size_of::<T>(),
+                    )
+                });
 
                 *key = Some(counter);
 
@@ -450,8 +507,8 @@ macro_rules! mint {
     }};
 }
 
-impl FullPrecision for f32 {
-    fn __expand_beam<'a>(
+impl FullPrecisionImpl for f32 {
+    fn make_expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f32>,
         query: &'a [f32],
@@ -480,7 +537,7 @@ impl FullPrecision for f32 {
         Ok(output)
     }
 
-    fn __prune<'a>(
+    fn make_prune<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
@@ -498,8 +555,8 @@ impl FullPrecision for f32 {
     }
 }
 
-impl FullPrecision for f16 {
-    fn __expand_beam<'a>(
+impl FullPrecisionImpl for f16 {
+    fn make_expand_beam<'a>(
         _: Hidden,
         full: &'a Full<f16>,
         query: &'a [f16],
@@ -528,7 +585,7 @@ impl FullPrecision for f16 {
         Ok(output)
     }
 
-    fn __prune<'a>(
+    fn make_prune<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
@@ -546,8 +603,8 @@ impl FullPrecision for f16 {
     }
 }
 
-impl FullPrecision for u8 {
-    fn __expand_beam<'a>(
+impl FullPrecisionImpl for u8 {
+    fn make_expand_beam<'a>(
         _: Hidden,
         full: &'a Full<u8>,
         query: &'a [u8],
@@ -574,7 +631,7 @@ impl FullPrecision for u8 {
         Ok(output)
     }
 
-    fn __prune<'a>(
+    fn make_prune<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
@@ -592,8 +649,8 @@ impl FullPrecision for u8 {
     }
 }
 
-impl FullPrecision for i8 {
-    fn __expand_beam<'a>(
+impl FullPrecisionImpl for i8 {
+    fn make_expand_beam<'a>(
         _: Hidden,
         full: &'a Full<i8>,
         query: &'a [i8],
@@ -614,7 +671,7 @@ impl FullPrecision for i8 {
         Ok(output)
     }
 
-    fn __prune<'a>(
+    fn make_prune<'a>(
         _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index b1c924fed0..c68b28e7fd 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -31,7 +31,7 @@ use std::num::NonZeroU16;
 use diskann::ANNResult;
 use thiserror::Error;
 
-use crate::{Hidden, num::Bytes, store::Store};
+use crate::{Hidden, counters::LocalCounters, num::Bytes, store::Store};
 
 mod full;
 pub use full::{Full, FullPrecision};
@@ -46,15 +46,6 @@ pub enum Status {
 
 /// Base layer for data representations.
 pub trait Layer: Send + Sync + 'static {
-    // /// Return the number of entries present in this layer.
-    // fn entries(&self) -> usize;
-
-    // /// Attempt to freeze entry `i`.
-    // unsafe fn freeze(&self, i: u32) -> ANNResult<()>;
-
-    // /// Attempt to delete entry `i`.
-    // unsafe fn delete(&self, i: u32) -> ANNResult<bool>;
-
     fn bytes(&self) -> Bytes;
 }
 
@@ -63,8 +54,7 @@ pub trait Set<T>: Layer {
 }
 
 // TODO: Try to hide?
-#[doc(hidden)]
-pub trait __ExpandBeam: Send + Sync + std::fmt::Debug {
+pub(crate) trait __ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Evaluate a raw distance against index `i`.
     fn __evaluate(&self, i: u32, _: Hidden) -> ANNResult<Option<f32>>;
 
@@ -121,17 +111,6 @@ struct Overflow;
 
 diskann::convert_error!(Overflow);
 
-// TODO: Try to hide?
-#[doc(hidden)]
-pub(crate) trait __Prune: Send + Sync + std::fmt::Debug {
-    fn __prepare(
-        &mut self,
-        items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
-    ) -> ANNResult<PruneKey>;
-
-    fn __evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
-}
-
 /// Enable search over vectors defined by a [`Layer`].
 pub trait Search: Send + Sync + 'static {
     /// The type of the query. This should be equivalent to the generic parameter in
@@ -140,12 +119,24 @@ pub trait Search: Send + Sync + 'static {
     type Query<'a>;
 
     #[doc(hidden)]
-    fn __search_expand_beam<'a>(
+    fn search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
         store: &'a Store,
-        _: Hidden,
-    ) -> ANNResult<Box<dyn __ExpandBeam + 'a>>;
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
+}
+
+// TODO: Try to hide?
+#[doc(hidden)]
+pub(crate) trait __Prune: Send + Sync + std::fmt::Debug {
+    fn __prepare(
+        &mut self,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
+    ) -> ANNResult<PruneKey>;
+
+    fn __evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
 }
 
 /// A insert-specific specialization of [`Search`].
@@ -153,17 +144,21 @@ pub trait Search: Send + Sync + 'static {
 /// Note that the bounds for this trait are unnecessarily complicated, but rely on changes
 /// to `diskann` to full resolve.
 pub trait Insert: Search + for<'a> Set<Self::Query<'a>> {
-    /// A specialization of [`Search::query_distance`] targeting vector insert specifically.
     #[doc(hidden)]
-    fn __insert_expand_beam<'a, V>(
+    fn insert_search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
         store: &'a Store,
-        _: Hidden,
-    ) -> ANNResult<Box<dyn __ExpandBeam + 'a>> {
-        self.__search_expand_beam(query, store, Hidden::new())
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
+        self.search_accessor(query, store, provider, counters)
     }
 
     #[doc(hidden)]
-    fn __prune<'a>(&'a self, store: &'a Store, _: Hidden) -> ANNResult<Box<dyn __Prune + 'a>>;
+    fn prune_accessor<'a>(
+        &'a self,
+        store: &'a Store,
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
 }
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 3c5aff07b1..4f88214b67 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -396,6 +396,26 @@ pub struct SearchAccessor<'a> {
     counters: LocalCounters<'a>,
 }
 
+impl<'a> SearchAccessor<'a> {
+    pub(crate) fn new(
+        neighbors: &'a Neighbors,
+        expand_beam: Box<dyn layers::__ExpandBeam + 'a>,
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        start_points: std::ops::Range<u32>,
+        counters: LocalCounters<'a>,
+    ) -> Self {
+        Self {
+            neighbors,
+            ids: AdjacencyList::with_capacity(neighbors.max_length()),
+            expand_beam,
+            buffer: vec![Default::default(); neighbors.max_length()],
+            provider,
+            start_points,
+            counters,
+        }
+    }
+}
+
 impl diskann::provider::HasId for SearchAccessor<'_> {
     type Id = u32;
 }
@@ -464,11 +484,6 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
                         .__expand_beam(&self.ids, &mut self.buffer, Hidden::new())
                 }?;
 
-                // let processed = unsafe {
-                //     self.expand_beam
-                //         .expand_beam(&self.ids, &self.reader, &mut self.buffer)
-                // }?;
-
                 self.counters.get_vector(processed as u64);
                 self.counters.query_distance(processed as u64);
 
@@ -712,6 +727,21 @@ pub struct PruneAccessor<'a> {
     counters: LocalCounters<'a>,
 }
 
+impl<'a> PruneAccessor<'a> {
+    pub(crate) fn new(
+        prune: Box<dyn layers::__Prune + 'a>,
+        neighbors: &'a Neighbors,
+        counters: LocalCounters<'a>,
+    ) -> Self {
+        Self {
+            prune,
+            keys: hashbrown::HashMap::new(),
+            neighbors,
+            counters,
+        }
+    }
+}
+
 /// The distance computer for [`PruneAccessor`].
 #[derive(Debug)]
 pub struct Distance<'a> {
@@ -869,24 +899,13 @@ where
         _context: &'a Context,
         query: L::Query<'a>,
     ) -> ANNResult<SearchAccessor<'a>> {
-        let reader = provider.store.reader()?;
-        let expand_beam = <L as layers::Search>::__search_expand_beam(
+        <L as layers::Search>::search_accessor(
             &provider.layer,
             query,
             &provider.store,
-            Hidden::new(),
-        )?;
-
-        let accessor = SearchAccessor {
-            neighbors: provider.store.temp_neighbors(),
-            ids: AdjacencyList::new(),
-            expand_beam,
-            buffer: vec![(0, 0.0); provider.max_degree()],
             provider,
-            start_points: provider.store.frozen(),
-            counters: provider.local_counters(),
-        };
-        Ok(accessor)
+            provider.local_counters(),
+        )
     }
 }
 
@@ -978,12 +997,11 @@ where
         _context: &'a Context,
         _capacity: usize,
     ) -> ANNResult<PruneAccessor<'a>> {
-        Ok(PruneAccessor {
-            prune: <L as layers::Insert>::__prune(&provider.layer, &provider.store, Hidden::new())?,
-            keys: hashbrown::HashMap::new(),
-            neighbors: provider.store.temp_neighbors(),
-            counters: provider.local_counters(),
-        })
+        <L as layers::Insert>::prune_accessor(
+            &provider.layer,
+            &provider.store,
+            provider.local_counters(),
+        )
     }
 }
 

From f54fbf5c2112e688be3681b65ea61f82e85a355f Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 18:03:38 -0700
Subject: [PATCH 10/34] Checkpoint cleanups.

---
 diskann-inmem/src/iter.rs        | 53 ----------------------
 diskann-inmem/src/layers/full.rs | 77 +++++++++++++-------------------
 diskann-inmem/src/layers/mod.rs  | 16 +++----
 diskann-inmem/src/lib.rs         | 17 -------
 diskann-inmem/src/provider.rs    | 21 +++++----
 5 files changed, 48 insertions(+), 136 deletions(-)
 delete mode 100644 diskann-inmem/src/iter.rs

diff --git a/diskann-inmem/src/iter.rs b/diskann-inmem/src/iter.rs
deleted file mode 100644
index e33418b58c..0000000000
--- a/diskann-inmem/src/iter.rs
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * Copyright (c) Microsoft Corporation.
- * Licensed under the MIT license.
- */
-
-use std::mem::MaybeUninit;
-
-#[derive(Debug)]
-pub(crate) struct StackBuffer<T: Copy, const N: usize>([MaybeUninit<T>; N]);
-
-impl<T, const N: usize> StackBuffer<T, N>
-where
-    T: Copy,
-{
-    pub(crate) fn new() -> Self {
-        Self(core::array::from_fn(|_| MaybeUninit::uninit()))
-    }
-
-    pub(crate) fn as_mut_slice(&mut self) -> StackSlice<'_, T> {
-        StackSlice(&mut self.0)
-    }
-}
-
-#[derive(Debug)]
-pub(crate) struct StackSlice<'a, T: Copy>(&'a mut [MaybeUninit<T>]);
-
-pub(crate) trait Chunked<T>: std::fmt::Debug
-where
-    T: Copy,
-{
-    fn next<'a>(&'a mut self, buffer: StackSlice<'a, T>) -> &'a [T];
-}
-
-#[derive(Debug)]
-pub(crate) struct Iter<I>(pub(crate) I);
-
-impl<I> Chunked<I::Item> for Iter<I>
-where
-    I: Iterator + std::fmt::Debug,
-    I::Item: Copy,
-{
-    fn next<'a>(&'a mut self, buffer: StackSlice<'a, I::Item>) -> &'a [I::Item] {
-        let raw = buffer.0;
-
-        let count = std::iter::zip(raw.iter_mut(), self.0.by_ref())
-            .map(|(dst, src)| {
-                dst.write(src);
-            })
-            .count();
-
-        unsafe { raw[..count].assume_init_ref() }
-    }
-}
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index c9f0630208..f34da91951 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -22,7 +22,6 @@ use half::f16;
 use thiserror::Error;
 
 use crate::{
-    Hidden,
     counters::LocalCounters,
     layers,
     num::Bytes,
@@ -30,29 +29,24 @@ use crate::{
     tag::AtomicTag,
 };
 
-/// The granularity for processing iterator chunks.
-const CHUNK_SIZE: usize = 16;
-
-/// A useful trait bound for types compatible with [`Full`].
-///
-/// This encompasses *everything* required for `Full: layers::Insert` and can be used as
-/// a single bound.
-pub(crate) trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
+trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_expand_beam<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         query: &'a [Self],
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>>;
+    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>>;
 
     #[doc(hidden)]
     fn make_prune<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>>;
+    ) -> ANNResult<Box<dyn layers::Prune + 'a>>;
 }
 
+/// A useful trait bound for types compatible with [`Full`].
+///
+/// This encompasses *everything* required for `Full: layers::Insert` and can be used as
+/// a single bound.
 pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     #[doc(hidden)]
     fn __search_accessor<'a>(
@@ -82,7 +76,7 @@ where
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        let expand_beam = T::make_expand_beam(Hidden::new(), layer, query, store)?;
+        let expand_beam = T::make_expand_beam(layer, query, store)?;
         Ok(crate::provider::SearchAccessor::new(
             store.temp_neighbors(),
             expand_beam,
@@ -97,7 +91,7 @@ where
         store: &'a Store,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
-        let prune = T::make_prune(Hidden::new(), layer, store)?;
+        let prune = T::make_prune(layer, store)?;
         Ok(crate::provider::PruneAccessor::new(
             prune,
             store.temp_neighbors(),
@@ -266,7 +260,7 @@ impl<'a, T, D> Prune<'a, T, D> {
     }
 }
 
-impl<T, D> layers::__Prune for Prune<'_, T, D>
+impl<T, D> layers::Prune for Prune<'_, T, D>
 where
     T: Send + Sync + 'static + Debug,
     D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, T>>
@@ -274,7 +268,7 @@ where
         + Sync
         + Debug,
 {
-    fn __prepare(
+    fn prepare(
         &mut self,
         items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
     ) -> ANNResult<layers::PruneKey> {
@@ -304,7 +298,7 @@ where
         Ok(counter)
     }
 
-    fn __evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
+    fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
         D::run(ARCH, self.buffer[a.index()], self.buffer[b.index()])
     }
 }
@@ -393,7 +387,7 @@ impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
 const LOOKAHEAD: usize = 8;
 const BYTES: usize = 0;
 
-impl<const PREFETCH: usize, T, U, D> layers::__ExpandBeam for QueryDistance<'_, PREFETCH, T, U, D>
+impl<const PREFETCH: usize, T, U, D> layers::ExpandBeam for QueryDistance<'_, PREFETCH, T, U, D>
 where
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
@@ -402,7 +396,7 @@ where
         + Sync
         + Debug,
 {
-    fn __evaluate(&self, i: u32, _: Hidden) -> ANNResult<Option<f32>> {
+    fn evaluate(&self, i: u32) -> ANNResult<Option<f32>> {
         if !self.reader.is_in_bounds(i.into_usize()) {
             return Err(ANNError::new(OutOfBounds(i)));
         } else {
@@ -413,11 +407,10 @@ where
         }
     }
 
-    unsafe fn __expand_beam(
+    unsafe fn expand_beam(
         &self,
         list: &[u32],
         buffer: &mut [(u32, f32)],
-        _: Hidden,
     ) -> ANNResult<usize> {
         let len = list.len();
         let lookahead = LOOKAHEAD.min(len);
@@ -509,17 +502,16 @@ macro_rules! mint {
 
 impl FullPrecisionImpl for f32 {
     fn make_expand_beam<'a>(
-        _: Hidden,
         full: &'a Full<f32>,
         query: &'a [f32],
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
         let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
+        let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 100 {
                     mint!(query, reader, f32 => { 100, SquaredL2 })
@@ -538,13 +530,12 @@ impl FullPrecisionImpl for f32 {
     }
 
     fn make_prune<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = store.temp_inner_reader()?;
 
-        let output: Box<dyn layers::__Prune> = match full.metric {
+        let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f32, SquaredL2>::new(reader)),
             Metric::InnerProduct => Box::new(Prune::<f32, InnerProduct>::new(reader)),
             Metric::Cosine => Box::new(Prune::<f32, Cosine>::new(reader)),
@@ -557,11 +548,10 @@ impl FullPrecisionImpl for f32 {
 
 impl FullPrecisionImpl for f16 {
     fn make_expand_beam<'a>(
-        _: Hidden,
         full: &'a Full<f16>,
         query: &'a [f16],
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
         let reader = store.temp_inner_reader()?;
 
@@ -569,7 +559,7 @@ impl FullPrecisionImpl for f16 {
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
         let query = Calf::Owned(as_f32);
 
-        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
+        let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 100 {
                     mint!(query, reader, { f32, f16 } => { 100, SquaredL2 })
@@ -586,13 +576,12 @@ impl FullPrecisionImpl for f16 {
     }
 
     fn make_prune<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = store.temp_inner_reader()?;
 
-        let output: Box<dyn layers::__Prune> = match full.metric {
+        let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f16, SquaredL2>::new(reader)),
             Metric::InnerProduct => Box::new(Prune::<f16, InnerProduct>::new(reader)),
             Metric::Cosine => Box::new(Prune::<f16, Cosine>::new(reader)),
@@ -605,17 +594,16 @@ impl FullPrecisionImpl for f16 {
 
 impl FullPrecisionImpl for u8 {
     fn make_expand_beam<'a>(
-        _: Hidden,
         full: &'a Full<u8>,
         query: &'a [u8],
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
         let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output: Box<dyn layers::__ExpandBeam> = match full.metric {
+        let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 128 {
                     mint!(query, reader, u8 => { 128, SquaredL2 })
@@ -632,13 +620,12 @@ impl FullPrecisionImpl for u8 {
     }
 
     fn make_prune<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = store.temp_inner_reader()?;
 
-        let output: Box<dyn layers::__Prune> = match full.metric {
+        let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<u8, SquaredL2>::new(reader)),
             Metric::InnerProduct => Box::new(Prune::<u8, InnerProduct>::new(reader)),
             Metric::Cosine => Box::new(Prune::<u8, Cosine>::new(reader)),
@@ -651,17 +638,16 @@ impl FullPrecisionImpl for u8 {
 
 impl FullPrecisionImpl for i8 {
     fn make_expand_beam<'a>(
-        _: Hidden,
         full: &'a Full<i8>,
         query: &'a [i8],
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
         let reader = store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
-        let output: Box<dyn layers::__ExpandBeam + 'a> = match full.metric {
+        let output: Box<dyn layers::ExpandBeam + 'a> = match full.metric {
             Metric::L2 => mint!(query, reader, i8 => SquaredL2),
             Metric::InnerProduct => mint!(query, reader, i8 => InnerProduct),
             Metric::Cosine => mint!(query, reader, i8 => Cosine),
@@ -672,13 +658,12 @@ impl FullPrecisionImpl for i8 {
     }
 
     fn make_prune<'a>(
-        _: Hidden,
         full: &'a Full<Self>,
         store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::__Prune + 'a>> {
+    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = store.temp_inner_reader()?;
 
-        let output: Box<dyn layers::__Prune> = match full.metric {
+        let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<i8, SquaredL2>::new(reader)),
             Metric::InnerProduct => Box::new(Prune::<i8, InnerProduct>::new(reader)),
             Metric::Cosine => Box::new(Prune::<i8, Cosine>::new(reader)),
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index c68b28e7fd..44a82c2b55 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -31,7 +31,7 @@ use std::num::NonZeroU16;
 use diskann::ANNResult;
 use thiserror::Error;
 
-use crate::{Hidden, counters::LocalCounters, num::Bytes, store::Store};
+use crate::{counters::LocalCounters, num::Bytes, store::Store};
 
 mod full;
 pub use full::{Full, FullPrecision};
@@ -53,10 +53,9 @@ pub trait Set<T>: Layer {
     fn set(&self, element: T, bytes: &mut [u8]) -> ANNResult<()>;
 }
 
-// TODO: Try to hide?
-pub(crate) trait __ExpandBeam: Send + Sync + std::fmt::Debug {
+pub(crate) trait ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Evaluate a raw distance against index `i`.
-    fn __evaluate(&self, i: u32, _: Hidden) -> ANNResult<Option<f32>>;
+    fn evaluate(&self, i: u32) -> ANNResult<Option<f32>>;
 
     /// Compute the distance between the query and each neighbor in `list`.
     ///
@@ -64,11 +63,10 @@ pub(crate) trait __ExpandBeam: Send + Sync + std::fmt::Debug {
     ///
     /// * All items in `list` must in-bounds with respect to `reader`.
     /// * `buffer.len() >= list.len()`.
-    unsafe fn __expand_beam(
+    unsafe fn expand_beam(
         &self,
         list: &[u32],
         buffer: &mut [(u32, f32)],
-        _: Hidden,
     ) -> ANNResult<usize>;
 }
 
@@ -130,13 +128,13 @@ pub trait Search: Send + Sync + 'static {
 
 // TODO: Try to hide?
 #[doc(hidden)]
-pub(crate) trait __Prune: Send + Sync + std::fmt::Debug {
-    fn __prepare(
+pub(crate) trait Prune: Send + Sync + std::fmt::Debug {
+    fn prepare(
         &mut self,
         items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
     ) -> ANNResult<PruneKey>;
 
-    fn __evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
+    fn evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
 }
 
 /// A insert-specific specialization of [`Search`].
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 4d36bfd9fa..db5ba75799 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -15,7 +15,6 @@ mod counters;
 mod epoch;
 mod freelist;
 mod ids;
-mod iter;
 mod neighbors;
 mod tag;
 
@@ -33,19 +32,3 @@ mod test;
 #[doc(hidden)]
 pub mod integration;
 
-//----------------//
-// Internal Tools //
-//----------------//
-
-/// A "public" type that can only be constructed by this crate.
-///
-/// This helps with public traits with internal methods that we don't want users to call.
-#[doc(hidden)]
-#[derive(Debug)]
-pub struct Hidden(());
-
-impl Hidden {
-    const fn new() -> Self {
-        Self(())
-    }
-}
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 4f88214b67..9ac3f781d9 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -50,7 +50,6 @@ use diskann_utils::views::Matrix;
 use thiserror::Error;
 
 use crate::{
-    Hidden,
     counters::{Counters, LocalCounters},
     ids::IdMap,
     layers,
@@ -387,7 +386,7 @@ where
 pub struct SearchAccessor<'a> {
     neighbors: &'a Neighbors,
     ids: AdjacencyList<u32>,
-    expand_beam: Box<dyn layers::__ExpandBeam + 'a>,
+    expand_beam: Box<dyn layers::ExpandBeam + 'a>,
     buffer: Vec<(u32, f32)>,
 
     // The parent provider for the accessor.
@@ -399,7 +398,7 @@ pub struct SearchAccessor<'a> {
 impl<'a> SearchAccessor<'a> {
     pub(crate) fn new(
         neighbors: &'a Neighbors,
-        expand_beam: Box<dyn layers::__ExpandBeam + 'a>,
+        expand_beam: Box<dyn layers::ExpandBeam + 'a>,
         provider: &'a (dyn std::any::Any + Send + Sync),
         start_points: std::ops::Range<u32>,
         counters: LocalCounters<'a>,
@@ -436,7 +435,7 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     {
         let work = move || {
             for p in self.start_points.clone() {
-                match self.expand_beam.__evaluate(p, Hidden::new())? {
+                match self.expand_beam.evaluate(p)? {
                     Some(distance) => {
                         // Counters are no-ops without `integration-test`.
                         self.counters.get_vector(1);
@@ -481,7 +480,7 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
                 // `self.buffer` is long enough to hold all the IDs.
                 let processed = unsafe {
                     self.expand_beam
-                        .__expand_beam(&self.ids, &mut self.buffer, Hidden::new())
+                        .expand_beam(&self.ids, &mut self.buffer)
                 }?;
 
                 self.counters.get_vector(processed as u64);
@@ -721,7 +720,7 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
 /// This type implements zero-copy access to the data within its parent provider during prunes.
 #[derive(Debug)]
 pub struct PruneAccessor<'a> {
-    prune: Box<dyn layers::__Prune + 'a>,
+    prune: Box<dyn layers::Prune + 'a>,
     keys: hashbrown::HashMap<u32, Option<layers::PruneKey>>,
     neighbors: &'a Neighbors,
     counters: LocalCounters<'a>,
@@ -729,7 +728,7 @@ pub struct PruneAccessor<'a> {
 
 impl<'a> PruneAccessor<'a> {
     pub(crate) fn new(
-        prune: Box<dyn layers::__Prune + 'a>,
+        prune: Box<dyn layers::Prune + 'a>,
         neighbors: &'a Neighbors,
         counters: LocalCounters<'a>,
     ) -> Self {
@@ -745,12 +744,12 @@ impl<'a> PruneAccessor<'a> {
 /// The distance computer for [`PruneAccessor`].
 #[derive(Debug)]
 pub struct Distance<'a> {
-    prune: &'a dyn layers::__Prune,
+    prune: &'a dyn layers::Prune,
     counters: LocalCounters<'a>,
 }
 
 impl<'a> Distance<'a> {
-    fn new(prune: &'a dyn layers::__Prune, counters: LocalCounters<'a>) -> Self {
+    fn new(prune: &'a dyn layers::Prune, counters: LocalCounters<'a>) -> Self {
         Self { prune, counters }
     }
 }
@@ -763,7 +762,7 @@ impl diskann_vector::DistanceFunction<layers::PruneKey, layers::PruneKey, f32> f
     #[inline]
     fn evaluate_similarity(&self, x: layers::PruneKey, y: layers::PruneKey) -> f32 {
         self.counters.distance_ref(1);
-        self.prune.__evaluate(x, y)
+        self.prune.evaluate(x, y)
     }
 }
 
@@ -802,7 +801,7 @@ impl glue::PruneAccessor for PruneAccessor<'_> {
     {
         self.keys.clear();
         self.keys.extend(itr.map(|i| (i, None)));
-        let count = self.prune.__prepare(self.keys.iter_mut())?;
+        let count = self.prune.prepare(self.keys.iter_mut())?;
         self.counters.get_vector(count.as_u64());
 
         Ok((self, Distance::new(&*self.prune, self.counters.fork())))

From 63b73635a8cb7a4c8f2d51ab27d329ccf1a89e03 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 17 Aug 2026 18:14:45 -0700
Subject: [PATCH 11/34] Turns out it's faster!

---
 diskann-inmem/src/layers/full.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index f34da91951..e1a24981e5 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -357,7 +357,6 @@ impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
         std::mem::size_of::<U>() * self.query.len()
     }
 
-    #[inline(never)]
     fn error(&self, len: usize) -> ANNResult<f32> {
         let error = QueryDistanceError {
             expected: self.bytes(),

From 49fddedd7db780c396f4e6d14e42f624db924a82 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 18 Aug 2026 12:04:35 -0700
Subject: [PATCH 12/34] Prepare for the great switcheroo.

---
 diskann-inmem/src/counters.rs          |  8 ++-
 diskann-inmem/src/integration/store.rs |  9 +--
 diskann-inmem/src/layers/full.rs       | 51 ++++++++++++--
 diskann-inmem/src/layers/mod.rs        | 20 ++++--
 diskann-inmem/src/lib.rs               |  1 -
 diskann-inmem/src/provider.rs          | 14 ++--
 diskann-inmem/src/store/mod.rs         | 94 +++++++++++++-------------
 7 files changed, 127 insertions(+), 70 deletions(-)

diff --git a/diskann-inmem/src/counters.rs b/diskann-inmem/src/counters.rs
index b53940dc2c..882ad4adda 100644
--- a/diskann-inmem/src/counters.rs
+++ b/diskann-inmem/src/counters.rs
@@ -22,8 +22,10 @@ mod inner {
         }
     }
 
+    // Must be public because it shows up in public APIs - but is unconstructable by
+    // downstream code.
     #[derive(Debug)]
-    pub(crate) struct LocalCounters<'a> {
+    pub struct LocalCounters<'a> {
         _marker: PhantomData<&'a ()>,
     }
 
@@ -86,8 +88,10 @@ mod inner {
         }
     }
 
+    // Must be public because it shows up in public APIs - but is unconstructable by
+    // downstream code.
     #[derive(Debug)]
-    pub(crate) struct LocalCounters<'a> {
+    pub struct LocalCounters<'a> {
         query_distance: u64,
         // This fields needs to be `AtomicU64` because we increment in some loops where we
         // have to increment it behind a shared reference.
diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 9a8b9865f7..78b61b7f2c 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -41,8 +41,9 @@ impl Store {
     /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
     /// the number of epoch guard slots are invalid (e.g. zero).
     pub fn new(config: Config) -> Self {
-        let mut store_config =
-            store::Config::new(config.capacity, Bytes::new(config.entry_bytes), 0);
+        let mut store_layout = store::Layout::new(config.capacity, 0);
+
+        let mut store_config = store::Config::default();
 
         store_config
             .epoch_guard_slots(
@@ -60,8 +61,8 @@ impl Store {
             );
 
         let data = Matrix::new(0u8, 1, config.entry_bytes);
-        let store =
-            store::Store::new(store_config, data.as_view()).expect("failed to construct store");
+        let store = store::Store::new(store_layout, store_config, data.as_view())
+            .expect("failed to construct store");
         Self { store }
     }
 
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index e1a24981e5..c4fffcf846 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -6,6 +6,7 @@
 use std::{fmt::Debug, marker::PhantomData};
 
 use diskann::{ANNError, ANNResult, utils::IntoUsize};
+use diskann_utils::views::Matrix;
 use diskann_vector::{
     UnalignedSlice,
     conversion::SliceCast,
@@ -29,6 +30,50 @@ use crate::{
     tag::AtomicTag,
 };
 
+#[derive(Debug)]
+pub struct Config<T> {
+    layout: store::Layout,
+    metric: Metric,
+    start_points: Matrix<T>,
+    store: store::Config,
+}
+
+impl<T> Config<T> {
+    pub fn new(
+        capacity: usize,
+        max_degree: usize,
+        metric: Metric,
+        start_points: Matrix<T>,
+    ) -> Self {
+        Self {
+            layout: store::Layout::new(capacity, max_degree),
+            metric,
+            start_points,
+            store: store::Config::default(),
+        }
+    }
+
+    pub fn store(mut self, config: store::Config) -> Self {
+        self.store = config;
+        self
+    }
+
+    fn dim(&self) -> usize {
+        self.start_points.ncols()
+    }
+}
+
+impl<T> layers::LayerConfig for Config<T>
+where
+    T: FullPrecision,
+{
+    type Layer = Full<T>;
+
+    fn build(self) -> ANNResult<Full<T>> {
+        Ok(Full::new(self.dim(), self.metric))
+    }
+}
+
 trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_expand_beam<'a>(
         full: &'a Full<Self>,
@@ -406,11 +451,7 @@ where
         }
     }
 
-    unsafe fn expand_beam(
-        &self,
-        list: &[u32],
-        buffer: &mut [(u32, f32)],
-    ) -> ANNResult<usize> {
+    unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize> {
         let len = list.len();
         let lookahead = LOOKAHEAD.min(len);
 
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 44a82c2b55..08e4e388c4 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -44,8 +44,22 @@ pub enum Status {
     Frozen,
 }
 
+impl Status {
+    #[must_use = "this function has no side-effects"]
+    pub fn is_readable(self) -> bool {
+        matches!(self, Self::Published | Self::Frozen)
+    }
+}
+
+pub trait LayerConfig {
+    type Layer: Layer;
+
+    fn build(self) -> ANNResult<Self::Layer>;
+}
+
 /// Base layer for data representations.
 pub trait Layer: Send + Sync + 'static {
+    // // Return the status of the item behind internal ID `i`.
     fn bytes(&self) -> Bytes;
 }
 
@@ -63,11 +77,7 @@ pub(crate) trait ExpandBeam: Send + Sync + std::fmt::Debug {
     ///
     /// * All items in `list` must in-bounds with respect to `reader`.
     /// * `buffer.len() >= list.len()`.
-    unsafe fn expand_beam(
-        &self,
-        list: &[u32],
-        buffer: &mut [(u32, f32)],
-    ) -> ANNResult<usize>;
+    unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize>;
 }
 
 #[derive(Debug, Clone, Copy)]
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index db5ba75799..1c63e613cb 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -31,4 +31,3 @@ mod test;
 #[cfg(feature = "integration-test")]
 #[doc(hidden)]
 pub mod integration;
-
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 9ac3f781d9..557c0ce882 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -107,7 +107,7 @@ where
             layers::Set::set(&layer, point, row)?;
         }
 
-        let mut store_config = store::Config::new(config.capacity(), bytes, config.max_degree());
+        let mut store_config = store::Config::default();
 
         if let Some(slots) = config.epoch_guard_slots {
             store_config.epoch_guard_slots(slots);
@@ -117,7 +117,9 @@ where
             store_config.freelist_recycle_capacity(capacity);
         }
 
-        let store = Store::new(store_config, data.as_view())
+        let mut store_layout = store::Layout::new(config.capacity(), config.max_degree());
+
+        let store = Store::new(store_layout, store_config, data.as_view())
             .map_err(|err| ProviderError::CreatingStore(Box::new(err)))?;
 
         let mapping = IdMap::new(config.capacity());
@@ -470,6 +472,8 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
                 self.counters.get_neighbors(1);
 
                 // Filter out unvisited IDs and ensure that all the IDs we are about
+                //
+                // TODO: We need safe provenance on the upper bound.
                 self.ids
                     .retain(|i| pred.eval_mut(i) && *i < self.neighbors.entries());
 
@@ -478,10 +482,8 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
 
                 // SAFETY: We've verified that each entry in `self.ids` is in-bounds and the
                 // `self.buffer` is long enough to hold all the IDs.
-                let processed = unsafe {
-                    self.expand_beam
-                        .expand_beam(&self.ids, &mut self.buffer)
-                }?;
+                let processed =
+                    unsafe { self.expand_beam.expand_beam(&self.ids, &mut self.buffer) }?;
 
                 self.counters.get_vector(processed as u64);
                 self.counters.query_distance(processed as u64);
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 35eb35b2bb..3a455d3c68 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -79,16 +79,7 @@ pub(crate) const TAG_SIZE: Bytes = AtomicTag::SIZE;
 
 /// Configuration for the concurrenct store.
 #[derive(Debug)]
-pub(crate) struct Config {
-    /// The number of non-frozen slots to create space for.
-    entries: usize,
-
-    /// The size of each slot.
-    bytes: Bytes,
-
-    /// The maximum number of neighbors in each adjacency list.
-    max_neighbors: usize,
-
+pub struct Config {
     /// The number of epoch guard slots.
     ///
     /// Increasing this number will increase the number of threads that can work concurrently
@@ -100,14 +91,10 @@ pub(crate) struct Config {
 }
 
 impl Config {
-    /// Create a new `Config` capable of holding `entries` non-frozen points each of size
-    /// `bytes`. All adjacency lists will have a maximum capacity of `max_neighbors`.
-    pub(crate) fn new(entries: usize, bytes: Bytes, max_neighbors: usize) -> Self {
+    /// Create a new [`Config`] with default concurrency parameters.
+    pub fn new() -> Self {
         const DEFAULT_FREELIST_RECYCLE_CAPACITY: NonZeroU32 = NonZeroU32::new(1024).unwrap();
         Self {
-            entries,
-            bytes,
-            max_neighbors,
             epoch_guard_slots: Registry::default_guard_slots(),
             freelist_recycle_capacity: DEFAULT_FREELIST_RECYCLE_CAPACITY,
         }
@@ -135,6 +122,33 @@ impl Config {
     }
 }
 
+impl Default for Config {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Layout {
+    /// The number of non-frozen slots to create space for.
+    entries: usize,
+
+    /// The maximum number of neighbors in each adjacency list.
+    max_neighbors: usize,
+}
+
+impl Layout {
+    /// Create a new [`Layout`] capable of holding `entries` non-frozen points.
+    ///
+    /// All adjacency lists will have a maximum capacity of `max_neighbors`.
+    pub(crate) fn new(entries: usize, max_neighbors: usize) -> Self {
+        Self {
+            entries,
+            max_neighbors,
+        }
+    }
+}
+
 /// A concurrent data and graph store.
 #[derive(Debug)]
 pub(crate) struct Store {
@@ -161,18 +175,22 @@ const RETRY_LIMIT: usize = 20;
 impl Store {
     /// Create a new [`Store`]. The entries within `init` will be used as frozen points
     /// within the store and must be compatible the the number of bytes in `config`.
-    pub(crate) fn new(config: Config, init: MatrixView<'_, u8>) -> Result<Self, StoreError> {
-        let Config {
+    pub(crate) fn new(
+        layout: Layout,
+        config: Config,
+        init: MatrixView<'_, u8>,
+    ) -> Result<Self, StoreError> {
+        let Layout {
             entries,
-            bytes,
             max_neighbors,
+        } = layout;
+
+        let Config {
             epoch_guard_slots,
             freelist_recycle_capacity,
         } = config;
 
-        if init.ncols() != bytes.value() {
-            return Err(StoreError::mismatched_frozen_point_dim(init.ncols(), bytes));
-        }
+        let bytes = Bytes::new(init.ncols());
 
         if init.nrows() == 0 {
             return Err(StoreError::need_frozen_point());
@@ -487,10 +505,6 @@ impl Store {
 pub(crate) struct StoreError(StoreErrorInner);
 
 impl StoreError {
-    fn mismatched_frozen_point_dim(dim: usize, bytes: Bytes) -> Self {
-        Self(StoreErrorInner::MismatchedFrozenPointDim { dim, bytes })
-    }
-
     fn need_frozen_point() -> Self {
         Self(StoreErrorInner::NeedFrozenPoint)
     }
@@ -518,12 +532,6 @@ impl From<NeighborsError> for StoreError {
 
 #[derive(Debug, Error)]
 enum StoreErrorInner {
-    #[error(
-        "frozen point dim ({}) must have the same dimensionality as requested bytes ({})",
-        dim,
-        bytes
-    )]
-    MismatchedFrozenPointDim { dim: usize, bytes: Bytes },
     #[error("at least one frozen point must be provided")]
     NeedFrozenPoint,
     #[error(
@@ -664,28 +672,18 @@ mod tests {
             base = base.wrapping_add(1);
         }
 
-        let mut config = Config::new(entries, Bytes::new(entry_bytes), 0);
+        let mut config = Config::new();
         config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
         config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
 
-        Store::new(config, data.as_view())
+        let layout = Layout::new(entries, 0);
+        Store::new(layout, config, data.as_view())
     }
 
     //------------------------//
     // Constructor validation //
     //------------------------//
 
-    #[test]
-    fn new_rejects_mismatched_frozen_dim() {
-        // Frozen point has 8 columns but the store is asked for 16-byte entries.
-        let data = Matrix::new(0u8, 1, 8);
-        let err = Store::new(Config::new(4, Bytes::new(16), 0), data.as_view()).unwrap_err();
-        assert!(matches!(
-            err.0,
-            StoreErrorInner::MismatchedFrozenPointDim { dim: 8, .. }
-        ));
-    }
-
     #[test]
     fn new_requires_a_frozen_point() {
         let err = store(4, 8, 0).unwrap_err();
@@ -697,7 +695,8 @@ mod tests {
         // `entries` alone fits in u32, but `entries + frozen` overflows it.
         let data = Matrix::new(0u8, 1, 8);
         let err = Store::new(
-            Config::new(u32::MAX as usize, Bytes::new(8), 0),
+            Layout::new(u32::MAX as usize, 0),
+            Config::default(),
             data.as_view(),
         )
         .unwrap_err();
@@ -708,7 +707,8 @@ mod tests {
     fn new_rejects_too_many_neighbors() {
         let data = Matrix::new(0u8, 1, 8);
         let err = Store::new(
-            Config::new(4, Bytes::new(8), u32::MAX.into_usize() + 1),
+            Layout::new(4, u32::MAX.into_usize() + 1),
+            Config::default(),
             data.as_view(),
         )
         .unwrap_err();

From 9e4e3852c6637e85799d5759b67a968d1ac9d5e4 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 18 Aug 2026 13:17:00 -0700
Subject: [PATCH 13/34] I actually did it?!?!

---
 diskann-inmem/integration/index/runner.rs |  37 +-
 diskann-inmem/src/layers/full.rs          | 195 ++++++----
 diskann-inmem/src/layers/mod.rs           |  43 +--
 diskann-inmem/src/num.rs                  |  25 ++
 diskann-inmem/src/provider.rs             | 421 +++-------------------
 diskann-inmem/src/store/mod.rs            |  12 +
 diskann-utils/src/views.rs                |  13 +
 7 files changed, 267 insertions(+), 479 deletions(-)

diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index 9b767f41cd..17391b556c 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -369,28 +369,49 @@ impl Test {
 
                 let dim = start_points.ncols();
                 let metric = self.data.metric;
-                let config = diskann_inmem::provider::Config::new(
-                    capacity,
-                    self.build.config.max_degree().get(),
-                );
+                let max_degree = self.build.config.max_degree().get();
+                // let config = diskann_inmem::provider::Config::new(
+                //     capacity,
+                //     self.build.config.max_degree().get(),
+                // );
 
                 let index_config = self.build.config.clone();
 
                 let index = match start_points {
                     DatasetView::F32(v) => finish(
-                        Provider::new(layers::Full::<f32>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(diskann_inmem::layers::full::Config::<f32>::new(
+                            capacity,
+                            max_degree,
+                            metric,
+                            v.to_owned(),
+                        ))?,
                         index_config,
                     ),
                     DatasetView::F16(v) => finish(
-                        Provider::new(layers::Full::<f16>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(diskann_inmem::layers::full::Config::<f16>::new(
+                            capacity,
+                            max_degree,
+                            metric,
+                            v.to_owned(),
+                        ))?,
                         index_config,
                     ),
                     DatasetView::U8(v) => finish(
-                        Provider::new(layers::Full::<u8>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(diskann_inmem::layers::full::Config::<u8>::new(
+                            capacity,
+                            max_degree,
+                            metric,
+                            v.to_owned(),
+                        ))?,
                         index_config,
                     ),
                     DatasetView::I8(v) => finish(
-                        Provider::new(layers::Full::<i8>::new(dim, metric), config, v.row_iter())?,
+                        Provider::new(diskann_inmem::layers::full::Config::<i8>::new(
+                            capacity,
+                            max_degree,
+                            metric,
+                            v.to_owned(),
+                        ))?,
                         index_config,
                     ),
                 };
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index c4fffcf846..218369e489 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -70,7 +70,7 @@ where
     type Layer = Full<T>;
 
     fn build(self) -> ANNResult<Full<T>> {
-        Ok(Full::new(self.dim(), self.metric))
+        Full::new(self)
     }
 }
 
@@ -78,14 +78,10 @@ trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_expand_beam<'a>(
         full: &'a Full<Self>,
         query: &'a [Self],
-        store: &'a Store,
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>>;
 
     #[doc(hidden)]
-    fn make_prune<'a>(
-        full: &'a Full<Self>,
-        store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::Prune + 'a>>;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>>;
 }
 
 /// A useful trait bound for types compatible with [`Full`].
@@ -97,7 +93,6 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn __search_accessor<'a>(
         layer: &'a Full<Self>,
         query: &'a [Self],
-        store: &'a Store,
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
@@ -105,7 +100,6 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     #[doc(hidden)]
     fn __prune_accessor<'a>(
         layer: &'a Full<Self>,
-        store: &'a Store,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
 }
@@ -117,29 +111,27 @@ where
     fn __search_accessor<'a>(
         layer: &'a Full<Self>,
         query: &'a [Self],
-        store: &'a Store,
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        let expand_beam = T::make_expand_beam(layer, query, store)?;
+        let expand_beam = T::make_expand_beam(layer, query)?;
         Ok(crate::provider::SearchAccessor::new(
-            store.temp_neighbors(),
+            layer.store.temp_neighbors(),
             expand_beam,
             provider,
-            store.frozen(),
+            layer.store.frozen(),
             counters,
         ))
     }
 
     fn __prune_accessor<'a>(
         layer: &'a Full<Self>,
-        store: &'a Store,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
-        let prune = T::make_prune(layer, store)?;
+        let prune = T::make_prune(layer)?;
         Ok(crate::provider::PruneAccessor::new(
             prune,
-            store.temp_neighbors(),
+            layer.store.temp_neighbors(),
             counters,
         ))
     }
@@ -153,6 +145,7 @@ where
 {
     dim: usize,
     metric: Metric,
+    store: Store,
     _type: PhantomData<T>,
 }
 
@@ -161,15 +154,25 @@ where
     T: 'static,
 {
     /// Create a new full-precision layer for data with the given `dim` and `metric`.
-    pub fn new(dim: usize, metric: Metric) -> Self
+    fn new(config: Config<T>) -> ANNResult<Self>
     where
         T: FullPrecision,
     {
-        Self {
-            dim,
+        let Config {
+            layout,
             metric,
+            start_points,
+            store,
+        } = config;
+
+        let store = Store::new(layout, store, start_points.as_bytes())?;
+
+        Ok(Self {
+            dim: start_points.ncols(),
+            metric,
+            store,
             _type: PhantomData,
-        }
+        })
     }
 
     /// Return the logical dimension of the data handled by this [`layers::Layer`].
@@ -194,12 +197,49 @@ where
     }
 }
 
+impl<T> Full<T>
+where
+    T: FullPrecision,
+{
+    pub(crate) fn get(&self, i: u32) -> ANNResult<Box<[T]>> {
+        let reader = self.store.reader()?;
+
+        let data = match reader.inner().read(i.into_usize()) {
+            Some(data) => data,
+            None => {
+                return Err(ANNError::message("item could not be read"));
+            }
+        };
+
+        let mut buf: Box<[_]> = std::iter::repeat_n(T::zeroed(), self.dim()).collect();
+
+        bytemuck::must_cast_slice_mut::<T, u8>(&mut buf).copy_from_slice(data);
+        Ok(buf)
+    }
+}
+
 impl<T> layers::Layer for Full<T>
 where
     T: FullPrecision,
 {
-    fn bytes(&self) -> Bytes {
-        <Full<T>>::bytes(self)
+    fn max_degree(&self) -> usize {
+        self.store.temp_neighbors().max_length()
+    }
+
+    fn retire(&self, i: u32) -> ANNResult<()> {
+        Ok(self.store.retire(i.into_usize())?)
+    }
+
+    fn is_readable(&self, i: u32) -> Option<bool> {
+        self.store.can_read_approximate(i.into_usize())
+    }
+
+    fn maximum(&self) -> u32 {
+        self.store.maximum()
+    }
+
+    fn capacity(&self) -> usize {
+        self.store.capacity()
     }
 }
 
@@ -207,38 +247,57 @@ impl<T> layers::Set<&[T]> for Full<T>
 where
     T: FullPrecision,
 {
-    fn set(&self, v: &[T], bytes: &mut [u8]) -> ANNResult<()> {
+    type Guard<'a> = Guard<'a>;
+
+    fn set(&self, v: &[T]) -> ANNResult<Guard<'_>> {
         if v.len() != self.dim() {
-            Err(ANNError::from(SetError::Dim {
+            return Err(ANNError::from(SetError {
                 got: v.len(),
                 expected: self.dim(),
-            }))
-        } else if bytes.len() != self.bytes().value() {
-            Err(ANNError::from(SetError::Bytes {
-                got: bytes.len(),
-                expected: self.bytes().value(),
-            }))
-        } else {
-            bytes.copy_from_slice(bytemuck::must_cast_slice::<T, u8>(v));
-            Ok(())
+            }));
         }
+
+        let mut slot = self
+            .store
+            .acquire()
+            .ok_or_else(|| ANNError::message("could not allocate a new slot"))?;
+
+        slot.as_mut_slice()
+            .copy_from_slice(bytemuck::must_cast_slice::<T, u8>(v));
+
+        Ok(Guard::new(slot))
+    }
+}
+
+#[derive(Debug)]
+pub struct Guard<'a> {
+    slot: store::Slot<'a>,
+}
+
+impl<'a> Guard<'a> {
+    fn new(slot: store::Slot<'a>) -> Self {
+        Self { slot }
+    }
+}
+
+impl layers::Guard for Guard<'_> {
+    fn publish(self) {
+        self.slot.publish();
+    }
+    fn id(&self) -> u32 {
+        self.slot.slot()
     }
 }
 
 #[derive(Debug, Error)]
-enum SetError {
-    #[error(
-        "data of dimension {} does not match full precision layer's dimension {}",
-        got,
-        expected
-    )]
-    Dim { got: usize, expected: usize },
-    #[error(
-        "raw byte slice of length {} does not match expected length {}",
-        got,
-        expected
-    )]
-    Bytes { got: usize, expected: usize },
+#[error(
+    "data of dimension {} does not match full precision layer's dimension {}",
+    self.got,
+    self.expected
+)]
+struct SetError {
+    got: usize,
+    expected: usize,
 }
 
 diskann::convert_error!(SetError);
@@ -252,11 +311,10 @@ where
     fn search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
-        store: &'a Store,
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        T::__search_accessor(self, query, store, provider, counters)
+        T::__search_accessor(self, query, provider, counters)
     }
 }
 
@@ -266,10 +324,9 @@ where
 {
     fn prune_accessor<'a>(
         &'a self,
-        store: &'a Store,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
-        T::__prune_accessor(self, store, counters)
+        T::__prune_accessor(self, counters)
     }
 }
 
@@ -544,10 +601,9 @@ impl FullPrecisionImpl for f32 {
     fn make_expand_beam<'a>(
         full: &'a Full<f32>,
         query: &'a [f32],
-        store: &'a Store,
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = store.temp_inner_reader()?;
+        let reader = full.store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
@@ -569,11 +625,8 @@ impl FullPrecisionImpl for f32 {
         Ok(output)
     }
 
-    fn make_prune<'a>(
-        full: &'a Full<Self>,
-        store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = store.temp_inner_reader()?;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+        let reader = full.store.temp_inner_reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f32, SquaredL2>::new(reader)),
@@ -590,10 +643,9 @@ impl FullPrecisionImpl for f16 {
     fn make_expand_beam<'a>(
         full: &'a Full<f16>,
         query: &'a [f16],
-        store: &'a Store,
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = store.temp_inner_reader()?;
+        let reader = full.store.temp_inner_reader()?;
 
         let mut as_f32: Box<[f32]> = std::iter::repeat_n(0.0, full.dim()).collect();
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
@@ -615,11 +667,8 @@ impl FullPrecisionImpl for f16 {
         Ok(output)
     }
 
-    fn make_prune<'a>(
-        full: &'a Full<Self>,
-        store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = store.temp_inner_reader()?;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+        let reader = full.store.temp_inner_reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f16, SquaredL2>::new(reader)),
@@ -636,10 +685,9 @@ impl FullPrecisionImpl for u8 {
     fn make_expand_beam<'a>(
         full: &'a Full<u8>,
         query: &'a [u8],
-        store: &'a Store,
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = store.temp_inner_reader()?;
+        let reader = full.store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
@@ -659,11 +707,8 @@ impl FullPrecisionImpl for u8 {
         Ok(output)
     }
 
-    fn make_prune<'a>(
-        full: &'a Full<Self>,
-        store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = store.temp_inner_reader()?;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+        let reader = full.store.temp_inner_reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<u8, SquaredL2>::new(reader)),
@@ -680,10 +725,9 @@ impl FullPrecisionImpl for i8 {
     fn make_expand_beam<'a>(
         full: &'a Full<i8>,
         query: &'a [i8],
-        store: &'a Store,
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = store.temp_inner_reader()?;
+        let reader = full.store.temp_inner_reader()?;
 
         let query = Calf::Borrowed(query);
 
@@ -697,11 +741,8 @@ impl FullPrecisionImpl for i8 {
         Ok(output)
     }
 
-    fn make_prune<'a>(
-        full: &'a Full<Self>,
-        store: &'a Store,
-    ) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = store.temp_inner_reader()?;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+        let reader = full.store.temp_inner_reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<i8, SquaredL2>::new(reader)),
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 08e4e388c4..9a2c54ebb1 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -31,26 +31,11 @@ use std::num::NonZeroU16;
 use diskann::ANNResult;
 use thiserror::Error;
 
-use crate::{counters::LocalCounters, num::Bytes, store::Store};
+use crate::{counters::LocalCounters, num::Bytes};
 
-mod full;
+pub mod full;
 pub use full::{Full, FullPrecision};
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum Status {
-    Available,
-    Published,
-    Retiring,
-    Frozen,
-}
-
-impl Status {
-    #[must_use = "this function has no side-effects"]
-    pub fn is_readable(self) -> bool {
-        matches!(self, Self::Published | Self::Frozen)
-    }
-}
-
 pub trait LayerConfig {
     type Layer: Layer;
 
@@ -59,12 +44,25 @@ pub trait LayerConfig {
 
 /// Base layer for data representations.
 pub trait Layer: Send + Sync + 'static {
-    // // Return the status of the item behind internal ID `i`.
-    fn bytes(&self) -> Bytes;
+    fn max_degree(&self) -> usize;
+
+    fn retire(&self, i: u32) -> ANNResult<()>;
+
+    fn is_readable(&self, i: u32) -> Option<bool>;
+
+    fn maximum(&self) -> u32;
+
+    fn capacity(&self) -> usize;
 }
 
 pub trait Set<T>: Layer {
-    fn set(&self, element: T, bytes: &mut [u8]) -> ANNResult<()>;
+    type Guard<'a>: Guard;
+    fn set(&self, element: T) -> ANNResult<Self::Guard<'_>>;
+}
+
+pub trait Guard {
+    fn id(&self) -> u32;
+    fn publish(self);
 }
 
 pub(crate) trait ExpandBeam: Send + Sync + std::fmt::Debug {
@@ -130,7 +128,6 @@ pub trait Search: Send + Sync + 'static {
     fn search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
-        store: &'a Store,
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
@@ -156,17 +153,15 @@ pub trait Insert: Search + for<'a> Set<Self::Query<'a>> {
     fn insert_search_accessor<'a>(
         &'a self,
         query: Self::Query<'a>,
-        store: &'a Store,
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        self.search_accessor(query, store, provider, counters)
+        self.search_accessor(query, provider, counters)
     }
 
     #[doc(hidden)]
     fn prune_accessor<'a>(
         &'a self,
-        store: &'a Store,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
 }
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 98c20d82be..2d6bda098c 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -161,6 +161,31 @@ impl std::fmt::Display for Align {
     }
 }
 
+//-------------------------//
+// General Number Wrappers //
+//-------------------------//
+
+macro_rules! typed_int {
+    ($name:ident, $T:ty) => {
+        #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+        pub struct $name($T);
+
+        impl $name {
+            pub fn new(value: $T) -> Self {
+                Self(value)
+            }
+
+            pub fn value(self) -> $T {
+                self.0
+            }
+        }
+    };
+}
+
+typed_int!(Capacity, usize);
+typed_int!(MaxDegree, usize);
+typed_int!(MaximumId, u32);
+
 ///////////
 // Tests //
 ///////////
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 557c0ce882..3ee62e62a1 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -73,15 +73,10 @@ pub struct Provider<L, M = u32>
 where
     M: Id,
 {
-    // The raw binary store
-    store: Store,
-    // Data representation.
+    // Data representation and storage.
     layer: L,
     // ID translation.
     mapping: IdMap<M>,
-    // Construction `Config`.
-    config: Config,
-
     // `Counters` is only non-trivial under the `integration-test` feature flag. Otherwise,
     // all counter related operations are no-ops.
     counters: Counters,
@@ -91,58 +86,11 @@ impl<L, M> Provider<L, M>
 where
     M: Id,
 {
-    /// Construct a new [`Provider`].
-    ///
-    /// The list of `start_points` must be must be compatible with `layer`.
-    pub fn new<I, T>(layer: L, config: Config, start_points: I) -> Result<Self, ProviderError>
-    where
-        I: IntoIterator<Item = T>,
-        L: layers::Set<T>,
-    {
-        let start_points: Vec<_> = start_points.into_iter().collect();
-        let bytes = layers::Layer::bytes(&layer);
-        let mut data = Matrix::new(0u8, start_points.len(), bytes.value());
-
-        for (row, point) in std::iter::zip(data.row_iter_mut(), start_points) {
-            layers::Set::set(&layer, point, row)?;
-        }
-
-        let mut store_config = store::Config::default();
-
-        if let Some(slots) = config.epoch_guard_slots {
-            store_config.epoch_guard_slots(slots);
-        }
-
-        if let Some(capacity) = config.freelist_recycle_capacity {
-            store_config.freelist_recycle_capacity(capacity);
-        }
-
-        let mut store_layout = store::Layout::new(config.capacity(), config.max_degree());
-
-        let store = Store::new(store_layout, store_config, data.as_view())
-            .map_err(|err| ProviderError::CreatingStore(Box::new(err)))?;
-
-        let mapping = IdMap::new(config.capacity());
-
-        Ok(Self {
-            store,
-            layer,
-            mapping,
-            config,
-            counters: Counters::new(),
-        })
-    }
-
     /// A local set of counters that update the provider-wide counters in bulk.
     fn local_counters(&self) -> LocalCounters<'_> {
         self.counters.local()
     }
 
-    /// Return the maximum number of neighbors that can be stored in the provider's graph.
-    pub fn max_degree(&self) -> usize {
-        self.store.max_degree()
-    }
-
     /// Return a snapshot of the current event counters.
     #[cfg(feature = "integration-test")]
     pub fn counters(&self) -> crate::integration::counters::CounterSnapshot {
@@ -150,72 +98,31 @@ where
     }
 }
 
-#[derive(Debug, Error)]
-pub enum ProviderError {
-    #[error("error when trying to set start points")]
-    SettingStartPoints(#[from] ANNError),
-    #[error("could not create data store")]
-    CreatingStore(#[source] Box<dyn std::error::Error + Send + Sync>),
-}
-
-/// Configuration for [`Provider`].
-#[derive(Debug)]
-pub struct Config {
-    capacity: usize,
-    max_degree: usize,
-    prefetch_lookahead: Option<NonZeroUsize>,
-    epoch_guard_slots: Option<NonZeroUsize>,
-    freelist_recycle_capacity: Option<NonZeroU32>,
-}
-
-impl Config {
-    const DEFAULT_PREFETCH_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(8).unwrap();
-
-    /// Construct a new [`Config`].
+impl<L, M> Provider<L, M>
+where
+    L: layers::Layer,
+    M: Id,
+{
+    /// Construct a new [`Provider`].
     ///
-    /// * `capacity`: The number of dynamic entries in the resulting provider.
-    /// * `max_degree`: The maximum degree of any adjacency list in the graph.
-    pub fn new(capacity: usize, max_degree: usize) -> Self {
-        Self {
-            capacity,
-            max_degree,
-            prefetch_lookahead: Some(Self::DEFAULT_PREFETCH_LOOKAHEAD),
-            epoch_guard_slots: None,
-            freelist_recycle_capacity: None,
-        }
-    }
+    /// The list of `start_points` must be must be compatible with `layer`.
+    pub fn new<C>(config: C) -> ANNResult<Self>
+    where
+        C: layers::LayerConfig<Layer = L>,
+    {
+        let layer = <_ as layers::LayerConfig>::build(config)?;
+        let mapping = IdMap::new(layer.capacity());
 
-    /// Return the number of dynamic entries in the resulting provider.
-    pub fn capacity(&self) -> usize {
-        self.capacity
+        Ok(Self {
+            layer,
+            mapping,
+            counters: Counters::new(),
+        })
     }
 
-    /// Return the maximum degree of any adjacency list.
+    /// Return the maximum number of neighbors that can be stored in the provider's graph.
     pub fn max_degree(&self) -> usize {
-        self.max_degree
-    }
-
-    /// Configure the prefetch lookahead.
-    ///
-    /// This is used during beam expansion to prefetch data into CPU caches.
-    pub fn set_prefetch_lookahead(&mut self, prefetch_lookahead: Option<NonZeroUsize>) {
-        self.prefetch_lookahead = prefetch_lookahead;
-    }
-
-    /// Configure the number of epoch guard slots.
-    ///
-    /// Increasing this number will increase the number of threads that can work concurrently
-    /// on the index at the cost of longer scan times for epoch advancement.
-    pub fn set_epoch_guard_slots(&mut self, slots: Option<NonZeroUsize>) {
-        self.epoch_guard_slots = slots;
-    }
-
-    /// Configure the capacity of the freelist recycle queue.
-    ///
-    /// Increasing the capacity of the queue will allow more recycled IDs to be retrieved
-    /// without triggering a scan, but will cost more memory.
-    pub fn set_freelist_recycle_capacity(&mut self, capacity: Option<NonZeroU32>) {
-        self.freelist_recycle_capacity = capacity;
+        self.layer.max_degree()
     }
 }
 
@@ -271,7 +178,7 @@ where
 // with this situation.
 impl<L, M> diskann::provider::Delete for Provider<L, M>
 where
-    L: Send + Sync + 'static,
+    L: layers::Layer,
     M: Id,
 {
     async fn delete(&self, _context: &Context, gid: &M) -> ANNResult<()> {
@@ -286,14 +193,14 @@ where
             Some(e) => e,
         };
 
-        match self.store.retire(entry.internal().into_usize()) {
-            Ok(()) => {
-                // Successfully retired the internal slot. We can safely release the ID mapping.
-                entry.delete();
-                Ok(())
-            }
-            Err(err) => Err(ANNError::new(err)),
-        }
+        // An early return here will cause `entry` to be dropped, which will *not* cause
+        // the delete to commit.
+        <L as layers::Layer>::retire(&self.layer, entry.internal())?;
+
+        // Successfully retired the internal slot. We can safely release the ID mapping.
+        entry.delete();
+
+        Ok(())
     }
 
     async fn release(&self, _context: &Context, _id: Self::InternalId) -> ANNResult<()> {
@@ -307,7 +214,7 @@ where
     ) -> ANNResult<diskann::provider::ElementStatus> {
         // Not that this check is approximate. A full check requires materialization of
         // a `reader`.
-        match self.store.can_read_approximate(id.into_usize()) {
+        match <L as layers::Layer>::is_readable(&self.layer, id) {
             Some(true) => Ok(diskann::provider::ElementStatus::Valid),
             Some(false) => Ok(diskann::provider::ElementStatus::Deleted),
             None => Err(ANNError::message("accessed invalid internal ID")),
@@ -348,19 +255,18 @@ where
         element: T,
     ) -> impl std::future::Future<Output = Result<Self::Guard, Self::SetError>> + Send {
         let work = move || {
-            let mut slot = self
-                .store
-                .acquire()
-                .ok_or_else(|| ANNError::message("could not allocate a new slot"))?;
-
             // TODO: Proper cleanup via `Guard` or some other mechanism on the event of
             // insert failure after `set_element` returns.
-            <L as layers::Set<T>>::set(&self.layer, element, slot.as_mut_slice())?;
-            self.mapping.insert(id.clone(), slot.slot())?;
+            //
+            // The internal `Guard` is sufficient for local rollback, but not after
+            // `set_element` returns.
+            let guard = <L as layers::Set<T>>::set(&self.layer, element)?;
+            let internal = <_ as layers::Guard>::id(&guard);
+            self.mapping.insert(id.clone(), internal)?;
 
             // Now that insert has succeeded - publish the slot. This method cannot fail, so
             // we do not need to worry about potentially unwinding the ID mapping.
-            let id = slot.publish();
+            <_ as layers::Guard>::publish(guard);
 
             // This is a rather expensive update.
             //
@@ -368,7 +274,7 @@ where
             // is not expected to be enabled for general use.
             self.local_counters().set_vector(1);
 
-            Ok(diskann::provider::NoopGuard::new(id))
+            Ok(diskann::provider::NoopGuard::new(internal))
         };
 
         ready(work)
@@ -501,218 +407,6 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
     }
 }
 
-// trait ExpandBeam: Send + Sync + std::fmt::Debug {
-//     /// Evaluate a raw distance function.
-//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32>;
-//
-//     /// Compute the distance between the query and each neighbor in `list`.
-//     ///
-//     /// # Safety
-//     ///
-//     /// * All items in `list` must in-bounds with respect to `reader`.
-//     /// * `buffer.len() >= list.len()`.
-//     unsafe fn expand_beam(
-//         &self,
-//         list: &[u32],
-//         reader: &store::Reader<'_>,
-//         buffer: &mut [(u32, f32)],
-//     ) -> ANNResult<usize>;
-// }
-//
-// #[derive(Debug)]
-// struct ExpandBeamImpl<T, const BYTES: usize> {
-//     inner: T,
-//     prefetch_lookahead: usize,
-// }
-//
-// impl<T, const BYTES: usize> ExpandBeamImpl<T, BYTES> {
-//     fn new(inner: T, prefetch_lookahead: usize) -> Self {
-//         Self {
-//             inner,
-//             prefetch_lookahead,
-//         }
-//     }
-// }
-//
-// impl<T, const BYTES: usize> ExpandBeam for ExpandBeamImpl<T, BYTES>
-// where
-//     T: layers::QueryDistance,
-// {
-//     fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
-//         self.inner.evaluate(x)
-//     }
-//
-//     unsafe fn expand_beam(
-//         &self,
-//         list: &[u32],
-//         reader: &store::Reader<'_>,
-//         buffer: &mut [(u32, f32)],
-//     ) -> ANNResult<usize> {
-//         // SAFETY: Inherited from caller.
-//         unsafe {
-//             expand_beam_inner::<T, BYTES>(
-//                 &self.inner,
-//                 list,
-//                 self.prefetch_lookahead,
-//                 reader,
-//                 buffer,
-//             )
-//         }
-//     }
-// }
-//
-// #[derive(Debug)]
-// struct ExpandBeamVisitor {
-//     bytes: Bytes,
-//     prefetch_lookahead: usize,
-// }
-//
-// impl<'a> layers::QueryVisitor<'a> for ExpandBeamVisitor {
-//     type Output = Box<dyn ExpandBeam + 'a>;
-//
-//     fn visit_sized<const BYTES: usize, T>(self, distance: T) -> Self::Output
-//     where
-//         T: QueryDistance + 'a,
-//     {
-//         // This is critical to ensure we emit the correct number of prefetches.
-//         assert!(Bytes::new(BYTES + store::TAG_SIZE.value()) <= self.bytes);
-//         Box::new(ExpandBeamImpl::<_, BYTES>::new(
-//             distance,
-//             self.prefetch_lookahead,
-//         ))
-//     }
-//
-//     fn visit<T>(self, distance: T) -> Self::Output
-//     where
-//         T: QueryDistance + 'a,
-//     {
-//         Box::new(ExpandBeamImpl::<_, 0>::new(
-//             distance,
-//             self.prefetch_lookahead,
-//         ))
-//     }
-// }
-//
-// /// Prefetch `len` bytes beginning at `ptr`.
-// ///
-// /// The last cache line prefetched first, followed by the rest in ascending order.
-// ///
-// /// # Safety
-// ///
-// /// The memory range `[ptr, ptr.add(len))` must be valid.
-// #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
-// #[inline(always)]
-// unsafe fn prefetch(ptr: *const u8, len: usize) {
-//     use std::arch::x86_64::*;
-//
-//     // Fetch the last cache line (the one with the tag) first.
-//     let stride = Bytes::CACHELINE.value();
-//     let ptr = ptr.cast::<i8>();
-//     let lines = len.div_ceil(stride);
-//     if lines == 0 {
-//         return;
-//     }
-//
-//     // SAFETY: Inherited from caller.
-//     unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
-//     for i in 0..(lines - 1) {
-//         // SAFETY: Inherited from caller.
-//         unsafe {
-//             _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
-//         }
-//     }
-// }
-//
-// /// Prefetch `len` bytes beginning at `ptr`.
-// ///
-// /// The last cache line prefetched first, followed by the rest in ascending order.
-// ///
-// /// # Safety
-// ///
-// /// The memory range `[ptr, ptr.add(len))` must be valid.
-// #[cfg(not(any(target_arch = "x86_64", target_feature = "avx2")))]
-// unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
-//
-// /// # Safety
-// ///
-// /// * All items in `list` must in-bounds with respect to `reader`.
-// /// * The number of bytes associated with `N` cache lines must "make sense".
-// /// * `buffer.len() >= list.len()`.
-// #[inline]
-// unsafe fn expand_beam_inner<T, const BYTES: usize>(
-//     distance: &T,
-//     list: &[u32],
-//     lookahead: usize,
-//     reader: &store::Reader<'_>,
-//     buffer: &mut [(u32, f32)],
-// ) -> ANNResult<usize>
-// where
-//     T: layers::QueryDistance,
-// {
-//     debug_assert!(
-//         BYTES + store::TAG_SIZE.value() <= reader.inner().bytes().value(),
-//         "we really rely on this: {}, bytes = {}",
-//         BYTES + store::TAG_SIZE.value(),
-//         reader.inner().bytes()
-//     );
-//
-//     debug_assert!(buffer.len() >= list.len());
-//
-//     let bytes = if BYTES == 0 {
-//         reader.inner().bytes().value()
-//     } else {
-//         BYTES + store::TAG_SIZE.value()
-//     };
-//
-//     let len = list.len();
-//     let lookahead = lookahead.min(len);
-//
-//     for j in 0..lookahead {
-//         // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
-//         // as the validity of the prefetch bounds.
-//         unsafe {
-//             prefetch(
-//                 reader
-//                     .inner()
-//                     .read_raw_unchecked(list.get_unchecked(j).into_usize())
-//                     .as_ptr()
-//                     .cast(),
-//                 bytes,
-//             )
-//         }
-//     }
-//
-//     // Disable prefetching if the lookahead is 0.
-//     let mut j = if lookahead == 0 { len } else { lookahead };
-//     let mut processed = 0;
-//     for &i in list.iter() {
-//         if j != len {
-//             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
-//             // well as the validity of the prefetch bounds.
-//             unsafe {
-//                 prefetch(
-//                     reader
-//                         .inner()
-//                         .read_raw_unchecked(list.get_unchecked(j).into_usize())
-//                         .as_ptr()
-//                         .cast(),
-//                     bytes,
-//                 )
-//             }
-//             j += 1;
-//         }
-//
-//         // SAFETY: Caller asserts that `i` is in-bounds.
-//         if let Some(data) = unsafe { reader.inner().read_in_bounds(i.into_usize()) } {
-//             // SAFETY: Inherited from caller.
-//             *unsafe { buffer.get_unchecked_mut(processed) } = (i, distance.evaluate(data)?);
-//             processed += 1;
-//         }
-//     }
-//
-//     Ok(processed)
-// }
-
 ////////////
 // Insert //
 ////////////
@@ -903,7 +597,6 @@ where
         <L as layers::Search>::search_accessor(
             &provider.layer,
             query,
-            &provider.store,
             provider,
             provider.local_counters(),
         )
@@ -998,11 +691,7 @@ where
         _context: &'a Context,
         _capacity: usize,
     ) -> ANNResult<PruneAccessor<'a>> {
-        <L as layers::Insert>::prune_accessor(
-            &provider.layer,
-            &provider.store,
-            provider.local_counters(),
-        )
+        <L as layers::Insert>::prune_accessor(&provider.layer, provider.local_counters())
     }
 }
 
@@ -1050,21 +739,7 @@ where
         id: u32,
     ) -> impl Future<Output = Result<Self::DeleteElementGuard, Self::DeleteElementError>> + Send
     {
-        let work = move || {
-            let reader = provider.store.reader()?;
-            let data = match reader.inner().read(id.into_usize()) {
-                Some(data) => data,
-                None => {
-                    return Err(ANNError::message("item could not be read"));
-                }
-            };
-
-            let mut buf: Box<[_]> =
-                std::iter::repeat_n(T::zeroed(), provider.layer.dim()).collect();
-
-            bytemuck::must_cast_slice_mut::<T, u8>(&mut buf).copy_from_slice(data);
-            Ok(buf)
-        };
+        let work = move || provider.layer.get(id);
         ready(work)
     }
 }
@@ -1112,12 +787,18 @@ mod tests {
         let start = grid.start_point(size);
         let degree = 6;
 
-        let full = Full::<f32>::new(grid.dim().into(), Metric::L2);
+        let config = layers::full::Config::new(
+            grid.num_points(size),
+            degree,
+            Metric::L2,
+            Matrix::row_vector(start.into()),
+        );
+
+        // let full = Full::<f32>::new(grid.dim().into(), Metric::L2);
 
-        let config = Config::new(grid.num_points(size), degree);
+        // let config = Config::new(grid.num_points(size), degree);
 
-        let provider =
-            Provider::<_, u64>::new(full, config, std::iter::once(start.as_slice())).unwrap();
+        let provider = Provider::<_, u64>::new(config).unwrap();
         assert_eq!(provider.max_degree(), degree);
 
         let config = diskann::graph::config::Builder::new(
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 3a455d3c68..34806b78c7 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -255,6 +255,14 @@ impl Store {
         self.neighbors.max_length()
     }
 
+    pub(crate) fn maximum(&self) -> u32 {
+        self.neighbors.entries()
+    }
+
+    pub(crate) fn capacity(&self) -> usize {
+        self.unfrozen
+    }
+
     pub(crate) fn temp_neighbors(&self) -> &Neighbors {
         &self.neighbors
     }
@@ -530,6 +538,8 @@ impl From<NeighborsError> for StoreError {
     }
 }
 
+diskann::convert_error!(StoreError);
+
 #[derive(Debug, Error)]
 enum StoreErrorInner {
     #[error("at least one frozen point must be provided")]
@@ -565,6 +575,8 @@ pub(crate) enum RetireError {
     CouldNotClaimSlot,
 }
 
+diskann::convert_error!(RetireError);
+
 /// An epoch protected reader into a [`Store`].
 ///
 /// Created via [`Store::reader`].
diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs
index a9352918c9..2d0cd49dbc 100644
--- a/diskann-utils/src/views.rs
+++ b/diskann-utils/src/views.rs
@@ -637,6 +637,19 @@ where
             ncols: self.ncols,
         }
     }
+
+    pub fn as_bytes(&self) -> MatrixView<'_, u8>
+    where
+        T::Elem: bytemuck::Pod,
+    {
+        let data = bytemuck::must_cast_slice::<T::Elem, u8>(self.as_slice());
+
+        MatrixView {
+            data,
+            nrows: self.nrows(),
+            ncols: self.ncols() * std::mem::size_of::<T::Elem>(),
+        }
+    }
 }
 
 /// Represents an owning, 2-dimensional view of a contiguous block of memory,

From 9d72c246fcb5d3f8fc88ee837a32544d28f004ee Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 18 Aug 2026 17:41:54 -0700
Subject: [PATCH 14/34] Stronger types.

---
 diskann-inmem/integration/index/runner.rs | 18 +++----
 diskann-inmem/src/ids.rs                  | 38 +++++++-------
 diskann-inmem/src/integration/store.rs    |  4 +-
 diskann-inmem/src/layers/full.rs          | 12 ++---
 diskann-inmem/src/layers/mod.rs           |  6 +--
 diskann-inmem/src/neighbors.rs            | 52 ++++++++++----------
 diskann-inmem/src/num.rs                  | 40 ++++++++++++---
 diskann-inmem/src/provider.rs             | 18 +++----
 diskann-inmem/src/store/mod.rs            | 60 +++++++++++------------
 9 files changed, 139 insertions(+), 109 deletions(-)

diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index 17391b556c..95b9b301a3 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -18,7 +18,7 @@ use diskann_vector::distance::Metric;
 use half::f16;
 use serde::{Deserialize, Serialize};
 
-use diskann_inmem::{Provider, layers};
+use diskann_inmem::{Provider, layers, num::{MaxDegree, Capacity}};
 
 use crate::{
     index::{Counters, Index},
@@ -380,8 +380,8 @@ impl Test {
                 let index = match start_points {
                     DatasetView::F32(v) => finish(
                         Provider::new(diskann_inmem::layers::full::Config::<f32>::new(
-                            capacity,
-                            max_degree,
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
                         ))?,
@@ -389,8 +389,8 @@ impl Test {
                     ),
                     DatasetView::F16(v) => finish(
                         Provider::new(diskann_inmem::layers::full::Config::<f16>::new(
-                            capacity,
-                            max_degree,
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
                         ))?,
@@ -398,8 +398,8 @@ impl Test {
                     ),
                     DatasetView::U8(v) => finish(
                         Provider::new(diskann_inmem::layers::full::Config::<u8>::new(
-                            capacity,
-                            max_degree,
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
                         ))?,
@@ -407,8 +407,8 @@ impl Test {
                     ),
                     DatasetView::I8(v) => finish(
                         Provider::new(diskann_inmem::layers::full::Config::<i8>::new(
-                            capacity,
-                            max_degree,
+                            Capacity::new(capacity),
+                            MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
                         ))?,
diff --git a/diskann-inmem/src/ids.rs b/diskann-inmem/src/ids.rs
index dcf41424a8..6d478ea914 100644
--- a/diskann-inmem/src/ids.rs
+++ b/diskann-inmem/src/ids.rs
@@ -13,6 +13,8 @@ use diskann::utils::IntoUsize;
 use parking_lot::{RwLock, RwLockWriteGuard};
 use thiserror::Error;
 
+use crate::num::Capacity;
+
 const SHARD_SIZE: usize = 1024;
 
 /// Bidirectional mapping between an external id `I` and a dense internal `u32` id.
@@ -27,19 +29,19 @@ where
     // Since we know the internal IDs are contiguous from `[0..self.capacity)`, we can
     // use a lighter-weight `backward` ID map than a full `DashMap`.
     backward: Vec<RwLock<Box<[Option<I>]>>>,
-    capacity: usize,
+    capacity: Capacity,
 }
 
 impl<I> IdMap<I>
 where
     I: Hash + Eq,
 {
-    pub(crate) fn new(capacity: usize) -> Self {
+    pub(crate) fn new(capacity: Capacity) -> Self {
         let backward = std::iter::repeat_with(|| {
             let shard = std::iter::repeat_with(|| None).take(SHARD_SIZE).collect();
             RwLock::new(shard)
         })
-        .take(capacity.div_ceil(SHARD_SIZE))
+        .take(capacity.value().div_ceil(SHARD_SIZE))
         .collect();
 
         Self {
@@ -60,7 +62,7 @@ where
     where
         I: Eq + Hash + Clone,
     {
-        if internal.into_usize() >= self.capacity {
+        if internal.into_usize() >= self.capacity.value() {
             return Err(InsertError::OutOfBounds);
         }
 
@@ -108,7 +110,7 @@ where
     where
         I: Clone,
     {
-        if internal.into_usize() >= self.capacity {
+        if internal.into_usize() >= self.capacity.value() {
             return None;
         }
 
@@ -154,7 +156,7 @@ where
     }
 
     #[cfg(test)]
-    fn capacity(&self) -> usize {
+    fn capacity(&self) -> Capacity {
         self.capacity
     }
 }
@@ -221,7 +223,7 @@ mod tests {
             SHARD_SIZE,
             SHARD_SIZE + 1,
             3 * SHARD_SIZE,
-        ] {
+        ].map(Capacity::new) {
             let map = IdMap::<u32>::new(capacity);
             assert_eq!(map.capacity(), capacity);
         }
@@ -229,7 +231,7 @@ mod tests {
 
     #[test]
     fn insert_round_trips() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         assert!(map.insert(100, 3).is_ok());
 
         assert_eq!(map.to_internal(&100), Some(3));
@@ -244,7 +246,7 @@ mod tests {
 
     #[test]
     fn insert_rejects_out_of_bounds_internal() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         assert!(matches!(map.insert(0, 16), Err(InsertError::OutOfBounds)));
         assert!(matches!(
             map.insert(0, u32::MAX),
@@ -257,7 +259,7 @@ mod tests {
 
     #[test]
     fn insert_rejects_duplicate_external_and_preserves_state() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         map.insert(7, 5).unwrap();
 
         assert!(matches!(map.insert(7, 6), Err(InsertError::ExternalExists)));
@@ -270,7 +272,7 @@ mod tests {
 
     #[test]
     fn insert_rejects_duplicate_internal_and_preserves_state() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         map.insert(7, 5).unwrap();
 
         assert!(matches!(map.insert(8, 5), Err(InsertError::InternalExists)));
@@ -283,7 +285,7 @@ mod tests {
 
     #[test]
     fn to_external_handles_bounds_and_empty_slots() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         // In-bounds but unmapped slot.
         assert_eq!(map.to_external(5), None);
         // Out-of-bounds slot.
@@ -292,7 +294,7 @@ mod tests {
 
     #[test]
     fn mappings_span_shard_boundaries() {
-        let capacity = 3 * SHARD_SIZE;
+        let capacity = Capacity::new(3 * SHARD_SIZE);
         let map = IdMap::<u32>::new(capacity);
 
         // Ids straddling every internal shard boundary.
@@ -302,7 +304,7 @@ mod tests {
             SHARD_SIZE as u32,
             (2 * SHARD_SIZE - 1) as u32,
             (2 * SHARD_SIZE) as u32,
-            (capacity - 1) as u32,
+            (capacity.value() - 1) as u32,
         ];
 
         for (external, &internal) in ids.iter().enumerate() {
@@ -317,7 +319,7 @@ mod tests {
 
     #[test]
     fn lookup_supports_borrowed_query() {
-        let map = IdMap::<String>::new(16);
+        let map = IdMap::<String>::new(Capacity::new(16));
         map.insert("alpha".to_string(), 1).unwrap();
 
         // Borrowed `&str` lookups against `String` keys.
@@ -329,7 +331,7 @@ mod tests {
 
     #[test]
     fn occupied_entry_exposes_mapping() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         map.insert(42, 9).unwrap();
 
         let entry = map.occupied_entry(42).expect("entry should exist");
@@ -339,13 +341,13 @@ mod tests {
 
     #[test]
     fn occupied_entry_absent_for_unmapped() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         assert!(map.occupied_entry(42).is_none());
     }
 
     #[test]
     fn entry_delete_clears_both_directions() {
-        let map = IdMap::<u32>::new(16);
+        let map = IdMap::<u32>::new(Capacity::new(16));
         map.insert(42, 9).unwrap();
 
         // Just creating and dropping an `occupied_entry` does not clear it.
diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 78b61b7f2c..877f29e7ec 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -12,7 +12,7 @@ use std::num::{NonZeroU32, NonZeroUsize};
 
 use diskann_utils::views::Matrix;
 
-use crate::{num::Bytes, store};
+use crate::{num::{Bytes, Capacity, MaxDegree}, store};
 
 #[derive(Debug)]
 pub struct Config {
@@ -41,7 +41,7 @@ impl Store {
     /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
     /// the number of epoch guard slots are invalid (e.g. zero).
     pub fn new(config: Config) -> Self {
-        let mut store_layout = store::Layout::new(config.capacity, 0);
+        let mut store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0));
 
         let mut store_config = store::Config::default();
 
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 218369e489..ee839e0edc 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -25,7 +25,7 @@ use thiserror::Error;
 use crate::{
     counters::LocalCounters,
     layers,
-    num::Bytes,
+    num::{Bytes, Capacity, MaxDegree},
     store::{self, Store},
     tag::AtomicTag,
 };
@@ -40,8 +40,8 @@ pub struct Config<T> {
 
 impl<T> Config<T> {
     pub fn new(
-        capacity: usize,
-        max_degree: usize,
+        capacity: Capacity,
+        max_degree: MaxDegree,
         metric: Metric,
         start_points: Matrix<T>,
     ) -> Self {
@@ -222,8 +222,8 @@ impl<T> layers::Layer for Full<T>
 where
     T: FullPrecision,
 {
-    fn max_degree(&self) -> usize {
-        self.store.temp_neighbors().max_length()
+    fn max_degree(&self) -> MaxDegree {
+        self.store.temp_neighbors().max_degree()
     }
 
     fn retire(&self, i: u32) -> ANNResult<()> {
@@ -238,7 +238,7 @@ where
         self.store.maximum()
     }
 
-    fn capacity(&self) -> usize {
+    fn capacity(&self) -> Capacity {
         self.store.capacity()
     }
 }
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index 9a2c54ebb1..f33028a7b7 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -31,7 +31,7 @@ use std::num::NonZeroU16;
 use diskann::ANNResult;
 use thiserror::Error;
 
-use crate::{counters::LocalCounters, num::Bytes};
+use crate::{counters::LocalCounters, num::{Capacity, Bytes, MaxDegree}};
 
 pub mod full;
 pub use full::{Full, FullPrecision};
@@ -44,7 +44,7 @@ pub trait LayerConfig {
 
 /// Base layer for data representations.
 pub trait Layer: Send + Sync + 'static {
-    fn max_degree(&self) -> usize;
+    fn max_degree(&self) -> MaxDegree;
 
     fn retire(&self, i: u32) -> ANNResult<()>;
 
@@ -52,7 +52,7 @@ pub trait Layer: Send + Sync + 'static {
 
     fn maximum(&self) -> u32;
 
-    fn capacity(&self) -> usize;
+    fn capacity(&self) -> Capacity;
 }
 
 pub trait Set<T>: Layer {
diff --git a/diskann-inmem/src/neighbors.rs b/diskann-inmem/src/neighbors.rs
index 88f9b786aa..c1840adcb2 100644
--- a/diskann-inmem/src/neighbors.rs
+++ b/diskann-inmem/src/neighbors.rs
@@ -31,7 +31,7 @@ use thiserror::Error;
 
 use crate::{
     buffer::{Buffer, BufferError},
-    num::{Align, Bytes},
+    num::{Align, Bytes, MaxDegree},
 };
 
 type Id = u32;
@@ -64,20 +64,20 @@ pub(crate) struct Neighbors {
 
 impl Neighbors {
     /// Construct a new [`Neighbors`] capable of holding `entries` adjacency lists with a
-    /// maximum length of `max_length`.
+    /// maximum length of `max_degree`.
     ///
     /// # Errors
     ///
-    /// Returns an error if `(max_length + 1) * size_of::<u32>()` overflows `usize`
+    /// Returns an error if `(max_degree + 1) * size_of::<u32>()` overflows `usize`
     /// (unreachable on 64-bit targets) or the resulting allocation would exceed
     /// `isize::MAX` bytes.
-    pub(crate) fn new(entries: u32, max_length: u32) -> Result<Self, NeighborsError> {
-        let bytes = max_length
+    pub(crate) fn new(entries: u32, max_degree: u32) -> Result<Self, NeighborsError> {
+        let bytes = max_degree
             .into_usize()
             .checked_add(1)
             .and_then(|len| len.checked_mul(std::mem::size_of::<Id>()))
             .map(Bytes::new)
-            .ok_or(NeighborsError::Overflow(max_length))?;
+            .ok_or(NeighborsError::Overflow(max_degree))?;
 
         // We materialize slices of `Id` into the raw byte buffers.
         //
@@ -101,15 +101,15 @@ impl Neighbors {
     }
 
     /// Return the maximum length for any adjacency list.
-    pub(crate) fn max_length(&self) -> usize {
+    pub(crate) fn max_degree(&self) -> MaxDegree {
         // We reserve 4 bytes at the beginning for the length of the adjacency list.
-        (self.neighbors.stride().value() - std::mem::size_of::<Id>()) / std::mem::size_of::<Id>()
+        MaxDegree::new((self.neighbors.stride().value() - std::mem::size_of::<Id>()) / std::mem::size_of::<Id>())
     }
 
     /// Return the maximum length for any adjacency list as a 32-bit integer.
-    pub(crate) fn max_length_u32(&self) -> u32 {
+    pub(crate) fn max_degree_u32(&self) -> u32 {
         // Lossless by the invariants on `Self::new`.
-        self.max_length() as u32
+        self.max_degree().value() as u32
     }
 
     /// Return the number of adjacency lists contained by this graph.
@@ -144,7 +144,7 @@ impl Neighbors {
         // SAFETY: We hold the read-lock, so reading is safe. From our bounds checks, we
         // know that this pointer is valid.
         let len: usize = unsafe { prefix.as_ptr().cast::<Id>().read() }
-            .min(self.max_length_u32())
+            .min(self.max_degree_u32())
             .into_usize();
 
         let mut resizer = neighbors.resize(len);
@@ -189,7 +189,7 @@ impl Neighbors {
 
         Lock {
             ptr: slice.as_non_null().cast::<Id>(),
-            capacity: self.max_length().into_usize(),
+            capacity: self.max_degree().value(),
             _lock: lock,
         }
     }
@@ -201,24 +201,24 @@ impl Neighbors {
     /// Returns an error if:
     ///
     /// * `i` exceeds [`Self::entries`].
-    /// * `neighbors.len()` exceeds [`Self::max_length_u32`].
+    /// * `neighbors.len()` exceeds [`Self::max_degree_u32`].
     ///
     /// If an error is returned, the graph is left unmodified.
     pub(crate) fn set(&self, i: u32, neighbors: &[u32]) -> Result<(), SetError> {
         self.check(i).map_err(SetError::OutOfBounds)?;
 
         // We can check the length of `neighbors` before acquiring any locks as an early exit.
-        if neighbors.len() > self.max_length().into_usize() {
+        if neighbors.len() > self.max_degree().value() {
             return Err(SetError::TooLong(TooLong {
                 got: neighbors.len(),
-                max: self.max_length_u32(),
+                max: self.max_degree_u32(),
             }));
         }
 
         // SAFETY: We've checked `i` is in-bounds.
         let lock = unsafe { self.lock_unchecked(i) };
 
-        // SAFETY: `neighbors.len() <= self.max_length()`.
+        // SAFETY: `neighbors.len() <= self.max_degree()`.
         unsafe { lock.write_unchecked(neighbors) };
         Ok(())
     }
@@ -235,7 +235,7 @@ impl Neighbors {
 /// Errors returned by [`Neighbors::new`].
 #[derive(Debug, Error)]
 pub(crate) enum NeighborsError {
-    /// Computing the per-list byte size `(max_length + 1) * size_of::<u32>()` overflowed
+    /// Computing the per-list byte size `(max_degree + 1) * size_of::<u32>()` overflowed
     /// `usize`.
     ///
     /// Unreachable on 64-bit targets.
@@ -566,7 +566,7 @@ mod tests {
     fn basic_test() {
         let mut neighbors = Neighbors::new(10, 4).unwrap();
         assert_eq!(neighbors.entries(), 10);
-        assert_eq!(neighbors.max_length(), 4);
+        assert_eq!(neighbors.max_degree(), MaxDegree::new(4));
 
         let mut list = AdjacencyList::new();
         for i in 0..neighbors.entries() {
@@ -576,7 +576,7 @@ mod tests {
             assert!(list.is_empty());
 
             let lock = neighbors.lock(i).unwrap();
-            assert_eq!(lock.capacity(), neighbors.max_length());
+            assert_eq!(lock.capacity(), neighbors.max_degree().value());
             assert_eq!(lock.len(), 0);
             assert!(lock.is_empty());
             assert_eq!(lock.as_slice(), &[]);
@@ -595,7 +595,7 @@ mod tests {
             |round: u32, entry: u32| -> Vec<u32> { (0..(round + 1)).map(|r| entry + r).collect() };
 
         // Test mutation via `Neighbors::set`.
-        for round in 0..neighbors.max_length_u32() {
+        for round in 0..neighbors.max_degree_u32() {
             for i in 0..neighbors.entries() {
                 let v = generate(round, i);
                 neighbors.set(i, &v).unwrap();
@@ -614,7 +614,7 @@ mod tests {
         clear(&mut neighbors);
 
         // Test mutation via `lock + write`.
-        for round in 0..neighbors.max_length_u32() {
+        for round in 0..neighbors.max_degree_u32() {
             for i in 0..neighbors.entries() {
                 let v = generate(round, i);
                 neighbors.lock(i).unwrap().write(&v).unwrap();
@@ -633,7 +633,7 @@ mod tests {
         clear(&mut neighbors);
 
         // Test mutation via `lock + append`.
-        for round in 0..neighbors.max_length_u32() {
+        for round in 0..neighbors.max_degree_u32() {
             for i in 0..neighbors.entries() {
                 neighbors.lock(i).unwrap().append(&[round + i]).unwrap();
             }
@@ -684,9 +684,9 @@ mod tests {
 
     #[test]
     fn many_appends() {
-        let max_length = if cfg!(miri) { 100 } else { 1000 };
+        let max_degree = if cfg!(miri) { 100 } else { 1000 };
 
-        let neighbors = Neighbors::new(1, max_length).unwrap();
+        let neighbors = Neighbors::new(1, max_degree).unwrap();
 
         let num_threads = 4;
         let barrier = std::sync::Barrier::new(num_threads);
@@ -699,7 +699,7 @@ mod tests {
                 s.spawn(move || {
                     barrier_ref.wait();
                     let mut i = thread_id as u32;
-                    let upper = neighbors_ref.max_length() as u32;
+                    let upper = neighbors_ref.max_degree_u32();
                     while i < upper {
                         neighbors_ref.lock(0).unwrap().append(&[i]).unwrap();
                         i += num_threads as u32;
@@ -709,7 +709,7 @@ mod tests {
         });
 
         let mut list = AdjacencyList::new();
-        let expected: Vec<_> = (0..neighbors.max_length()).map(|i| i as u32).collect();
+        let expected: Vec<_> = (0..neighbors.max_degree().value()).map(|i| i as u32).collect();
         neighbors.get(0, &mut list).unwrap();
         list.sort();
 
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 2d6bda098c..ca09fae071 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -166,25 +166,53 @@ impl std::fmt::Display for Align {
 //-------------------------//
 
 macro_rules! typed_int {
-    ($name:ident, $T:ty) => {
+    ($(#[$doc:meta])* $name:ident, $T:ty $(,)?) => {
+        $(#[$doc])*
         #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
         pub struct $name($T);
 
         impl $name {
-            pub fn new(value: $T) -> Self {
+            pub const fn new(value: $T) -> Self {
                 Self(value)
             }
 
-            pub fn value(self) -> $T {
+            pub const fn value(self) -> $T {
                 self.0
             }
         }
+
+        impl std::fmt::Display for $name {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                write!(f, concat!(stringify!($name), "({})"), self.value())
+            }
+        }
     };
 }
 
-typed_int!(Capacity, usize);
-typed_int!(MaxDegree, usize);
-typed_int!(MaximumId, u32);
+// TODO: Provide a linkable reference for "immutable" points.
+
+typed_int!(
+    /// The number of distinct slots a [`crate::Provider`] or [`crate::Layer`] has capacity
+    /// for. This is logically distinct from [`MaximumId`], which may be greater due too
+    /// immutable points within a storage container.
+    Capacity,
+    usize,
+);
+
+typed_int!(
+    /// The maximum degree of an adjacency list.
+    MaxDegree,
+    usize
+);
+
+typed_int!(
+    /// The **inclusive** maximum ID that a [`crate::Provider`], [`crate::Layer`], or other
+    /// such store in this crate can access in-bounds.
+    ///
+    /// This is related to [`Capacity`] but is often larger due to immutable points.
+    MaximumIds,
+    u32
+);
 
 ///////////
 // Tests //
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 3ee62e62a1..53b8aa8e6e 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -54,7 +54,7 @@ use crate::{
     ids::IdMap,
     layers,
     neighbors::Neighbors,
-    num::Bytes,
+    num::{Bytes, MaxDegree},
     store::{self, Store},
 };
 
@@ -121,7 +121,7 @@ where
     }
 
     /// Return the maximum number of neighbors that can be stored in the provider's graph.
-    pub fn max_degree(&self) -> usize {
+    pub fn max_degree(&self) -> MaxDegree {
         self.layer.max_degree()
     }
 }
@@ -313,9 +313,9 @@ impl<'a> SearchAccessor<'a> {
     ) -> Self {
         Self {
             neighbors,
-            ids: AdjacencyList::with_capacity(neighbors.max_length()),
+            ids: AdjacencyList::with_capacity(neighbors.max_degree().value()),
             expand_beam,
-            buffer: vec![Default::default(); neighbors.max_length()],
+            buffer: vec![Default::default(); neighbors.max_degree().value()],
             provider,
             start_points,
             counters,
@@ -759,7 +759,7 @@ mod tests {
     };
     use diskann_vector::distance::Metric;
 
-    use crate::layers::Full;
+    use crate::{num::Capacity, layers::Full};
 
     /// The true tests live in the integration tests for this repo.
     ///
@@ -788,8 +788,8 @@ mod tests {
         let degree = 6;
 
         let config = layers::full::Config::new(
-            grid.num_points(size),
-            degree,
+            Capacity::new(grid.num_points(size)),
+            MaxDegree::new(degree),
             Metric::L2,
             Matrix::row_vector(start.into()),
         );
@@ -799,11 +799,11 @@ mod tests {
         // let config = Config::new(grid.num_points(size), degree);
 
         let provider = Provider::<_, u64>::new(config).unwrap();
-        assert_eq!(provider.max_degree(), degree);
+        assert_eq!(provider.max_degree(), MaxDegree::new(degree));
 
         let config = diskann::graph::config::Builder::new(
             2 * (grid.dim() as usize),
-            diskann::graph::config::MaxDegree::new(provider.max_degree()),
+            diskann::graph::config::MaxDegree::new(provider.max_degree().value()),
             10,
             (Metric::L2).into(),
         )
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 34806b78c7..7073397853 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -67,7 +67,7 @@ use crate::{
     epoch::{self, Registry},
     freelist::{self, Freelist},
     neighbors::{Neighbors, NeighborsError},
-    num::Bytes,
+    num::{Bytes, Capacity, MaxDegree},
     tag::{AtomicTag, Tag},
 };
 
@@ -131,20 +131,20 @@ impl Default for Config {
 #[derive(Debug)]
 pub(crate) struct Layout {
     /// The number of non-frozen slots to create space for.
-    entries: usize,
+    capacity: Capacity,
 
     /// The maximum number of neighbors in each adjacency list.
-    max_neighbors: usize,
+    max_degree: MaxDegree,
 }
 
 impl Layout {
-    /// Create a new [`Layout`] capable of holding `entries` non-frozen points.
+    /// Create a new [`Layout`] capable of holding `capacity` non-frozen points.
     ///
-    /// All adjacency lists will have a maximum capacity of `max_neighbors`.
-    pub(crate) fn new(entries: usize, max_neighbors: usize) -> Self {
+    /// All adjacency lists will have a maximum capacity of `max_degree`.
+    pub(crate) fn new(capacity: Capacity, max_degree: MaxDegree) -> Self {
         Self {
-            entries,
-            max_neighbors,
+            capacity,
+            max_degree,
         }
     }
 }
@@ -156,7 +156,7 @@ pub(crate) struct Store {
     plugin: invasive::Invasive,
 
     // The number of unfrozen points. This is guaranteed to be less than `buffer`.
-    unfrozen: usize,
+    unfrozen: Capacity,
 
     // The authoritative source of truth for the state of each slot.
     tags: Vec<AtomicTag>,
@@ -181,8 +181,8 @@ impl Store {
         init: MatrixView<'_, u8>,
     ) -> Result<Self, StoreError> {
         let Layout {
-            entries,
-            max_neighbors,
+            capacity,
+            max_degree,
         } = layout;
 
         let Config {
@@ -196,24 +196,24 @@ impl Store {
             return Err(StoreError::need_frozen_point());
         }
 
-        let too_many_entries = || StoreError::too_many_entries(entries, init.nrows());
+        let too_many_entries = || StoreError::too_many_entries(capacity, init.nrows());
 
         // We have a hard upper-bound of `u32::MAX` total slots.
         //
         // This enforces that bound.
-        let entries: u32 = entries.try_into().map_err(|_| too_many_entries())?;
+        let entries: u32 = capacity.value().try_into().map_err(|_| too_many_entries())?;
 
         let frozen: u32 = init.nrows().try_into().map_err(|_| too_many_entries())?;
 
         let total: u32 = entries.checked_add(frozen).ok_or_else(too_many_entries)?;
 
-        let max_neighbors: u32 = max_neighbors
+        let max_degree: u32 = max_degree.value()
             .try_into()
-            .map_err(|_| StoreError::too_many_neighbors(max_neighbors))?;
+            .map_err(|_| StoreError::too_many_neighbors(max_degree))?;
 
         let me = Self {
             plugin: invasive::Invasive::new(total.into_usize(), bytes),
-            unfrozen: entries.into_usize(),
+            unfrozen: capacity,
             tags: repeat_n(Tag::AVAILABLE, total.into_usize())
                 .map(AtomicTag::new)
                 .collect(),
@@ -222,7 +222,7 @@ impl Store {
             // we do not want it to release frozen IDs.
             freelist: Freelist::new(entries, freelist_recycle_capacity),
             registry: Registry::with_capacity(epoch_guard_slots),
-            neighbors: Neighbors::new(total, max_neighbors)?,
+            neighbors: Neighbors::new(total, max_degree)?,
         };
 
         // Populate frozen points.
@@ -247,19 +247,19 @@ impl Store {
 
     /// Return the range of slots containing frozen items in `self`.
     pub(crate) fn frozen(&self) -> std::ops::Range<u32> {
-        (self.unfrozen as u32)..self.neighbors.entries()
+        (self.unfrozen.value() as u32)..self.neighbors.entries()
     }
 
     /// Return the maximum degree that can be stored in the graph.
-    pub(crate) fn max_degree(&self) -> usize {
-        self.neighbors.max_length()
+    pub(crate) fn max_degree(&self) -> MaxDegree {
+        self.neighbors.max_degree()
     }
 
     pub(crate) fn maximum(&self) -> u32 {
         self.neighbors.entries()
     }
 
-    pub(crate) fn capacity(&self) -> usize {
+    pub(crate) fn capacity(&self) -> Capacity {
         self.unfrozen
     }
 
@@ -405,7 +405,7 @@ impl Store {
     fn scan_acquire(&self) -> Option<Slot<'_>> {
         // This is potentially quite slow - but stop if we've scanned the entire range
         // without finding anything.
-        let mut remaining = self.unfrozen.div_ceil(RETRY_LIMIT);
+        let mut remaining = self.unfrozen.value().div_ceil(RETRY_LIMIT);
         let mut chunks_since_freelist_check = 0;
         let mut acquired: Option<Slot<'_>> = None;
 
@@ -503,7 +503,7 @@ impl Store {
 
     #[cfg(test)]
     fn writable(&self) -> std::ops::Range<u32> {
-        0..self.unfrozen as u32
+        0..self.unfrozen.value() as u32
     }
 }
 
@@ -517,12 +517,12 @@ impl StoreError {
         Self(StoreErrorInner::NeedFrozenPoint)
     }
 
-    fn too_many_entries(entries: usize, frozen: usize) -> Self {
-        Self(StoreErrorInner::TooManyEntries { entries, frozen })
+    fn too_many_entries(capacity: Capacity, frozen: usize) -> Self {
+        Self(StoreErrorInner::TooManyEntries { entries: capacity.value(), frozen })
     }
 
-    fn too_many_neighbors(neighbors: usize) -> Self {
-        Self(StoreErrorInner::TooManyNeighbors { neighbors })
+    fn too_many_neighbors(neighbors: MaxDegree) -> Self {
+        Self(StoreErrorInner::TooManyNeighbors { neighbors: neighbors.value() })
     }
 }
 
@@ -688,7 +688,7 @@ mod tests {
         config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
         config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
 
-        let layout = Layout::new(entries, 0);
+        let layout = Layout::new(Capacity::new(entries), MaxDegree::new(0));
         Store::new(layout, config, data.as_view())
     }
 
@@ -707,7 +707,7 @@ mod tests {
         // `entries` alone fits in u32, but `entries + frozen` overflows it.
         let data = Matrix::new(0u8, 1, 8);
         let err = Store::new(
-            Layout::new(u32::MAX as usize, 0),
+            Layout::new(Capacity::new(u32::MAX as usize), MaxDegree::new(0)),
             Config::default(),
             data.as_view(),
         )
@@ -719,7 +719,7 @@ mod tests {
     fn new_rejects_too_many_neighbors() {
         let data = Matrix::new(0u8, 1, 8);
         let err = Store::new(
-            Layout::new(4, u32::MAX.into_usize() + 1),
+            Layout::new(Capacity::new(4), MaxDegree::new(u32::MAX.into_usize() + 1)),
             Config::default(),
             data.as_view(),
         )

From 0583c31346f5e110f464b6ac462f106f3efa4cb3 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Wed, 19 Aug 2026 10:36:17 -0700
Subject: [PATCH 15/34] More progress.

---
 diskann-benchmark/src/index/inmem2.rs     | 24 +++++++-----
 diskann-inmem/integration/index/runner.rs | 20 +++++-----
 diskann-inmem/src/counters.rs             |  5 ---
 diskann-inmem/src/ids.rs                  |  4 +-
 diskann-inmem/src/integration/store.rs    |  7 +++-
 diskann-inmem/src/layers/full.rs          | 37 +++++++++++++-----
 diskann-inmem/src/layers/mod.rs           | 40 +++++++++++++++----
 diskann-inmem/src/neighbors.rs            | 47 +++++++++++++----------
 diskann-inmem/src/num.rs                  | 28 ++++++++++++--
 diskann-inmem/src/provider.rs             | 30 +++++++--------
 diskann-inmem/src/store/invasive.rs       | 25 ++++++++++--
 diskann-inmem/src/store/mod.rs            | 40 +++++++++++--------
 diskann-inmem/src/store/plugin.rs         |  4 ++
 13 files changed, 207 insertions(+), 104 deletions(-)

diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs
index 8622afe243..f8013ec767 100644
--- a/diskann-benchmark/src/index/inmem2.rs
+++ b/diskann-benchmark/src/index/inmem2.rs
@@ -30,6 +30,7 @@ use diskann_benchmark_runner::{
 };
 use diskann_inmem::{
     layers::{Full, FullPrecision},
+    num::{Capacity, MaxDegree},
     Provider, Strategy,
 };
 use diskann_utils::views::{Matrix, MatrixView};
@@ -474,11 +475,14 @@ where
 
         // Compute the medoid of the dataset as the single start point.
         let start = StartPointStrategy::Medoid.compute(data.as_view())?;
-        let layer = Full::<T>::new(dim, input.data.distance);
-        let config =
-            diskann_inmem::provider::Config::new(num_points, input.build.config.max_degree().get());
-        let provider = Provider::<_, u32>::new(layer, config, start.row_iter())?;
+        let config = Full::config(
+            Capacity::new(num_points),
+            MaxDegree::new(input.build.config.max_degree().get()),
+            input.data.distance,
+            start,
+        );
 
+        let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(
             input.build.config.clone(),
             provider,
@@ -844,12 +848,14 @@ where
         // Compute the medoid of the dataset as the single start point.
         let start = StartPointStrategy::Medoid.compute(dataset.as_view())?;
         let index_config = input.build.config.clone();
-        let layer = Full::<T>::new(dim, input.data.distance);
-
-        let config =
-            diskann_inmem::provider::Config::new(max_points, index_config.max_degree().get());
-        let provider = Provider::<_, u32>::new(layer, config, start.row_iter())?;
+        let config = Full::config(
+            Capacity::new(max_points),
+            MaxDegree::new(index_config.max_degree().get()),
+            input.data.distance,
+            start,
+        );
 
+        let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(index_config, provider, None));
 
         let num_threads = input.build.num_threads;
diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index 95b9b301a3..eb6a4f2fa1 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -18,7 +18,11 @@ use diskann_vector::distance::Metric;
 use half::f16;
 use serde::{Deserialize, Serialize};
 
-use diskann_inmem::{Provider, layers, num::{MaxDegree, Capacity}};
+use diskann_inmem::{
+    Provider,
+    layers::Full,
+    num::{Capacity, MaxDegree},
+};
 
 use crate::{
     index::{Counters, Index},
@@ -367,19 +371,13 @@ impl Test {
                     );
                 }
 
-                let dim = start_points.ncols();
                 let metric = self.data.metric;
                 let max_degree = self.build.config.max_degree().get();
-                // let config = diskann_inmem::provider::Config::new(
-                //     capacity,
-                //     self.build.config.max_degree().get(),
-                // );
-
                 let index_config = self.build.config.clone();
 
                 let index = match start_points {
                     DatasetView::F32(v) => finish(
-                        Provider::new(diskann_inmem::layers::full::Config::<f32>::new(
+                        Provider::new(Full::config(
                             Capacity::new(capacity),
                             MaxDegree::new(max_degree),
                             metric,
@@ -388,7 +386,7 @@ impl Test {
                         index_config,
                     ),
                     DatasetView::F16(v) => finish(
-                        Provider::new(diskann_inmem::layers::full::Config::<f16>::new(
+                        Provider::new(Full::config(
                             Capacity::new(capacity),
                             MaxDegree::new(max_degree),
                             metric,
@@ -397,7 +395,7 @@ impl Test {
                         index_config,
                     ),
                     DatasetView::U8(v) => finish(
-                        Provider::new(diskann_inmem::layers::full::Config::<u8>::new(
+                        Provider::new(Full::config(
                             Capacity::new(capacity),
                             MaxDegree::new(max_degree),
                             metric,
@@ -406,7 +404,7 @@ impl Test {
                         index_config,
                     ),
                     DatasetView::I8(v) => finish(
-                        Provider::new(diskann_inmem::layers::full::Config::<i8>::new(
+                        Provider::new(Full::config(
                             Capacity::new(capacity),
                             MaxDegree::new(max_degree),
                             metric,
diff --git a/diskann-inmem/src/counters.rs b/diskann-inmem/src/counters.rs
index 882ad4adda..5c10954268 100644
--- a/diskann-inmem/src/counters.rs
+++ b/diskann-inmem/src/counters.rs
@@ -43,7 +43,6 @@ mod inner {
         pub(crate) fn query_distance(&mut self, _i: u64) {}
         pub(crate) fn distance_ref(&self, _i: u64) {}
         pub(crate) fn get_vector(&mut self, _i: u64) {}
-        pub(crate) fn get_vector_ref(&self, _i: u64) {}
         pub(crate) fn set_vector(&mut self, _i: u64) {}
         pub(crate) fn get_neighbors(&mut self, _i: u64) {}
         pub(crate) fn set_neighbors(&mut self, _i: u64) {}
@@ -136,10 +135,6 @@ mod inner {
             *self.get_vector.get_mut() += i;
         }
 
-        pub(crate) fn get_vector_ref(&self, i: u64) {
-            self.get_vector.fetch_add(i, Relaxed);
-        }
-
         pub(crate) fn set_vector(&mut self, i: u64) {
             self.set_vector += i;
         }
diff --git a/diskann-inmem/src/ids.rs b/diskann-inmem/src/ids.rs
index 6d478ea914..7fe6dee4be 100644
--- a/diskann-inmem/src/ids.rs
+++ b/diskann-inmem/src/ids.rs
@@ -223,7 +223,9 @@ mod tests {
             SHARD_SIZE,
             SHARD_SIZE + 1,
             3 * SHARD_SIZE,
-        ].map(Capacity::new) {
+        ]
+        .map(Capacity::new)
+        {
             let map = IdMap::<u32>::new(capacity);
             assert_eq!(map.capacity(), capacity);
         }
diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 877f29e7ec..8aacda28f2 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -12,7 +12,10 @@ use std::num::{NonZeroU32, NonZeroUsize};
 
 use diskann_utils::views::Matrix;
 
-use crate::{num::{Bytes, Capacity, MaxDegree}, store};
+use crate::{
+    num::{Capacity, MaxDegree},
+    store,
+};
 
 #[derive(Debug)]
 pub struct Config {
@@ -41,7 +44,7 @@ impl Store {
     /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
     /// the number of epoch guard slots are invalid (e.g. zero).
     pub fn new(config: Config) -> Self {
-        let mut store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0));
+        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0));
 
         let mut store_config = store::Config::default();
 
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index ee839e0edc..67ade7ecc7 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -10,10 +10,7 @@ use diskann_utils::views::Matrix;
 use diskann_vector::{
     UnalignedSlice,
     conversion::SliceCast,
-    distance::{
-        self, Cosine, CosineNormalized, DistanceProvider, InnerProduct, Metric, Specialize,
-        SquaredL2,
-    },
+    distance::{Cosine, CosineNormalized, InnerProduct, Metric, Specialize, SquaredL2},
 };
 use diskann_wide::{
     ARCH,
@@ -25,12 +22,12 @@ use thiserror::Error;
 use crate::{
     counters::LocalCounters,
     layers,
-    num::{Bytes, Capacity, MaxDegree},
+    num::{Bytes, Capacity, IdLimit, MaxDegree},
     store::{self, Store},
     tag::AtomicTag,
 };
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Config<T> {
     layout: store::Layout,
     metric: Metric,
@@ -58,7 +55,8 @@ impl<T> Config<T> {
         self
     }
 
-    fn dim(&self) -> usize {
+    /// Return the vector dimension of this configuration and the resulting [`Full`].
+    pub fn dim(&self) -> usize {
         self.start_points.ncols()
     }
 }
@@ -153,7 +151,21 @@ impl<T> Full<T>
 where
     T: 'static,
 {
+    /// Initialize a [`Config`] for this layer.
+    ///
+    /// See also: [`Config::new`].
+    pub fn config(
+        capacity: Capacity,
+        max_degree: MaxDegree,
+        metric: Metric,
+        start_points: Matrix<T>,
+    ) -> Config<T> {
+        Config::new(capacity, max_degree, metric, start_points)
+    }
+
     /// Create a new full-precision layer for data with the given `dim` and `metric`.
+    ///
+    /// See: [`Config::build`].
     fn new(config: Config<T>) -> ANNResult<Self>
     where
         T: FullPrecision,
@@ -234,8 +246,8 @@ where
         self.store.can_read_approximate(i.into_usize())
     }
 
-    fn maximum(&self) -> u32 {
-        self.store.maximum()
+    fn id_limit(&self) -> IdLimit {
+        self.store.id_limit()
     }
 
     fn capacity(&self) -> Capacity {
@@ -488,7 +500,8 @@ impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
 const LOOKAHEAD: usize = 8;
 const BYTES: usize = 0;
 
-impl<const PREFETCH: usize, T, U, D> layers::ExpandBeam for QueryDistance<'_, PREFETCH, T, U, D>
+unsafe impl<const PREFETCH: usize, T, U, D> layers::ExpandBeam
+    for QueryDistance<'_, PREFETCH, T, U, D>
 where
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
@@ -508,6 +521,10 @@ where
         }
     }
 
+    fn id_limit(&self) -> IdLimit {
+        self.reader.id_limit()
+    }
+
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize> {
         let len = list.len();
         let lookahead = LOOKAHEAD.min(len);
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index f33028a7b7..a6f9dd6edd 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -31,7 +31,10 @@ use std::num::NonZeroU16;
 use diskann::ANNResult;
 use thiserror::Error;
 
-use crate::{counters::LocalCounters, num::{Capacity, Bytes, MaxDegree}};
+use crate::{
+    counters::LocalCounters,
+    num::{Capacity, IdLimit, MaxDegree},
+};
 
 pub mod full;
 pub use full::{Full, FullPrecision};
@@ -46,13 +49,13 @@ pub trait LayerConfig {
 pub trait Layer: Send + Sync + 'static {
     fn max_degree(&self) -> MaxDegree;
 
-    fn retire(&self, i: u32) -> ANNResult<()>;
+    fn id_limit(&self) -> IdLimit;
 
-    fn is_readable(&self, i: u32) -> Option<bool>;
+    fn capacity(&self) -> Capacity;
 
-    fn maximum(&self) -> u32;
+    fn retire(&self, i: u32) -> ANNResult<()>;
 
-    fn capacity(&self) -> Capacity;
+    fn is_readable(&self, i: u32) -> Option<bool>;
 }
 
 pub trait Set<T>: Layer {
@@ -65,15 +68,36 @@ pub trait Guard {
     fn publish(self);
 }
 
-pub(crate) trait ExpandBeam: Send + Sync + std::fmt::Debug {
+/// Trait object based implementation of [`diskann::graph::glue::SearchAccessor::expand_beam`].
+///
+/// Dynamic dispatch is used to enable aggressive specialization of this primitive without
+/// monomorphizing the entire search algorithm. Examples specializations include:
+///
+/// * Optimizing for certain fixed dimensions.
+/// * Inlining metric specific distance functions.
+/// * Tailoring prefetching to the dimension.
+///
+/// # Safety
+///
+/// This trait is `unsafe` because [`Self::id_limit`] **must** work for [`Self::expand_beam`]'s
+/// safety pre-conditions.
+pub(crate) unsafe trait ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Evaluate a raw distance against index `i`.
     fn evaluate(&self, i: u32) -> ANNResult<Option<f32>>;
 
+    /// Return an [`IdLimit`] for this primitive.
+    ///
+    /// Callers must be able to use this limit to satisfy the safety pre-conditions for
+    /// [`Self::expand_beam`].
+    ///
+    /// See also: [`IdLimit::is_in_bound`].
+    fn id_limit(&self) -> IdLimit;
+
     /// Compute the distance between the query and each neighbor in `list`.
     ///
     /// # Safety
     ///
-    /// * All items in `list` must in-bounds with respect to `reader`.
+    /// * All items in `list` must in-bounds with respect to [`Self::id_limit`].
     /// * `buffer.len() >= list.len()`.
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize>;
 }
@@ -113,7 +137,7 @@ impl<'a> diskann_utils::Reborrow<'a> for PruneKey {
 
 #[derive(Debug, Error)]
 #[error("prune list exceeded u16::MAX")]
-struct Overflow;
+pub(crate) struct Overflow;
 
 diskann::convert_error!(Overflow);
 
diff --git a/diskann-inmem/src/neighbors.rs b/diskann-inmem/src/neighbors.rs
index c1840adcb2..60f651acc6 100644
--- a/diskann-inmem/src/neighbors.rs
+++ b/diskann-inmem/src/neighbors.rs
@@ -31,7 +31,7 @@ use thiserror::Error;
 
 use crate::{
     buffer::{Buffer, BufferError},
-    num::{Align, Bytes, MaxDegree},
+    num::{Align, Bytes, MaxDegree, IdLimit},
 };
 
 type Id = u32;
@@ -63,7 +63,7 @@ pub(crate) struct Neighbors {
 }
 
 impl Neighbors {
-    /// Construct a new [`Neighbors`] capable of holding `entries` adjacency lists with a
+    /// Construct a new [`Neighbors`] capable of holding `id_limit` adjacency lists with a
     /// maximum length of `max_degree`.
     ///
     /// # Errors
@@ -71,7 +71,7 @@ impl Neighbors {
     /// Returns an error if `(max_degree + 1) * size_of::<u32>()` overflows `usize`
     /// (unreachable on 64-bit targets) or the resulting allocation would exceed
     /// `isize::MAX` bytes.
-    pub(crate) fn new(entries: u32, max_degree: u32) -> Result<Self, NeighborsError> {
+    pub(crate) fn new(id_limit: IdLimit, max_degree: u32) -> Result<Self, NeighborsError> {
         let bytes = max_degree
             .into_usize()
             .checked_add(1)
@@ -91,10 +91,10 @@ impl Neighbors {
             );
         }
 
-        let neighbors = Buffer::new(entries.into_usize(), bytes, ALIGN)?;
+        let neighbors = Buffer::new(id_limit.as_usize(), bytes, ALIGN)?;
 
         let locks = std::iter::repeat_with(|| RwLock::new(()))
-            .take(entries.into_usize().div_ceil(LOCK_GRANULARITY))
+            .take(id_limit.as_usize().div_ceil(LOCK_GRANULARITY))
             .collect();
 
         Ok(Self { neighbors, locks })
@@ -103,7 +103,10 @@ impl Neighbors {
     /// Return the maximum length for any adjacency list.
     pub(crate) fn max_degree(&self) -> MaxDegree {
         // We reserve 4 bytes at the beginning for the length of the adjacency list.
-        MaxDegree::new((self.neighbors.stride().value() - std::mem::size_of::<Id>()) / std::mem::size_of::<Id>())
+        MaxDegree::new(
+            (self.neighbors.stride().value() - std::mem::size_of::<Id>())
+                / std::mem::size_of::<Id>(),
+        )
     }
 
     /// Return the maximum length for any adjacency list as a 32-bit integer.
@@ -420,7 +423,7 @@ mod tests {
 
     #[test]
     fn out_of_bounds_rejects_indices_beyond_entries() {
-        let n = Neighbors::new(4, 4).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 4).unwrap();
         // entries == 4, so valid indices are 0..=3.
         // Regression test: a buggy `check` using `i == entries()` would let
         // `entries+1`, `entries+2`, ... slip through to UB.
@@ -434,7 +437,7 @@ mod tests {
 
     #[test]
     fn empty_neighbors_rejects_all_access() {
-        let n = Neighbors::new(0, 4).unwrap();
+        let n = Neighbors::new(IdLimit::new(0), 4).unwrap();
         let mut out = AdjacencyList::with_capacity(4);
         for i in [0u32, 1, u32::MAX] {
             assert!(matches!(n.get(i, &mut out), Err(OutOfBounds(_))));
@@ -447,21 +450,21 @@ mod tests {
 
     #[test]
     fn set_rejects_oversized_neighbors() {
-        let n = Neighbors::new(4, 3).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 3).unwrap();
         let too_many = &[1, 2, 3, 4];
         assert!(matches!(n.set(0, too_many), Err(SetError::TooLong(_))));
     }
 
     #[test]
     fn lock_write_rejects_oversized_neighbors() {
-        let n = Neighbors::new(4, 3).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 3).unwrap();
         let lock = n.lock(0).unwrap();
         assert!(lock.write(&[1, 2, 3, 4]).is_err());
     }
 
     #[test]
     fn lock_append_rejects_overflow() {
-        let n = Neighbors::new(4, 3).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 3).unwrap();
         n.set(0, &[1, 2]).unwrap();
         let lock = n.lock(0).unwrap();
         assert!(lock.append(&[3, 4]).is_err());
@@ -469,7 +472,7 @@ mod tests {
 
     #[test]
     fn lock_implements_debug() {
-        let n = Neighbors::new(4, 3).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 3).unwrap();
         let lock = n.lock(0).unwrap();
         let _ = format!("{:?}", lock);
     }
@@ -478,7 +481,7 @@ mod tests {
 
     #[test]
     fn append_preserves_existing_and_adds_new() {
-        let n = Neighbors::new(4, 6).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 6).unwrap();
         n.set(0, &[10, 20]).unwrap();
 
         let lock = n.lock(0).unwrap();
@@ -492,7 +495,7 @@ mod tests {
 
     #[test]
     fn append_to_empty() {
-        let n = Neighbors::new(4, 4).unwrap();
+        let n = Neighbors::new(IdLimit::new(4), 4).unwrap();
 
         let lock = n.lock(0).unwrap();
         assert_eq!(lock.as_slice(), &[]);
@@ -505,7 +508,7 @@ mod tests {
 
     #[test]
     fn append_fills_to_capacity() {
-        let n = Neighbors::new(1, 3).unwrap();
+        let n = Neighbors::new(IdLimit::new(1), 3).unwrap();
         n.set(0, &[1]).unwrap();
 
         let lock = n.lock(0).unwrap();
@@ -518,7 +521,7 @@ mod tests {
 
     #[test]
     fn append_empty_slice_is_noop() {
-        let n = Neighbors::new(1, 4).unwrap();
+        let n = Neighbors::new(IdLimit::new(1), 4).unwrap();
         n.set(0, &[10, 20]).unwrap();
 
         let lock = n.lock(0).unwrap();
@@ -531,7 +534,7 @@ mod tests {
 
     #[test]
     fn write_overwrites_longer_list() {
-        let n = Neighbors::new(1, 5).unwrap();
+        let n = Neighbors::new(IdLimit::new(1), 5).unwrap();
         n.set(0, &[1, 2, 3, 4, 5]).unwrap();
 
         // Overwrite with a shorter list.
@@ -564,7 +567,7 @@ mod tests {
 
     #[test]
     fn basic_test() {
-        let mut neighbors = Neighbors::new(10, 4).unwrap();
+        let mut neighbors = Neighbors::new(IdLimit::new(10), 4).unwrap();
         assert_eq!(neighbors.entries(), 10);
         assert_eq!(neighbors.max_degree(), MaxDegree::new(4));
 
@@ -660,7 +663,7 @@ mod tests {
     #[test]
     fn lock_blocks_get() {
         for _ in 0..10 {
-            let neighbors = Neighbors::new(3, 4).unwrap();
+            let neighbors = Neighbors::new(IdLimit::new(3), 4).unwrap();
             let seq = Sequencer::new();
 
             std::thread::scope(|s| {
@@ -686,7 +689,7 @@ mod tests {
     fn many_appends() {
         let max_degree = if cfg!(miri) { 100 } else { 1000 };
 
-        let neighbors = Neighbors::new(1, max_degree).unwrap();
+        let neighbors = Neighbors::new(IdLimit::new(1), max_degree).unwrap();
 
         let num_threads = 4;
         let barrier = std::sync::Barrier::new(num_threads);
@@ -709,7 +712,9 @@ mod tests {
         });
 
         let mut list = AdjacencyList::new();
-        let expected: Vec<_> = (0..neighbors.max_degree().value()).map(|i| i as u32).collect();
+        let expected: Vec<_> = (0..neighbors.max_degree().value())
+            .map(|i| i as u32)
+            .collect();
         neighbors.get(0, &mut list).unwrap();
         list.sort();
 
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index ca09fae071..da11ac0bd5 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -206,14 +206,34 @@ typed_int!(
 );
 
 typed_int!(
-    /// The **inclusive** maximum ID that a [`crate::Provider`], [`crate::Layer`], or other
-    /// such store in this crate can access in-bounds.
+    /// One larger than the maximum ID that a [`crate::Provider`], [`crate::Layer`], or
+    /// other such store in this crate can access in-bounds.
     ///
-    /// This is related to [`Capacity`] but is often larger due to immutable points.
-    MaximumIds,
+    /// This implies that access to ids `[0..self)` are in-bounds.
+    ///
+    /// [`Capacity`] is related, but the [`IdLimit`] for a collection may be larger due to
+    /// immutable points.
+    IdLimit,
     u32
 );
 
+impl IdLimit {
+    /// Return `true` if `i` is within `[0..self)`.
+    pub const fn is_in_bounds(self, i: u32) -> bool {
+        i < self.value()
+    }
+
+    /// Return the [`Self::value`] as a [`usize`].
+    pub const fn as_usize(self) -> usize {
+        // We cannot use the `IntoUsize` trait in a const function unfortunately.
+        //
+        // Instead, we need to re-create the check that makes this conversion safe.
+        const { assert!(std::mem::size_of::<u32>() <= std::mem::size_of::<usize>()); }
+
+        self.value() as usize
+    }
+}
+
 ///////////
 // Tests //
 ///////////
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 53b8aa8e6e..ac26c35db8 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -30,10 +30,7 @@
 //! * Lack of save/load support: The index is currently ephemeral, but there are plans to
 //!   address this gap.
 
-use std::{
-    hash::Hash,
-    num::{NonZeroU32, NonZeroUsize},
-};
+use std::hash::Hash;
 
 use diskann::{
     ANNError, ANNResult,
@@ -44,18 +41,14 @@ use diskann::{
     },
     neighbor::Neighbor,
     provider,
-    utils::IntoUsize,
 };
-use diskann_utils::views::Matrix;
-use thiserror::Error;
 
 use crate::{
     counters::{Counters, LocalCounters},
     ids::IdMap,
     layers,
     neighbors::Neighbors,
-    num::{Bytes, MaxDegree},
-    store::{self, Store},
+    num::{IdLimit, MaxDegree},
 };
 
 /// Aggregate trait for the external ID type of [`Provider`].
@@ -295,6 +288,7 @@ pub struct SearchAccessor<'a> {
     neighbors: &'a Neighbors,
     ids: AdjacencyList<u32>,
     expand_beam: Box<dyn layers::ExpandBeam + 'a>,
+    id_limit: IdLimit,
     buffer: Vec<(u32, f32)>,
 
     // The parent provider for the accessor.
@@ -311,10 +305,12 @@ impl<'a> SearchAccessor<'a> {
         start_points: std::ops::Range<u32>,
         counters: LocalCounters<'a>,
     ) -> Self {
+        let id_limit = expand_beam.id_limit();
         Self {
             neighbors,
             ids: AdjacencyList::with_capacity(neighbors.max_degree().value()),
             expand_beam,
+            id_limit,
             buffer: vec![Default::default(); neighbors.max_degree().value()],
             provider,
             start_points,
@@ -378,16 +374,19 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
                 self.counters.get_neighbors(1);
 
                 // Filter out unvisited IDs and ensure that all the IDs we are about
-                //
-                // TODO: We need safe provenance on the upper bound.
                 self.ids
-                    .retain(|i| pred.eval_mut(i) && *i < self.neighbors.entries());
+                    .retain(|i| pred.eval_mut(i) && self.id_limit.is_in_bounds(*i));
 
                 // This should always hold, but let's double check.
                 assert!(self.buffer.len() >= self.ids.len());
 
-                // SAFETY: We've verified that each entry in `self.ids` is in-bounds and the
-                // `self.buffer` is long enough to hold all the IDs.
+                // SAFETY:
+                //
+                // 1. We've verified that each entry in `self.ids` is in-bounds using
+                //   the `IdLimit` derived from this `ExpandBeam` object (enforced in the
+                //   constructor).
+                //
+                // 2. `self.buffer` is long enough to hold all the IDs.
                 let processed =
                     unsafe { self.expand_beam.expand_beam(&self.ids, &mut self.buffer) }?;
 
@@ -757,9 +756,10 @@ mod tests {
         neighbor::Neighbor,
         provider::{DataProvider, Delete},
     };
+    use diskann_utils::views::Matrix;
     use diskann_vector::distance::Metric;
 
-    use crate::{num::Capacity, layers::Full};
+    use crate::num::Capacity;
 
     /// The true tests live in the integration tests for this repo.
     ///
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index 5015aaac0b..b8a2f5517c 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -10,7 +10,7 @@ use diskann::utils::IntoUsize;
 use crate::{
     buffer::{Buffer, RawSlice},
     epoch,
-    num::{Align, Bytes},
+    num::{Align, Bytes, Capacity, IdLimit},
     tag::{AtomicTag, Tag},
 };
 
@@ -28,18 +28,24 @@ pub(crate) struct Invasive {
 const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
 
 impl Invasive {
-    pub(crate) fn new(entries: usize, bytes: Bytes) -> Self {
+    pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Self {
         let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
         let padded_bytes = unpadded
             .checked_next_multiple_of(Bytes::CACHELINE.div(TWO))
             .unwrap();
 
         Self {
-            buffer: Buffer::new(entries, padded_bytes, Align::_128).unwrap(),
+            buffer: Buffer::new(id_limit.as_usize(), padded_bytes, Align::_128).unwrap(),
             unpadded,
         }
     }
 
+    pub(crate) fn id_limit(&self) -> IdLimit {
+        // The numeric cast is save because `Invasive::new` takes an `IdLimit` in its
+        // constructor, and thus `self.buffer.len()` cannot exceed `u32::MAX`.
+        IdLimit::new(self.buffer.len() as u32)
+    }
+
     pub(crate) fn bytes(&self) -> Bytes {
         self.unpadded
     }
@@ -82,6 +88,10 @@ impl Invasive {
 impl super::plugin::Plugin for Invasive {
     type Slot<'a> = Slot<'a>;
 
+    fn id_limit(&self) -> IdLimit {
+        <Invasive>::id_limit(self)
+    }
+
     unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
         let Some((tag, data)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
@@ -148,6 +158,15 @@ impl<'a> Reader<'a> {
         i < self.buffer.len()
     }
 
+    /// Return the [`IdLimit`] for this collection.
+    #[inline]
+    #[must_use = "this function has no side-effects"]
+    pub(crate) fn id_limit(&self) -> IdLimit {
+        // Like `Invasive::id_limit`, the numberic cast is safe because by construction,
+        // the underlying buffer is limited to `u32::MAX`.
+        IdLimit::new(self.buffer.len() as u32)
+    }
+
     /// Return `true` if it is safe to read the data at position `i`.
     ///
     /// This guarantee only holds while `self` is alive. Construction of a new [`Reader`]
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 7073397853..ec512d516c 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -67,18 +67,18 @@ use crate::{
     epoch::{self, Registry},
     freelist::{self, Freelist},
     neighbors::{Neighbors, NeighborsError},
-    num::{Bytes, Capacity, MaxDegree},
+    num::{Bytes, Capacity, IdLimit, MaxDegree},
     tag::{AtomicTag, Tag},
 };
 
 pub(crate) mod invasive;
 pub(crate) mod plugin;
-pub(crate) mod stacked;
 
-pub(crate) const TAG_SIZE: Bytes = AtomicTag::SIZE;
+// TODO: Remove?
+// pub(crate) mod stacked;
 
 /// Configuration for the concurrenct store.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Config {
     /// The number of epoch guard slots.
     ///
@@ -128,7 +128,7 @@ impl Default for Config {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct Layout {
     /// The number of non-frozen slots to create space for.
     capacity: Capacity,
@@ -201,20 +201,24 @@ impl Store {
         // We have a hard upper-bound of `u32::MAX` total slots.
         //
         // This enforces that bound.
-        let entries: u32 = capacity.value().try_into().map_err(|_| too_many_entries())?;
+        let entries: u32 = capacity
+            .value()
+            .try_into()
+            .map_err(|_| too_many_entries())?;
 
         let frozen: u32 = init.nrows().try_into().map_err(|_| too_many_entries())?;
 
-        let total: u32 = entries.checked_add(frozen).ok_or_else(too_many_entries)?;
+        let id_limit = IdLimit::new(entries.checked_add(frozen).ok_or_else(too_many_entries)?);
 
-        let max_degree: u32 = max_degree.value()
+        let max_degree: u32 = max_degree
+            .value()
             .try_into()
             .map_err(|_| StoreError::too_many_neighbors(max_degree))?;
 
         let me = Self {
-            plugin: invasive::Invasive::new(total.into_usize(), bytes),
+            plugin: invasive::Invasive::new(id_limit, bytes),
             unfrozen: capacity,
-            tags: repeat_n(Tag::AVAILABLE, total.into_usize())
+            tags: repeat_n(Tag::AVAILABLE, id_limit.as_usize())
                 .map(AtomicTag::new)
                 .collect(),
 
@@ -222,7 +226,7 @@ impl Store {
             // we do not want it to release frozen IDs.
             freelist: Freelist::new(entries, freelist_recycle_capacity),
             registry: Registry::with_capacity(epoch_guard_slots),
-            neighbors: Neighbors::new(total, max_degree)?,
+            neighbors: Neighbors::new(id_limit, max_degree)?,
         };
 
         // Populate frozen points.
@@ -255,8 +259,9 @@ impl Store {
         self.neighbors.max_degree()
     }
 
-    pub(crate) fn maximum(&self) -> u32 {
-        self.neighbors.entries()
+    pub(crate) fn id_limit(&self) -> IdLimit {
+        // TODO: Figure out how to justify this once plugins are a thing.
+        IdLimit::new(self.neighbors.entries())
     }
 
     pub(crate) fn capacity(&self) -> Capacity {
@@ -518,11 +523,16 @@ impl StoreError {
     }
 
     fn too_many_entries(capacity: Capacity, frozen: usize) -> Self {
-        Self(StoreErrorInner::TooManyEntries { entries: capacity.value(), frozen })
+        Self(StoreErrorInner::TooManyEntries {
+            entries: capacity.value(),
+            frozen,
+        })
     }
 
     fn too_many_neighbors(neighbors: MaxDegree) -> Self {
-        Self(StoreErrorInner::TooManyNeighbors { neighbors: neighbors.value() })
+        Self(StoreErrorInner::TooManyNeighbors {
+            neighbors: neighbors.value(),
+        })
     }
 }
 
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index 4e243cc817..f0c0458348 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -5,9 +5,13 @@
 
 use std::fmt::Debug;
 
+use crate::num::IdLimit;
+
 pub(crate) trait Plugin: Debug + 'static {
     type Slot<'a>: Slot;
 
+    fn id_limit(&self) -> IdLimit;
+
     unsafe fn acquire(&self, i: u32) -> Self::Slot<'_>;
     fn reclaim(&self, i: u32);
     fn retire(&self, i: u32);

From d2b6f4a9b193864269df216d95bf12f001f7a6ba Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Wed, 19 Aug 2026 12:20:18 -0700
Subject: [PATCH 16/34] Over the main hurdle.

---
 diskann-inmem/src/integration/store.rs |  28 +-
 diskann-inmem/src/layers/full.rs       |  62 ++-
 diskann-inmem/src/neighbors.rs         |   2 +-
 diskann-inmem/src/num.rs               |   4 +-
 diskann-inmem/src/store/invasive.rs    |  45 +-
 diskann-inmem/src/store/mod.rs         | 737 +++++++++++++------------
 diskann-inmem/src/store/plugin.rs      |  22 +-
 diskann-inmem/src/store/stacked.rs     |  76 ---
 8 files changed, 498 insertions(+), 478 deletions(-)
 delete mode 100644 diskann-inmem/src/store/stacked.rs

diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 8aacda28f2..47f0526e06 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -13,7 +13,7 @@ use std::num::{NonZeroU32, NonZeroUsize};
 use diskann_utils::views::Matrix;
 
 use crate::{
-    num::{Capacity, MaxDegree},
+    num::{Bytes, Capacity, MaxDegree},
     store,
 };
 
@@ -27,7 +27,7 @@ pub struct Config {
 
 #[derive(Debug)]
 pub struct Store {
-    store: store::Store,
+    store: store::Store<store::invasive::Invasive>,
 }
 
 impl Store {
@@ -44,7 +44,7 @@ impl Store {
     /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
     /// the number of epoch guard slots are invalid (e.g. zero).
     pub fn new(config: Config) -> Self {
-        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0));
+        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 1);
 
         let mut store_config = store::Config::default();
 
@@ -63,9 +63,10 @@ impl Store {
                 .expect("`freelist_recycle_capacity` must be non-zero"),
             );
 
-        let data = Matrix::new(0u8, 1, config.entry_bytes);
-        let store = store::Store::new(store_layout, store_config, data.as_view())
+        let plugin_config = store::invasive::Invasive::config(Bytes::new(config.entry_bytes));
+        let store = store::Store::new(store_layout, store_config, plugin_config)
             .expect("failed to construct store");
+
         Self { store }
     }
 
@@ -94,7 +95,10 @@ impl Store {
     }
 
     pub fn reader(&self) -> Option<Reader<'_>> {
-        match self.store.reader() {
+        match self
+            .store
+            .guard(|plugin, guard| unsafe { plugin.reader(guard) })
+        {
             Ok(reader) => Some(Reader::new(reader)),
             Err(crate::epoch::Unavailable) => None,
         }
@@ -102,25 +106,25 @@ impl Store {
 }
 
 pub struct Reader<'a> {
-    reader: store::Reader<'a>,
+    reader: store::invasive::Reader<'a>,
 }
 
 impl<'a> Reader<'a> {
-    fn new(reader: store::Reader<'a>) -> Self {
+    fn new(reader: store::invasive::Reader<'a>) -> Self {
         Self { reader }
     }
 
     pub fn read(&self, i: usize) -> Option<&[u8]> {
-        self.reader.inner().read(i)
+        self.reader.read(i)
     }
 }
 
 pub struct Writer<'a> {
-    slot: store::Slot<'a>,
+    slot: store::Slot<'a, store::invasive::Slot<'a>>,
 }
 
 impl<'a> Writer<'a> {
-    fn new(slot: store::Slot<'a>) -> Self {
+    fn new(slot: store::Slot<'a, store::invasive::Slot<'a>>) -> Self {
         Self { slot }
     }
 
@@ -129,6 +133,6 @@ impl<'a> Writer<'a> {
     }
 
     pub fn as_mut_slice(&mut self) -> &mut [u8] {
-        self.slot.as_mut_slice()
+        self.slot.data().as_mut_slice()
     }
 }
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 67ade7ecc7..4f15ccc917 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -23,7 +23,10 @@ use crate::{
     counters::LocalCounters,
     layers,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
-    store::{self, Store},
+    store::{
+        self, Store,
+        invasive::{self, Invasive},
+    },
     tag::AtomicTag,
 };
 
@@ -43,7 +46,11 @@ impl<T> Config<T> {
         start_points: Matrix<T>,
     ) -> Self {
         Self {
-            layout: store::Layout::new(capacity, max_degree),
+            layout: store::Layout::new(
+                capacity,
+                max_degree,
+                start_points.nrows().try_into().unwrap(),
+            ),
             metric,
             start_points,
             store: store::Config::default(),
@@ -143,7 +150,7 @@ where
 {
     dim: usize,
     metric: Metric,
-    store: Store,
+    store: Store<Invasive>,
     _type: PhantomData<T>,
 }
 
@@ -177,7 +184,19 @@ where
             store,
         } = config;
 
-        let store = Store::new(layout, store, start_points.as_bytes())?;
+        let bytes = Bytes::new(start_points.ncols() * std::mem::size_of::<T>());
+        let invasive = Invasive::config(bytes);
+        let store = Store::new(layout, store, invasive)?;
+
+        // Initialize start points.
+        for (i, row) in std::iter::zip(store.frozen(), start_points.row_iter()) {
+            let mut slot = store.slot(i).unwrap();
+            slot.data()
+                .as_mut_slice()
+                .copy_from_slice(bytemuck::must_cast_slice::<T, u8>(row));
+
+            slot.freeze();
+        }
 
         Ok(Self {
             dim: start_points.ncols(),
@@ -207,6 +226,12 @@ where
             Ok(())
         }
     }
+
+    fn reader(&self) -> ANNResult<invasive::Reader<'_>> {
+        Ok(self
+            .store
+            .guard(|invasive, guard| unsafe { invasive.reader(guard) })?)
+    }
 }
 
 impl<T> Full<T>
@@ -214,9 +239,8 @@ where
     T: FullPrecision,
 {
     pub(crate) fn get(&self, i: u32) -> ANNResult<Box<[T]>> {
-        let reader = self.store.reader()?;
-
-        let data = match reader.inner().read(i.into_usize()) {
+        let reader = self.reader()?;
+        let data = match reader.read(i.into_usize()) {
             Some(data) => data,
             None => {
                 return Err(ANNError::message("item could not be read"));
@@ -274,7 +298,8 @@ where
             .acquire()
             .ok_or_else(|| ANNError::message("could not allocate a new slot"))?;
 
-        slot.as_mut_slice()
+        slot.data()
+            .as_mut_slice()
             .copy_from_slice(bytemuck::must_cast_slice::<T, u8>(v));
 
         Ok(Guard::new(slot))
@@ -283,11 +308,11 @@ where
 
 #[derive(Debug)]
 pub struct Guard<'a> {
-    slot: store::Slot<'a>,
+    slot: store::Slot<'a, invasive::Slot<'a>>,
 }
 
 impl<'a> Guard<'a> {
-    fn new(slot: store::Slot<'a>) -> Self {
+    fn new(slot: store::Slot<'a, invasive::Slot<'a>>) -> Self {
         Self { slot }
     }
 }
@@ -620,8 +645,7 @@ impl FullPrecisionImpl for f32 {
         query: &'a [f32],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = full.store.temp_inner_reader()?;
-
+        let reader = full.reader()?;
         let query = Calf::Borrowed(query);
 
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
@@ -643,7 +667,7 @@ impl FullPrecisionImpl for f32 {
     }
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f32, SquaredL2>::new(reader)),
@@ -662,7 +686,7 @@ impl FullPrecisionImpl for f16 {
         query: &'a [f16],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let mut as_f32: Box<[f32]> = std::iter::repeat_n(0.0, full.dim()).collect();
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
@@ -685,7 +709,7 @@ impl FullPrecisionImpl for f16 {
     }
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<f16, SquaredL2>::new(reader)),
@@ -704,7 +728,7 @@ impl FullPrecisionImpl for u8 {
         query: &'a [u8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let query = Calf::Borrowed(query);
 
@@ -725,7 +749,7 @@ impl FullPrecisionImpl for u8 {
     }
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<u8, SquaredL2>::new(reader)),
@@ -744,7 +768,7 @@ impl FullPrecisionImpl for i8 {
         query: &'a [i8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         full.check_dim(query.len())?;
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let query = Calf::Borrowed(query);
 
@@ -759,7 +783,7 @@ impl FullPrecisionImpl for i8 {
     }
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
-        let reader = full.store.temp_inner_reader()?;
+        let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => Box::new(Prune::<i8, SquaredL2>::new(reader)),
diff --git a/diskann-inmem/src/neighbors.rs b/diskann-inmem/src/neighbors.rs
index 60f651acc6..f8ad42b0f1 100644
--- a/diskann-inmem/src/neighbors.rs
+++ b/diskann-inmem/src/neighbors.rs
@@ -31,7 +31,7 @@ use thiserror::Error;
 
 use crate::{
     buffer::{Buffer, BufferError},
-    num::{Align, Bytes, MaxDegree, IdLimit},
+    num::{Align, Bytes, IdLimit, MaxDegree},
 };
 
 type Id = u32;
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index da11ac0bd5..2ce7fc5bc1 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -228,7 +228,9 @@ impl IdLimit {
         // We cannot use the `IntoUsize` trait in a const function unfortunately.
         //
         // Instead, we need to re-create the check that makes this conversion safe.
-        const { assert!(std::mem::size_of::<u32>() <= std::mem::size_of::<usize>()); }
+        const {
+            assert!(std::mem::size_of::<u32>() <= std::mem::size_of::<usize>());
+        }
 
         self.value() as usize
     }
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index b8a2f5517c..9cf1b96b53 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -5,15 +5,40 @@
 
 use std::{num::NonZeroUsize, sync::atomic::Ordering};
 
-use diskann::utils::IntoUsize;
+use diskann::{ANNResult, utils::IntoUsize};
 
 use crate::{
     buffer::{Buffer, RawSlice},
     epoch,
     num::{Align, Bytes, Capacity, IdLimit},
     tag::{AtomicTag, Tag},
+    store::Lifecycle,
 };
 
+#[derive(Debug, Clone)]
+pub(crate) struct Config {
+    bytes: Bytes,
+}
+
+impl Config {
+    pub(crate) fn new(bytes: Bytes) -> Self {
+        Self { bytes }
+    }
+
+    pub(crate) fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
+        let Self { bytes } = self;
+
+        Ok(Invasive::new(id_limit, bytes))
+    }
+}
+
+impl super::plugin::PluginConfig for Config {
+    type Plugin = Invasive;
+    fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
+        <Config>::build(self, id_limit)
+    }
+}
+
 /// The invasive store where concurrency tags are stored inline with the data.
 #[derive(Debug)]
 pub(crate) struct Invasive {
@@ -28,6 +53,10 @@ pub(crate) struct Invasive {
 const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
 
 impl Invasive {
+    pub(crate) fn config(bytes: Bytes) -> Config {
+        Config::new(bytes)
+    }
+
     pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Self {
         let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
         let padded_bytes = unpadded
@@ -92,7 +121,7 @@ impl super::plugin::Plugin for Invasive {
         <Invasive>::id_limit(self)
     }
 
-    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
+    unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_> {
         let Some((tag, data)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
         };
@@ -111,7 +140,7 @@ impl super::plugin::Plugin for Invasive {
         Slot { tag, data }
     }
 
-    fn reclaim(&self, i: u32) {
+    unsafe fn reclaim(&self, i: u32, _: Lifecycle) {
         let Some((tag, _)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
         };
@@ -119,7 +148,7 @@ impl super::plugin::Plugin for Invasive {
         tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 
-    fn retire(&self, i: u32) {
+    unsafe fn retire(&self, i: u32, _: Lifecycle) {
         let Some((tag, _)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
         };
@@ -265,19 +294,19 @@ pub(crate) struct Slot<'a> {
 }
 
 impl<'a> Slot<'a> {
-    pub(crate) unsafe fn as_mut_slice(&mut self) -> &mut [u8] {
+    pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
         unsafe { self.data.as_mut_slice() }
     }
 }
 
 impl super::plugin::Slot for Slot<'_> {
-    fn publish(self) {
+    fn publish(self, _: Lifecycle) {
         self.tag.store(Tag::PUBLISHED, Ordering::Release);
     }
-    fn freeze(self) {
+    fn freeze(self, _: Lifecycle) {
         self.tag.store(Tag::FROZEN, Ordering::Release);
     }
-    fn abort(self) {
+    fn abort(self, _: Lifecycle) {
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index ec512d516c..5174541d01 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -58,7 +58,7 @@ use std::{
     sync::atomic::Ordering,
 };
 
-use diskann::utils::IntoUsize;
+use diskann::{ANNError, utils::IntoUsize};
 use diskann_utils::views::MatrixView;
 use thiserror::Error;
 
@@ -74,8 +74,21 @@ use crate::{
 pub(crate) mod invasive;
 pub(crate) mod plugin;
 
-// TODO: Remove?
-// pub(crate) mod stacked;
+/// To make extra sure that [`plugin::Plugin`] life-cycle arguments are not callable outside
+/// of this module (i.e., elsewhere in this crate), this [`Lifecycle`] marker type is used
+/// that is only constructible in this module.
+#[derive(Debug)]
+pub(crate) struct Lifecycle(());
+
+impl Lifecycle {
+    /// Construct a new [`Lifecycle`].
+    ///
+    /// DO NOT MAKE THIS `pub(anything)`. It helps prevent accidentally interacting with
+    /// plugins when all uses should be managed in this file instead.
+    const fn new() -> Self {
+        Self(())
+    }
+}
 
 /// Configuration for the concurrenct store.
 #[derive(Debug, Clone)]
@@ -135,25 +148,29 @@ pub(crate) struct Layout {
 
     /// The maximum number of neighbors in each adjacency list.
     max_degree: MaxDegree,
+
+    /// The number of immutable points to reserve at the end of the [`Store`].
+    frozen: u32,
 }
 
 impl Layout {
     /// Create a new [`Layout`] capable of holding `capacity` non-frozen points.
     ///
     /// All adjacency lists will have a maximum capacity of `max_degree`.
-    pub(crate) fn new(capacity: Capacity, max_degree: MaxDegree) -> Self {
+    pub(crate) fn new(capacity: Capacity, max_degree: MaxDegree, frozen: u32) -> Self {
         Self {
             capacity,
             max_degree,
+            frozen,
         }
     }
 }
 
 /// A concurrent data and graph store.
 #[derive(Debug)]
-pub(crate) struct Store {
-    // This is a temporary concrete type until [`Store`] is properly parameterized by its plugin.
-    plugin: invasive::Invasive,
+pub(crate) struct Store<P> {
+    // The [`plugin::Plugin`] managed by this [`Store`].
+    plugin: P,
 
     // The number of unfrozen points. This is guaranteed to be less than `buffer`.
     unfrozen: Capacity,
@@ -172,17 +189,20 @@ pub(crate) struct Store {
 // TODO: This is a guess and probably needs tuning.
 const RETRY_LIMIT: usize = 20;
 
-impl Store {
+impl<P> Store<P>
+where
+    P: plugin::Plugin,
+{
     /// Create a new [`Store`]. The entries within `init` will be used as frozen points
     /// within the store and must be compatible the the number of bytes in `config`.
-    pub(crate) fn new(
-        layout: Layout,
-        config: Config,
-        init: MatrixView<'_, u8>,
-    ) -> Result<Self, StoreError> {
+    pub(crate) fn new<C>(layout: Layout, config: Config, plugin: C) -> Result<Self, StoreError>
+    where
+        C: plugin::PluginConfig<Plugin = P>,
+    {
         let Layout {
             capacity,
             max_degree,
+            frozen,
         } = layout;
 
         let Config {
@@ -190,13 +210,7 @@ impl Store {
             freelist_recycle_capacity,
         } = config;
 
-        let bytes = Bytes::new(init.ncols());
-
-        if init.nrows() == 0 {
-            return Err(StoreError::need_frozen_point());
-        }
-
-        let too_many_entries = || StoreError::too_many_entries(capacity, init.nrows());
+        let too_many_entries = || StoreError::too_many_entries(capacity, frozen);
 
         // We have a hard upper-bound of `u32::MAX` total slots.
         //
@@ -206,8 +220,6 @@ impl Store {
             .try_into()
             .map_err(|_| too_many_entries())?;
 
-        let frozen: u32 = init.nrows().try_into().map_err(|_| too_many_entries())?;
-
         let id_limit = IdLimit::new(entries.checked_add(frozen).ok_or_else(too_many_entries)?);
 
         let max_degree: u32 = max_degree
@@ -215,8 +227,10 @@ impl Store {
             .try_into()
             .map_err(|_| StoreError::too_many_neighbors(max_degree))?;
 
+        let plugin = plugin::PluginConfig::build(plugin, id_limit).map_err(StoreError::plugin)?;
+
         let me = Self {
-            plugin: invasive::Invasive::new(id_limit, bytes),
+            plugin,
             unfrozen: capacity,
             tags: repeat_n(Tag::AVAILABLE, id_limit.as_usize())
                 .map(AtomicTag::new)
@@ -229,23 +243,10 @@ impl Store {
             neighbors: Neighbors::new(id_limit, max_degree)?,
         };
 
-        // Populate frozen points.
-        for (i, data) in init.row_iter().enumerate() {
-            // We have checked that the total number of entries fits in `u32`, so this
-            // arithmetic cannot overflow.
-            #[expect(clippy::expect_used, reason = "this should always succeed")]
-            let mut slot = me
-                .slot(entries + (i as u32))
-                .expect("store was just created - claiming the slot must succeed");
-
-            slot.as_mut_slice().copy_from_slice(data);
-            slot.freeze();
-        }
-
         Ok(me)
     }
 
-    pub(crate) fn plugin(&self) -> &invasive::Invasive {
+    pub(crate) fn plugin(&self) -> &P {
         &self.plugin
     }
 
@@ -260,8 +261,7 @@ impl Store {
     }
 
     pub(crate) fn id_limit(&self) -> IdLimit {
-        // TODO: Figure out how to justify this once plugins are a thing.
-        IdLimit::new(self.neighbors.entries())
+        plugin::Plugin::id_limit(&self.plugin)
     }
 
     pub(crate) fn capacity(&self) -> Capacity {
@@ -289,7 +289,7 @@ impl Store {
 
             // We release the plugin before the main tag. The other direction would
             // prematurely advertise availability.
-            plugin::Plugin::reclaim(self.plugin(), i);
+            unsafe { plugin::Plugin::reclaim(self.plugin(), i, Lifecycle::new()) };
 
             // Use `Release` ordering to ensure that the store to the mirror cannot get moved
             // after the store to the authoritative list.
@@ -311,28 +311,36 @@ impl Store {
         Some(items)
     }
 
-    /// Return a [`Reader`] into the store.
-    ///
-    /// # Errors
-    ///
-    /// Returns [`epoch::Unavailable`] if there are too many active readers.
-    pub(crate) fn reader(&self) -> Result<Reader<'_>, epoch::Unavailable> {
-        Ok(Reader {
-            inner: unsafe { self.plugin().reader(self.registry.guard()?) },
-            neighbors: &self.neighbors,
-        })
+    pub(crate) fn guard<'a, F, R>(&'a self, f: F) -> Result<R, epoch::Unavailable>
+    where
+        F: FnOnce(&'a P, epoch::Guard<'a>) -> R,
+    {
+        let guard = self.registry.guard()?;
+        Ok(f(self.plugin(), guard))
     }
 
-    // TODO: Rework neighbor storage.
-    pub(crate) fn temp_inner_reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
-        Ok(unsafe { self.plugin().reader(self.registry.guard()?) })
-    }
+    // /// Return a [`Reader`] into the store.
+    // ///
+    // /// # Errors
+    // ///
+    // /// Returns [`epoch::Unavailable`] if there are too many active readers.
+    // pub(crate) fn reader(&self) -> Result<Reader<'_>, epoch::Unavailable> {
+    //     Ok(Reader {
+    //         inner: unsafe { self.plugin().reader(self.registry.guard()?) },
+    //         neighbors: &self.neighbors,
+    //     })
+    // }
+
+    // // TODO: Rework neighbor storage.
+    // pub(crate) fn temp_inner_reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
+    //     Ok(unsafe { self.plugin().reader(self.registry.guard()?) })
+    // }
 
     /// Attempt to acquire a new [`Slot`] for writing.
     ///
     /// This method first consults the freelist and falls back to scanning the tags list
     /// if no ID is available from the fast path.
-    pub(crate) fn acquire(&self) -> Option<Slot<'_>> {
+    pub(crate) fn acquire(&self) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
         for _ in 0..RETRY_LIMIT {
             match self.freelist.pop() {
                 freelist::Id::Found(id) => {
@@ -388,7 +396,7 @@ impl Store {
         match tag.compare_exchange(current, retiring, Ordering::Relaxed, Ordering::Relaxed) {
             Ok(_) => {
                 // Set the metadata in the mirror as well.
-                plugin::Plugin::retire(self.plugin(), i.try_into().unwrap());
+                unsafe { plugin::Plugin::retire(self.plugin(), i.try_into().unwrap(), Lifecycle::new()) };
                 guard.retire(i as u32);
                 Ok(())
             }
@@ -407,12 +415,12 @@ impl Store {
     ///
     /// Periodically, the freelist is checked to see if another thread has found an available
     /// slot for us.
-    fn scan_acquire(&self) -> Option<Slot<'_>> {
+    fn scan_acquire(&self) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
         // This is potentially quite slow - but stop if we've scanned the entire range
         // without finding anything.
         let mut remaining = self.unfrozen.value().div_ceil(RETRY_LIMIT);
         let mut chunks_since_freelist_check = 0;
-        let mut acquired: Option<Slot<'_>> = None;
+        let mut acquired: Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> = None;
 
         while remaining != 0 {
             let chunk = self.freelist.scan();
@@ -458,7 +466,7 @@ impl Store {
         None
     }
 
-    fn slot(&self, i: u32) -> Option<Slot<'_>> {
+    pub(crate) fn slot(&self, i: u32) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
         let tag = &self.tags.get(i.into_usize())?;
 
         // SAFETY: We've guaranteed that `tag` belongs to `slot`.
@@ -471,7 +479,11 @@ impl Store {
     ///
     /// Caller asserts that `tag` was obtained from `self.tags[slot]`. This is meant as
     /// a performance optimization where `tag` is first queried for potential availability.
-    unsafe fn try_acquire<'a>(&'a self, tag: &'a AtomicTag, slot: u32) -> Option<Slot<'a>> {
+    unsafe fn try_acquire<'a>(
+        &'a self,
+        tag: &'a AtomicTag,
+        slot: u32,
+    ) -> Option<Slot<'a, <P as plugin::Plugin>::Slot<'a>>> {
         if tag.load(Ordering::Relaxed) != Tag::AVAILABLE {
             return None;
         }
@@ -483,7 +495,7 @@ impl Store {
             Ordering::Relaxed,
         ) {
             Ok(_) => {
-                let data = unsafe { plugin::Plugin::acquire(self.plugin(), slot) };
+                let data = unsafe { plugin::Plugin::acquire(self.plugin(), slot, Lifecycle::new()) };
                 Some(Slot {
                     tag,
                     data: ManuallyDrop::new(data),
@@ -522,7 +534,7 @@ impl StoreError {
         Self(StoreErrorInner::NeedFrozenPoint)
     }
 
-    fn too_many_entries(capacity: Capacity, frozen: usize) -> Self {
+    fn too_many_entries(capacity: Capacity, frozen: u32) -> Self {
         Self(StoreErrorInner::TooManyEntries {
             entries: capacity.value(),
             frozen,
@@ -534,6 +546,10 @@ impl StoreError {
             neighbors: neighbors.value(),
         })
     }
+
+    fn plugin(err: ANNError) -> Self {
+        Self(StoreErrorInner::PluginError(err))
+    }
 }
 
 impl From<BufferError> for StoreError {
@@ -559,13 +575,15 @@ enum StoreErrorInner {
         entries,
         frozen
     )]
-    TooManyEntries { entries: usize, frozen: usize },
+    TooManyEntries { entries: usize, frozen: u32 },
     #[error("number of neighbors ({}) may not exceed `u32::MAX`", neighbors)]
     TooManyNeighbors { neighbors: usize },
     #[error(transparent)]
     BufferError(#[from] BufferError),
     #[error(transparent)]
     NeighborsError(#[from] NeighborsError),
+    #[error("error creating plugin")]
+    PluginError(ANNError),
 }
 
 /// Error conditions for [`Store::retire`].
@@ -587,47 +605,52 @@ pub(crate) enum RetireError {
 
 diskann::convert_error!(RetireError);
 
-/// An epoch protected reader into a [`Store`].
-///
-/// Created via [`Store::reader`].
-#[derive(Debug)]
-pub(crate) struct Reader<'a> {
-    inner: invasive::Reader<'a>,
-    neighbors: &'a Neighbors,
-}
-
-impl<'a> Reader<'a> {
-    /// Return `true` if the index `i` is in-bounds.
-    #[inline]
-    #[must_use = "this function has no side-effects"]
-    pub(crate) fn is_in_bounds(&self, i: usize) -> bool {
-        i < self.neighbors.entries().into_usize()
-    }
-
-    #[inline]
-    pub(crate) fn inner(&self) -> &invasive::Reader<'_> {
-        &self.inner
-    }
-
-    /// Return [`Neighbors`].
-    pub(crate) fn neighbors(&self) -> &Neighbors {
-        self.neighbors
-    }
-}
+// /// An epoch protected reader into a [`Store`].
+// ///
+// /// Created via [`Store::reader`].
+// #[derive(Debug)]
+// pub(crate) struct Reader<'a> {
+//     inner: invasive::Reader<'a>,
+//     neighbors: &'a Neighbors,
+// }
+//
+// impl<'a> Reader<'a> {
+//     /// Return `true` if the index `i` is in-bounds.
+//     #[inline]
+//     #[must_use = "this function has no side-effects"]
+//     pub(crate) fn is_in_bounds(&self, i: usize) -> bool {
+//         i < self.neighbors.entries().into_usize()
+//     }
+//
+//     #[inline]
+//     pub(crate) fn inner(&self) -> &invasive::Reader<'_> {
+//         &self.inner
+//     }
+//
+//     /// Return [`Neighbors`].
+//     pub(crate) fn neighbors(&self) -> &Neighbors {
+//         self.neighbors
+//     }
+// }
 
 /// A writable buffer into the data managed by a [`Store`], obtained from [`Store::acquire`].
 #[derive(Debug)]
-pub(crate) struct Slot<'a> {
+pub(crate) struct Slot<'a, S>
+where
+    S: plugin::Slot,
+{
     tag: &'a AtomicTag,
-    data: ManuallyDrop<invasive::Slot<'a>>,
+    data: ManuallyDrop<S>,
     slot: u32,
 }
 
-impl<'a> Slot<'a> {
-    /// View the managed data as a mutable slice.
-    pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
-        // SAFETY: The slot guarantees exclusive access to its corresponding data.
-        unsafe { self.data.as_mut_slice() }
+impl<'a, S> Slot<'a, S>
+where
+    S: plugin::Slot,
+{
+    /// View the raw inner slot.
+    pub(crate) fn data(&mut self) -> &mut S {
+        &mut self.data
     }
 
     /// Return the slot associated with this write.
@@ -635,12 +658,12 @@ impl<'a> Slot<'a> {
         self.slot
     }
 
-    fn freeze(self) {
+    pub(crate) fn freeze(self) {
         // Suppress normal `Drop`.
         let mut me = ManuallyDrop::new(self);
 
         // Freeze the inner slot.
-        plugin::Slot::freeze(unsafe { ManuallyDrop::take(&mut me.data) });
+        plugin::Slot::freeze(unsafe { ManuallyDrop::take(&mut me.data) }, Lifecycle::new());
 
         // Update the authoritative store.
         me.tag.store(Tag::FROZEN, Ordering::Release);
@@ -656,7 +679,7 @@ impl<'a> Slot<'a> {
         let mut me = ManuallyDrop::new(self);
 
         // Publish the inner slot.
-        plugin::Slot::publish(unsafe { ManuallyDrop::take(&mut me.data) });
+        plugin::Slot::publish(unsafe { ManuallyDrop::take(&mut me.data) }, Lifecycle::new());
 
         // Update the authoritative store.
         me.tag.store(Tag::PUBLISHED, Ordering::Release);
@@ -664,262 +687,268 @@ impl<'a> Slot<'a> {
     }
 }
 
-impl Drop for Slot<'_> {
+impl<S> Drop for Slot<'_, S>
+where
+    S: plugin::Slot,
+{
     fn drop(&mut self) {
-        plugin::Slot::abort(unsafe { ManuallyDrop::take(&mut self.data) });
+        plugin::Slot::abort(unsafe { ManuallyDrop::take(&mut self.data) }, Lifecycle::new());
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
 
-///////////
-// Tests //
-///////////
-
-/// These tests are basic functionality tests for the store.
-///
-/// Longer running conurrency tests are in the integration test suite.
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    use diskann_utils::views::Matrix;
-
-    // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
-    // zeroed frozen points. The frozen points occupy the highest slot indices.
-    fn store(entries: usize, entry_bytes: usize, frozen: usize) -> Result<Store, StoreError> {
-        let mut data = Matrix::new(0u8, frozen, entry_bytes);
-        let mut base = 0u8;
-        for row in data.row_iter_mut() {
-            row.fill(base);
-            base = base.wrapping_add(1);
-        }
-
-        let mut config = Config::new();
-        config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
-        config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
-
-        let layout = Layout::new(Capacity::new(entries), MaxDegree::new(0));
-        Store::new(layout, config, data.as_view())
-    }
-
-    //------------------------//
-    // Constructor validation //
-    //------------------------//
-
-    #[test]
-    fn new_requires_a_frozen_point() {
-        let err = store(4, 8, 0).unwrap_err();
-        assert!(matches!(err.0, StoreErrorInner::NeedFrozenPoint));
-    }
-
-    #[test]
-    fn new_rejects_total_slot_overflow() {
-        // `entries` alone fits in u32, but `entries + frozen` overflows it.
-        let data = Matrix::new(0u8, 1, 8);
-        let err = Store::new(
-            Layout::new(Capacity::new(u32::MAX as usize), MaxDegree::new(0)),
-            Config::default(),
-            data.as_view(),
-        )
-        .unwrap_err();
-        assert!(matches!(err.0, StoreErrorInner::TooManyEntries { .. }));
-    }
-
-    #[test]
-    fn new_rejects_too_many_neighbors() {
-        let data = Matrix::new(0u8, 1, 8);
-        let err = Store::new(
-            Layout::new(Capacity::new(4), MaxDegree::new(u32::MAX.into_usize() + 1)),
-            Config::default(),
-            data.as_view(),
-        )
-        .unwrap_err();
-        assert!(matches!(err.0, StoreErrorInner::TooManyNeighbors { .. }));
-    }
-
-    //--------//
-    // Layout //
-    //--------//
-
-    #[test]
-    fn frozen_range_follows_writable_slots() {
-        let s = store(4, 8, 2).unwrap();
-
-        // Writable slots are [0, 4); frozen points occupy [4, 6).
-        assert_eq!(s.frozen(), 4..6);
-
-        let reader = s.reader().unwrap();
-        for i in 0..4 {
-            assert!(!s.can_read_approximate(i).unwrap());
-            assert!(!reader.inner().can_read(i).unwrap());
-            assert!(reader.inner().read(i).is_none());
-        }
-
-        assert!(s.can_read_approximate(4).unwrap());
-        assert!(reader.inner().can_read(4).unwrap());
-        assert_eq!(reader.inner().read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
-
-        assert!(s.can_read_approximate(5).unwrap());
-        assert!(reader.inner().can_read(5).unwrap());
-        assert_eq!(reader.inner().read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
-
-        assert!(s.can_read_approximate(6).is_none());
-        assert!(reader.inner().can_read(6).is_none());
-        assert!(reader.inner().read(6).is_none());
-    }
-
-    ///////////////
-    // Lifecycle //
-    ///////////////
-
-    #[test]
-    fn acquire_write_publish_read_roundtrip() {
-        let s = store(4, 8, 1).unwrap();
-
-        let reader = s.reader().expect("reader guard available");
-
-        let idx = {
-            let mut slot = s.acquire().expect("a fresh store has free slots");
-            let idx = slot.slot() as usize;
-            slot.as_mut_slice()
-                .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
-
-            // Before the slot is dropped - we should not be able to read it.
-            assert!(reader.inner().read(idx).is_none());
-            assert!(!s.can_read_approximate(idx).unwrap());
-            slot.publish();
-            idx
-        };
-
-        assert_eq!(
-            reader.inner().read(idx),
-            Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice())
-        );
-        assert!(s.can_read_approximate(idx).unwrap());
-    }
-
-    #[test]
-    fn unpublished_slots_are_immediately_available() {
-        let s = store(4, 8, 1).unwrap();
-
-        let reader = s.reader().expect("reader guard available");
-
-        let idx = {
-            let mut slot = s.acquire().expect("a fresh store has free slots");
-            let idx = slot.slot() as usize;
-            slot.as_mut_slice()
-                .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
-
-            // Before the slot is dropped - we should not be able to read it.
-            assert!(reader.inner().read(idx).is_none());
-            assert!(!s.can_read_approximate(idx).unwrap());
-
-            // NOTE: We do not explicitly publish the slot.
-            idx
-        };
-
-        assert!(reader.inner().read(idx).is_none());
-        assert!(!s.can_read_approximate(idx).unwrap());
-    }
-
-    #[test]
-    fn acquire_exhausts_then_reports_none() {
-        let s = store(2, 8, 1).unwrap();
-        // Hold the guards so the slots stay owned.
-        let _a = s.acquire().expect("first writable slot");
-        let _b = s.acquire().expect("second writable slot");
-        assert!(
-            s.acquire().is_none(),
-            "all writable slots are owned, so acquire must fail"
-        );
-    }
-
-    //--------//
-    // Retire //
-    //--------//
-
-    #[test]
-    fn retire_out_of_bounds() {
-        let s = store(4, 8, 1).unwrap();
-        assert!(matches!(s.retire(999), Err(RetireError::OutOfBounds)));
-    }
-
-    #[test]
-    fn retire_rejects_reserved_slots() {
-        let s = store(4, 8, 1).unwrap();
-        // An untouched writable slot is AVAILABLE, which is a reserved state.
-        assert!(matches!(
-            s.retire(0),
-            Err(RetireError::SlotIsReserved { .. })
-        ));
-        // A frozen slot is likewise reserved.
-        let frozen = s.frozen().start as usize;
-        assert!(matches!(
-            s.retire(frozen),
-            Err(RetireError::SlotIsReserved { .. })
-        ));
-        // An owned slot is not retirable.
-        let slot = s.acquire().unwrap();
-        assert!(matches!(
-            s.retire(slot.slot() as usize),
-            Err(RetireError::SlotIsReserved { .. })
-        ));
-    }
-
-    #[test]
-    fn retire_published_slot_then_unreadable() {
-        let s = store(4, 8, 1).unwrap();
-
-        let idx = {
-            let slot = s.acquire().unwrap();
-            slot.publish() as usize
-        };
-
-        assert!(s.retire(idx).is_ok());
-
-        // A reader opened after retirement must not observe the retired slot.
-        let reader = s.reader().unwrap();
-        assert_eq!(reader.inner().read(idx), None);
-        assert_eq!(reader.inner().can_read(idx), Some(false));
-
-        // The slot can also not be retired again.
-        assert!(matches!(
-            s.retire(idx),
-            Err(RetireError::SlotIsReserved { .. })
-        ));
-    }
-
-    //---------//
-    // Recycle //
-    //---------//
-
-    #[test]
-    fn test_recycling() {
-        let entries = if cfg!(miri) { 16 } else { 2048 };
-
-        let s = store(entries, 4, 2).unwrap();
-
-        // Claim all slots.
-        let mut count = 0;
-        while let Some(slot) = s.acquire() {
-            slot.publish();
-            count += 1;
-        }
-
-        assert_eq!(count, s.writable().len());
-
-        // Now that all slots are claimed - retire all slots.
-        for i in s.writable() {
-            s.retire(i.into_usize()).unwrap();
-        }
-
-        // Verify that we can claim all slots again.
-        let mut count = 0;
-        while let Some(slot) = s.acquire() {
-            slot.publish();
-            count += 1;
-        }
-
-        assert_eq!(count, s.writable().len());
-    }
-}
+// ///////////
+// // Tests //
+// ///////////
+//
+// /// These tests are basic functionality tests for the store.
+// ///
+// /// Longer running conurrency tests are in the integration test suite.
+// #[cfg(test)]
+// mod tests {
+//     use super::*;
+//
+//     use diskann_utils::views::Matrix;
+//
+//     // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
+//     // zeroed frozen points. The frozen points occupy the highest slot indices.
+//     fn store(entries: usize, entry_bytes: usize, frozen: usize) -> Result<Store, StoreError> {
+//         let mut data = Matrix::new(0u8, frozen, entry_bytes);
+//         let mut base = 0u8;
+//         for row in data.row_iter_mut() {
+//             row.fill(base);
+//             base = base.wrapping_add(1);
+//         }
+//
+//         let mut config = Config::new();
+//         config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
+//         config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
+//
+//         let layout = Layout::new(
+//             Capacity::new(entries),
+//             MaxDegree::new(0),
+//         );
+//         Store::new(layout, config, data.as_view())
+//     }
+//
+//     //------------------------//
+//     // Constructor validation //
+//     //------------------------//
+//
+//     #[test]
+//     fn new_requires_a_frozen_point() {
+//         let err = store(4, 8, 0).unwrap_err();
+//         assert!(matches!(err.0, StoreErrorInner::NeedFrozenPoint));
+//     }
+//
+//     #[test]
+//     fn new_rejects_total_slot_overflow() {
+//         // `entries` alone fits in u32, but `entries + frozen` overflows it.
+//         let data = Matrix::new(0u8, 1, 8);
+//         let err = Store::new(
+//             Layout::new(Capacity::new(u32::MAX as usize), MaxDegree::new(0)),
+//             Config::default(),
+//             data.as_view(),
+//         )
+//         .unwrap_err();
+//         assert!(matches!(err.0, StoreErrorInner::TooManyEntries { .. }));
+//     }
+//
+//     #[test]
+//     fn new_rejects_too_many_neighbors() {
+//         let data = Matrix::new(0u8, 1, 8);
+//         let err = Store::new(
+//             Layout::new(Capacity::new(4), MaxDegree::new(u32::MAX.into_usize() + 1)),
+//             Config::default(),
+//             data.as_view(),
+//         )
+//         .unwrap_err();
+//         assert!(matches!(err.0, StoreErrorInner::TooManyNeighbors { .. }));
+//     }
+//
+//     //--------//
+//     // Layout //
+//     //--------//
+//
+//     #[test]
+//     fn frozen_range_follows_writable_slots() {
+//         let s = store(4, 8, 2).unwrap();
+//
+//         // Writable slots are [0, 4); frozen points occupy [4, 6).
+//         assert_eq!(s.frozen(), 4..6);
+//
+//         let reader = s.reader().unwrap();
+//         for i in 0..4 {
+//             assert!(!s.can_read_approximate(i).unwrap());
+//             assert!(!reader.inner().can_read(i).unwrap());
+//             assert!(reader.inner().read(i).is_none());
+//         }
+//
+//         assert!(s.can_read_approximate(4).unwrap());
+//         assert!(reader.inner().can_read(4).unwrap());
+//         assert_eq!(reader.inner().read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
+//
+//         assert!(s.can_read_approximate(5).unwrap());
+//         assert!(reader.inner().can_read(5).unwrap());
+//         assert_eq!(reader.inner().read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
+//
+//         assert!(s.can_read_approximate(6).is_none());
+//         assert!(reader.inner().can_read(6).is_none());
+//         assert!(reader.inner().read(6).is_none());
+//     }
+//
+//     ///////////////
+//     // Lifecycle //
+//     ///////////////
+//
+//     #[test]
+//     fn acquire_write_publish_read_roundtrip() {
+//         let s = store(4, 8, 1).unwrap();
+//
+//         let reader = s.reader().expect("reader guard available");
+//
+//         let idx = {
+//             let mut slot = s.acquire().expect("a fresh store has free slots");
+//             let idx = slot.slot() as usize;
+//             slot.as_mut_slice()
+//                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
+//
+//             // Before the slot is dropped - we should not be able to read it.
+//             assert!(reader.inner().read(idx).is_none());
+//             assert!(!s.can_read_approximate(idx).unwrap());
+//             slot.publish();
+//             idx
+//         };
+//
+//         assert_eq!(
+//             reader.inner().read(idx),
+//             Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice())
+//         );
+//         assert!(s.can_read_approximate(idx).unwrap());
+//     }
+//
+//     #[test]
+//     fn unpublished_slots_are_immediately_available() {
+//         let s = store(4, 8, 1).unwrap();
+//
+//         let reader = s.reader().expect("reader guard available");
+//
+//         let idx = {
+//             let mut slot = s.acquire().expect("a fresh store has free slots");
+//             let idx = slot.slot() as usize;
+//             slot.as_mut_slice()
+//                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
+//
+//             // Before the slot is dropped - we should not be able to read it.
+//             assert!(reader.inner().read(idx).is_none());
+//             assert!(!s.can_read_approximate(idx).unwrap());
+//
+//             // NOTE: We do not explicitly publish the slot.
+//             idx
+//         };
+//
+//         assert!(reader.inner().read(idx).is_none());
+//         assert!(!s.can_read_approximate(idx).unwrap());
+//     }
+//
+//     #[test]
+//     fn acquire_exhausts_then_reports_none() {
+//         let s = store(2, 8, 1).unwrap();
+//         // Hold the guards so the slots stay owned.
+//         let _a = s.acquire().expect("first writable slot");
+//         let _b = s.acquire().expect("second writable slot");
+//         assert!(
+//             s.acquire().is_none(),
+//             "all writable slots are owned, so acquire must fail"
+//         );
+//     }
+//
+//     //--------//
+//     // Retire //
+//     //--------//
+//
+//     #[test]
+//     fn retire_out_of_bounds() {
+//         let s = store(4, 8, 1).unwrap();
+//         assert!(matches!(s.retire(999), Err(RetireError::OutOfBounds)));
+//     }
+//
+//     #[test]
+//     fn retire_rejects_reserved_slots() {
+//         let s = store(4, 8, 1).unwrap();
+//         // An untouched writable slot is AVAILABLE, which is a reserved state.
+//         assert!(matches!(
+//             s.retire(0),
+//             Err(RetireError::SlotIsReserved { .. })
+//         ));
+//         // A frozen slot is likewise reserved.
+//         let frozen = s.frozen().start as usize;
+//         assert!(matches!(
+//             s.retire(frozen),
+//             Err(RetireError::SlotIsReserved { .. })
+//         ));
+//         // An owned slot is not retirable.
+//         let slot = s.acquire().unwrap();
+//         assert!(matches!(
+//             s.retire(slot.slot() as usize),
+//             Err(RetireError::SlotIsReserved { .. })
+//         ));
+//     }
+//
+//     #[test]
+//     fn retire_published_slot_then_unreadable() {
+//         let s = store(4, 8, 1).unwrap();
+//
+//         let idx = {
+//             let slot = s.acquire().unwrap();
+//             slot.publish() as usize
+//         };
+//
+//         assert!(s.retire(idx).is_ok());
+//
+//         // A reader opened after retirement must not observe the retired slot.
+//         let reader = s.reader().unwrap();
+//         assert_eq!(reader.inner().read(idx), None);
+//         assert_eq!(reader.inner().can_read(idx), Some(false));
+//
+//         // The slot can also not be retired again.
+//         assert!(matches!(
+//             s.retire(idx),
+//             Err(RetireError::SlotIsReserved { .. })
+//         ));
+//     }
+//
+//     //---------//
+//     // Recycle //
+//     //---------//
+//
+//     #[test]
+//     fn test_recycling() {
+//         let entries = if cfg!(miri) { 16 } else { 2048 };
+//
+//         let s = store(entries, 4, 2).unwrap();
+//
+//         // Claim all slots.
+//         let mut count = 0;
+//         while let Some(slot) = s.acquire() {
+//             slot.publish();
+//             count += 1;
+//         }
+//
+//         assert_eq!(count, s.writable().len());
+//
+//         // Now that all slots are claimed - retire all slots.
+//         for i in s.writable() {
+//             s.retire(i.into_usize()).unwrap();
+//         }
+//
+//         // Verify that we can claim all slots again.
+//         let mut count = 0;
+//         while let Some(slot) = s.acquire() {
+//             slot.publish();
+//             count += 1;
+//         }
+//
+//         assert_eq!(count, s.writable().len());
+//     }
+// }
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index f0c0458348..b3afb1fe3d 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -3,22 +3,30 @@
  * Licensed under the MIT license.
  */
 
+use diskann::ANNResult;
 use std::fmt::Debug;
 
 use crate::num::IdLimit;
 
+use super::Lifecycle;
+
+pub(crate) trait PluginConfig: Debug {
+    type Plugin: Plugin;
+    fn build(self, id_limit: IdLimit) -> ANNResult<Self::Plugin>;
+}
+
 pub(crate) trait Plugin: Debug + 'static {
     type Slot<'a>: Slot;
-
     fn id_limit(&self) -> IdLimit;
 
-    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_>;
-    fn reclaim(&self, i: u32);
-    fn retire(&self, i: u32);
+    unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_>;
+    unsafe fn reclaim(&self, i: u32, _: Lifecycle);
+    unsafe fn retire(&self, i: u32, _: Lifecycle);
 }
 
 pub(crate) trait Slot: Debug {
-    fn publish(self);
-    fn freeze(self);
-    fn abort(self);
+    fn publish(self, _: Lifecycle);
+    fn freeze(self, _: Lifecycle);
+    fn abort(self, _: Lifecycle);
 }
+
diff --git a/diskann-inmem/src/store/stacked.rs b/diskann-inmem/src/store/stacked.rs
deleted file mode 100644
index 65a3a8b56c..0000000000
--- a/diskann-inmem/src/store/stacked.rs
+++ /dev/null
@@ -1,76 +0,0 @@
-/*
- * Copyright (c) Microsoft Corporation.
- * Licensed under the MIT license.
- */
-
-use super::plugin::{self, Plugin};
-
-/// A [`super::Plugin`] for cascading multiple plugins together.
-#[derive(Debug)]
-pub(crate) struct Stacked<T, U> {
-    first: T,
-    rest: U,
-}
-
-impl<T, U> Plugin for Stacked<T, U>
-where
-    T: Plugin,
-    U: Plugin,
-{
-    type Slot<'a> = Slot<T::Slot<'a>, U::Slot<'a>>;
-
-    unsafe fn acquire(&self, i: u32) -> Self::Slot<'_> {
-        Slot::new(unsafe { self.first.acquire(i) }, unsafe {
-            self.rest.acquire(i)
-        })
-    }
-
-    fn reclaim(&self, i: u32) {
-        self.first.reclaim(i);
-        self.rest.reclaim(i);
-    }
-
-    fn retire(&self, i: u32) {
-        self.first.retire(i);
-        self.rest.retire(i);
-    }
-}
-
-#[derive(Debug)]
-pub(crate) struct Slot<T, U> {
-    first: T,
-    rest: U,
-}
-
-impl<T, U> Slot<T, U> {
-    fn new(first: T, rest: U) -> Self {
-        Self { first, rest }
-    }
-
-    pub(crate) fn first(&self) -> &T {
-        &self.first
-    }
-
-    pub(crate) fn rest(&self) -> &U {
-        &self.rest
-    }
-}
-
-impl<T, U> plugin::Slot for Slot<T, U>
-where
-    T: plugin::Slot,
-    U: plugin::Slot,
-{
-    fn publish(self) {
-        self.first.publish();
-        self.rest.publish();
-    }
-    fn freeze(self) {
-        self.first.freeze();
-        self.rest.freeze();
-    }
-    fn abort(self) {
-        self.first.abort();
-        self.rest.abort();
-    }
-}

From a58db48391a9d5c2731b44590be58d2c123bd8e2 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Wed, 19 Aug 2026 16:20:30 -0700
Subject: [PATCH 17/34] checkpoint

---
 diskann-inmem/integration/index/runner.rs |   1 -
 diskann-inmem/src/integration/store.rs    |   7 +-
 diskann-inmem/src/layers/full.rs          |  10 +-
 diskann-inmem/src/num.rs                  |   6 -
 diskann-inmem/src/store/checked.rs        | 243 +++++++++
 diskann-inmem/src/store/invasive.rs       |  73 ++-
 diskann-inmem/src/store/mod.rs            | 602 ++++++++++------------
 diskann-inmem/src/store/plugin.rs         | 128 ++++-
 8 files changed, 701 insertions(+), 369 deletions(-)
 create mode 100644 diskann-inmem/src/store/checked.rs

diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index eb6a4f2fa1..74994f3cc6 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -15,7 +15,6 @@ use diskann_benchmark_runner::{
 };
 use diskann_utils::views::Matrix;
 use diskann_vector::distance::Metric;
-use half::f16;
 use serde::{Deserialize, Serialize};
 
 use diskann_inmem::{
diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
index 47f0526e06..a05df5840c 100644
--- a/diskann-inmem/src/integration/store.rs
+++ b/diskann-inmem/src/integration/store.rs
@@ -10,8 +10,6 @@
 
 use std::num::{NonZeroU32, NonZeroUsize};
 
-use diskann_utils::views::Matrix;
-
 use crate::{
     num::{Bytes, Capacity, MaxDegree},
     store,
@@ -95,10 +93,7 @@ impl Store {
     }
 
     pub fn reader(&self) -> Option<Reader<'_>> {
-        match self
-            .store
-            .guard(|plugin, guard| unsafe { plugin.reader(guard) })
-        {
+        match store::invasive::Invasive::reader(&self.store) {
             Ok(reader) => Some(Reader::new(reader)),
             Err(crate::epoch::Unavailable) => None,
         }
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 4f15ccc917..b0c885aa3b 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -121,7 +121,7 @@ where
     ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
         let expand_beam = T::make_expand_beam(layer, query)?;
         Ok(crate::provider::SearchAccessor::new(
-            layer.store.temp_neighbors(),
+            layer.store.neighbors(),
             expand_beam,
             provider,
             layer.store.frozen(),
@@ -136,7 +136,7 @@ where
         let prune = T::make_prune(layer)?;
         Ok(crate::provider::PruneAccessor::new(
             prune,
-            layer.store.temp_neighbors(),
+            layer.store.neighbors(),
             counters,
         ))
     }
@@ -228,9 +228,7 @@ where
     }
 
     fn reader(&self) -> ANNResult<invasive::Reader<'_>> {
-        Ok(self
-            .store
-            .guard(|invasive, guard| unsafe { invasive.reader(guard) })?)
+        Ok(Invasive::reader(&self.store)?)
     }
 }
 
@@ -259,7 +257,7 @@ where
     T: FullPrecision,
 {
     fn max_degree(&self) -> MaxDegree {
-        self.store.temp_neighbors().max_degree()
+        self.store.neighbors().max_degree()
     }
 
     fn retire(&self, i: u32) -> ANNResult<()> {
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 2ce7fc5bc1..4a5ab74f9a 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -46,12 +46,6 @@ impl Bytes {
         }
     }
 
-    /// Perform integer division of `self` by `other`.
-    #[inline]
-    pub(crate) const fn div(self, other: NonZeroUsize) -> Bytes {
-        Bytes::new(self.value() / other.get())
-    }
-
     /// Subtract `other` from `self` without checking for underflow.
     #[inline]
     pub(crate) const fn unchecked_sub(self, other: Bytes) -> Bytes {
diff --git a/diskann-inmem/src/store/checked.rs b/diskann-inmem/src/store/checked.rs
new file mode 100644
index 0000000000..838b9c0a67
--- /dev/null
+++ b/diskann-inmem/src/store/checked.rs
@@ -0,0 +1,243 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{
+    assert_matches,
+    sync::atomic::{AtomicBool, Ordering},
+};
+
+use diskann::utils::IntoUsize;
+use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+use crate::{epoch, num::IdLimit};
+
+use super::{Lifecycle, plugin};
+
+#[derive(Debug, Default)]
+enum State {
+    #[default]
+    Available,
+    Readable {
+        value: u64,
+        retired: AtomicBool,
+    },
+    Frozen {
+        value: u64,
+    },
+}
+
+#[derive(Debug)]
+pub(crate) struct Config(());
+
+impl Config {
+    pub(crate) fn new() -> Self {
+        Self(())
+    }
+}
+
+impl plugin::PluginConfig for Config {
+    type Plugin = Checked;
+
+    fn build(self, id_limit: IdLimit) -> diskann::ANNResult<Checked> {
+        Ok(Checked::new(id_limit))
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Checked {
+    states: Vec<RwLock<State>>,
+}
+
+impl Checked {
+    pub(crate) fn config() -> Config {
+        Config::new()
+    }
+
+    pub(crate) fn new(id_limit: IdLimit) -> Self {
+        Self {
+            states: std::iter::repeat_with(|| RwLock::new(State::default()))
+                .take(id_limit.as_usize())
+                .collect(),
+        }
+    }
+
+    pub(crate) fn id_limit(&self) -> IdLimit {
+        IdLimit::new(self.states.len().try_into().unwrap())
+    }
+
+    fn expect_write(&self, i: u32) -> RwLockWriteGuard<'_, State> {
+        let i = i.into_usize();
+
+        // Note: this will panic if `i` is out-of-bounds.
+        let entry = &self.states[i];
+
+        // Correct usage of the concurrency protocol means that this `try_write` failing
+        // is a bug.
+        let Some(guard) = entry.try_write() else {
+            panic!("concurrency violation when acquiring write guard");
+        };
+
+        guard
+    }
+
+    fn expect_read(&self, i: u32) -> RwLockReadGuard<'_, State> {
+        let i = i.into_usize();
+
+        // Note: this will panic if `i` is out-of-bounds.
+        let entry = &self.states[i];
+
+        // Correct usage of the concurrency protocol means that this `try_read` failing
+        // is a bug.
+        let Some(guard) = entry.try_read() else {
+            panic!("concurrency violation when acquiring read guard");
+        };
+
+        guard
+    }
+
+    pub(crate) fn reader<'a>(&'a self, guard: epoch::Guard<'a>) -> Reader<'a> {
+        Reader {
+            parent: self,
+            _guard: guard,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Value<'a> {
+    value: u64,
+    _guard: RwLockReadGuard<'a, State>,
+}
+
+impl Value<'_> {
+    pub(crate) fn get(&self) -> u64 {
+        self.value
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Reader<'a> {
+    parent: &'a Checked,
+    _guard: epoch::Guard<'a>,
+}
+
+impl Reader<'_> {
+    pub(crate) fn read(&self, i: u32) -> Option<Value<'_>> {
+        // This is kind of messy. The overall summary is this:
+        //
+        // 1. `i` has to be inbounds.
+        // 2. We have to be able to read the slot (if a write guard is active, then we
+        //    clearly should not be reading).
+        // 3a. If the state is frozen, then we can read it.
+        // 3b. If the state is readable and not retired, we can read it.
+        //
+        //     What happens if we transition to retired just after reading?
+        //
+        //     Fortunately, the EBR guard will keep the slot from being reclaimed until
+        //     the current `Reader` goes out-of-scope. Holding onto the
+        //     `RwLockReadGuard` allows us to detect bugs in the EBR protocol as
+        //     `Checked::expect_write` will fail on reclamation if a returned `Value`
+        //     is still active.
+        if let Some(state) = self.parent.states.get(i.into_usize())
+            && let Some(guard) = state.try_read()
+        {
+            let value = match &*guard {
+                State::Frozen { value } => *value,
+                State::Readable { value, retired } => {
+                    if retired.load(Ordering::Relaxed) {
+                        return None;
+                    }
+                    *value
+                }
+                _ => return None,
+            };
+            Some(Value {
+                value,
+                _guard: guard,
+            })
+        } else {
+            None
+        }
+    }
+}
+
+impl plugin::Plugin for Checked {
+    type Slot<'a> = Slot<'a>;
+
+    fn id_limit(&self) -> IdLimit {
+        <Checked>::id_limit(self)
+    }
+
+    unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_> {
+        let guard = self.expect_write(i);
+        assert_matches!(*guard, State::Available, "slot is in an invalid state");
+        Slot::new(guard)
+    }
+
+    unsafe fn retire(&self, i: u32, _: Lifecycle) {
+        let guard = self.expect_read(i);
+        match &*guard {
+            State::Available => panic!("invalid \"Available\" state"),
+            State::Readable { retired, .. } => {
+                let old = retired.swap(true, Ordering::Relaxed);
+                if old {
+                    panic!("slot {i} was retired multiple times");
+                }
+            }
+            State::Frozen { .. } => panic!("tried to retire frozen point {i}"),
+        }
+    }
+
+    unsafe fn reclaim(&self, i: u32, _: Lifecycle) {
+        let mut guard = self.expect_write(i);
+        match &*guard {
+            State::Available => panic!("invalid \"Available\" state"),
+            State::Readable { retired, .. } => {
+                assert!(
+                    retired.load(Ordering::Relaxed),
+                    "tried to reclaim {i} before it has been retired!",
+                );
+            }
+            State::Frozen { .. } => panic!("tried to reclaim frozen point {i}"),
+        }
+
+        *guard = State::Available;
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Slot<'a> {
+    guard: RwLockWriteGuard<'a, State>,
+    value: Option<u64>,
+}
+
+impl<'a> Slot<'a> {
+    fn new(guard: RwLockWriteGuard<'a, State>) -> Self {
+        Self { guard, value: None }
+    }
+
+    pub(crate) fn set(&mut self, value: u64) {
+        self.value = Some(value)
+    }
+}
+
+impl plugin::Slot for Slot<'_> {
+    fn publish(mut self, _: Lifecycle) {
+        let value = self.value.expect("`value` was not set");
+        *self.guard = State::Readable {
+            value,
+            retired: AtomicBool::new(false),
+        };
+    }
+
+    fn freeze(mut self, _: Lifecycle) {
+        let value = self.value.expect("`value` was not set");
+        *self.guard = State::Frozen { value };
+    }
+
+    fn abort(mut self, _: Lifecycle) {
+        *self.guard = State::Available;
+    }
+}
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index 9cf1b96b53..0b948e0a56 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -3,43 +3,52 @@
  * Licensed under the MIT license.
  */
 
+//! A store [`plugin::Plugin`] that maintains in invasive slot state where the data in each
+//! slot is a contiguous slice of memory.
+//!
+//! Slot state is stored as an [`AtomicTag`] immediately after the slot data.
+
 use std::{num::NonZeroUsize, sync::atomic::Ordering};
 
 use diskann::{ANNResult, utils::IntoUsize};
+use thiserror::Error;
 
 use crate::{
     buffer::{Buffer, RawSlice},
     epoch,
-    num::{Align, Bytes, Capacity, IdLimit},
+    num::{Align, Bytes, IdLimit},
+    store::{Lifecycle, Store, plugin},
     tag::{AtomicTag, Tag},
-    store::Lifecycle,
 };
 
+/// A [`plugin::PluginConfig`] for [`Invasive`].
 #[derive(Debug, Clone)]
 pub(crate) struct Config {
+    /// The number of bytes held in each slot.
     bytes: Bytes,
 }
 
 impl Config {
+    /// Create a new [`Config`] for [`Invasive`] reserving `bytes` bytes for each slot.
     pub(crate) fn new(bytes: Bytes) -> Self {
         Self { bytes }
     }
 
+    /// Build an [`Invasive`] store holding `id_limit` slots.
     pub(crate) fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
         let Self { bytes } = self;
-
         Ok(Invasive::new(id_limit, bytes))
     }
 }
 
-impl super::plugin::PluginConfig for Config {
+impl plugin::PluginConfig for Config {
     type Plugin = Invasive;
     fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
         <Config>::build(self, id_limit)
     }
 }
 
-/// The invasive store where concurrency tags are stored inline with the data.
+/// The invasive store where concurrency tags are stored inline just after the data.
 #[derive(Debug)]
 pub(crate) struct Invasive {
     // The inline tags are `AtomicTag`s stored after the data.
@@ -50,17 +59,19 @@ pub(crate) struct Invasive {
     unpadded: Bytes,
 }
 
-const TWO: NonZeroUsize = NonZeroUsize::new(2).unwrap();
-
 impl Invasive {
+    /// Construct the [`Config`] for [`Self`].
+    ///
+    /// See also: [`Config::new`].
     pub(crate) fn config(bytes: Bytes) -> Config {
         Config::new(bytes)
     }
 
+    /// Create a new [`Invasive`] with capacity for `id_limit` slots of `bytes`.
     pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Self {
         let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
         let padded_bytes = unpadded
-            .checked_next_multiple_of(Bytes::CACHELINE.div(TWO))
+            .checked_next_multiple_of(Bytes::CACHELINE)
             .unwrap();
 
         Self {
@@ -69,22 +80,22 @@ impl Invasive {
         }
     }
 
+    /// Return the [`IdLimit`] for this store.
     pub(crate) fn id_limit(&self) -> IdLimit {
         // The numeric cast is save because `Invasive::new` takes an `IdLimit` in its
         // constructor, and thus `self.buffer.len()` cannot exceed `u32::MAX`.
         IdLimit::new(self.buffer.len() as u32)
     }
 
-    pub(crate) fn bytes(&self) -> Bytes {
-        self.unpadded
-    }
-
-    pub(crate) unsafe fn reader<'a>(&'a self, guard: epoch::Guard<'a>) -> Reader<'a> {
-        Reader {
-            buffer: &self.buffer,
-            unpadded: self.unpadded,
-            _guard: guard,
-        }
+    /// Return a [`Reader`] over [`Self`] inside `store`.
+    pub(crate) fn reader<'a>(store: &'a Store<Self>) -> Result<Reader<'a>, epoch::Unavailable> {
+        store.guard(|this, guard: epoch::Guard<'a>| {
+            Reader {
+                buffer: &this.buffer,
+                unpadded: this.unpadded,
+                _guard: guard
+            }
+        })
     }
 
     /// Return the data at position `i` without bound-checking.
@@ -114,7 +125,7 @@ impl Invasive {
     }
 }
 
-impl super::plugin::Plugin for Invasive {
+impl plugin::Plugin for Invasive {
     type Slot<'a> = Slot<'a>;
 
     fn id_limit(&self) -> IdLimit {
@@ -128,7 +139,7 @@ impl super::plugin::Plugin for Invasive {
 
         // This is a pessimistic check to ensure that the caller is correctly using the
         // `plugin` API.
-        assert_eq!(
+        debug_assert_eq!(
             tag.load(Ordering::Relaxed),
             Tag::AVAILABLE,
             "concurrency violation",
@@ -157,6 +168,7 @@ impl super::plugin::Plugin for Invasive {
     }
 }
 
+/// A reader into an [`Invasive`] store.
 #[derive(Debug)]
 pub(crate) struct Reader<'a> {
     buffer: &'a Buffer,
@@ -200,12 +212,9 @@ impl<'a> Reader<'a> {
     ///
     /// This guarantee only holds while `self` is alive. Construction of a new [`Reader`]
     /// requires a separate check.
-    #[cfg_attr(
-        not(test),
-        expect(
-            dead_code,
-            reason = "this is non-trivial method that likely be used in the future"
-        )
+    #[expect(
+        dead_code,
+        reason = "this is non-trivial method that likely be used in the future"
     )]
     pub(crate) fn can_read(&self, i: usize) -> Option<bool> {
         if !self.is_in_bounds(i) {
@@ -287,6 +296,7 @@ impl<'a> Reader<'a> {
     }
 }
 
+/// A [`plugin::Slot`] for [`Invasive`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
     tag: &'a AtomicTag,
@@ -294,12 +304,20 @@ pub(crate) struct Slot<'a> {
 }
 
 impl<'a> Slot<'a> {
+    /// Return the data within this slot as a mutable slice.
+    ///
+    /// The length of this slice is guaranteed to be the number of bytes passed to
+    /// [`Invasive::new`] or [`Config::new`].
     pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
+        // SAFETY: Users of the `plugin::Slot` are obligated to ensure exclusivity.
+        //
+        // Since `Reader` obeys the plugin life-cycle requirements, a concurrent reader
+        // of this data should not be possible.
         unsafe { self.data.as_mut_slice() }
     }
 }
 
-impl super::plugin::Slot for Slot<'_> {
+impl plugin::Slot for Slot<'_> {
     fn publish(self, _: Lifecycle) {
         self.tag.store(Tag::PUBLISHED, Ordering::Release);
     }
@@ -310,3 +328,4 @@ impl super::plugin::Slot for Slot<'_> {
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
+
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 5174541d01..f080ac06ee 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -59,7 +59,6 @@ use std::{
 };
 
 use diskann::{ANNError, utils::IntoUsize};
-use diskann_utils::views::MatrixView;
 use thiserror::Error;
 
 use crate::{
@@ -67,13 +66,16 @@ use crate::{
     epoch::{self, Registry},
     freelist::{self, Freelist},
     neighbors::{Neighbors, NeighborsError},
-    num::{Bytes, Capacity, IdLimit, MaxDegree},
+    num::{Capacity, IdLimit, MaxDegree},
     tag::{AtomicTag, Tag},
 };
 
 pub(crate) mod invasive;
 pub(crate) mod plugin;
 
+#[cfg(test)]
+mod checked;
+
 /// To make extra sure that [`plugin::Plugin`] life-cycle arguments are not callable outside
 /// of this module (i.e., elsewhere in this crate), this [`Lifecycle`] marker type is used
 /// that is only constructible in this module.
@@ -117,7 +119,7 @@ impl Config {
     ///
     /// Increasing this number will increase the number of threads that can work concurrently
     /// on the index at the cost of longer scan times for epoch advancement.
-    pub(crate) fn epoch_guard_slots(&mut self, epoch_guard_slots: NonZeroUsize) -> &mut Self {
+    pub fn epoch_guard_slots(&mut self, epoch_guard_slots: NonZeroUsize) -> &mut Self {
         self.epoch_guard_slots = epoch_guard_slots;
         self
     }
@@ -126,7 +128,7 @@ impl Config {
     ///
     /// Increasing the capacity of the queue will allow more recycled IDs to be retrieved
     /// without triggering a scan, but will cost more memory.
-    pub(crate) fn freelist_recycle_capacity(
+    pub fn freelist_recycle_capacity(
         &mut self,
         freelist_recycle_capacity: NonZeroU32,
     ) -> &mut Self {
@@ -255,11 +257,6 @@ where
         (self.unfrozen.value() as u32)..self.neighbors.entries()
     }
 
-    /// Return the maximum degree that can be stored in the graph.
-    pub(crate) fn max_degree(&self) -> MaxDegree {
-        self.neighbors.max_degree()
-    }
-
     pub(crate) fn id_limit(&self) -> IdLimit {
         plugin::Plugin::id_limit(&self.plugin)
     }
@@ -268,7 +265,7 @@ where
         self.unfrozen
     }
 
-    pub(crate) fn temp_neighbors(&self) -> &Neighbors {
+    pub(crate) fn neighbors(&self) -> &Neighbors {
         &self.neighbors
     }
 
@@ -319,23 +316,6 @@ where
         Ok(f(self.plugin(), guard))
     }
 
-    // /// Return a [`Reader`] into the store.
-    // ///
-    // /// # Errors
-    // ///
-    // /// Returns [`epoch::Unavailable`] if there are too many active readers.
-    // pub(crate) fn reader(&self) -> Result<Reader<'_>, epoch::Unavailable> {
-    //     Ok(Reader {
-    //         inner: unsafe { self.plugin().reader(self.registry.guard()?) },
-    //         neighbors: &self.neighbors,
-    //     })
-    // }
-
-    // // TODO: Rework neighbor storage.
-    // pub(crate) fn temp_inner_reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
-    //     Ok(unsafe { self.plugin().reader(self.registry.guard()?) })
-    // }
-
     /// Attempt to acquire a new [`Slot`] for writing.
     ///
     /// This method first consults the freelist and falls back to scanning the tags list
@@ -396,7 +376,9 @@ where
         match tag.compare_exchange(current, retiring, Ordering::Relaxed, Ordering::Relaxed) {
             Ok(_) => {
                 // Set the metadata in the mirror as well.
-                unsafe { plugin::Plugin::retire(self.plugin(), i.try_into().unwrap(), Lifecycle::new()) };
+                unsafe {
+                    plugin::Plugin::retire(self.plugin(), i.try_into().unwrap(), Lifecycle::new())
+                };
                 guard.retire(i as u32);
                 Ok(())
             }
@@ -495,7 +477,8 @@ where
             Ordering::Relaxed,
         ) {
             Ok(_) => {
-                let data = unsafe { plugin::Plugin::acquire(self.plugin(), slot, Lifecycle::new()) };
+                let data =
+                    unsafe { plugin::Plugin::acquire(self.plugin(), slot, Lifecycle::new()) };
                 Some(Slot {
                     tag,
                     data: ManuallyDrop::new(data),
@@ -530,10 +513,6 @@ where
 pub(crate) struct StoreError(StoreErrorInner);
 
 impl StoreError {
-    fn need_frozen_point() -> Self {
-        Self(StoreErrorInner::NeedFrozenPoint)
-    }
-
     fn too_many_entries(capacity: Capacity, frozen: u32) -> Self {
         Self(StoreErrorInner::TooManyEntries {
             entries: capacity.value(),
@@ -568,8 +547,6 @@ diskann::convert_error!(StoreError);
 
 #[derive(Debug, Error)]
 enum StoreErrorInner {
-    #[error("at least one frozen point must be provided")]
-    NeedFrozenPoint,
     #[error(
         "total points ({} + {} frozen) must not exceed `u32::MAX`",
         entries,
@@ -605,34 +582,6 @@ pub(crate) enum RetireError {
 
 diskann::convert_error!(RetireError);
 
-// /// An epoch protected reader into a [`Store`].
-// ///
-// /// Created via [`Store::reader`].
-// #[derive(Debug)]
-// pub(crate) struct Reader<'a> {
-//     inner: invasive::Reader<'a>,
-//     neighbors: &'a Neighbors,
-// }
-//
-// impl<'a> Reader<'a> {
-//     /// Return `true` if the index `i` is in-bounds.
-//     #[inline]
-//     #[must_use = "this function has no side-effects"]
-//     pub(crate) fn is_in_bounds(&self, i: usize) -> bool {
-//         i < self.neighbors.entries().into_usize()
-//     }
-//
-//     #[inline]
-//     pub(crate) fn inner(&self) -> &invasive::Reader<'_> {
-//         &self.inner
-//     }
-//
-//     /// Return [`Neighbors`].
-//     pub(crate) fn neighbors(&self) -> &Neighbors {
-//         self.neighbors
-//     }
-// }
-
 /// A writable buffer into the data managed by a [`Store`], obtained from [`Store::acquire`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a, S>
@@ -663,7 +612,10 @@ where
         let mut me = ManuallyDrop::new(self);
 
         // Freeze the inner slot.
-        plugin::Slot::freeze(unsafe { ManuallyDrop::take(&mut me.data) }, Lifecycle::new());
+        plugin::Slot::freeze(
+            unsafe { ManuallyDrop::take(&mut me.data) },
+            Lifecycle::new(),
+        );
 
         // Update the authoritative store.
         me.tag.store(Tag::FROZEN, Ordering::Release);
@@ -679,7 +631,10 @@ where
         let mut me = ManuallyDrop::new(self);
 
         // Publish the inner slot.
-        plugin::Slot::publish(unsafe { ManuallyDrop::take(&mut me.data) }, Lifecycle::new());
+        plugin::Slot::publish(
+            unsafe { ManuallyDrop::take(&mut me.data) },
+            Lifecycle::new(),
+        );
 
         // Update the authoritative store.
         me.tag.store(Tag::PUBLISHED, Ordering::Release);
@@ -692,263 +647,268 @@ where
     S: plugin::Slot,
 {
     fn drop(&mut self) {
-        plugin::Slot::abort(unsafe { ManuallyDrop::take(&mut self.data) }, Lifecycle::new());
+        plugin::Slot::abort(
+            unsafe { ManuallyDrop::take(&mut self.data) },
+            Lifecycle::new(),
+        );
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
 
-// ///////////
-// // Tests //
-// ///////////
-//
-// /// These tests are basic functionality tests for the store.
-// ///
-// /// Longer running conurrency tests are in the integration test suite.
-// #[cfg(test)]
-// mod tests {
-//     use super::*;
-//
-//     use diskann_utils::views::Matrix;
-//
-//     // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
-//     // zeroed frozen points. The frozen points occupy the highest slot indices.
-//     fn store(entries: usize, entry_bytes: usize, frozen: usize) -> Result<Store, StoreError> {
-//         let mut data = Matrix::new(0u8, frozen, entry_bytes);
-//         let mut base = 0u8;
-//         for row in data.row_iter_mut() {
-//             row.fill(base);
-//             base = base.wrapping_add(1);
-//         }
-//
-//         let mut config = Config::new();
-//         config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
-//         config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
-//
-//         let layout = Layout::new(
-//             Capacity::new(entries),
-//             MaxDegree::new(0),
-//         );
-//         Store::new(layout, config, data.as_view())
-//     }
-//
-//     //------------------------//
-//     // Constructor validation //
-//     //------------------------//
-//
-//     #[test]
-//     fn new_requires_a_frozen_point() {
-//         let err = store(4, 8, 0).unwrap_err();
-//         assert!(matches!(err.0, StoreErrorInner::NeedFrozenPoint));
-//     }
-//
-//     #[test]
-//     fn new_rejects_total_slot_overflow() {
-//         // `entries` alone fits in u32, but `entries + frozen` overflows it.
-//         let data = Matrix::new(0u8, 1, 8);
-//         let err = Store::new(
-//             Layout::new(Capacity::new(u32::MAX as usize), MaxDegree::new(0)),
-//             Config::default(),
-//             data.as_view(),
-//         )
-//         .unwrap_err();
-//         assert!(matches!(err.0, StoreErrorInner::TooManyEntries { .. }));
-//     }
-//
-//     #[test]
-//     fn new_rejects_too_many_neighbors() {
-//         let data = Matrix::new(0u8, 1, 8);
-//         let err = Store::new(
-//             Layout::new(Capacity::new(4), MaxDegree::new(u32::MAX.into_usize() + 1)),
-//             Config::default(),
-//             data.as_view(),
-//         )
-//         .unwrap_err();
-//         assert!(matches!(err.0, StoreErrorInner::TooManyNeighbors { .. }));
-//     }
-//
-//     //--------//
-//     // Layout //
-//     //--------//
-//
-//     #[test]
-//     fn frozen_range_follows_writable_slots() {
-//         let s = store(4, 8, 2).unwrap();
-//
-//         // Writable slots are [0, 4); frozen points occupy [4, 6).
-//         assert_eq!(s.frozen(), 4..6);
-//
-//         let reader = s.reader().unwrap();
-//         for i in 0..4 {
-//             assert!(!s.can_read_approximate(i).unwrap());
-//             assert!(!reader.inner().can_read(i).unwrap());
-//             assert!(reader.inner().read(i).is_none());
-//         }
-//
-//         assert!(s.can_read_approximate(4).unwrap());
-//         assert!(reader.inner().can_read(4).unwrap());
-//         assert_eq!(reader.inner().read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
-//
-//         assert!(s.can_read_approximate(5).unwrap());
-//         assert!(reader.inner().can_read(5).unwrap());
-//         assert_eq!(reader.inner().read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
-//
-//         assert!(s.can_read_approximate(6).is_none());
-//         assert!(reader.inner().can_read(6).is_none());
-//         assert!(reader.inner().read(6).is_none());
-//     }
-//
-//     ///////////////
-//     // Lifecycle //
-//     ///////////////
-//
-//     #[test]
-//     fn acquire_write_publish_read_roundtrip() {
-//         let s = store(4, 8, 1).unwrap();
-//
-//         let reader = s.reader().expect("reader guard available");
-//
-//         let idx = {
-//             let mut slot = s.acquire().expect("a fresh store has free slots");
-//             let idx = slot.slot() as usize;
-//             slot.as_mut_slice()
-//                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
-//
-//             // Before the slot is dropped - we should not be able to read it.
-//             assert!(reader.inner().read(idx).is_none());
-//             assert!(!s.can_read_approximate(idx).unwrap());
-//             slot.publish();
-//             idx
-//         };
-//
-//         assert_eq!(
-//             reader.inner().read(idx),
-//             Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice())
-//         );
-//         assert!(s.can_read_approximate(idx).unwrap());
-//     }
-//
-//     #[test]
-//     fn unpublished_slots_are_immediately_available() {
-//         let s = store(4, 8, 1).unwrap();
-//
-//         let reader = s.reader().expect("reader guard available");
-//
-//         let idx = {
-//             let mut slot = s.acquire().expect("a fresh store has free slots");
-//             let idx = slot.slot() as usize;
-//             slot.as_mut_slice()
-//                 .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
-//
-//             // Before the slot is dropped - we should not be able to read it.
-//             assert!(reader.inner().read(idx).is_none());
-//             assert!(!s.can_read_approximate(idx).unwrap());
-//
-//             // NOTE: We do not explicitly publish the slot.
-//             idx
-//         };
-//
-//         assert!(reader.inner().read(idx).is_none());
-//         assert!(!s.can_read_approximate(idx).unwrap());
-//     }
-//
-//     #[test]
-//     fn acquire_exhausts_then_reports_none() {
-//         let s = store(2, 8, 1).unwrap();
-//         // Hold the guards so the slots stay owned.
-//         let _a = s.acquire().expect("first writable slot");
-//         let _b = s.acquire().expect("second writable slot");
-//         assert!(
-//             s.acquire().is_none(),
-//             "all writable slots are owned, so acquire must fail"
-//         );
-//     }
-//
-//     //--------//
-//     // Retire //
-//     //--------//
-//
-//     #[test]
-//     fn retire_out_of_bounds() {
-//         let s = store(4, 8, 1).unwrap();
-//         assert!(matches!(s.retire(999), Err(RetireError::OutOfBounds)));
-//     }
-//
-//     #[test]
-//     fn retire_rejects_reserved_slots() {
-//         let s = store(4, 8, 1).unwrap();
-//         // An untouched writable slot is AVAILABLE, which is a reserved state.
-//         assert!(matches!(
-//             s.retire(0),
-//             Err(RetireError::SlotIsReserved { .. })
-//         ));
-//         // A frozen slot is likewise reserved.
-//         let frozen = s.frozen().start as usize;
-//         assert!(matches!(
-//             s.retire(frozen),
-//             Err(RetireError::SlotIsReserved { .. })
-//         ));
-//         // An owned slot is not retirable.
-//         let slot = s.acquire().unwrap();
-//         assert!(matches!(
-//             s.retire(slot.slot() as usize),
-//             Err(RetireError::SlotIsReserved { .. })
-//         ));
-//     }
-//
-//     #[test]
-//     fn retire_published_slot_then_unreadable() {
-//         let s = store(4, 8, 1).unwrap();
-//
-//         let idx = {
-//             let slot = s.acquire().unwrap();
-//             slot.publish() as usize
-//         };
-//
-//         assert!(s.retire(idx).is_ok());
-//
-//         // A reader opened after retirement must not observe the retired slot.
-//         let reader = s.reader().unwrap();
-//         assert_eq!(reader.inner().read(idx), None);
-//         assert_eq!(reader.inner().can_read(idx), Some(false));
-//
-//         // The slot can also not be retired again.
-//         assert!(matches!(
-//             s.retire(idx),
-//             Err(RetireError::SlotIsReserved { .. })
-//         ));
-//     }
-//
-//     //---------//
-//     // Recycle //
-//     //---------//
-//
-//     #[test]
-//     fn test_recycling() {
-//         let entries = if cfg!(miri) { 16 } else { 2048 };
-//
-//         let s = store(entries, 4, 2).unwrap();
-//
-//         // Claim all slots.
-//         let mut count = 0;
-//         while let Some(slot) = s.acquire() {
-//             slot.publish();
-//             count += 1;
-//         }
-//
-//         assert_eq!(count, s.writable().len());
-//
-//         // Now that all slots are claimed - retire all slots.
-//         for i in s.writable() {
-//             s.retire(i.into_usize()).unwrap();
-//         }
-//
-//         // Verify that we can claim all slots again.
-//         let mut count = 0;
-//         while let Some(slot) = s.acquire() {
-//             slot.publish();
-//             count += 1;
-//         }
-//
-//         assert_eq!(count, s.writable().len());
-//     }
-// }
+///////////
+// Tests //
+///////////
+
+/// These tests are basic functionality tests for the store.
+///
+/// Longer running conurrency tests are in the integration test suite.
+#[cfg(test)]
+mod tests {
+    use super::{checked::Checked, *};
+
+    use std::assert_matches;
+
+    // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
+    // zeroed frozen points. The frozen points occupy the highest slot indices.
+    fn store(entries: usize, frozen: u32) -> Result<Store<Checked>, StoreError> {
+        let mut config = Config::new();
+        config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
+        config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
+
+        let layout = Layout::new(Capacity::new(entries), MaxDegree::new(0), frozen);
+        let store = Store::new(layout, config, Checked::config())?;
+
+        assert_eq!(store.frozen().len(), frozen.into_usize());
+
+        for (i, id) in store.frozen().enumerate() {
+            let mut slot = store.slot(id).unwrap();
+            slot.data().set(i as u64);
+            slot.freeze();
+        }
+
+        Ok(store)
+    }
+
+    fn reader(store: &Store<Checked>) -> checked::Reader<'_> {
+        store.guard(|checked, guard| checked.reader(guard)).unwrap()
+    }
+
+    //------------------------//
+    // Constructor validation //
+    //------------------------//
+
+    #[test]
+    fn new_rejects_total_slot_overflow() {
+        // `entries` alone fits in u32, but `entries + frozen` overflows it.
+        let err = Store::new(
+            Layout::new(Capacity::new(u32::MAX as usize), MaxDegree::new(0), 1),
+            Config::default(),
+            Checked::config(),
+        )
+        .unwrap_err();
+        assert!(matches!(err.0, StoreErrorInner::TooManyEntries { .. }));
+    }
+
+    #[test]
+    fn new_rejects_too_many_neighbors() {
+        let err = Store::new(
+            Layout::new(
+                Capacity::new(4),
+                MaxDegree::new(u32::MAX.into_usize() + 1),
+                0,
+            ),
+            Config::default(),
+            Checked::config(),
+        )
+        .unwrap_err();
+        assert!(matches!(err.0, StoreErrorInner::TooManyNeighbors { .. }));
+    }
+
+    //--------//
+    // Layout //
+    //--------//
+
+    #[test]
+    fn frozen_range_follows_writable_slots() {
+        let s = store(4, 2).unwrap();
+
+        // Writable slots are [0, 4); frozen points occupy [4, 6).
+        assert_eq!(s.frozen(), 4..6);
+
+        let reader = reader(&s);
+        for i in 0u32..4 {
+            assert!(!s.can_read_approximate(i.into_usize()).unwrap());
+            assert!(reader.read(i).is_none());
+        }
+
+        assert!(s.can_read_approximate(4).unwrap());
+        assert_eq!(reader.read(4).unwrap().get(), 0);
+
+        assert!(s.can_read_approximate(5).unwrap());
+        assert_eq!(reader.read(5).unwrap().get(), 1);
+
+        assert!(s.can_read_approximate(6).is_none());
+        assert!(reader.read(6).is_none());
+    }
+
+    ///////////////
+    // Lifecycle //
+    ///////////////
+
+    #[test]
+    fn acquire_write_publish_read_roundtrip() {
+        let s = store(4, 1).unwrap();
+
+        let reader = reader(&s);
+
+        let idx = {
+            let mut slot = s.acquire().expect("a fresh store has free slots");
+            let idx = slot.slot();
+            slot.data().set(10);
+
+            // Before the slot is dropped - we should not be able to read it.
+            assert!(reader.read(idx).is_none());
+            assert!(!s.can_read_approximate(idx.into_usize()).unwrap());
+            slot.publish();
+            idx
+        };
+
+        assert_eq!(reader.read(idx).unwrap().get(), 10,);
+        assert!(s.can_read_approximate(idx.into_usize()).unwrap());
+    }
+
+    #[test]
+    fn unpublished_slots_are_immediately_available() {
+        let s = store(4, 1).unwrap();
+
+        let reader = reader(&s);
+
+        let idx = {
+            let mut slot = s.acquire().expect("a fresh store has free slots");
+            let idx = slot.slot();
+            slot.data().set(100);
+
+            // Before the slot is dropped - we should not be able to read it.
+            assert!(reader.read(idx).is_none());
+            assert!(!s.can_read_approximate(idx.into_usize()).unwrap());
+
+            // NOTE: We do not explicitly publish the slot.
+            idx
+        };
+
+        assert!(reader.read(idx).is_none());
+        assert!(!s.can_read_approximate(idx.into_usize()).unwrap());
+    }
+
+    #[test]
+    fn acquire_exhausts_then_reports_none() {
+        let s = store(2, 1).unwrap();
+        // Hold the guards so the slots stay owned.
+        let _a = s.acquire().expect("first writable slot");
+        let _b = s.acquire().expect("second writable slot");
+        assert!(
+            s.acquire().is_none(),
+            "all writable slots are owned, so acquire must fail"
+        );
+    }
+
+    //--------//
+    // Retire //
+    //--------//
+
+    #[test]
+    fn retire_out_of_bounds() {
+        let s = store(4, 1).unwrap();
+        assert!(matches!(s.retire(999), Err(RetireError::OutOfBounds)));
+    }
+
+    #[test]
+    fn retire_rejects_reserved_slots() {
+        let s = store(4, 1).unwrap();
+        // An untouched writable slot is AVAILABLE, which is a reserved state.
+        assert!(matches!(
+            s.retire(0),
+            Err(RetireError::SlotIsReserved { .. })
+        ));
+        // A frozen slot is likewise reserved.
+        let frozen = s.frozen().start as usize;
+        assert!(matches!(
+            s.retire(frozen),
+            Err(RetireError::SlotIsReserved { .. })
+        ));
+        // An owned slot is not retirable.
+        let slot = s.acquire().unwrap();
+        assert!(matches!(
+            s.retire(slot.slot() as usize),
+            Err(RetireError::SlotIsReserved { .. })
+        ));
+    }
+
+    #[test]
+    fn retire_published_slot_then_unreadable() {
+        let s = store(4, 1).unwrap();
+
+        let idx = {
+            let mut slot = s.acquire().unwrap();
+            slot.data().set(101);
+            slot.publish()
+        };
+
+        assert!(s.retire(idx.into_usize()).is_ok());
+
+        // A reader opened after retirement must not observe the retired slot.
+        let reader = reader(&s);
+        assert_matches!(reader.read(idx), None);
+
+        // The slot can also not be retired again.
+        assert!(matches!(
+            s.retire(idx.into_usize()),
+            Err(RetireError::SlotIsReserved { .. })
+        ));
+    }
+
+    //---------//
+    // Recycle //
+    //---------//
+
+    #[test]
+    fn test_recycling() {
+        let entries = if cfg!(miri) { 16 } else { 2048 };
+
+        let s = store(entries, 2).unwrap();
+
+        assert_eq!(s.writable().len(), entries);
+
+        // Claim all slots.
+        let mut count = 0;
+        while let Some(mut slot) = s.acquire() {
+            slot.data().set(count as u64);
+            slot.publish();
+            count += 1;
+        }
+
+        assert_eq!(count, entries);
+
+        {
+            let reader = reader(&s);
+            for i in 0..entries {
+                assert_eq!(reader.read(i as u32).unwrap().get(), i as u64);
+            }
+        }
+
+        // Now that all slots are claimed - retire all slots.
+        for i in s.writable() {
+            s.retire(i.into_usize()).unwrap();
+        }
+
+        // Verify that we can claim all slots again.
+        let mut count = 0;
+        while let Some(mut slot) = s.acquire() {
+            slot.data().set(count as u64);
+            slot.publish();
+            count += 1;
+        }
+
+        assert_eq!(count, entries);
+    }
+}
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index b3afb1fe3d..e700bee8eb 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -3,6 +3,71 @@
  * Licensed under the MIT license.
  */
 
+//! # EBR lifecycle hooks for [`super::Store`].
+//!
+//! Please read this section carefully - the protocol is not difficult, but it *is* subtle.
+//!
+//! The transitions are a simplified version of the protocol described in [`crate::tag`]
+//! that storage plugins need to implement to be compatible. A state diagram is shown below:
+//!
+//! ```text
+//!         +--------------- `reclaim` ----------------+
+//!         |                                          |
+//!         V                                          |
+//!   +-----------+                               +----------+
+//!   | Available |<---+                          | Retiring |
+//!   +-----------+    |                          +----------+
+//!         |          |                               ^
+//!         |          |                               |
+//!         |       `abort`                         `retire`
+//!     `acquire`      |                               |
+//!         |          |                               |
+//!         |      +----------+                   +-----------+
+//!         +----->| Slot<'_> |---- `publish` --->| Published |
+//!                +----------+                   +-----------+
+//!                    |
+//!                 `freeze`
+//!                    |
+//!        +-----------+
+//!        |
+//!        V
+//!    +--------+
+//!    | Frozen |
+//!    +--------+
+//! ```
+//!
+//! ## Readable States
+//!
+//! * `published`: **New** references to slots may be given out in the "published" state. It is
+//!   possible for a transition to go from "published" to "retiring" while references are lent
+//!   out. This is fine as long as the lifetime of these references is bounded by a
+//!   [`crate::epoch::Guard`]. Using [`super::Store::guard`] will provide such a guard.
+//!
+//! * `frozen`: Since "frozen" is a terminal state, it is safe to give out references to
+//!   frozen slots.
+//!
+//! ## Writable States
+//!
+//! * [`Slot`]: Slots are a little spooky. Plugins can assume that a [`Slot`] for an index
+//!   `i` is exclusive for its duration. This means that [`Slot`] implementations can lend
+//!   out mutable references to its contents (for example, [`invasive::Slot::as_mut_slice`]).
+//!
+//!   Code in [`super`] is very careful to maintain this invariant and all users of [`Slot`]
+//!   must carefully maintain this as well.
+//!
+//! * `reclaim`: On a call to [`Plugin::reclaim`], it can be assumed that the plugin has
+//!   exclusive access to the indicated slot for the duration of the function call.
+//!
+//! ## Contracts
+//!
+//! Users of [`Plugin`] must ensure that the lifecycle shown above is strictly observed.
+//! Furthermore, for [`Slot`]s, exactly one of the terminal methods **must** be called.
+//!
+//! State transitions are driven by the authoritative [`super::Store`]. Before invoking a
+//! plugin transition, the store ensures the slot is not externally available in its previous
+//! state. Further, the store commits the destination state only after the plugin API call
+//! completes.
+
 use diskann::ANNResult;
 use std::fmt::Debug;
 
@@ -10,23 +75,82 @@ use crate::num::IdLimit;
 
 use super::Lifecycle;
 
+/// A configuration for a [`Plugin`].
 pub(crate) trait PluginConfig: Debug {
+    /// The type of the resulting [`Plugin`].
     type Plugin: Plugin;
+
+    /// Build the associated [`Plugin`] from self with the [`IdLimit`].
     fn build(self, id_limit: IdLimit) -> ANNResult<Self::Plugin>;
 }
 
+/// A lifecycle backend for [`super::Store`]'s EBR scheme.
+///
+/// See the [module level documentation](self) for details.
 pub(crate) trait Plugin: Debug + 'static {
+    /// The writable [`Slot`] for this plugin.
     type Slot<'a>: Slot;
+
+    /// Return the exclusive upper bound for indices provided to this API.
+    ///
+    /// Callers should ensure that indices are in the range `[0..self.id_limit())`.
     fn id_limit(&self) -> IdLimit;
 
+    /// Immediately transition slot `i` from the "available" state to the "slot" state.
+    ///
+    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
+    /// `i < Self::id_limit` for memory safety.
+    ///
+    /// # Safety
+    ///
+    /// Callers must ensure **all** of the following:
+    ///
+    /// 1. The plugin is in the implicit "available" state according to the [module docs](self).
+    ///
+    /// 2. Access to slot `i` is exclusive before invoking this method and that exclusivity
+    ///    is maintained until the returned [`Slot`] is consumed by a terminal method.
+    ///
+    /// 3. Exactly one of the [`Slot`] terminal methods is called. The [`Slot`] **may not**
+    ///    be dropped or forgotten without one of these methods being called.
     unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_>;
-    unsafe fn reclaim(&self, i: u32, _: Lifecycle);
+
+    /// Transition slot `i` from the "published" state to the "retiring" state.
+    ///
+    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
+    /// `i < Self::id_limit` for memory safety.
+    ///
+    /// # Safety
+    ///
+    /// The plugin is in the implicit "published" state.
     unsafe fn retire(&self, i: u32, _: Lifecycle);
+
+    /// Transition slot `i` from the "retiring" state to the "available" state.
+    ///
+    /// Implementations may panic when `i` is out-of-bounds, but must not rely on
+    /// `i < Self::id_limit` for memory safety.
+    ///
+    /// # Safety
+    ///
+    /// Callers must ensure **all** of the following:
+    ///
+    /// 1. The plugin is in the implicit "retiring" state.
+    ///
+    /// 2. All [`crate::epoch::Guard`]s for this [`Plugin`] that could have obtained a
+    ///    reference while this slot was in the "published" state have been dropped.
+    unsafe fn reclaim(&self, i: u32, _: Lifecycle);
 }
 
+/// A writable slot for [`Plugin`].
+///
+/// [`Slot`]s may assume that they have exclusive ownership of their plugin slots for their
+/// duration in accordance with [`Plugin::acquire`].
 pub(crate) trait Slot: Debug {
+    /// Mark this slot as readable, transition it to the "published" state.
     fn publish(self, _: Lifecycle);
+
+    /// Mark this slot as "frozen".
     fn freeze(self, _: Lifecycle);
+
+    /// Abort any action, returning the slot to "available".
     fn abort(self, _: Lifecycle);
 }
-

From 8f717d738f9825d30273b8a5d86f7c4a3bbe92ec Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 20 Aug 2026 17:26:50 -0700
Subject: [PATCH 18/34] checkpoint

---
 .../integration/jsons/store-stress-test.json  |  24 +-
 .../integration/jsons/store-stress.json       |  41 +-
 diskann-inmem/integration/main.rs             |   2 +-
 diskann-inmem/integration/store.rs            | 597 ------------------
 diskann-inmem/integration/store/checked.rs    | 222 +++++++
 diskann-inmem/integration/store/invasive.rs   | 284 +++++++++
 diskann-inmem/integration/store/mod.rs        | 492 +++++++++++++++
 diskann-inmem/src/integration/store.rs        | 133 ----
 .../src/integration/store/checked.rs          |  93 +++
 .../src/integration/store/invasive.rs         |  80 +++
 diskann-inmem/src/integration/store/mod.rs    |  87 +++
 diskann-inmem/src/store/checked.rs            | 252 ++++----
 diskann-inmem/src/store/invasive.rs           |  17 +-
 diskann-inmem/src/store/mod.rs                |  23 +-
 14 files changed, 1472 insertions(+), 875 deletions(-)
 delete mode 100644 diskann-inmem/integration/store.rs
 create mode 100644 diskann-inmem/integration/store/checked.rs
 create mode 100644 diskann-inmem/integration/store/invasive.rs
 create mode 100644 diskann-inmem/integration/store/mod.rs
 delete mode 100644 diskann-inmem/src/integration/store.rs
 create mode 100644 diskann-inmem/src/integration/store/checked.rs
 create mode 100644 diskann-inmem/src/integration/store/invasive.rs
 create mode 100644 diskann-inmem/src/integration/store/mod.rs

diff --git a/diskann-inmem/integration/jsons/store-stress-test.json b/diskann-inmem/integration/jsons/store-stress-test.json
index 5a1e8351f4..fed64e095f 100644
--- a/diskann-inmem/integration/jsons/store-stress-test.json
+++ b/diskann-inmem/integration/jsons/store-stress-test.json
@@ -2,19 +2,21 @@
   "search_directories": [],
   "jobs": [
     {
-      "type": "store-stress",
+      "type": "store-stress-invasive",
       "content": {
-        "readers": 4,
-        "writers": 2,
-        "retirers": 1,
-        "capacity": 512,
         "entry_bytes": 64,
-        "epoch_guard_slots": 64,
-        "freelist_recycle_capacity": 32,
-        "low_watermark": 128,
-        "duration_secs": 2,
-        "max_ops": 2000000,
-        "seed": 11939873485092837375
+        "setup": {
+          "readers": 4,
+          "writers": 2,
+          "retirers": 1,
+          "capacity": 512,
+          "epoch_guard_slots": 64,
+          "freelist_recycle_capacity": 32,
+          "low_watermark": 128,
+          "duration_secs": 2,
+          "max_ops": 2000000,
+          "seed": 11939873485092837375
+        }
       }
     }
   ]
diff --git a/diskann-inmem/integration/jsons/store-stress.json b/diskann-inmem/integration/jsons/store-stress.json
index 89584764b5..4fdce96604 100644
--- a/diskann-inmem/integration/jsons/store-stress.json
+++ b/diskann-inmem/integration/jsons/store-stress.json
@@ -2,19 +2,38 @@
   "search_directories": [],
   "jobs": [
     {
-      "type": "store-stress",
+      "type": "store-stress-checked",
+      "content": {
+        "setup": {
+          "readers": 8,
+          "writers": 4,
+          "retirers": 2,
+          "capacity": 4096,
+          "epoch_guard_slots": 64,
+          "freelist_recycle_capacity": 512,
+          "low_watermark": 1024,
+          "duration_secs": 10,
+          "max_ops": 50000000,
+          "seed": 11939873485092837375
+        }
+      }
+    },
+    {
+      "type": "store-stress-invasive",
       "content": {
-        "readers": 8,
-        "writers": 4,
-        "retirers": 2,
-        "capacity": 4096,
         "entry_bytes": 128,
-        "epoch_guard_slots": 64,
-        "freelist_recycle_capacity": 512,
-        "low_watermark": 1024,
-        "duration_secs": 10,
-        "max_ops": 50000000,
-        "seed": 11939873485092837375
+        "setup": {
+          "readers": 8,
+          "writers": 4,
+          "retirers": 2,
+          "capacity": 4096,
+          "epoch_guard_slots": 64,
+          "freelist_recycle_capacity": 512,
+          "low_watermark": 1024,
+          "duration_secs": 10,
+          "max_ops": 50000000,
+          "seed": 11939873485092837375
+        }
       }
     }
   ]
diff --git a/diskann-inmem/integration/main.rs b/diskann-inmem/integration/main.rs
index d3667b6842..9a9d0be453 100644
--- a/diskann-inmem/integration/main.rs
+++ b/diskann-inmem/integration/main.rs
@@ -12,7 +12,7 @@ use diskann_benchmark_runner::{App, Registry, output};
 /// Build a [`Registry`] with all integration benchmarks registered.
 fn registry() -> anyhow::Result<Registry> {
     let mut registry = Registry::new();
-    registry.register("store-stress", store::StoreStress)?;
+    store::register(&mut registry)?;
     index::register(&mut registry)?;
     Ok(registry)
 }
diff --git a/diskann-inmem/integration/store.rs b/diskann-inmem/integration/store.rs
deleted file mode 100644
index ec4e31f4c8..0000000000
--- a/diskann-inmem/integration/store.rs
+++ /dev/null
@@ -1,597 +0,0 @@
-/*
- * Copyright (c) Microsoft Corporation.
- * Licensed under the MIT license.
- */
-
-//! Concurrency stress test for the in-memory [`Store`](diskann_inmem::integration::store::Store).
-//!
-//! Reader, writer, and retirer threads hammer the epoch-based store concurrently while a
-//! per-guard invariant checker verifies the store's safety guarantees:
-//!
-//! 1. Reads are never torn.
-//! 2. A readable value is stable for the lifetime of a single reader guard.
-//! 3. A slot never resurrects (`readable -> unreadable -> readable`) within one guard.
-
-#![expect(
-    clippy::unwrap_used,
-    reason = "this code works mainly as an integration test"
-)]
-
-use std::{
-    collections::HashMap,
-    io::Write,
-    sync::{
-        Mutex,
-        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering::Relaxed},
-    },
-    time::{Duration, Instant},
-};
-
-use diskann_benchmark_runner::{
-    Benchmark, Checker, Checkpoint, Input, Output,
-    benchmark::{MatchContext, Score},
-    utils::fmt::KeyValue,
-};
-use rand::{Rng, SeedableRng, distr::Uniform, rngs::StdRng};
-use serde::{Deserialize, Serialize};
-
-use diskann_inmem::integration::store::{Config, Store};
-
-/// Maximum number of concurrent reader guards supported by the epoch registry.
-const GUARD_CAPACITY: usize = 256;
-
-/// Number of slots a reader inspects per guard. Kept small so guards are short-lived,
-/// allowing the epoch to advance and reclamation to make progress.
-const READER_WINDOW: usize = 64;
-
-/// Number of times a reader re-reads its window within a single guard. Re-reading is what
-/// exercises the value-stability and no-resurrection invariants.
-const READER_PASSES: usize = 4;
-
-/// How often (in retirer iterations) a retirer attempts to reclaim retired slots.
-const RECLAIM_EVERY: u64 = 16;
-
-///////////
-// Input //
-///////////
-
-/// Configuration for a [`StoreStress`] run.
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct StoreStressInput {
-    /// Number of reader threads. Must be below [`GUARD_CAPACITY`].
-    readers: usize,
-    /// Number of writer threads.
-    writers: usize,
-    /// Number of retirer threads.
-    retirers: usize,
-    /// Number of writable (non-frozen) slots.
-    capacity: usize,
-    /// Bytes per entry. Must be a non-zero multiple of 8 (the stamp lane width).
-    entry_bytes: usize,
-    /// The number of epoch guard slots.
-    epoch_guard_slots: usize,
-    /// The capacity of the freelist recycle queue capacity.
-    freelist_recycle_capacity: usize,
-    /// Retirers only retire while the live published population exceeds this watermark.
-    low_watermark: usize,
-    /// Wall-clock cap for the run, in seconds. Zero means unbounded (rely on `max_ops`).
-    duration_secs: u64,
-    /// Total-operation cap across all worker threads. Zero means unbounded (rely on
-    /// `duration_secs`).
-    max_ops: u64,
-    /// Seed for the worker pseudo-random number generators.
-    seed: u64,
-}
-
-impl StoreStressInput {
-    fn check(self) -> anyhow::Result<Self> {
-        if self.readers == 0 || self.writers == 0 {
-            anyhow::bail!("`readers` and `writers` must be non-zero");
-        }
-        if self.readers >= GUARD_CAPACITY {
-            anyhow::bail!(
-                "`readers` ({}) must be below the epoch guard capacity ({GUARD_CAPACITY})",
-                self.readers,
-            );
-        }
-        if self.capacity == 0 {
-            anyhow::bail!("`capacity` must be non-zero");
-        }
-        if self.entry_bytes == 0 || !self.entry_bytes.is_multiple_of(8) {
-            anyhow::bail!(
-                "`entry_bytes` ({}) must be a non-zero multiple of 8",
-                self.entry_bytes,
-            );
-        }
-        if self.low_watermark > self.capacity {
-            anyhow::bail!(
-                "`low_watermark` ({}) must not exceed `capacity` ({})",
-                self.low_watermark,
-                self.capacity,
-            );
-        }
-        if self.duration_secs == 0 && self.max_ops == 0 {
-            anyhow::bail!("at least one of `duration_secs` or `max_ops` must be non-zero");
-        }
-        Ok(self)
-    }
-}
-
-impl Input for StoreStressInput {
-    type Raw = Self;
-
-    fn tag() -> &'static str {
-        "store-stress"
-    }
-
-    fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result<Self> {
-        Self::check(raw)
-    }
-
-    fn serialize(&self) -> anyhow::Result<serde_json::Value> {
-        Ok(serde_json::to_value(self)?)
-    }
-
-    fn example() -> Self::Raw {
-        StoreStressInput {
-            readers: 8,
-            writers: 4,
-            retirers: 2,
-            capacity: 4096,
-            entry_bytes: 128,
-            epoch_guard_slots: 256,
-            freelist_recycle_capacity: 1024,
-            low_watermark: 1024,
-            duration_secs: 5,
-            max_ops: 50_000_000,
-            seed: 0xA5A5_1234_DEAD_BEEF,
-        }
-    }
-}
-
-impl std::fmt::Display for StoreStressInput {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let mut kv = KeyValue::new();
-        kv.push("readers", &self.readers);
-        kv.push("writers", &self.writers);
-        kv.push("retirers", &self.retirers);
-        kv.push("capacity", &self.capacity);
-        kv.push("entry_bytes", &self.entry_bytes);
-        kv.push("epoch_guard_slots", &self.epoch_guard_slots);
-        kv.push("freelist_recycle_capacity", &self.freelist_recycle_capacity);
-        kv.push("low_watermark", &self.low_watermark);
-        kv.push("duration_secs", &self.duration_secs);
-        kv.push("max_ops", &self.max_ops);
-        kv.push("seed", &self.seed);
-        write!(f, "{}", kv)
-    }
-}
-
-////////////
-// Output //
-////////////
-
-/// Summary statistics produced by a [`StoreStress`] run.
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct StoreStressStats {
-    elapsed_secs: f64,
-    reads: u64,
-    acquires_ok: u64,
-    acquires_fail: u64,
-    retires_ok: u64,
-    retires_fail: u64,
-    reclaims: u64,
-    /// Observed `readable -> unreadable` transitions across all reader guards.
-    transitions: u64,
-    /// Peak observed live (published, not-yet-retired) population.
-    peak_live: usize,
-}
-
-impl std::fmt::Display for StoreStressStats {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let mut kv = KeyValue::new();
-        kv.push("elapsed_secs", &self.elapsed_secs);
-        kv.push("reads", &self.reads);
-        kv.push("acquires_ok", &self.acquires_ok);
-        kv.push("acquires_fail", &self.acquires_fail);
-        kv.push("retires_ok", &self.retires_ok);
-        kv.push("retires_fail", &self.retires_fail);
-        kv.push("reclaims", &self.reclaims);
-        kv.push("transitions", &self.transitions);
-        kv.push("peak_live", &self.peak_live);
-        write!(f, "{}", kv)
-    }
-}
-
-/////////////
-// Payload //
-/////////////
-
-/// Fill `buf` with `stamp` replicated across every 8-byte lane.
-fn write_stamp(buf: &mut [u8], stamp: u64) {
-    let bytes = stamp.to_ne_bytes();
-    for lane in buf.chunks_exact_mut(8) {
-        lane.copy_from_slice(&bytes);
-    }
-}
-
-/// Read the stamp from `buf`, returning `Err` if any 8-byte lane disagrees (a torn read).
-fn read_stamp(buf: &[u8]) -> Result<u64, ()> {
-    let (lanes, _) = buf.as_chunks::<8>();
-    let mut lanes = lanes.iter();
-    let first = u64::from_ne_bytes(*lanes.next().ok_or(())?);
-    for lane in lanes {
-        if u64::from_ne_bytes(*lane) != first {
-            return Err(());
-        }
-    }
-    Ok(first)
-}
-
-////////////////
-// Invariants //
-////////////////
-
-/// Per-guard observation of a single slot.
-#[derive(Debug, Clone, Copy)]
-enum SlotObservations {
-    /// The slot was observed readable with the given stamp.
-    Readable(u64),
-    /// The slot was observed readable and then became unreadable (retired).
-    Retired,
-}
-
-/// Feed a single observation of slot `i` into the per-guard checker, recording a violation
-/// on the shared state if a safety invariant is broken.
-fn observe(
-    shared: &Shared,
-    observed: &mut HashMap<usize, SlotObservations>,
-    i: usize,
-    read: Option<&[u8]>,
-) {
-    match (observed.get(&i).copied(), read) {
-        // Not yet observed readable; an unreadable slot tells us nothing actionable.
-        (None, None) => {}
-        // First readable observation: record the stamp (after a tearing check).
-        (None, Some(bytes)) => match read_stamp(bytes) {
-            Ok(stamp) => {
-                observed.insert(i, SlotObservations::Readable(stamp));
-            }
-            Err(()) => record_violation(shared, format!("torn read at slot {i}")),
-        },
-        // Still readable: the value must be identical and untorn.
-        (Some(SlotObservations::Readable(prev)), Some(bytes)) => match read_stamp(bytes) {
-            Ok(stamp) if stamp != prev => record_violation(
-                shared,
-                format!("slot {i} value changed within guard: {prev} -> {stamp}"),
-            ),
-            Ok(_) => {}
-            Err(()) => record_violation(shared, format!("torn read at slot {i}")),
-        },
-        // Readable -> unreadable: an allowed, terminal transition.
-        (Some(SlotObservations::Readable(_)), None) => {
-            observed.insert(i, SlotObservations::Retired);
-            shared.transitions.fetch_add(1, Relaxed);
-        }
-        // Resurrection: a slot that retired came back to life within the same guard.
-        (Some(SlotObservations::Retired), Some(_)) => record_violation(
-            shared,
-            format!("resurrection at slot {i}: unreadable -> readable within one guard"),
-        ),
-        (Some(SlotObservations::Retired), None) => {}
-    }
-}
-
-////////////
-// Shared //
-////////////
-
-struct Local<'a> {
-    counter: u64,
-    parent: &'a AtomicU64,
-}
-
-impl<'a> Local<'a> {
-    fn new(parent: &'a AtomicU64) -> Self {
-        Self { counter: 0, parent }
-    }
-
-    fn add(&mut self, by: u64) {
-        self.counter += by;
-
-        if self.counter >= 2048 {
-            self.parent.fetch_add(self.counter, Relaxed);
-            self.counter = 0;
-        }
-    }
-}
-
-impl Drop for Local<'_> {
-    fn drop(&mut self) {
-        self.parent.fetch_add(self.counter, Relaxed);
-    }
-}
-
-struct LocalMax<'a> {
-    max: usize,
-    parent: &'a AtomicUsize,
-}
-
-impl<'a> LocalMax<'a> {
-    fn new(parent: &'a AtomicUsize) -> Self {
-        Self { max: 0, parent }
-    }
-
-    fn max(&mut self, m: usize) {
-        self.max = self.max.max(m);
-    }
-}
-
-impl Drop for LocalMax<'_> {
-    fn drop(&mut self) {
-        self.parent.fetch_max(self.max, Relaxed);
-    }
-}
-
-/// State shared by all worker threads for the duration of a run.
-struct Shared {
-    store: Store,
-    slots: usize,
-    readable: Uniform<usize>,
-    writable: Uniform<usize>,
-    low_watermark: usize,
-    max_ops: u64,
-    deadline: Instant,
-
-    stop: AtomicBool,
-    violation: Mutex<Vec<String>>,
-
-    stamp: AtomicU64,
-    live: AtomicUsize,
-    peak_live: AtomicUsize,
-
-    ops: AtomicU64,
-    reads: AtomicU64,
-    acquires_ok: AtomicU64,
-    acquires_fail: AtomicU64,
-    retires_ok: AtomicU64,
-    retires_fail: AtomicU64,
-    reclaims: AtomicU64,
-    transitions: AtomicU64,
-}
-
-/// Record an observed invariant violation and signal all workers to stop.
-fn record_violation(shared: &Shared, message: String) {
-    let mut slot = shared.violation.lock().unwrap();
-    slot.push(message);
-    shared.stop.store(true, Relaxed);
-}
-
-/// Return `true` once any termination condition is met.
-fn should_stop(shared: &Shared) -> bool {
-    shared.stop.load(Relaxed)
-        || shared.ops.load(Relaxed) >= shared.max_ops
-        || Instant::now() >= shared.deadline
-}
-
-/////////////
-// Workers //
-/////////////
-
-fn writer(shared: &Shared) {
-    let mut ops = Local::new(&shared.ops);
-    let mut acquires_ok = Local::new(&shared.acquires_ok);
-    let mut acquires_fail = Local::new(&shared.acquires_fail);
-
-    let mut peak_live = LocalMax::new(&shared.peak_live);
-
-    while !should_stop(shared) {
-        ops.add(1);
-        match shared.store.acquire() {
-            Some(mut writer) => {
-                let stamp = shared.stamp.fetch_add(1, Relaxed);
-                write_stamp(writer.as_mut_slice(), stamp);
-                writer.publish();
-
-                let live = shared.live.fetch_add(1, Relaxed) + 1;
-                peak_live.max(live);
-                acquires_ok.add(1);
-            }
-            None => {
-                acquires_fail.add(1);
-                std::thread::yield_now();
-            }
-        }
-    }
-}
-
-fn retirer(shared: &Shared, seed: u64) {
-    let mut rng = StdRng::seed_from_u64(seed);
-    let mut iteration: u64 = 0;
-
-    let mut ops = Local::new(&shared.ops);
-    let mut retires_ok = Local::new(&shared.retires_ok);
-    let mut retires_fail = Local::new(&shared.retires_fail);
-    let mut reclaims = Local::new(&shared.reclaims);
-
-    while !should_stop(shared) {
-        ops.add(1);
-        iteration += 1;
-
-        // Flow control: keep a steady readable population.
-        if shared.live.load(Relaxed) > shared.low_watermark {
-            let i = rng.sample(shared.writable);
-            if shared.store.retire(i) {
-                shared.live.fetch_sub(1, Relaxed);
-                retires_ok.add(1);
-            } else {
-                retires_fail.add(1);
-            }
-        }
-
-        if iteration.is_multiple_of(RECLAIM_EVERY)
-            && let Some(reclaimed) = shared.store.reclaim()
-        {
-            reclaims.add(reclaimed as u64);
-        }
-
-        std::thread::yield_now();
-    }
-}
-
-fn reader(shared: &Shared, seed: u64) {
-    let mut rng = StdRng::seed_from_u64(seed);
-    let slots = shared.slots;
-    let window = READER_WINDOW.min(slots);
-    let mut observations = HashMap::with_capacity(window);
-
-    let mut ops = Local::new(&shared.ops);
-    let mut reads = Local::new(&shared.reads);
-
-    while !should_stop(shared) {
-        ops.add(1);
-        let Some(guard) = shared.store.reader() else {
-            // All guard slots are occupied; back off and retry.
-            std::thread::yield_now();
-            continue;
-        };
-
-        observations.clear();
-        let start = rng.sample(shared.readable);
-        for _ in 0..READER_PASSES {
-            for k in 0..window {
-                let i = (start + k) % slots;
-                observe(shared, &mut observations, i, guard.read(i));
-                reads.add(1);
-            }
-        }
-    }
-}
-
-///////////////
-// Benchmark //
-///////////////
-
-/// The store concurrency stress benchmark.
-#[derive(Debug)]
-pub struct StoreStress;
-
-impl Benchmark for StoreStress {
-    type Input = StoreStressInput;
-    type Output = StoreStressStats;
-
-    fn try_match(&self, _input: &StoreStressInput, context: &MatchContext) -> Score {
-        context.success(0)
-    }
-
-    fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "concurrency stress test for the in-memory store (readers/writers/retirers)"
-        )
-    }
-
-    fn run(
-        &self,
-        input: &StoreStressInput,
-        _checkpoint: Checkpoint<'_>,
-        mut output: &mut dyn Output,
-    ) -> anyhow::Result<StoreStressStats> {
-        let config = Config {
-            capacity: input.capacity,
-            entry_bytes: input.entry_bytes,
-            epoch_guard_slots: input.epoch_guard_slots,
-            freelist_recycle_capacity: input.freelist_recycle_capacity,
-        };
-
-        let store = Store::new(config);
-        let writable = store.writable();
-        let slots = store.slots();
-        let start = Instant::now();
-
-        let shared = Shared {
-            store,
-            slots,
-            readable: Uniform::new(0, slots)?,
-            writable: Uniform::try_from(writable)?,
-            low_watermark: input.low_watermark,
-            max_ops: if input.max_ops == 0 {
-                u64::MAX
-            } else {
-                input.max_ops
-            },
-            deadline: if input.duration_secs == 0 {
-                // Effectively unbounded; the op cap terminates the run.
-                start + Duration::from_secs(u64::from(u32::MAX))
-            } else {
-                start + Duration::from_secs(input.duration_secs)
-            },
-            stop: AtomicBool::new(false),
-            violation: Mutex::new(Vec::new()),
-            // Stamp 0 is reserved for the zeroed frozen point.
-            stamp: AtomicU64::new(1),
-            live: AtomicUsize::new(0),
-            peak_live: AtomicUsize::new(0),
-            ops: AtomicU64::new(0),
-            reads: AtomicU64::new(0),
-            acquires_ok: AtomicU64::new(0),
-            acquires_fail: AtomicU64::new(0),
-            retires_ok: AtomicU64::new(0),
-            retires_fail: AtomicU64::new(0),
-            reclaims: AtomicU64::new(0),
-            transitions: AtomicU64::new(0),
-        };
-
-        writeln!(output, "{}", input)?;
-
-        std::thread::scope(|scope| {
-            let shared = &shared;
-            for _ in 0..input.writers {
-                scope.spawn(move || writer(shared));
-            }
-            for t in 0..input.retirers {
-                let seed = input.seed ^ (0x2000_0000 + t as u64);
-                scope.spawn(move || retirer(shared, seed));
-            }
-            for t in 0..input.readers {
-                let seed = input.seed ^ (0x4000_0000 + t as u64);
-                scope.spawn(move || reader(shared, seed));
-            }
-        });
-
-        let errors: Vec<_> = std::mem::take(&mut *shared.violation.lock().unwrap());
-        if !errors.is_empty() {
-            anyhow::bail!("invariants violated: {:?}", errors);
-        }
-
-        let elapsed = start.elapsed();
-        let stats = StoreStressStats {
-            elapsed_secs: elapsed.as_secs_f64(),
-            reads: shared.reads.load(Relaxed),
-            acquires_ok: shared.acquires_ok.load(Relaxed),
-            acquires_fail: shared.acquires_fail.load(Relaxed),
-            retires_ok: shared.retires_ok.load(Relaxed),
-            retires_fail: shared.retires_fail.load(Relaxed),
-            reclaims: shared.reclaims.load(Relaxed),
-            transitions: shared.transitions.load(Relaxed),
-            peak_live: shared.peak_live.load(Relaxed),
-        };
-
-        writeln!(output, "{}", stats)?;
-        Ok(stats)
-    }
-}
-
-///////////
-// Tests //
-///////////
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn make_sure_example_parses() {
-        let _ = StoreStressInput::check(StoreStressInput::example()).unwrap();
-    }
-}
diff --git a/diskann-inmem/integration/store/checked.rs b/diskann-inmem/integration/store/checked.rs
new file mode 100644
index 0000000000..9f074f9625
--- /dev/null
+++ b/diskann-inmem/integration/store/checked.rs
@@ -0,0 +1,222 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{collections::HashMap, io::Write, sync::atomic::Ordering::Relaxed};
+
+use diskann_benchmark_runner as dbr;
+use diskann_inmem::integration::store::checked;
+use serde::{Deserialize, Serialize};
+
+pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::RegistryError> {
+    registry.register("store-stress-test-checked", Stress)
+}
+
+/// Configuration for a [`StoreStress`] run.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Input {
+    /// Shared stress test setup.
+    setup: super::Setup,
+}
+
+impl Input {
+    fn check(self) -> anyhow::Result<Self> {
+        self.setup.check()?;
+        Ok(self)
+    }
+}
+
+impl dbr::Input for Input {
+    type Raw = Self;
+
+    fn tag() -> &'static str {
+        "store-stress-checked"
+    }
+
+    fn from_raw(raw: Self::Raw, _checker: &mut dbr::Checker) -> anyhow::Result<Self> {
+        raw.check()
+    }
+
+    fn serialize(&self) -> anyhow::Result<serde_json::Value> {
+        Ok(serde_json::to_value(self)?)
+    }
+
+    fn example() -> Self::Raw {
+        Input {
+            setup: super::Setup::example(),
+        }
+    }
+}
+
+impl std::fmt::Display for Input {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Input { setup } = self;
+
+        let mut kv = dbr::utils::fmt::KeyValue::new();
+        kv.push("setup", &setup);
+        write!(f, "{}", kv)
+    }
+}
+
+#[derive(Debug)]
+struct Stress;
+
+impl dbr::Benchmark for Stress {
+    type Input = Input;
+    type Output = super::StoreStressStats;
+
+    fn try_match(
+        &self,
+        _input: &Input,
+        context: &dbr::benchmark::MatchContext,
+    ) -> dbr::benchmark::Score {
+        context.success(0)
+    }
+
+    fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "concurrency stress test for the checked in-memory store")
+    }
+
+    fn run(
+        &self,
+        input: &Input,
+        _checkpoint: dbr::Checkpoint<'_>,
+        mut output: &mut dyn dbr::Output,
+    ) -> anyhow::Result<Self::Output> {
+        let config = checked::Config {
+            capacity: input.setup.capacity,
+            epoch_guard_slots: input.setup.epoch_guard_slots,
+            freelist_recycle_capacity: input.setup.freelist_recycle_capacity,
+        };
+
+        writeln!(output, "{}", input)?;
+        let stats = super::run_benchmark(checked::Store::new(config), &input.setup)?;
+        writeln!(output, "{}", stats)?;
+        Ok(stats)
+    }
+}
+
+impl super::Testable for checked::Store {
+    type Writer<'a> = checked::Writer<'a>;
+    type ReaderState<'a> = ReaderState<'a>;
+
+    fn writer(&self) -> Option<Self::Writer<'_>> {
+        <checked::Store>::acquire(self)
+    }
+
+    fn reader_state<'a>(
+        &'a self,
+        capacity_hint: usize,
+        shared: &'a super::Shared<Self>,
+    ) -> Self::ReaderState<'a> {
+        ReaderState {
+            store: self,
+            shared,
+            capacity_hint,
+        }
+    }
+
+    fn retire(&self, i: usize) -> bool {
+        <checked::Store>::retire(self, i)
+    }
+
+    fn reclaim(&self) -> Option<usize> {
+        <checked::Store>::reclaim(self)
+    }
+
+    fn readable_slots(&self) -> usize {
+        <checked::Store>::slots(self)
+    }
+
+    fn writable_slots(&self) -> usize {
+        <checked::Store>::writable(self)
+    }
+}
+
+impl super::Writer for checked::Writer<'_> {
+    fn write(mut self, stamp: u64) {
+        self.set(stamp);
+        self.publish();
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct ReaderState<'a> {
+    store: &'a checked::Store,
+    shared: &'a super::Shared<checked::Store>,
+    capacity_hint: usize,
+}
+
+impl super::ReaderState for ReaderState<'_> {
+    type Reader<'a> = Reader<'a>;
+
+    fn try_with_reader<F>(&mut self, f: F) -> bool
+    where
+        F: FnOnce(Self::Reader<'_>),
+    {
+        let Some(reader) = self.store.reader() else { return false; };
+        let observed = HashMap::with_capacity(self.capacity_hint);
+
+        f(Reader {
+            reader: &reader,
+            observed,
+            shared: self.shared,
+        });
+
+        true
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct Reader<'a> {
+    observed: HashMap<usize, checked::Value<'a>>,
+    reader: &'a checked::Reader<'a>,
+    shared: &'a super::Shared<checked::Store>,
+}
+
+impl super::Reader for Reader<'_> {
+    /// Feed a single observation of slot `i` into the per-guard checker, recording a
+    /// violation on the shared state if a safety invariant is broken.
+    fn observe(&mut self, i: usize) {
+        let read = self.reader.read(i);
+        let observed = self.observed.get(&i).map(|v| v.get());
+
+        match (observed, read) {
+            // Not yet observed readable; an unreadable slot tells us nothing actionable.
+            (None, None) => {}
+            // First readable observation: record the stamp (after a tearing check).
+            (None, Some(value)) => {
+                self.observed.insert(i, value);
+            }
+            // Still readable: the value must be identical and untorn.
+            (Some(previous), Some(value)) => {
+                if previous != value.get() {
+                    self.shared.record_violation(format!(
+                        "slot {i} value changed within guard: {} -> {}",
+                        previous,
+                        value.get(),
+                    ))
+                }
+            }
+            // Readable -> unreadable: an allowed, terminal transition.
+            (Some(previous), None) => {
+                self.shared.transitions.fetch_add(1, Relaxed);
+            }
+        }
+    }
+}
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn make_sure_example_parses() {
+        let _ = Input::check(<Input as dbr::Input>::example()).unwrap();
+    }
+}
diff --git a/diskann-inmem/integration/store/invasive.rs b/diskann-inmem/integration/store/invasive.rs
new file mode 100644
index 0000000000..71e529af95
--- /dev/null
+++ b/diskann-inmem/integration/store/invasive.rs
@@ -0,0 +1,284 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use std::{collections::HashMap, io::Write, sync::atomic::Ordering::Relaxed};
+
+use diskann_benchmark_runner as dbr;
+use diskann_inmem::integration::store::invasive;
+use serde::{Deserialize, Serialize};
+
+pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::RegistryError> {
+    registry.register("invasive-store-stress-test", Stress)
+}
+
+/// Configuration for a [`StoreStress`] run.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Input {
+    /// Shared stress test setup.
+    setup: super::Setup,
+
+    /// Bytes per entry. Must be a non-zero multiple of 8 (the stamp lane width).
+    entry_bytes: usize,
+}
+
+impl Input {
+    fn check(self) -> anyhow::Result<Self> {
+        self.setup.check()?;
+
+        if self.entry_bytes == 0 || !self.entry_bytes.is_multiple_of(8) {
+            anyhow::bail!(
+                "`entry_bytes` ({}) must be a non-zero multiple of 8",
+                self.entry_bytes,
+            );
+        }
+
+        Ok(self)
+    }
+}
+
+impl dbr::Input for Input {
+    type Raw = Self;
+
+    fn tag() -> &'static str {
+        "store-stress-invasive"
+    }
+
+    fn from_raw(raw: Self::Raw, _checker: &mut dbr::Checker) -> anyhow::Result<Self> {
+        raw.check()
+    }
+
+    fn serialize(&self) -> anyhow::Result<serde_json::Value> {
+        Ok(serde_json::to_value(self)?)
+    }
+
+    fn example() -> Self::Raw {
+        Input {
+            setup: super::Setup::example(),
+            entry_bytes: 128,
+        }
+    }
+}
+
+impl std::fmt::Display for Input {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let Input { setup, entry_bytes } = self;
+
+        let mut kv = dbr::utils::fmt::KeyValue::new();
+        kv.push("setup", &setup);
+        kv.push("entry_bytes", &entry_bytes);
+        write!(f, "{}", kv)
+    }
+}
+
+#[derive(Debug)]
+struct Stress;
+
+impl dbr::Benchmark for Stress {
+    type Input = Input;
+    type Output = super::StoreStressStats;
+
+    fn try_match(
+        &self,
+        _input: &Input,
+        context: &dbr::benchmark::MatchContext,
+    ) -> dbr::benchmark::Score {
+        context.success(0)
+    }
+
+    fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "concurrency stress test for the invasive in-memory store"
+        )
+    }
+
+    fn run(
+        &self,
+        input: &Input,
+        _checkpoint: dbr::Checkpoint<'_>,
+        mut output: &mut dyn dbr::Output,
+    ) -> anyhow::Result<Self::Output> {
+        let config = invasive::Config {
+            capacity: input.setup.capacity,
+            entry_bytes: input.entry_bytes,
+            epoch_guard_slots: input.setup.epoch_guard_slots,
+            freelist_recycle_capacity: input.setup.freelist_recycle_capacity,
+        };
+
+        writeln!(output, "{}", input)?;
+        let stats = super::run_benchmark(invasive::Store::new(config), &input.setup)?;
+        writeln!(output, "{}", stats)?;
+        Ok(stats)
+    }
+}
+
+/// Per-guard observation of a single slot.
+#[derive(Debug, Clone, Copy)]
+enum SlotObservations {
+    /// The slot was observed readable with the given stamp.
+    Readable(u64),
+    /// The slot was observed readable and then became unreadable (retired).
+    Retired,
+}
+
+/// Fill `buf` with `stamp` replicated across every 8-byte lane.
+fn write_stamp(buf: &mut [u8], stamp: u64) {
+    let bytes = stamp.to_ne_bytes();
+    for lane in buf.chunks_exact_mut(8) {
+        lane.copy_from_slice(&bytes);
+    }
+}
+
+/// Read the stamp from `buf`, returning `Err` if any 8-byte lane disagrees (a torn read).
+fn read_stamp(buf: &[u8]) -> Result<u64, ()> {
+    let (lanes, _) = buf.as_chunks::<8>();
+    let mut lanes = lanes.iter();
+    let first = u64::from_ne_bytes(*lanes.next().ok_or(())?);
+    for lane in lanes {
+        if u64::from_ne_bytes(*lane) != first {
+            return Err(());
+        }
+    }
+    Ok(first)
+}
+
+impl super::Testable for invasive::Store {
+    type Writer<'a> = invasive::Writer<'a>;
+    type ReaderState<'a> = ReaderState<'a>;
+
+    fn writer(&self) -> Option<Self::Writer<'_>> {
+        <invasive::Store>::acquire(self)
+    }
+
+    fn reader_state<'a>(
+        &'a self,
+        capacity_hint: usize,
+        shared: &'a super::Shared<Self>,
+    ) -> Self::ReaderState<'a> {
+        let observed = HashMap::with_capacity(capacity_hint);
+        ReaderState {
+            store: self,
+            observed,
+            shared,
+        }
+    }
+
+    fn retire(&self, i: usize) -> bool {
+        <invasive::Store>::retire(self, i)
+    }
+
+    fn reclaim(&self) -> Option<usize> {
+        <invasive::Store>::reclaim(self)
+    }
+
+    fn readable_slots(&self) -> usize {
+        <invasive::Store>::slots(self)
+    }
+
+    fn writable_slots(&self) -> usize {
+        <invasive::Store>::writable(self)
+    }
+}
+
+impl super::Writer for invasive::Writer<'_> {
+    fn write(mut self, stamp: u64) {
+        write_stamp(self.as_mut_slice(), stamp);
+        self.publish();
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct ReaderState<'a> {
+    store: &'a invasive::Store,
+    observed: HashMap<usize, SlotObservations>,
+    shared: &'a super::Shared<invasive::Store>,
+}
+
+impl super::ReaderState for ReaderState<'_> {
+    type Reader<'a>
+        = Reader<'a>;
+
+    fn try_with_reader<F>(&mut self, f: F) -> bool
+    where
+        F: FnOnce(Self::Reader<'_>),
+    {
+        let Some(reader) = self.store.reader() else {
+            return false;
+        };
+        self.observed.clear();
+
+        f(Reader {
+            reader,
+            observed: &mut self.observed,
+            shared: self.shared,
+        });
+
+        true
+    }
+}
+
+#[derive(Debug)]
+pub(super) struct Reader<'a> {
+    reader: invasive::Reader<'a>,
+    observed: &'a mut HashMap<usize, SlotObservations>,
+    shared: &'a super::Shared<invasive::Store>,
+}
+
+impl super::Reader for Reader<'_> {
+    /// Feed a single observation of slot `i` into the per-guard checker, recording a
+    /// violation on the shared state if a safety invariant is broken.
+    fn observe(&mut self, i: usize) {
+        let read = self.reader.read(i);
+        let observed = self.observed.get(&i).copied();
+
+        match (observed, read) {
+            // Not yet observed readable; an unreadable slot tells us nothing actionable.
+            (None, None) => {}
+            // First readable observation: record the stamp (after a tearing check).
+            (None, Some(bytes)) => match read_stamp(bytes) {
+                Ok(stamp) => {
+                    self.observed.insert(i, SlotObservations::Readable(stamp));
+                }
+                Err(()) => self
+                    .shared
+                    .record_violation(format!("torn read at slot {i}")),
+            },
+            // Still readable: the value must be identical and untorn.
+            (Some(SlotObservations::Readable(prev)), Some(bytes)) => match read_stamp(bytes) {
+                Ok(stamp) if stamp != prev => self.shared.record_violation(format!(
+                    "slot {i} value changed within guard: {prev} -> {stamp}"
+                )),
+                Ok(_) => {}
+                Err(()) => self
+                    .shared
+                    .record_violation(format!("torn read at slot {i}")),
+            },
+            // Readable -> unreadable: an allowed, terminal transition.
+            (Some(SlotObservations::Readable(_)), None) => {
+                self.observed.insert(i, SlotObservations::Retired);
+                self.shared.transitions.fetch_add(1, Relaxed);
+            }
+            // Resurrection: a slot that retired came back to life within the same guard.
+            (Some(SlotObservations::Retired), Some(_)) => self.shared.record_violation(format!(
+                "resurrection at slot {i}: unreadable -> readable within one guard"
+            )),
+            (Some(SlotObservations::Retired), None) => {}
+        }
+    }
+}
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn make_sure_example_parses() {
+        let _ = Input::check(<Input as dbr::Input>::example()).unwrap();
+    }
+}
diff --git a/diskann-inmem/integration/store/mod.rs b/diskann-inmem/integration/store/mod.rs
new file mode 100644
index 0000000000..589ba99add
--- /dev/null
+++ b/diskann-inmem/integration/store/mod.rs
@@ -0,0 +1,492 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+//! Concurrency stress test for the in-memory [`Store`](diskann_inmem::integration::store::Store).
+//!
+//! Reader, writer, and retirer threads hammer the epoch-based store concurrently while a
+//! per-guard invariant checker verifies the store's safety guarantees:
+//!
+//! 1. Reads are never torn.
+//! 2. A readable value is stable for the lifetime of a single reader guard.
+//! 3. A slot never resurrects (`readable -> unreadable -> readable`) within one guard.
+
+#![expect(
+    clippy::unwrap_used,
+    reason = "this code works mainly as an integration test"
+)]
+
+use std::{
+    sync::{
+        Mutex,
+        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering::Relaxed},
+    },
+    time::{Duration, Instant},
+};
+
+use diskann_benchmark_runner::{Registry, RegistryError, utils::fmt::KeyValue};
+use rand::{Rng, SeedableRng, distr::Uniform, rngs::StdRng};
+use serde::{Deserialize, Serialize};
+
+// use diskann_inmem::integration::store::invasive::{self, Config, Store};
+
+/// Number of slots a reader inspects per guard. Kept small so guards are short-lived,
+/// allowing the epoch to advance and reclamation to make progress.
+const READER_WINDOW: usize = 64;
+
+/// Number of times a reader re-reads its window within a single guard. Re-reading is what
+/// exercises the value-stability and no-resurrection invariants.
+const READER_PASSES: usize = 4;
+
+/// How often (in retirer iterations) a retirer attempts to reclaim retired slots.
+const RECLAIM_EVERY: u64 = 16;
+
+mod checked;
+mod invasive;
+
+pub(super) fn register(registry: &mut Registry) -> Result<(), RegistryError> {
+    invasive::register(registry)?;
+    checked::register(registry)?;
+
+    Ok(())
+}
+
+///////////
+// Input //
+///////////
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Setup {
+    /// Number of reader threads. Must be below `epoch_guard_slots`.
+    readers: usize,
+    /// Number of writer threads.
+    writers: usize,
+    /// Number of retirer threads.
+    retirers: usize,
+    /// Number of writable (non-frozen) slots.
+    capacity: usize,
+    /// The number of epoch guard slots.
+    epoch_guard_slots: usize,
+    /// The capacity of the freelist recycle queue capacity.
+    freelist_recycle_capacity: usize,
+    /// Retirers only retire while the live published population exceeds this watermark.
+    low_watermark: usize,
+    /// Wall-clock cap for the run, in seconds. Zero means unbounded (rely on `max_ops`).
+    duration_secs: u64,
+    /// Total-operation cap across all worker threads. Zero means unbounded (rely on
+    /// `duration_secs`).
+    max_ops: u64,
+    /// Seed for the worker pseudo-random number generators.
+    seed: u64,
+}
+
+impl Setup {
+    fn check(&self) -> anyhow::Result<()> {
+        if self.readers == 0 || self.writers == 0 {
+            anyhow::bail!("`readers` and `writers` must be non-zero");
+        }
+        if self.readers >= self.epoch_guard_slots {
+            anyhow::bail!(
+                "`readers` ({}) must be below the epoch guard capacity ({})",
+                self.readers,
+                self.epoch_guard_slots,
+            );
+        }
+        if self.capacity == 0 {
+            anyhow::bail!("`capacity` must be non-zero");
+        }
+        if self.low_watermark > self.capacity {
+            anyhow::bail!(
+                "`low_watermark` ({}) must not exceed `capacity` ({})",
+                self.low_watermark,
+                self.capacity,
+            );
+        }
+        if self.duration_secs == 0 && self.max_ops == 0 {
+            anyhow::bail!("at least one of `duration_secs` or `max_ops` must be non-zero");
+        }
+
+        Ok(())
+    }
+
+    fn example() -> Self {
+        Setup {
+            readers: 8,
+            writers: 4,
+            retirers: 2,
+            capacity: 4096,
+            epoch_guard_slots: 256,
+            freelist_recycle_capacity: 1024,
+            low_watermark: 1024,
+            duration_secs: 5,
+            max_ops: 50_000_000,
+            seed: 0xA5A5_1234_DEAD_BEEF,
+        }
+    }
+}
+
+impl std::fmt::Display for Setup {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut kv = KeyValue::new();
+        kv.push("readers", &self.readers);
+        kv.push("writers", &self.writers);
+        kv.push("retirers", &self.retirers);
+        kv.push("capacity", &self.capacity);
+        kv.push("epoch_guard_slots", &self.epoch_guard_slots);
+        kv.push("freelist_recycle_capacity", &self.freelist_recycle_capacity);
+        kv.push("low_watermark", &self.low_watermark);
+        kv.push("duration_secs", &self.duration_secs);
+        kv.push("max_ops", &self.max_ops);
+        kv.push("seed", &self.seed);
+        write!(f, "{}", kv)
+    }
+}
+
+trait Testable: std::fmt::Debug + Sized + Sync {
+    type Writer<'a>: Writer
+    where
+        Self: 'a;
+
+    type ReaderState<'a>: ReaderState
+    where
+        Self: 'a;
+
+    fn writer(&self) -> Option<Self::Writer<'_>>;
+
+    fn reader_state<'a>(
+        &'a self,
+        capacity_hint: usize,
+        shared: &'a Shared<Self>,
+    ) -> Self::ReaderState<'a>;
+
+    fn retire(&self, i: usize) -> bool;
+
+    fn reclaim(&self) -> Option<usize>;
+
+    fn readable_slots(&self) -> usize;
+    fn writable_slots(&self) -> usize;
+}
+
+trait Writer: std::fmt::Debug {
+    fn write(self, stamp: u64);
+}
+
+trait ReaderState: std::fmt::Debug {
+    type Reader<'a>: Reader;
+
+    #[must_use]
+    fn try_with_reader<F>(&mut self, f: F) -> bool
+    where
+        F: FnOnce(Self::Reader<'_>);
+}
+
+trait Reader: std::fmt::Debug {
+    fn observe(&mut self, i: usize);
+}
+
+////////////
+// Output //
+////////////
+
+/// Summary statistics produced by a [`StoreStress`] run.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StoreStressStats {
+    elapsed_secs: f64,
+    reads: u64,
+    acquires_ok: u64,
+    acquires_fail: u64,
+    retires_ok: u64,
+    retires_fail: u64,
+    reclaims: u64,
+    /// Observed `readable -> unreadable` transitions across all reader guards.
+    transitions: u64,
+    /// Peak observed live (published, not-yet-retired) population.
+    peak_live: usize,
+}
+
+impl std::fmt::Display for StoreStressStats {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut kv = KeyValue::new();
+        kv.push("elapsed_secs", &self.elapsed_secs);
+        kv.push("reads", &self.reads);
+        kv.push("acquires_ok", &self.acquires_ok);
+        kv.push("acquires_fail", &self.acquires_fail);
+        kv.push("retires_ok", &self.retires_ok);
+        kv.push("retires_fail", &self.retires_fail);
+        kv.push("reclaims", &self.reclaims);
+        kv.push("transitions", &self.transitions);
+        kv.push("peak_live", &self.peak_live);
+        write!(f, "{}", kv)
+    }
+}
+
+////////////
+// Shared //
+////////////
+
+struct Local<'a> {
+    counter: u64,
+    parent: &'a AtomicU64,
+}
+
+impl<'a> Local<'a> {
+    fn new(parent: &'a AtomicU64) -> Self {
+        Self { counter: 0, parent }
+    }
+
+    fn add(&mut self, by: u64) {
+        self.counter += by;
+
+        if self.counter >= 2048 {
+            self.parent.fetch_add(self.counter, Relaxed);
+            self.counter = 0;
+        }
+    }
+}
+
+impl Drop for Local<'_> {
+    fn drop(&mut self) {
+        self.parent.fetch_add(self.counter, Relaxed);
+    }
+}
+
+struct LocalMax<'a> {
+    max: usize,
+    parent: &'a AtomicUsize,
+}
+
+impl<'a> LocalMax<'a> {
+    fn new(parent: &'a AtomicUsize) -> Self {
+        Self { max: 0, parent }
+    }
+
+    fn max(&mut self, m: usize) {
+        self.max = self.max.max(m);
+    }
+}
+
+impl Drop for LocalMax<'_> {
+    fn drop(&mut self) {
+        self.parent.fetch_max(self.max, Relaxed);
+    }
+}
+
+fn run_benchmark<T>(store: T, setup: &Setup) -> anyhow::Result<StoreStressStats>
+where
+    T: Testable,
+{
+    let writable = store.writable_slots();
+    let readable = store.readable_slots();
+    let start = Instant::now();
+
+    let shared = Shared {
+        store,
+        slots: readable,
+        readable: Uniform::new(0, readable)?,
+        writable: Uniform::new(0, writable)?,
+        low_watermark: setup.low_watermark,
+        max_ops: if setup.max_ops == 0 {
+            u64::MAX
+        } else {
+            setup.max_ops
+        },
+        deadline: if setup.duration_secs == 0 {
+            // Effectively unbounded; the op cap terminates the run.
+            start + Duration::from_secs(u64::from(u32::MAX))
+        } else {
+            start + Duration::from_secs(setup.duration_secs)
+        },
+        stop: AtomicBool::new(false),
+        violation: Mutex::new(Vec::new()),
+        // Stamp 0 is reserved for the zeroed frozen point.
+        stamp: AtomicU64::new(1),
+        live: AtomicUsize::new(0),
+        peak_live: AtomicUsize::new(0),
+        ops: AtomicU64::new(0),
+        reads: AtomicU64::new(0),
+        acquires_ok: AtomicU64::new(0),
+        acquires_fail: AtomicU64::new(0),
+        retires_ok: AtomicU64::new(0),
+        retires_fail: AtomicU64::new(0),
+        reclaims: AtomicU64::new(0),
+        transitions: AtomicU64::new(0),
+    };
+
+    std::thread::scope(|scope| {
+        let shared = &shared;
+        for _ in 0..setup.writers {
+            scope.spawn(move || shared.writer());
+        }
+        for t in 0..setup.retirers {
+            let seed = setup.seed ^ (0x2000_0000 + t as u64);
+            scope.spawn(move || shared.retirer(seed));
+        }
+        for t in 0..setup.readers {
+            let seed = setup.seed ^ (0x4000_0000 + t as u64);
+            scope.spawn(move || shared.reader(seed));
+        }
+    });
+
+    let errors: Vec<_> = std::mem::take(&mut *shared.violation.lock().unwrap());
+    if !errors.is_empty() {
+        anyhow::bail!("invariants violated: {:?}", errors);
+    }
+
+    let elapsed = start.elapsed();
+    let stats = StoreStressStats {
+        elapsed_secs: elapsed.as_secs_f64(),
+        reads: shared.reads.load(Relaxed),
+        acquires_ok: shared.acquires_ok.load(Relaxed),
+        acquires_fail: shared.acquires_fail.load(Relaxed),
+        retires_ok: shared.retires_ok.load(Relaxed),
+        retires_fail: shared.retires_fail.load(Relaxed),
+        reclaims: shared.reclaims.load(Relaxed),
+        transitions: shared.transitions.load(Relaxed),
+        peak_live: shared.peak_live.load(Relaxed),
+    };
+
+    Ok(stats)
+}
+
+/// State shared by all worker threads for the duration of a run.
+#[derive(Debug)]
+struct Shared<T> {
+    store: T,
+    slots: usize,
+    readable: Uniform<usize>,
+    writable: Uniform<usize>,
+    low_watermark: usize,
+    max_ops: u64,
+    deadline: Instant,
+
+    stop: AtomicBool,
+    violation: Mutex<Vec<String>>,
+
+    stamp: AtomicU64,
+    live: AtomicUsize,
+    peak_live: AtomicUsize,
+
+    ops: AtomicU64,
+    reads: AtomicU64,
+    acquires_ok: AtomicU64,
+    acquires_fail: AtomicU64,
+    retires_ok: AtomicU64,
+    retires_fail: AtomicU64,
+    reclaims: AtomicU64,
+    transitions: AtomicU64,
+}
+
+impl<T> Shared<T>
+where
+    T: Testable,
+{
+    /// Record an observed invariant violation and signal all workers to stop.
+    fn record_violation(&self, message: String) {
+        let mut slot = self.violation.lock().unwrap();
+        slot.push(message);
+        self.stop.store(true, Relaxed);
+    }
+
+    /// Return `true` once any termination condition is met.
+    fn should_stop(&self) -> bool {
+        self.stop.load(Relaxed)
+            || self.ops.load(Relaxed) >= self.max_ops
+            || Instant::now() >= self.deadline
+    }
+
+    //---------//
+    // Workers //
+    //---------//
+
+    fn writer(&self) {
+        let mut ops = Local::new(&self.ops);
+        let mut acquires_ok = Local::new(&self.acquires_ok);
+        let mut acquires_fail = Local::new(&self.acquires_fail);
+
+        let mut peak_live = LocalMax::new(&self.peak_live);
+
+        while !self.should_stop() {
+            ops.add(1);
+            match self.store.writer() {
+                Some(writer) => {
+                    let stamp = self.stamp.fetch_add(1, Relaxed);
+                    writer.write(stamp);
+
+                    let live = self.live.fetch_add(1, Relaxed) + 1;
+                    peak_live.max(live);
+                    acquires_ok.add(1);
+                }
+                None => {
+                    acquires_fail.add(1);
+                    std::thread::yield_now();
+                }
+            }
+        }
+    }
+
+    fn retirer(&self, seed: u64) {
+        let mut rng = StdRng::seed_from_u64(seed);
+        let mut iteration: u64 = 0;
+
+        let mut ops = Local::new(&self.ops);
+        let mut retires_ok = Local::new(&self.retires_ok);
+        let mut retires_fail = Local::new(&self.retires_fail);
+        let mut reclaims = Local::new(&self.reclaims);
+
+        while !self.should_stop() {
+            ops.add(1);
+            iteration += 1;
+
+            // Flow control: keep a steady readable population.
+            if self.live.load(Relaxed) > self.low_watermark {
+                let i = rng.sample(self.writable);
+                if self.store.retire(i) {
+                    self.live.fetch_sub(1, Relaxed);
+                    retires_ok.add(1);
+                } else {
+                    retires_fail.add(1);
+                }
+            }
+
+            if iteration.is_multiple_of(RECLAIM_EVERY)
+                && let Some(reclaimed) = self.store.reclaim()
+            {
+                reclaims.add(reclaimed as u64);
+            }
+
+            std::thread::yield_now();
+        }
+    }
+
+    fn reader(&self, seed: u64) {
+        let mut rng = StdRng::seed_from_u64(seed);
+        let slots = self.slots;
+        let window = READER_WINDOW.min(slots);
+
+        let mut ops = Local::new(&self.ops);
+        let mut reads = Local::new(&self.reads);
+
+        let mut reader_state = self.store.reader_state(window, &self);
+
+        while !self.should_stop() {
+            ops.add(1);
+
+            let succeeded = reader_state.try_with_reader(|mut reader| {
+                let start = rng.sample(self.readable);
+                for _ in 0..READER_PASSES {
+                    for k in 0..window {
+                        let i = (start + k) % slots;
+                        reader.observe(i);
+                        reads.add(1);
+                    }
+                }
+            });
+
+            // All guard slots are occupied; back off and retry.
+            if !succeeded {
+                std::thread::yield_now();
+            };
+        }
+    }
+}
diff --git a/diskann-inmem/src/integration/store.rs b/diskann-inmem/src/integration/store.rs
deleted file mode 100644
index a05df5840c..0000000000
--- a/diskann-inmem/src/integration/store.rs
+++ /dev/null
@@ -1,133 +0,0 @@
-/*
- * Copyright (c) Microsoft Corporation.
- * Licensed under the MIT license.
- */
-
-#![expect(
-    clippy::expect_used,
-    reason = "integration test tools are not production code"
-)]
-
-use std::num::{NonZeroU32, NonZeroUsize};
-
-use crate::{
-    num::{Bytes, Capacity, MaxDegree},
-    store,
-};
-
-#[derive(Debug)]
-pub struct Config {
-    pub capacity: usize,
-    pub entry_bytes: usize,
-    pub epoch_guard_slots: usize,
-    pub freelist_recycle_capacity: usize,
-}
-
-#[derive(Debug)]
-pub struct Store {
-    store: store::Store<store::invasive::Invasive>,
-}
-
-impl Store {
-    /// Construct a store with `config.capacity` writable slots, each holding
-    /// `config.entry_bytes` bytes.
-    ///
-    /// A single zeroed frozen point is created internally to satisfy the underlying
-    /// store's requirement of at least one frozen entry; it occupies the highest slot
-    /// index and is always readable.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the underlying store could not be constructed (e.g. `config.capacity` plus
-    /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
-    /// the number of epoch guard slots are invalid (e.g. zero).
-    pub fn new(config: Config) -> Self {
-        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 1);
-
-        let mut store_config = store::Config::default();
-
-        store_config
-            .epoch_guard_slots(
-                NonZeroUsize::new(config.epoch_guard_slots)
-                    .expect("`epoch_guard_slots` must be non-zero"),
-            )
-            .freelist_recycle_capacity(
-                NonZeroU32::new(
-                    config
-                        .freelist_recycle_capacity
-                        .try_into()
-                        .expect("`freelist_recycle_capacity` must fit within 32-bits"),
-                )
-                .expect("`freelist_recycle_capacity` must be non-zero"),
-            );
-
-        let plugin_config = store::invasive::Invasive::config(Bytes::new(config.entry_bytes));
-        let store = store::Store::new(store_layout, store_config, plugin_config)
-            .expect("failed to construct store");
-
-        Self { store }
-    }
-
-    /// Return the total number of slots, including the frozen point.
-    pub fn slots(&self) -> usize {
-        self.store.frozen().end as usize
-    }
-
-    /// Return the range of writable (non-frozen) slot indices.
-    pub fn writable(&self) -> std::ops::Range<usize> {
-        0..(self.store.frozen().start as usize)
-    }
-
-    /// Attempt to reclaim retired slots, returning the number reclaimed if any.
-    pub fn reclaim(&self) -> Option<usize> {
-        self.store.try_drain()
-    }
-
-    pub fn acquire(&self) -> Option<Writer<'_>> {
-        self.store.acquire().map(Writer::new)
-    }
-
-    #[must_use = "result indicates success or failure"]
-    pub fn retire(&self, i: usize) -> bool {
-        self.store.retire(i).is_ok()
-    }
-
-    pub fn reader(&self) -> Option<Reader<'_>> {
-        match store::invasive::Invasive::reader(&self.store) {
-            Ok(reader) => Some(Reader::new(reader)),
-            Err(crate::epoch::Unavailable) => None,
-        }
-    }
-}
-
-pub struct Reader<'a> {
-    reader: store::invasive::Reader<'a>,
-}
-
-impl<'a> Reader<'a> {
-    fn new(reader: store::invasive::Reader<'a>) -> Self {
-        Self { reader }
-    }
-
-    pub fn read(&self, i: usize) -> Option<&[u8]> {
-        self.reader.read(i)
-    }
-}
-
-pub struct Writer<'a> {
-    slot: store::Slot<'a, store::invasive::Slot<'a>>,
-}
-
-impl<'a> Writer<'a> {
-    fn new(slot: store::Slot<'a, store::invasive::Slot<'a>>) -> Self {
-        Self { slot }
-    }
-
-    pub fn publish(self) {
-        self.slot.publish();
-    }
-
-    pub fn as_mut_slice(&mut self) -> &mut [u8] {
-        self.slot.data().as_mut_slice()
-    }
-}
diff --git a/diskann-inmem/src/integration/store/checked.rs b/diskann-inmem/src/integration/store/checked.rs
new file mode 100644
index 0000000000..ff8f1a6eb0
--- /dev/null
+++ b/diskann-inmem/src/integration/store/checked.rs
@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+#![expect(
+    clippy::expect_used,
+    reason = "integration test tools are not production code"
+)]
+
+use std::num::{NonZeroU32, NonZeroUsize};
+
+use crate::{
+    num::{Bytes, Capacity, MaxDegree},
+    store::{self, checked},
+};
+
+use super::boilerplate;
+
+#[derive(Debug)]
+pub struct Config {
+    pub capacity: usize,
+    pub epoch_guard_slots: usize,
+    pub freelist_recycle_capacity: usize,
+}
+
+boilerplate!(
+    checked::Checked => Store,
+    for<'a> checked::Reader<'a> => Reader,
+    for<'a> checked::Slot<'a> => Writer,
+);
+
+impl Store {
+    /// Construct a store with `config.capacity` writable slots.
+    ///
+    /// A single zeroed frozen point is created internally to satisfy the underlying
+    /// store's requirement of at least one frozen entry; it occupies the highest slot
+    /// index and is always readable.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the underlying store could not be constructed (e.g. `config.capacity`
+    /// exceeds `u32::MAX`) or if other configuration parameters such as the number of
+    /// epoch guard slots are invalid (e.g. zero).
+    pub fn new(config: Config) -> Self {
+        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 0);
+
+        let store_config = store::Config::__exhaustive(
+            NonZeroUsize::new(config.epoch_guard_slots)
+                .expect("`epoch_guard_slots` must be non-zero"),
+            NonZeroU32::new(
+                config
+                    .freelist_recycle_capacity
+                    .try_into()
+                    .expect("`freelist_recycle_capacity` must fit within 32-bits"),
+            )
+            .expect("`freelist_recycle_capacity` must be non-zero"),
+        );
+
+        let plugin_config = checked::Checked::config();
+        let store = store::Store::new(store_layout, store_config, plugin_config)
+            .expect("failed to construct store");
+
+        Self { store }
+    }
+}
+
+#[derive(Debug)]
+pub struct Value<'a> {
+    value: checked::Value<'a>,
+}
+
+impl<'a> Value<'a> {
+    fn new(value: checked::Value<'a>) -> Self {
+        Self { value }
+    }
+
+    pub fn get(&self) -> u64 {
+        self.value.get()
+    }
+}
+
+impl Reader<'_> {
+    pub fn read(&self, i: usize) -> Option<Value<'_>> {
+        self.reader.read(i.try_into().unwrap()).map(Value::new)
+    }
+}
+
+impl<'a> Writer<'a> {
+    pub fn set(&mut self, v: u64) {
+        self.slot.data().set(v)
+    }
+}
diff --git a/diskann-inmem/src/integration/store/invasive.rs b/diskann-inmem/src/integration/store/invasive.rs
new file mode 100644
index 0000000000..b21300dda9
--- /dev/null
+++ b/diskann-inmem/src/integration/store/invasive.rs
@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+#![expect(
+    clippy::expect_used,
+    reason = "integration test tools are not production code"
+)]
+
+use std::num::{NonZeroU32, NonZeroUsize};
+
+use crate::{
+    num::{Bytes, Capacity, MaxDegree},
+    store,
+};
+
+use super::boilerplate;
+
+#[derive(Debug)]
+pub struct Config {
+    pub capacity: usize,
+    pub entry_bytes: usize,
+    pub epoch_guard_slots: usize,
+    pub freelist_recycle_capacity: usize,
+}
+
+boilerplate!(
+    store::invasive::Invasive => Store,
+    for<'a> store::invasive::Reader<'a> => Reader,
+    for<'a> store::invasive::Slot<'a> => Writer,
+);
+
+impl Store {
+    /// Construct a store with `config.capacity` writable slots, each holding
+    /// `config.entry_bytes` bytes.
+    ///
+    /// A single zeroed frozen point is created internally to satisfy the underlying
+    /// store's requirement of at least one frozen entry; it occupies the highest slot
+    /// index and is always readable.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the underlying store could not be constructed (e.g. `config.capacity` plus
+    /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
+    /// the number of epoch guard slots are invalid (e.g. zero).
+    pub fn new(config: Config) -> Self {
+        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 1);
+
+        let store_config = store::Config::__exhaustive(
+            NonZeroUsize::new(config.epoch_guard_slots)
+                .expect("`epoch_guard_slots` must be non-zero"),
+            NonZeroU32::new(
+                config
+                    .freelist_recycle_capacity
+                    .try_into()
+                    .expect("`freelist_recycle_capacity` must fit within 32-bits"),
+            )
+            .expect("`freelist_recycle_capacity` must be non-zero"),
+        );
+
+        let plugin_config = store::invasive::Invasive::config(Bytes::new(config.entry_bytes));
+        let store = store::Store::new(store_layout, store_config, plugin_config)
+            .expect("failed to construct store");
+
+        Self { store }
+    }
+}
+
+impl<'a> Reader<'a> {
+    pub fn read(&self, i: usize) -> Option<&[u8]> {
+        self.reader.read(i)
+    }
+}
+
+impl<'a> Writer<'a> {
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        self.slot.data().as_mut_slice()
+    }
+}
diff --git a/diskann-inmem/src/integration/store/mod.rs b/diskann-inmem/src/integration/store/mod.rs
new file mode 100644
index 0000000000..1d39b97cd2
--- /dev/null
+++ b/diskann-inmem/src/integration/store/mod.rs
@@ -0,0 +1,87 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+pub mod checked;
+pub mod invasive;
+
+macro_rules! boilerplate {
+    (
+        $plugin:ty => $store:ident,
+        for<$read_lt:lifetime> $read:ty => $reader:ident,
+        for<$slot_lt:lifetime> $slot:ty => $writer:ident,
+    ) => {
+        #[derive(Debug)]
+        pub struct $store {
+            store: $crate::store::Store<$plugin>,
+        }
+
+        impl $store {
+            /// Return the total number of slots, including the frozen point.
+            pub fn slots(&self) -> usize {
+                self.store.frozen().end as usize
+            }
+
+            /// Return the range of writable (non-frozen) slot indices.
+            pub fn writable(&self) -> usize {
+                self.store.frozen().start as usize
+            }
+
+            /// Attempt to reclaim retired slots, returning the number reclaimed if any.
+            pub fn reclaim(&self) -> Option<usize> {
+                self.store.try_drain()
+            }
+
+            /// Acquire a slot, returning a [`Writer`]. Returns `None` if no slot is
+            /// available for writing.
+            pub fn acquire(&self) -> Option<$writer<'_>> {
+                self.store.acquire().map(Writer::new)
+            }
+
+            /// Attempt to retire slot `i`. Returns `true` only if the slot was successfully
+            /// retired.
+            #[must_use = "result indicates success or failure"]
+            pub fn retire(&self, i: usize) -> bool {
+                self.store.retire(i).is_ok()
+            }
+
+            /// Attain a reader into the store. Returns `None` if all epoch guard slots
+            /// are used.
+            pub fn reader(&self) -> Option<$reader<'_>> {
+                match <$plugin>::reader(&self.store) {
+                    Ok(reader) => Some($reader::new(reader)),
+                    Err($crate::epoch::Unavailable) => None,
+                }
+            }
+        }
+
+        #[derive(Debug)]
+        pub struct $reader<$read_lt> {
+            reader: $read,
+        }
+
+        impl<$read_lt> $reader<$read_lt> {
+            fn new(reader: $read) -> Self {
+                Self { reader }
+            }
+        }
+
+        #[derive(Debug)]
+        pub struct $writer<$slot_lt> {
+            slot: $crate::store::Slot<$slot_lt, $slot>,
+        }
+
+        impl<$slot_lt> $writer<$slot_lt> {
+            fn new(slot: $crate::store::Slot<$slot_lt, $slot>) -> Self {
+                Self { slot }
+            }
+
+            pub fn publish(self) {
+                self.slot.publish();
+            }
+        }
+    };
+}
+
+use boilerplate;
diff --git a/diskann-inmem/src/store/checked.rs b/diskann-inmem/src/store/checked.rs
index 838b9c0a67..29d2f963e7 100644
--- a/diskann-inmem/src/store/checked.rs
+++ b/diskann-inmem/src/store/checked.rs
@@ -11,7 +11,7 @@ use std::{
 use diskann::utils::IntoUsize;
 use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
-use crate::{epoch, num::IdLimit};
+use crate::{epoch, num::IdLimit, store::Store};
 
 use super::{Lifecycle, plugin};
 
@@ -21,13 +21,125 @@ enum State {
     Available,
     Readable {
         value: u64,
-        retired: AtomicBool,
     },
     Frozen {
         value: u64,
     },
 }
 
+#[derive(Debug, Default)]
+struct Entry {
+    readable: AtomicBool,
+    state: RwLock<State>,
+}
+
+impl Entry {
+    #[must_use]
+    fn is_readable(&self) -> bool {
+        self.readable.load(Ordering::Acquire)
+    }
+
+    fn try_read(&self) -> Option<ReadEntry<'_>> {
+        if self.is_readable() {
+            Some(self.expect_read())
+        } else {
+            None
+        }
+    }
+
+    fn expect_read(&self) -> ReadEntry<'_> {
+        // NOTE: we *DO NOT* check for `entry.is_readable()` because there is a race where
+        // the slot is retired after checking the readable state but before this function
+        // is called. We still expect to acquire the `RwLockReadGuard` in this situation.
+
+        let Some(guard) = self.state.try_read() else {
+            panic!("concurrency violation when acquiring read guard");
+        };
+
+        ReadEntry {
+            readable: &self.readable,
+            guard,
+        }
+    }
+
+    fn expect_write(&self) -> WriteEntry<'_> {
+        assert!(
+            !self.is_readable(),
+            "concurrency violation - entry should not be readable"
+        );
+
+        // Correct usage of the concurrency protocol means that this `try_write` failing
+        // is a bug.
+        let Some(guard) = self.state.try_write() else {
+            panic!("concurrency violation when acquiring write guard");
+        };
+
+        WriteEntry {
+            readable: &self.readable,
+            guard,
+        }
+    }
+}
+
+#[derive(Debug)]
+struct ReadEntry<'a> {
+    readable: &'a AtomicBool,
+    guard: RwLockReadGuard<'a, State>,
+}
+
+impl ReadEntry<'_> {
+    fn retire(self) {
+        assert_matches!(*self.guard, State::Readable { .. });
+
+        // TODO: Document the slightly weird order.
+        drop(self.guard);
+        self.readable.store(false, Ordering::Release);
+    }
+
+    fn state(&self) -> &State {
+        &self.guard
+    }
+}
+
+#[derive(Debug)]
+struct WriteEntry<'a> {
+    readable: &'a AtomicBool,
+    guard: RwLockWriteGuard<'a, State>,
+}
+
+impl WriteEntry<'_> {
+    fn publish(mut self, value: u64) {
+        let old = self.replace(State::Readable { value });
+        assert_matches!(old, State::Available);
+
+        drop(self.guard);
+        self.readable.store(true, Ordering::Release);
+    }
+
+    fn freeze(mut self, value: u64) {
+        let old = self.replace(State::Frozen { value });
+        assert_matches!(old, State::Available);
+
+        drop(self.guard);
+        self.readable.store(true, Ordering::Release);
+    }
+
+    fn reclaim(mut self) {
+        let old = self.replace(State::Available);
+        assert_matches!(old, State::Readable { .. });
+    }
+
+    /// Replace the proctected state with `state`, returning the old state.
+    fn replace(&mut self, mut state: State) -> State {
+        std::mem::swap(&mut *self.guard, &mut state);
+        state
+    }
+
+    fn state(&self) -> &State {
+        &self.guard
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct Config(());
 
@@ -47,7 +159,7 @@ impl plugin::PluginConfig for Config {
 
 #[derive(Debug)]
 pub(crate) struct Checked {
-    states: Vec<RwLock<State>>,
+    entries: Vec<Entry>,
 }
 
 impl Checked {
@@ -57,58 +169,28 @@ impl Checked {
 
     pub(crate) fn new(id_limit: IdLimit) -> Self {
         Self {
-            states: std::iter::repeat_with(|| RwLock::new(State::default()))
+            entries: std::iter::repeat_with(|| Entry::default())
                 .take(id_limit.as_usize())
                 .collect(),
         }
     }
 
     pub(crate) fn id_limit(&self) -> IdLimit {
-        IdLimit::new(self.states.len().try_into().unwrap())
+        IdLimit::new(self.entries.len().try_into().unwrap())
     }
 
-    fn expect_write(&self, i: u32) -> RwLockWriteGuard<'_, State> {
-        let i = i.into_usize();
-
-        // Note: this will panic if `i` is out-of-bounds.
-        let entry = &self.states[i];
-
-        // Correct usage of the concurrency protocol means that this `try_write` failing
-        // is a bug.
-        let Some(guard) = entry.try_write() else {
-            panic!("concurrency violation when acquiring write guard");
-        };
-
-        guard
-    }
-
-    fn expect_read(&self, i: u32) -> RwLockReadGuard<'_, State> {
-        let i = i.into_usize();
-
-        // Note: this will panic if `i` is out-of-bounds.
-        let entry = &self.states[i];
-
-        // Correct usage of the concurrency protocol means that this `try_read` failing
-        // is a bug.
-        let Some(guard) = entry.try_read() else {
-            panic!("concurrency violation when acquiring read guard");
-        };
-
-        guard
-    }
-
-    pub(crate) fn reader<'a>(&'a self, guard: epoch::Guard<'a>) -> Reader<'a> {
-        Reader {
-            parent: self,
+    pub(crate) fn reader(store: &Store<Self>) -> Result<Reader<'_>, epoch::Unavailable> {
+        store.guard(|this, guard: epoch::Guard<'_>| Reader {
+            parent: this,
             _guard: guard,
-        }
+        })
     }
 }
 
 #[derive(Debug)]
 pub(crate) struct Value<'a> {
     value: u64,
-    _guard: RwLockReadGuard<'a, State>,
+    _entry: ReadEntry<'a>,
 }
 
 impl Value<'_> {
@@ -125,37 +207,15 @@ pub(crate) struct Reader<'a> {
 
 impl Reader<'_> {
     pub(crate) fn read(&self, i: u32) -> Option<Value<'_>> {
-        // This is kind of messy. The overall summary is this:
-        //
-        // 1. `i` has to be inbounds.
-        // 2. We have to be able to read the slot (if a write guard is active, then we
-        //    clearly should not be reading).
-        // 3a. If the state is frozen, then we can read it.
-        // 3b. If the state is readable and not retired, we can read it.
-        //
-        //     What happens if we transition to retired just after reading?
-        //
-        //     Fortunately, the EBR guard will keep the slot from being reclaimed until
-        //     the current `Reader` goes out-of-scope. Holding onto the
-        //     `RwLockReadGuard` allows us to detect bugs in the EBR protocol as
-        //     `Checked::expect_write` will fail on reclamation if a returned `Value`
-        //     is still active.
-        if let Some(state) = self.parent.states.get(i.into_usize())
-            && let Some(guard) = state.try_read()
-        {
-            let value = match &*guard {
-                State::Frozen { value } => *value,
-                State::Readable { value, retired } => {
-                    if retired.load(Ordering::Relaxed) {
-                        return None;
-                    }
-                    *value
-                }
-                _ => return None,
+        if let Some(entry) = self.parent.entries.get(i.into_usize())?.try_read() {
+            let value = match entry.state() {
+                State::Frozen { value } | State::Readable { value } => value,
+                State::Available => panic!("concurrency violation"),
             };
+
             Some(Value {
-                value,
-                _guard: guard,
+                value: *value,
+                _entry: entry,
             })
         } else {
             None
@@ -171,51 +231,30 @@ impl plugin::Plugin for Checked {
     }
 
     unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_> {
-        let guard = self.expect_write(i);
-        assert_matches!(*guard, State::Available, "slot is in an invalid state");
-        Slot::new(guard)
+        Slot::new(self.entries[i.into_usize()].expect_write())
     }
 
     unsafe fn retire(&self, i: u32, _: Lifecycle) {
-        let guard = self.expect_read(i);
-        match &*guard {
-            State::Available => panic!("invalid \"Available\" state"),
-            State::Readable { retired, .. } => {
-                let old = retired.swap(true, Ordering::Relaxed);
-                if old {
-                    panic!("slot {i} was retired multiple times");
-                }
-            }
-            State::Frozen { .. } => panic!("tried to retire frozen point {i}"),
-        }
+        self.entries[i.into_usize()].expect_read().retire();
     }
 
     unsafe fn reclaim(&self, i: u32, _: Lifecycle) {
-        let mut guard = self.expect_write(i);
-        match &*guard {
-            State::Available => panic!("invalid \"Available\" state"),
-            State::Readable { retired, .. } => {
-                assert!(
-                    retired.load(Ordering::Relaxed),
-                    "tried to reclaim {i} before it has been retired!",
-                );
-            }
-            State::Frozen { .. } => panic!("tried to reclaim frozen point {i}"),
-        }
-
-        *guard = State::Available;
+        self.entries[i.into_usize()].expect_write().reclaim();
     }
 }
 
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
-    guard: RwLockWriteGuard<'a, State>,
+    entry: WriteEntry<'a>,
     value: Option<u64>,
 }
 
 impl<'a> Slot<'a> {
-    fn new(guard: RwLockWriteGuard<'a, State>) -> Self {
-        Self { guard, value: None }
+    fn new(entry: WriteEntry<'a>) -> Self {
+        Self {
+            entry,
+            value: None,
+        }
     }
 
     pub(crate) fn set(&mut self, value: u64) {
@@ -224,20 +263,17 @@ impl<'a> Slot<'a> {
 }
 
 impl plugin::Slot for Slot<'_> {
-    fn publish(mut self, _: Lifecycle) {
+    fn publish(self, _: Lifecycle) {
         let value = self.value.expect("`value` was not set");
-        *self.guard = State::Readable {
-            value,
-            retired: AtomicBool::new(false),
-        };
+        self.entry.publish(value);
     }
 
-    fn freeze(mut self, _: Lifecycle) {
+    fn freeze(self, _: Lifecycle) {
         let value = self.value.expect("`value` was not set");
-        *self.guard = State::Frozen { value };
+        self.entry.freeze(value);
     }
 
-    fn abort(mut self, _: Lifecycle) {
-        *self.guard = State::Available;
+    fn abort(self, _: Lifecycle) {
+        assert_matches!(self.entry.state(), State::Available);
     }
 }
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index 0b948e0a56..bcbdd41f3d 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -70,9 +70,7 @@ impl Invasive {
     /// Create a new [`Invasive`] with capacity for `id_limit` slots of `bytes`.
     pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Self {
         let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
-        let padded_bytes = unpadded
-            .checked_next_multiple_of(Bytes::CACHELINE)
-            .unwrap();
+        let padded_bytes = unpadded.checked_next_multiple_of(Bytes::CACHELINE).unwrap();
 
         Self {
             buffer: Buffer::new(id_limit.as_usize(), padded_bytes, Align::_128).unwrap(),
@@ -88,13 +86,11 @@ impl Invasive {
     }
 
     /// Return a [`Reader`] over [`Self`] inside `store`.
-    pub(crate) fn reader<'a>(store: &'a Store<Self>) -> Result<Reader<'a>, epoch::Unavailable> {
-        store.guard(|this, guard: epoch::Guard<'a>| {
-            Reader {
-                buffer: &this.buffer,
-                unpadded: this.unpadded,
-                _guard: guard
-            }
+    pub(crate) fn reader(store: &Store<Self>) -> Result<Reader<'_>, epoch::Unavailable> {
+        store.guard(|this, guard: epoch::Guard<'_>| Reader {
+            buffer: &this.buffer,
+            unpadded: this.unpadded,
+            _guard: guard,
         })
     }
 
@@ -328,4 +324,3 @@ impl plugin::Slot for Slot<'_> {
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
-
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index f080ac06ee..bc60c7e9ba 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -73,8 +73,8 @@ use crate::{
 pub(crate) mod invasive;
 pub(crate) mod plugin;
 
-#[cfg(test)]
-mod checked;
+#[cfg(any(test, feature = "integration-test"))]
+pub(crate) mod checked;
 
 /// To make extra sure that [`plugin::Plugin`] life-cycle arguments are not callable outside
 /// of this module (i.e., elsewhere in this crate), this [`Lifecycle`] marker type is used
@@ -135,6 +135,23 @@ impl Config {
         self.freelist_recycle_capacity = freelist_recycle_capacity;
         self
     }
+
+    /// An exhaustive constructor initializing everye element.
+    ///
+    /// This is under the "integration-test" since it will change to reflect the state
+    /// of the underlying data structure, potentially causing more churn for users if it
+    /// were unconditionally exposed.
+    #[cfg(feature = "integration-test")]
+    #[doc(hidden)]
+    pub fn __exhaustive(
+        epoch_guard_slots: NonZeroUsize,
+        freelist_recycle_capacity: NonZeroU32,
+    ) -> Self {
+        Self {
+            epoch_guard_slots,
+            freelist_recycle_capacity,
+        }
+    }
 }
 
 impl Default for Config {
@@ -690,7 +707,7 @@ mod tests {
     }
 
     fn reader(store: &Store<Checked>) -> checked::Reader<'_> {
-        store.guard(|checked, guard| checked.reader(guard)).unwrap()
+        Checked::reader(store).unwrap()
     }
 
     //------------------------//

From adc1c2c7c2cb218078c568706ff00249368c564e Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Fri, 21 Aug 2026 12:30:41 -0700
Subject: [PATCH 19/34] Harden the storage layer.

---
 .../integration/jsons/store-stress-test.json  |  17 +
 diskann-inmem/integration/store/checked.rs    |   4 +-
 diskann-inmem/integration/store/invasive.rs   |   3 +-
 diskann-inmem/integration/store/mod.rs        |   2 +-
 diskann-inmem/integration/support/datatype.rs |   2 +-
 .../src/integration/store/checked.rs          |   2 +-
 diskann-inmem/src/layers/full.rs              |   6 +-
 diskann-inmem/src/provider.rs                 |   4 -
 diskann-inmem/src/store/checked.rs            | 201 ++++++++--
 diskann-inmem/src/store/invasive.rs           | 357 ++++++++++++++++--
 diskann-inmem/src/store/mod.rs                |  41 +-
 diskann-inmem/src/store/plugin.rs             |   8 +-
 12 files changed, 578 insertions(+), 69 deletions(-)

diff --git a/diskann-inmem/integration/jsons/store-stress-test.json b/diskann-inmem/integration/jsons/store-stress-test.json
index fed64e095f..3404133d8d 100644
--- a/diskann-inmem/integration/jsons/store-stress-test.json
+++ b/diskann-inmem/integration/jsons/store-stress-test.json
@@ -18,6 +18,23 @@
           "seed": 11939873485092837375
         }
       }
+    },
+    {
+      "type": "store-stress-checked",
+      "content": {
+        "setup": {
+          "readers": 4,
+          "writers": 2,
+          "retirers": 1,
+          "capacity": 512,
+          "epoch_guard_slots": 64,
+          "freelist_recycle_capacity": 32,
+          "low_watermark": 128,
+          "duration_secs": 2,
+          "max_ops": 2000000,
+          "seed": 11939873485092837375
+        }
+      }
     }
   ]
 }
diff --git a/diskann-inmem/integration/store/checked.rs b/diskann-inmem/integration/store/checked.rs
index 9f074f9625..2078ac81c8 100644
--- a/diskann-inmem/integration/store/checked.rs
+++ b/diskann-inmem/integration/store/checked.rs
@@ -155,7 +155,9 @@ impl super::ReaderState for ReaderState<'_> {
     where
         F: FnOnce(Self::Reader<'_>),
     {
-        let Some(reader) = self.store.reader() else { return false; };
+        let Some(reader) = self.store.reader() else {
+            return false;
+        };
         let observed = HashMap::with_capacity(self.capacity_hint);
 
         f(Reader {
diff --git a/diskann-inmem/integration/store/invasive.rs b/diskann-inmem/integration/store/invasive.rs
index 71e529af95..8f826381c1 100644
--- a/diskann-inmem/integration/store/invasive.rs
+++ b/diskann-inmem/integration/store/invasive.rs
@@ -197,8 +197,7 @@ pub(super) struct ReaderState<'a> {
 }
 
 impl super::ReaderState for ReaderState<'_> {
-    type Reader<'a>
-        = Reader<'a>;
+    type Reader<'a> = Reader<'a>;
 
     fn try_with_reader<F>(&mut self, f: F) -> bool
     where
diff --git a/diskann-inmem/integration/store/mod.rs b/diskann-inmem/integration/store/mod.rs
index 589ba99add..04b2cf6884 100644
--- a/diskann-inmem/integration/store/mod.rs
+++ b/diskann-inmem/integration/store/mod.rs
@@ -467,7 +467,7 @@ where
         let mut ops = Local::new(&self.ops);
         let mut reads = Local::new(&self.reads);
 
-        let mut reader_state = self.store.reader_state(window, &self);
+        let mut reader_state = self.store.reader_state(window, self);
 
         while !self.should_stop() {
             ops.add(1);
diff --git a/diskann-inmem/integration/support/datatype.rs b/diskann-inmem/integration/support/datatype.rs
index f0e63debe4..34a60190aa 100644
--- a/diskann-inmem/integration/support/datatype.rs
+++ b/diskann-inmem/integration/support/datatype.rs
@@ -193,7 +193,7 @@ impl<'a> SliceMut<'a> {
 
         match (self, rhs) {
             (SliceMut::F32(dst), Slice::F32(src)) => dst.copy_from_slice(src),
-            (SliceMut::F32(dst), Slice::F16(src)) => map(dst, src, |x| cast_f16_to_f32(x)),
+            (SliceMut::F32(dst), Slice::F16(src)) => map(dst, src, cast_f16_to_f32),
             (SliceMut::F32(dst), Slice::U8(src)) => map(dst, src, |x| x.into()),
             (SliceMut::F32(dst), Slice::I8(src)) => map(dst, src, |x| x.into()),
 
diff --git a/diskann-inmem/src/integration/store/checked.rs b/diskann-inmem/src/integration/store/checked.rs
index ff8f1a6eb0..fb84a13ad9 100644
--- a/diskann-inmem/src/integration/store/checked.rs
+++ b/diskann-inmem/src/integration/store/checked.rs
@@ -11,7 +11,7 @@
 use std::num::{NonZeroU32, NonZeroUsize};
 
 use crate::{
-    num::{Bytes, Capacity, MaxDegree},
+    num::{Capacity, MaxDegree},
     store::{self, checked},
 };
 
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index b0c885aa3b..bbb3604ce6 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -535,7 +535,7 @@ where
 {
     fn evaluate(&self, i: u32) -> ANNResult<Option<f32>> {
         if !self.reader.is_in_bounds(i.into_usize()) {
-            return Err(ANNError::new(OutOfBounds(i)));
+            Err(ANNError::new(OutOfBounds(i)))
         } else {
             match unsafe { self.reader.read_in_bounds(i.into_usize()) } {
                 Some(data) => Ok(Some(self.run(data)?)),
@@ -561,6 +561,8 @@ where
         for j in 0..lookahead {
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
             // as the validity of the prefetch bounds.
+            //
+            // We do not materialize the `RawSlice` as a reference.
             unsafe {
                 crate::arch::prefetch(
                     self.reader
@@ -579,6 +581,8 @@ where
             if j != len {
                 // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
                 // well as the validity of the prefetch bounds.
+                //
+                // We do not materialize the `RawSlice` as a reference.
                 unsafe {
                     crate::arch::prefetch(
                         self.reader
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index ac26c35db8..dcd03d7264 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -449,10 +449,6 @@ impl<'a> Distance<'a> {
     }
 }
 
-#[expect(
-    clippy::unwrap_used,
-    reason = "prune does not allow fallible distance functions yet"
-)]
 impl diskann_vector::DistanceFunction<layers::PruneKey, layers::PruneKey, f32> for Distance<'_> {
     #[inline]
     fn evaluate_similarity(&self, x: layers::PruneKey, y: layers::PruneKey) -> f32 {
diff --git a/diskann-inmem/src/store/checked.rs b/diskann-inmem/src/store/checked.rs
index 29d2f963e7..607cbd669d 100644
--- a/diskann-inmem/src/store/checked.rs
+++ b/diskann-inmem/src/store/checked.rs
@@ -3,6 +3,75 @@
  * Licensed under the MIT license.
  */
 
+//! # A [`Store`] plugin for pedantically testing the EBR protocol
+//!
+//! The goal here is to detect violations of the state machine outlined in [`plugin`].
+//! This is accomplished by using a [`RwLock`] to protect internal [`State`] with guards
+//! **only** acquired via [`RwLock::try_read`] and [`RwLock::try_write`]. A correct EBR
+//! protocol should ensure that:
+//!
+//! 1. States where exclusive access is needed have no readers or other concurrent exclusive
+//!    accesses.
+//!
+//! 2. States where read access is allowed have no attempts at exclusive access.
+//!
+//! The "try" interfaces provide us with these checks:
+//!
+//! 1. [`RwLock::try_write`] will fail if there is any concurrent reader or writer.
+//! 2. [`RwLock::try_read`] will fail if there is a writer.
+//!
+//! Importantly, we *only* use these non-blocking APIs.
+//!
+//! Reads from the store return a [`Value`], which contains a [`RwLockReadGuard`] for the
+//! corresponding slot. Many such guards can coexist. Providing long-lived guards like this
+//! improves our chances of catching a bug in the EBR scheme where exclusive access is
+//! attempted too early. [`Reader::read`] constrains each [`Value`] to the borrow of the
+//! [`Reader`]. Since [`Reader`]s own an [`epoch::Guard`], [`Value`]s are guaranteed to be
+//! dropped before their protecting [`epoch::Guard`] is dropped.
+//!
+//! ## Lifecycle Details
+//!
+//! The plugin lifecycle is carefully designed to allow readers of the plugin to avoid any
+//! accesses to the authoritative [`Store`] for read-only operations. The [`Checked`] test
+//! code follows this pattern, but this does introduce a subtle detail that is worth
+//! highlighting. [`Reader::read`] needs to be able to check a slot for readability
+//! **without** trying to acquire a [`RwLockReadGuard`] for that slot. Doing so even briefly
+//! will cause a [`RwLock::try_write`] on an otherwise correct [`plugin::Plugin`] state
+//! transition to fail.
+//!
+//! To circumvent this, an additional [`AtomicBool`] is bundled with [`Entry`] to broadcast
+//! readability. This must be checked first on an optimistic read before attempting to
+//! acquire the [`RwLockReadGuard`]. This boolean flag is toggled on the following transitions:
+//!
+//! * "slot" -> "published"/"frozen": Publish as readable. Note that this transition goes
+//!   from an "exclusive" owned (with a [`RwLockWriteGuard`]) to a shared readable state.
+//!
+//!   As such, it's necessary to release the guard before toggling the "readable" flag to
+//!   ensure that if [`Reader::read`] observes the "published" state it is guaranteed to
+//!   succeed in [`RwLockReadGuard`] acquisition.
+//!
+//! * "published" -> "retiring": This switches to non-readable. Note that there is a race
+//!   condition where:
+//!
+//!   - (Thread A) [`Reader::read`] observes "readable" and decides to acquire the read guard.
+//!   - (Thread B) Toggles the "readable" flag.
+//!   - (Thread A) Finishes acquiring the read guard, even though the "readable" state is no
+//!     longer broadcasted.
+//!
+//!   This is **exactly** what the EBR protocol makes safe. These races are perfectly fine
+//!   because the protected [`State::Published`] value remains valid for existing readers
+//!   while the lifecycle state is "retiring". The EBR protocol ensures that a slot does
+//!   **not** transition away from "retiring" until thread `A` (and all other concurrent
+//!   threads that could have observed the slot while it was published) have dropped their
+//!   read guards, and thus any data accessed while under the guard, since these are
+//!   outlived by a proper [`epoch::Guard`].
+
+#![expect(
+    clippy::panic,
+    clippy::expect_used,
+    reason = "integration-test code is not production code"
+)]
+
 use std::{
     assert_matches,
     sync::atomic::{AtomicBool, Ordering},
@@ -15,11 +84,12 @@ use crate::{epoch, num::IdLimit, store::Store};
 
 use super::{Lifecycle, plugin};
 
+/// The state of a slot.
 #[derive(Debug, Default)]
 enum State {
     #[default]
     Available,
-    Readable {
+    Published {
         value: u64,
     },
     Frozen {
@@ -27,6 +97,24 @@ enum State {
     },
 }
 
+/// A slot entry. See the [module level docs](self) for a discussion on the contents of this
+/// struct. The table below describes how the combination of fields in this struct maps to
+/// the plugin lifecycle states.
+/// ```text
+/// +-----------------+----------+-----------+------------------------+
+/// | Lifecycle state | readable | State     | Lock Expectation       |
+/// +=================+==========+===========+========================+
+/// |       Available |  false   | Available | unlocked               |
+/// +-----------------+----------+-----------+------------------------+
+/// |            Slot |  false   | Available | write-locked           |
+/// +-----------------+----------+-----------+------------------------+
+/// |       Published |  true    | Published | shared reads allowed   |
+/// +-----------------+----------+-----------+------------------------+
+/// |        Retiring |  false   | Published | existing reads allowed |
+/// +-----------------+----------+-----------+------------------------+
+/// |          Frozen |  true    | Frozen    | shared reads allowed   |
+/// +-----------------+----------+-----------+------------------------+
+/// ```
 #[derive(Debug, Default)]
 struct Entry {
     readable: AtomicBool,
@@ -34,19 +122,33 @@ struct Entry {
 }
 
 impl Entry {
+    /// Return whether or not this [`Entry`] is broadcasted as readable.
     #[must_use]
     fn is_readable(&self) -> bool {
         self.readable.load(Ordering::Acquire)
     }
 
+    /// Attempt to acquire a [`ReadEntry`], failing if the [`Entry`] is not marked as
+    /// readable without attempting to acquire any locks.
     fn try_read(&self) -> Option<ReadEntry<'_>> {
         if self.is_readable() {
+            // We are relying on EBR to avoid problems with the TOCTOU/ABA race between
+            // checking the `bool` and attempting to acquire the read guard.
+            //
+            // EBR prevents the slot from being reclaimed and reused during the race interval.
+            //
+            // See the module level docs.
             Some(self.expect_read())
         } else {
             None
         }
     }
 
+    /// Acquire a [`ReadEntry`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if the [`RwLockReadGuard`] cannot be immediately acquired.
     fn expect_read(&self) -> ReadEntry<'_> {
         // NOTE: we *DO NOT* check for `entry.is_readable()` because there is a race where
         // the slot is retired after checking the readable state but before this function
@@ -62,6 +164,11 @@ impl Entry {
         }
     }
 
+    /// Acquire a [`WriteEntry`].
+    ///
+    /// # Panics
+    ///
+    /// Panics if the [`RwLockWriteGuard`] cannot be immediately acquired.
     fn expect_write(&self) -> WriteEntry<'_> {
         assert!(
             !self.is_readable(),
@@ -81,6 +188,7 @@ impl Entry {
     }
 }
 
+/// A readable version of [`Entry`].
 #[derive(Debug)]
 struct ReadEntry<'a> {
     readable: &'a AtomicBool,
@@ -88,12 +196,22 @@ struct ReadEntry<'a> {
 }
 
 impl ReadEntry<'_> {
+    /// Mark this slot as "retired".
     fn retire(self) {
-        assert_matches!(*self.guard, State::Readable { .. });
+        assert_matches!(
+            *self.guard,
+            State::Published { .. },
+            "\"retire\" should transition out of the \"published\" state",
+        );
 
-        // TODO: Document the slightly weird order.
-        drop(self.guard);
-        self.readable.store(false, Ordering::Release);
+        // The ordering of dropping the guard vs clearing the "readable" flag doesn't
+        // really matter because the EBR protocol ensures that final reclamation is
+        // sufficiently deferred.
+        let old = self.readable.swap(false, Ordering::Release);
+        assert!(
+            old,
+            "\"retire\" should transition out of the \"published\" state"
+        );
     }
 
     fn state(&self) -> &State {
@@ -101,6 +219,7 @@ impl ReadEntry<'_> {
     }
 }
 
+/// A writable version of [`Entry`].
 #[derive(Debug)]
 struct WriteEntry<'a> {
     readable: &'a AtomicBool,
@@ -108,28 +227,54 @@ struct WriteEntry<'a> {
 }
 
 impl WriteEntry<'_> {
+    /// Transition this slot to "published".
     fn publish(mut self, value: u64) {
-        let old = self.replace(State::Readable { value });
-        assert_matches!(old, State::Available);
+        let old = self.replace(State::Published { value });
+        assert_matches!(
+            old,
+            State::Available,
+            "\"publish\" must transition out of the \"available\" state",
+        );
 
+        // The ordering here matters. `Reader` must be guaranteed to acquire the read guard
+        // if it observes `readable = true`. We must drop the guard before broadcasting.
         drop(self.guard);
         self.readable.store(true, Ordering::Release);
     }
 
+    /// Transition this slot to "frozen".
     fn freeze(mut self, value: u64) {
         let old = self.replace(State::Frozen { value });
-        assert_matches!(old, State::Available);
+        assert_matches!(
+            old,
+            State::Available,
+            "\"freeze\" must transition out of the \"available\" state",
+        );
 
+        // See `Self::publish` for ordering details.
         drop(self.guard);
         self.readable.store(true, Ordering::Release);
     }
 
+    /// Transition this slot to "available".
     fn reclaim(mut self) {
         let old = self.replace(State::Available);
-        assert_matches!(old, State::Readable { .. });
+
+        // Note: The combination of `!readable` and `State::Published` implies "retiring".
+        // We check that here.
+        assert!(
+            !self.readable.load(Ordering::Relaxed),
+            "\"reclaim\" must transition out of \"retired\"",
+        );
+
+        assert_matches!(
+            old,
+            State::Published { .. },
+            "\"reclaim\" must transition out of \"retired\"",
+        );
     }
 
-    /// Replace the proctected state with `state`, returning the old state.
+    /// Replace the protected state with `state`, returning the old state.
     fn replace(&mut self, mut state: State) -> State {
         std::mem::swap(&mut *self.guard, &mut state);
         state
@@ -140,6 +285,7 @@ impl WriteEntry<'_> {
     }
 }
 
+/// A [`plugin::PluginConfig`] for [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Config(());
 
@@ -151,34 +297,40 @@ impl Config {
 
 impl plugin::PluginConfig for Config {
     type Plugin = Checked;
+    type Error = diskann::error::Infallible;
 
-    fn build(self, id_limit: IdLimit) -> diskann::ANNResult<Checked> {
+    fn build(self, id_limit: IdLimit) -> Result<Checked, diskann::error::Infallible> {
         Ok(Checked::new(id_limit))
     }
 }
 
+/// A correctness checking [`plugin::Plugin`]. See the [module level docs](self) for details.
 #[derive(Debug)]
 pub(crate) struct Checked {
     entries: Vec<Entry>,
 }
 
 impl Checked {
-    pub(crate) fn config() -> Config {
-        Config::new()
-    }
-
+    /// Create a new [`Checked`] with `id_limit` slots.
     pub(crate) fn new(id_limit: IdLimit) -> Self {
         Self {
-            entries: std::iter::repeat_with(|| Entry::default())
+            entries: std::iter::repeat_with(Entry::default)
                 .take(id_limit.as_usize())
                 .collect(),
         }
     }
 
+    /// Return the [`plugin::PluginConfig`] for [`Self`].
+    pub(crate) fn config() -> Config {
+        Config::new()
+    }
+
+    /// Return the [`IdLimit`] for this store.
     pub(crate) fn id_limit(&self) -> IdLimit {
-        IdLimit::new(self.entries.len().try_into().unwrap())
+        IdLimit::new(self.entries.len() as u32)
     }
 
+    /// Return an epoch-protected [`Reader`] into [`Self`].
     pub(crate) fn reader(store: &Store<Self>) -> Result<Reader<'_>, epoch::Unavailable> {
         store.guard(|this, guard: epoch::Guard<'_>| Reader {
             parent: this,
@@ -187,6 +339,7 @@ impl Checked {
     }
 }
 
+/// A valid, readable entry of [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Value<'a> {
     value: u64,
@@ -194,11 +347,13 @@ pub(crate) struct Value<'a> {
 }
 
 impl Value<'_> {
+    /// Return the payload for this [`Value`].
     pub(crate) fn get(&self) -> u64 {
         self.value
     }
 }
 
+/// A reader for [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Reader<'a> {
     parent: &'a Checked,
@@ -206,10 +361,13 @@ pub(crate) struct Reader<'a> {
 }
 
 impl Reader<'_> {
+    /// Attempt to read the value at slot `i`.
+    ///
+    /// Fails if `i` is out-of-bounds, or the slot is not in a readable state.
     pub(crate) fn read(&self, i: u32) -> Option<Value<'_>> {
         if let Some(entry) = self.parent.entries.get(i.into_usize())?.try_read() {
             let value = match entry.state() {
-                State::Frozen { value } | State::Readable { value } => value,
+                State::Frozen { value } | State::Published { value } => value,
                 State::Available => panic!("concurrency violation"),
             };
 
@@ -243,6 +401,7 @@ impl plugin::Plugin for Checked {
     }
 }
 
+/// A writable [`plugin::Slot`] for [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
     entry: WriteEntry<'a>,
@@ -251,12 +410,10 @@ pub(crate) struct Slot<'a> {
 
 impl<'a> Slot<'a> {
     fn new(entry: WriteEntry<'a>) -> Self {
-        Self {
-            entry,
-            value: None,
-        }
+        Self { entry, value: None }
     }
 
+    /// Write `value` into the slot on a successful state transition.
     pub(crate) fn set(&mut self, value: u64) {
         self.value = Some(value)
     }
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index bcbdd41f3d..a48fbcb64e 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -3,18 +3,46 @@
  * Licensed under the MIT license.
  */
 
-//! A store [`plugin::Plugin`] that maintains in invasive slot state where the data in each
+//! A store [`plugin::Plugin`] that maintains an invasive slot state where the data in each
 //! slot is a contiguous slice of memory.
 //!
 //! Slot state is stored as an [`AtomicTag`] immediately after the slot data.
+//!
+//! ## Lifecycle Details
+//!
+//! The plugin lifecycle details are relatively straightforward. The invasive [`AtomicTag`]
+//! mostly follows the transitions made by the [`Store`]. A [`Reader`] checks the tag for
+//! readability before creating a shared reference to the data payload.
+//!
+//! The problematic transition from "published" to "retiring" is made safe because
+//!
+//! 1. Each [`Reader`] stores an [`epoch::Guard`].
+//! 2. Slices yielded by [`Reader`] are scoped to the **borrow** of [`Reader`].
+//!
+//! This ensures that slices cannot outlive the [`epoch::Guard`] protecting their slot. EBR
+//! prevents the slot from transitioning from "retiring" back to "available" and being
+//! reused while the guard remains active.
+//!
+//! The transitions [`Slot::publish`] and [`Slot::freeze`] use release stores. Since these
+//! are terminal slot operations, their release stores occur after all payload writes. The
+//! acquire load in [`Reader::read`] makes those writes visible before creating a shared
+//! slice.
+//!
+//! ## Safety
+//!
+//! The safety of this module depends on [`Invasive`] being embedded in a [`Store`] that
+//! observes the plugin lifecycle. Every lifecycle operation requires a [`Lifecycle`] token,
+//! which is constructible only by the parent store module. The unsafe [`plugin::Plugin`]
+//! methods additionally rely on [`Store`] to satisfy their documented state and exclusivity
+//! preconditions.
 
-use std::{num::NonZeroUsize, sync::atomic::Ordering};
+use std::sync::atomic::Ordering;
 
-use diskann::{ANNResult, utils::IntoUsize};
+use diskann::utils::IntoUsize;
 use thiserror::Error;
 
 use crate::{
-    buffer::{Buffer, RawSlice},
+    buffer::{Buffer, BufferError, RawSlice},
     epoch,
     num::{Align, Bytes, IdLimit},
     store::{Lifecycle, Store, plugin},
@@ -35,15 +63,16 @@ impl Config {
     }
 
     /// Build an [`Invasive`] store holding `id_limit` slots.
-    pub(crate) fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
+    pub(crate) fn build(self, id_limit: IdLimit) -> Result<Invasive, InvasiveError> {
         let Self { bytes } = self;
-        Ok(Invasive::new(id_limit, bytes))
+        Invasive::new(id_limit, bytes)
     }
 }
 
 impl plugin::PluginConfig for Config {
     type Plugin = Invasive;
-    fn build(self, id_limit: IdLimit) -> ANNResult<Invasive> {
+    type Error = InvasiveError;
+    fn build(self, id_limit: IdLimit) -> Result<Invasive, InvasiveError> {
         <Config>::build(self, id_limit)
     }
 }
@@ -68,19 +97,30 @@ impl Invasive {
     }
 
     /// Create a new [`Invasive`] with capacity for `id_limit` slots of `bytes`.
-    pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Self {
-        let unpadded = bytes.checked_add(AtomicTag::SIZE).unwrap();
-        let padded_bytes = unpadded.checked_next_multiple_of(Bytes::CACHELINE).unwrap();
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the internal buffer allocation exceeds `isize::MAX` or
+    /// computation of the padded, invasive bytes exceeds `usize::MAX`.
+    pub(crate) fn new(id_limit: IdLimit, bytes: Bytes) -> Result<Self, InvasiveError> {
+        let Some(unpadded) = bytes.checked_add(AtomicTag::SIZE) else {
+            return Err(InvasiveError::bytes_overflowed());
+        };
+        let Some(padded_bytes) = unpadded.checked_next_multiple_of(Bytes::CACHELINE) else {
+            return Err(InvasiveError::bytes_overflowed());
+        };
 
-        Self {
-            buffer: Buffer::new(id_limit.as_usize(), padded_bytes, Align::_128).unwrap(),
-            unpadded,
-        }
+        let buffer = match Buffer::new(id_limit.as_usize(), padded_bytes, Align::_128) {
+            Ok(buffer) => buffer,
+            Err(err) => return Err(InvasiveError::buffer_error(err)),
+        };
+
+        Ok(Self { buffer, unpadded })
     }
 
     /// Return the [`IdLimit`] for this store.
     pub(crate) fn id_limit(&self) -> IdLimit {
-        // The numeric cast is save because `Invasive::new` takes an `IdLimit` in its
+        // The numeric cast is safe because `Invasive::new` takes an `IdLimit` in its
         // constructor, and thus `self.buffer.len()` cannot exceed `u32::MAX`.
         IdLimit::new(self.buffer.len() as u32)
     }
@@ -98,15 +138,15 @@ impl Invasive {
     ///
     /// # Safety
     ///
-    /// The index `i` must be less then `self.buffer.len()`.
+    /// The index `i` must be less than `self.buffer.len()`.
     unsafe fn data_unchecked(&self, i: usize) -> (&AtomicTag, RawSlice<'_>) {
         // SAFETY: inherited from caller.
         let (data, mirror) = unsafe { self.buffer.get_unchecked(i) }
             .truncate(self.unpadded)
             .split(self.unpadded.unchecked_sub(AtomicTag::SIZE));
         (
-            // SAFETY: We're careful in this module to ensure the inline tags are only
-            // ever accessed atomically.
+            // SAFETY: The tag byte lies within the zero-initialized row, is sufficiently
+            // aligned for `AtomicTag`, and is only accessed through atomic operations.
             unsafe { AtomicTag::from_ptr(mirror.as_mut_ptr().cast()) },
             data,
         )
@@ -116,11 +156,34 @@ impl Invasive {
         if i >= self.buffer.len() {
             None
         } else {
+            // SAFETY: We've checked that `i` is in-bounds.
             Some(unsafe { self.data_unchecked(i) })
         }
     }
 }
 
+#[derive(Debug, Error)]
+#[error(transparent)]
+pub(crate) struct InvasiveError(InvasiveErrorInner);
+
+impl InvasiveError {
+    fn bytes_overflowed() -> Self {
+        Self(InvasiveErrorInner::BytesOverflowed)
+    }
+
+    fn buffer_error(err: BufferError) -> Self {
+        Self(InvasiveErrorInner::BufferError(err))
+    }
+}
+
+#[derive(Debug, Error)]
+enum InvasiveErrorInner {
+    #[error("computation of the bytes per slot overflowed")]
+    BytesOverflowed,
+    #[error(transparent)]
+    BufferError(BufferError),
+}
+
 impl plugin::Plugin for Invasive {
     type Slot<'a> = Slot<'a>;
 
@@ -128,6 +191,7 @@ impl plugin::Plugin for Invasive {
         <Invasive>::id_limit(self)
     }
 
+    #[expect(clippy::panic, reason = "out-of-bounds is a hard program bug")]
     unsafe fn acquire(&self, i: u32, _: Lifecycle) -> Self::Slot<'_> {
         let Some((tag, data)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
@@ -147,6 +211,7 @@ impl plugin::Plugin for Invasive {
         Slot { tag, data }
     }
 
+    #[expect(clippy::panic, reason = "out-of-bounds is a hard program bug")]
     unsafe fn reclaim(&self, i: u32, _: Lifecycle) {
         let Some((tag, _)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
@@ -155,6 +220,7 @@ impl plugin::Plugin for Invasive {
         tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 
+    #[expect(clippy::panic, reason = "out-of-bounds is a hard program bug")]
     unsafe fn retire(&self, i: u32, _: Lifecycle) {
         let Some((tag, _)) = self.data(i.into_usize()) else {
             panic!("index {i} is out-of-bounds");
@@ -199,7 +265,7 @@ impl<'a> Reader<'a> {
     #[inline]
     #[must_use = "this function has no side-effects"]
     pub(crate) fn id_limit(&self) -> IdLimit {
-        // Like `Invasive::id_limit`, the numberic cast is safe because by construction,
+        // Like `Invasive::id_limit`, the numeric cast is safe because by construction,
         // the underlying buffer is limited to `u32::MAX`.
         IdLimit::new(self.buffer.len() as u32)
     }
@@ -208,9 +274,12 @@ impl<'a> Reader<'a> {
     ///
     /// This guarantee only holds while `self` is alive. Construction of a new [`Reader`]
     /// requires a separate check.
-    #[expect(
-        dead_code,
-        reason = "this is non-trivial method that likely be used in the future"
+    #[cfg_attr(
+        not(test),
+        expect(
+            dead_code,
+            reason = "this is non-trivial method that is likely to be used in the future"
+        )
     )]
     pub(crate) fn can_read(&self, i: usize) -> Option<bool> {
         if !self.is_in_bounds(i) {
@@ -219,7 +288,7 @@ impl<'a> Reader<'a> {
 
         // SAFETY: We've checked that `i` is in-bounds.
         //
-        // Further, we guarantee that `self.unpadded >= AtomicTag:::SIZE`, so the pointer
+        // Further, we guarantee that `self.unpadded >= AtomicTag::SIZE`, so the pointer
         // arithmetic is in-bounds.
         let tag_ptr = unsafe {
             self.buffer
@@ -249,7 +318,7 @@ impl<'a> Reader<'a> {
         // SAFETY:
         //
         // * The caller asserts `i` is in-bounds.
-        // * We maintain an internal invariant that `self.buffer.stride() <= self.unpadded`.
+        // * We maintain the internal invariant that `self.unpadded <= self.buffer.stride()`.
         // * Further, we maintain that `self.unpadded >= AtomicTag::SIZE`.
         let (data, tag_ptr) = unsafe {
             self.buffer
@@ -277,12 +346,17 @@ impl<'a> Reader<'a> {
 
     /// Return the raw data slice for index `i` without any race guarantees.
     ///
+    /// This includes both the data **and** the invasive tag.
+    ///
     /// # Safety
     ///
-    /// The index `i` must be satisfy [`Self::is_in_bounds`].
+    /// The index `i` must satisfy [`Self::is_in_bounds`].
+    ///
+    /// The returned [`RawSlice`] may only be used for prefetching. Callers must never
+    /// materialize it as a proper slice or reference.
     #[inline]
     pub(crate) unsafe fn read_raw_unchecked(&self, i: usize) -> RawSlice<'_> {
-        // SAFETY: Inherited from caller: `i` is inbounds.
+        // SAFETY: Inherited from caller: `i` is in bounds.
         unsafe { self.buffer.get_unchecked(i) }.truncate(self.unpadded)
     }
 
@@ -295,6 +369,7 @@ impl<'a> Reader<'a> {
 /// A [`plugin::Slot`] for [`Invasive`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
+    // NOTE: `tag` and `data` must belong to the same slot.
     tag: &'a AtomicTag,
     data: RawSlice<'a>,
 }
@@ -324,3 +399,233 @@ impl plugin::Slot for Slot<'_> {
         self.tag.store(Tag::AVAILABLE, Ordering::Release);
     }
 }
+
+///////////
+// Tests //
+///////////
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use std::{
+        assert_matches,
+        num::{NonZeroU32, NonZeroUsize},
+    };
+
+    use crate::{
+        num::{Capacity, MaxDegree},
+        store,
+    };
+
+    // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
+    // zeroed frozen points. The frozen points occupy the highest slot indices.
+    fn store(
+        entries: usize,
+        entry_bytes: usize,
+        frozen: usize,
+    ) -> Result<Store<Invasive>, store::StoreError> {
+        let store = Store::new(
+            store::Layout::new(
+                Capacity::new(entries),
+                MaxDegree::new(0),
+                frozen.try_into().unwrap(),
+            ),
+            store::Config::__exhaustive(
+                NonZeroUsize::new(10).unwrap(),
+                NonZeroU32::new(16).unwrap(),
+            ),
+            Config::new(Bytes::new(entry_bytes)),
+        )?;
+
+        for (base, id) in store.frozen().enumerate() {
+            let mut slot = store.slot(id).unwrap();
+            slot.data().as_mut_slice().fill(base as u8);
+            slot.freeze();
+        }
+
+        Ok(store)
+    }
+
+    //--------//
+    // Layout //
+    //--------//
+
+    #[test]
+    fn frozen_range_follows_writable_slots() {
+        let s = store(4, 8, 2).unwrap();
+
+        // Writable slots are [0, 4); frozen points occupy [4, 6).
+        assert_eq!(s.frozen(), 4..6);
+
+        let reader = Invasive::reader(&s).unwrap();
+        for i in 0..4 {
+            assert!(!s.can_read_approximate(i).unwrap());
+            assert!(!reader.can_read(i).unwrap());
+            assert!(reader.read(i).is_none());
+        }
+
+        assert!(s.can_read_approximate(4).unwrap());
+        assert!(reader.can_read(4).unwrap());
+        assert_eq!(reader.read(4).unwrap(), &[0, 0, 0, 0, 0, 0, 0, 0]);
+
+        assert!(s.can_read_approximate(5).unwrap());
+        assert!(reader.can_read(5).unwrap());
+        assert_eq!(reader.read(5).unwrap(), &[1, 1, 1, 1, 1, 1, 1, 1]);
+
+        assert!(s.can_read_approximate(6).is_none());
+        assert!(reader.can_read(6).is_none());
+        assert!(reader.read(6).is_none());
+    }
+
+    ///////////////
+    // Lifecycle //
+    ///////////////
+
+    #[test]
+    fn acquire_write_publish_read_roundtrip() {
+        let s = store(4, 8, 1).unwrap();
+
+        let reader = Invasive::reader(&s).expect("reader guard available");
+
+        let idx = {
+            let mut slot = s.acquire().expect("a fresh store has free slots");
+            let idx = slot.slot() as usize;
+            slot.data()
+                .as_mut_slice()
+                .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
+
+            // Before the slot is dropped - we should not be able to read it.
+            assert!(reader.read(idx).is_none());
+            assert!(!s.can_read_approximate(idx).unwrap());
+            slot.publish();
+            idx
+        };
+
+        assert_eq!(reader.read(idx), Some([1, 2, 3, 4, 5, 6, 7, 8].as_slice()));
+        assert!(s.can_read_approximate(idx).unwrap());
+    }
+
+    #[test]
+    fn unpublished_slots_are_immediately_available() {
+        let s = store(4, 8, 1).unwrap();
+
+        let reader = Invasive::reader(&s).expect("reader guard available");
+
+        let idx = {
+            let mut slot = s.acquire().expect("a fresh store has free slots");
+            let idx = slot.slot() as usize;
+            slot.data()
+                .as_mut_slice()
+                .copy_from_slice(&[1, 2, 3, 4, 5, 6, 7, 8]);
+
+            // Before the slot is dropped - we should not be able to read it.
+            assert!(reader.read(idx).is_none());
+            assert!(!s.can_read_approximate(idx).unwrap());
+
+            // NOTE: We do not explicitly publish the slot.
+            idx
+        };
+
+        assert!(reader.read(idx).is_none());
+        assert!(!s.can_read_approximate(idx).unwrap());
+    }
+
+    #[test]
+    fn acquire_exhausts_then_reports_none() {
+        let s = store(2, 8, 1).unwrap();
+        // Hold the guards so the slots stay owned.
+        let _a = s.acquire().expect("first writable slot");
+        let _b = s.acquire().expect("second writable slot");
+        assert!(
+            s.acquire().is_none(),
+            "all writable slots are owned, so acquire must fail"
+        );
+    }
+
+    //--------//
+    // Retire //
+    //--------//
+
+    #[test]
+    fn retire_out_of_bounds() {
+        let s = store(4, 8, 1).unwrap();
+        assert_matches!(s.retire(999), Err(store::RetireError::OutOfBounds));
+    }
+
+    #[test]
+    fn retire_rejects_reserved_slots() {
+        let s = store(4, 8, 1).unwrap();
+        // An untouched writable slot is AVAILABLE, which is a reserved state.
+        assert_matches!(s.retire(0), Err(store::RetireError::SlotIsReserved { .. }));
+        // A frozen slot is likewise reserved.
+        let frozen = s.frozen().start as usize;
+        assert_matches!(
+            s.retire(frozen),
+            Err(store::RetireError::SlotIsReserved { .. })
+        );
+        // An owned slot is not retirable.
+        let slot = s.acquire().unwrap();
+        assert_matches!(
+            s.retire(slot.slot() as usize),
+            Err(store::RetireError::SlotIsReserved { .. })
+        );
+    }
+
+    #[test]
+    fn retire_published_slot_then_unreadable() {
+        let s = store(4, 8, 1).unwrap();
+
+        let idx = {
+            let slot = s.acquire().unwrap();
+            slot.publish() as usize
+        };
+
+        assert!(s.retire(idx).is_ok());
+
+        // A reader opened after retirement must not observe the retired slot.
+        let reader = Invasive::reader(&s).unwrap();
+        assert_eq!(reader.read(idx), None);
+        assert_eq!(reader.can_read(idx), Some(false));
+
+        // The slot can also not be retired again.
+        assert_matches!(
+            s.retire(idx),
+            Err(store::RetireError::SlotIsReserved { .. })
+        );
+    }
+
+    //---------//
+    // Recycle //
+    //---------//
+
+    #[test]
+    fn test_recycling() {
+        let entries = if cfg!(miri) { 16 } else { 2048 };
+
+        let s = store(entries, 4, 2).unwrap();
+
+        // Claim all slots.
+        let mut count = 0;
+        while let Some(slot) = s.acquire() {
+            slot.publish();
+            count += 1;
+        }
+
+        assert_eq!(count, s.writable().len());
+
+        // Now that all slots are claimed - retire all slots.
+        for i in s.writable() {
+            s.retire(i.into_usize()).unwrap();
+        }
+
+        // Verify that we can claim all slots again.
+        let mut count = 0;
+        while let Some(slot) = s.acquire() {
+            slot.publish();
+            count += 1;
+        }
+
+        assert_eq!(count, s.writable().len());
+    }
+}
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index bc60c7e9ba..1b5a4ce355 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -141,7 +141,7 @@ impl Config {
     /// This is under the "integration-test" since it will change to reflect the state
     /// of the underlying data structure, potentially causing more churn for users if it
     /// were unconditionally exposed.
-    #[cfg(feature = "integration-test")]
+    #[cfg(any(test, feature = "integration-test"))]
     #[doc(hidden)]
     pub fn __exhaustive(
         epoch_guard_slots: NonZeroUsize,
@@ -293,6 +293,7 @@ where
         let drain = self.registry.try_advance()?;
         let items = drain.len();
         for i in drain {
+            #[expect(clippy::panic, reason = "this is an unrecoverable program bug")]
             let Some(tag) = self.tags.get(i.into_usize()) else {
                 panic!(
                     "received an invalid ID ({}) while reclaiming slots - max allowed is {}",
@@ -303,6 +304,12 @@ where
 
             // We release the plugin before the main tag. The other direction would
             // prematurely advertise availability.
+            //
+            // SAFETY: IDs only get added to the `epoch::Registry` only upon retiring, and
+            // are not released until the registry confirms that all guards active then the
+            // id was retired have been dropped.
+            //
+            // Therefore, this slot has no accessors and is ready to be reclaimed.
             unsafe { plugin::Plugin::reclaim(self.plugin(), i, Lifecycle::new()) };
 
             // Use `Release` ordering to ensure that the store to the mirror cannot get moved
@@ -393,8 +400,12 @@ where
         match tag.compare_exchange(current, retiring, Ordering::Relaxed, Ordering::Relaxed) {
             Ok(_) => {
                 // Set the metadata in the mirror as well.
+                //
+                // SAFETY: The above compare-exchange ensures that we transitioned the
+                // authoritative state from "published" to "retired" and prevents other
+                // threads from attempting the same transition.
                 unsafe {
-                    plugin::Plugin::retire(self.plugin(), i.try_into().unwrap(), Lifecycle::new())
+                    plugin::Plugin::retire(self.plugin(), i as u32, Lifecycle::new())
                 };
                 guard.retire(i as u32);
                 Ok(())
@@ -494,8 +505,14 @@ where
             Ordering::Relaxed,
         ) {
             Ok(_) => {
+                // SAFETY: The above compare-exchange ensures that this slot was previously
+                // "available" and prevents other threads from trying acquire this slot.
+                //
+                // The `Slot` data structure ensures that exactly one of the terminal methods
+                // for `plugin::Slot` is called.
                 let data =
                     unsafe { plugin::Plugin::acquire(self.plugin(), slot, Lifecycle::new()) };
+
                 Some(Slot {
                     tag,
                     data: ManuallyDrop::new(data),
@@ -543,8 +560,12 @@ impl StoreError {
         })
     }
 
-    fn plugin(err: ANNError) -> Self {
-        Self(StoreErrorInner::PluginError(err))
+    #[track_caller]
+    fn plugin<E>(err: E) -> Self
+    where
+        E: std::error::Error + Send + Sync + 'static,
+    {
+        Self(StoreErrorInner::PluginError(ANNError::new(err)))
     }
 }
 
@@ -600,6 +621,10 @@ pub(crate) enum RetireError {
 diskann::convert_error!(RetireError);
 
 /// A writable buffer into the data managed by a [`Store`], obtained from [`Store::acquire`].
+///
+/// This is the only safe way to interace with a [`plugin::Slot`] since this ensure that one
+/// of the terminal methods is called. Dropping a [`Slot`] without calling [`Slot::publish`]
+/// or [`Slot::freeze`] automatically invokes [`plugin::Slot::abort`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a, S>
 where
@@ -630,6 +655,7 @@ where
 
         // Freeze the inner slot.
         plugin::Slot::freeze(
+            // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut me.data) },
             Lifecycle::new(),
         );
@@ -649,6 +675,7 @@ where
 
         // Publish the inner slot.
         plugin::Slot::publish(
+            // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut me.data) },
             Lifecycle::new(),
         );
@@ -665,6 +692,7 @@ where
 {
     fn drop(&mut self) {
         plugin::Slot::abort(
+            // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut self.data) },
             Lifecycle::new(),
         );
@@ -688,9 +716,8 @@ mod tests {
     // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
     // zeroed frozen points. The frozen points occupy the highest slot indices.
     fn store(entries: usize, frozen: u32) -> Result<Store<Checked>, StoreError> {
-        let mut config = Config::new();
-        config.epoch_guard_slots(NonZeroUsize::new(10).unwrap());
-        config.freelist_recycle_capacity(NonZeroU32::new(16).unwrap());
+        let config =
+            Config::__exhaustive(NonZeroUsize::new(10).unwrap(), NonZeroU32::new(16).unwrap());
 
         let layout = Layout::new(Capacity::new(entries), MaxDegree::new(0), frozen);
         let store = Store::new(layout, config, Checked::config())?;
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index e700bee8eb..ac3977026f 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-//! # EBR lifecycle hooks for [`super::Store`].
+//! # EBR lifecycle hooks for [`super::Store`]
 //!
 //! Please read this section carefully - the protocol is not difficult, but it *is* subtle.
 //!
@@ -68,7 +68,6 @@
 //! state. Further, the store commits the destination state only after the plugin API call
 //! completes.
 
-use diskann::ANNResult;
 use std::fmt::Debug;
 
 use crate::num::IdLimit;
@@ -80,8 +79,11 @@ pub(crate) trait PluginConfig: Debug {
     /// The type of the resulting [`Plugin`].
     type Plugin: Plugin;
 
+    /// Construction errors.
+    type Error: std::error::Error + Send + Sync + 'static;
+
     /// Build the associated [`Plugin`] from self with the [`IdLimit`].
-    fn build(self, id_limit: IdLimit) -> ANNResult<Self::Plugin>;
+    fn build(self, id_limit: IdLimit) -> Result<Self::Plugin, Self::Error>;
 }
 
 /// A lifecycle backend for [`super::Store`]'s EBR scheme.

From ba34c05ff36fbb2e3a5d745ea1a4c8dbae24a579 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Fri, 21 Aug 2026 13:16:43 -0700
Subject: [PATCH 20/34] Test for mismatched `id_limit`.

---
 diskann-inmem/src/store/mod.rs | 151 +++++++++++++++++++++++----------
 1 file changed, 107 insertions(+), 44 deletions(-)

diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 1b5a4ce355..14d62b9775 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-//! A concurrent in-memory data store for uniformly sized data.
+//! A concurrent in-memory data store for driving [`plugin::Plugin`]s.
 //!
 //! This supports concurrent data access, deletes, and inserts through a safe interface.
 //! Data is stored internally in slots indexed from `[0..N)` with `K` points reserved at the
@@ -11,14 +11,19 @@
 //!
 //! ## Reading
 //!
-//! Read access requires a [`Reader`] produced by [`Store::reader`]. [`Reader::read`]
-//! provides read-only access to data at slot `i` if the data is valid for reads.
+//! A [`Store`] provides no direct way of reading data. Instead, the [`plugin::Plugin`] is
+//! responsible for exposing an appropriate reader (e.g., [`invasive::Invasive::reader`]) in
+//! accordance with its lifecycle implementation. [`Store::guard`] can be used for
+//! this purpose by acquiring an [`epoch::Guard`] for a [`Store`].
 //!
 //! ## Writing
 //!
 //! [`Store::acquire`] is used to find and claim an unused internal [`Slot`]. A [`Slot`]
-//! provides write access to its corresponding data which is published when the [`Slot`] is
-//! dropped.
+//! provides write access to its corresponding [`Slot::data`]. Either [`Slot::publish`] or
+//! [`Slot::freeze`] can be used to make data readable.
+//!
+//! If a [`Slot`] is dropped, its corresponding slot is returned to the [`Store`] without
+//! publishing its contents.
 //!
 //! The index of the slot chosen may be obtained via [`Slot::slot`].
 //!
@@ -26,7 +31,7 @@
 //!
 //! Data is deleted via [`Store::retire`]. This immediately marks the corresponding slot as
 //! unavailable for future readers. However, the retired slot will not be reused until the
-//! [`Store`] can guarantee that no [`Reader`]s that could be using the data are active.
+//! [`Store`] can guarantee that no readers that could be using the data are active.
 //!
 //! Slots are automatically reclaimed as part of slot acquisition in the "writing" phase.
 //!
@@ -39,17 +44,12 @@
 //! # Details
 //!
 //! This uses an implementation of the epoch-based reclamation (EBR) provided by [`Registry`].
-//! Concurrency tags are mirrored inline with the stored data (just after the data payload)
-//! to keep memory access localized. As such, high-performance implementations will want to
-//! fetch the last cache line of data first to ensure the tag is resident in cache for faster
-//! data checks.
+//! Plugins follow the lifecycle process defined in the [plugin module docs](plugin).
 //!
 //! The EBR scheme allows readers to safely access data while only generating read traffic to
 //! the CPU caches. The cost is that there is a delay between when slots are retired and when
-//! they can be reused, with a long lived [`Reader`] blocking this reclamation. As such,
-//! users of this data structure should ensure that [`Reader`]s are reasonably short lived.
-//!
-//! Internally, the data belongs to a single allocation.
+//! they can be reused, with a long lived reader blocking this reclamation. As such, users of
+//! this data structure should ensure that readers are reasonably short lived.
 
 use std::{
     iter::repeat_n,
@@ -92,7 +92,7 @@ impl Lifecycle {
     }
 }
 
-/// Configuration for the concurrenct store.
+/// Configuration for the concurrent store.
 #[derive(Debug, Clone)]
 pub struct Config {
     /// The number of epoch guard slots.
@@ -136,11 +136,11 @@ impl Config {
         self
     }
 
-    /// An exhaustive constructor initializing everye element.
+    /// An exhaustive constructor initializing every element.
     ///
-    /// This is under the "integration-test" since it will change to reflect the state
-    /// of the underlying data structure, potentially causing more churn for users if it
-    /// were unconditionally exposed.
+    /// This is exposed under "integration-test" since it will change to reflect the state
+    /// of the underlying data structure, potentially causing more churn for users if it were
+    /// unconditionally exposed.
     #[cfg(any(test, feature = "integration-test"))]
     #[doc(hidden)]
     pub fn __exhaustive(
@@ -160,6 +160,7 @@ impl Default for Config {
     }
 }
 
+/// Layout parameters for [`Store`] and the corresponding [`plugin::Plugin`].
 #[derive(Debug, Clone)]
 pub(crate) struct Layout {
     /// The number of non-frozen slots to create space for.
@@ -173,7 +174,13 @@ pub(crate) struct Layout {
 }
 
 impl Layout {
-    /// Create a new [`Layout`] capable of holding `capacity` non-frozen points.
+    /// Create a new [`Layout`] capable of holding `capacity` non-frozen points and `frozen`
+    /// reserved points at the end.
+    ///
+    /// After construction, the only distinction between non-frozen and frozen points is that
+    /// IDs for frozen points will not be selected through [`Store::acquire`]. Frozen slots
+    /// IDs can be obtained via [`Store::frozen`], and direct acquisition can be done with
+    /// [`Store::slot`].
     ///
     /// All adjacency lists will have a maximum capacity of `max_degree`.
     pub(crate) fn new(capacity: Capacity, max_degree: MaxDegree, frozen: u32) -> Self {
@@ -191,11 +198,13 @@ pub(crate) struct Store<P> {
     // The [`plugin::Plugin`] managed by this [`Store`].
     plugin: P,
 
-    // The number of unfrozen points. This is guaranteed to be less than `buffer`.
+    // The number of unfrozen points.
     unfrozen: Capacity,
 
     // The authoritative source of truth for the state of each slot.
     tags: Vec<AtomicTag>,
+
+    // Acceleration of finding free slot IDs.
     freelist: Freelist,
 
     // EBR registry.
@@ -212,8 +221,7 @@ impl<P> Store<P>
 where
     P: plugin::Plugin,
 {
-    /// Create a new [`Store`]. The entries within `init` will be used as frozen points
-    /// within the store and must be compatible the the number of bytes in `config`.
+    /// Create a new [`Store`].
     pub(crate) fn new<C>(layout: Layout, config: Config, plugin: C) -> Result<Self, StoreError>
     where
         C: plugin::PluginConfig<Plugin = P>,
@@ -248,6 +256,11 @@ where
 
         let plugin = plugin::PluginConfig::build(plugin, id_limit).map_err(StoreError::plugin)?;
 
+        let plugin_id_limit = plugin.id_limit();
+        if plugin_id_limit != id_limit {
+            return Err(StoreError::invalid_construction(plugin_id_limit, id_limit));
+        }
+
         let me = Self {
             plugin,
             unfrozen: capacity,
@@ -265,6 +278,7 @@ where
         Ok(me)
     }
 
+    /// Return the [`plugin::Plugin`] for this store.
     pub(crate) fn plugin(&self) -> &P {
         &self.plugin
     }
@@ -274,14 +288,19 @@ where
         (self.unfrozen.value() as u32)..self.neighbors.entries()
     }
 
+    /// Return the [`IdLimit`] for this store.
     pub(crate) fn id_limit(&self) -> IdLimit {
-        plugin::Plugin::id_limit(&self.plugin)
+        // The numeric cast is safe: We verify during construction that `self.tags.len()` fits
+        // in a `u32`.
+        IdLimit::new(self.tags.len() as u32)
     }
 
+    /// Return the [`Capacity`] for this store.
     pub(crate) fn capacity(&self) -> Capacity {
         self.unfrozen
     }
 
+    /// Return the [`Neighbors`] for this store.
     pub(crate) fn neighbors(&self) -> &Neighbors {
         &self.neighbors
     }
@@ -305,9 +324,9 @@ where
             // We release the plugin before the main tag. The other direction would
             // prematurely advertise availability.
             //
-            // SAFETY: IDs only get added to the `epoch::Registry` only upon retiring, and
-            // are not released until the registry confirms that all guards active then the
-            // id was retired have been dropped.
+            // SAFETY: IDs only get added to the `epoch::Registry` upon retiring, and are
+            // not released until the registry confirms that all guards active when the id
+            // was retired have been dropped.
             //
             // Therefore, this slot has no accessors and is ready to be reclaimed.
             unsafe { plugin::Plugin::reclaim(self.plugin(), i, Lifecycle::new()) };
@@ -318,8 +337,9 @@ where
             // The `load + check` is just runtime validation. The calling thread is expected
             // to have exclusive ownership of this tag.
             //
-            // Using a load + store avoids using a CAS style loop, which can be cheaper for
-            // the bulk styl operations we're going here at the cost of precision.
+            // Using a load followed by a store avoids a CAS loop, which may be cheaper for
+            // bulk reclamation at the cost of detecting concurrent modification only through
+            // the assertion.
             assert_eq!(
                 tag.load(Ordering::Relaxed),
                 Tag::RETIRING,
@@ -332,6 +352,11 @@ where
         Some(items)
     }
 
+    /// Create an [`epoch::Guard`] for the [`epoch::Registry`] within `self` and invoke `f`
+    /// with that guard, returning the result.
+    ///
+    /// This can be used by [`plugin::Plugin`] readers to establish a verifiable chain of
+    /// custody for an [`epoch::Guard`] over the plugin.
     pub(crate) fn guard<'a, F, R>(&'a self, f: F) -> Result<R, epoch::Unavailable>
     where
         F: FnOnce(&'a P, epoch::Guard<'a>) -> R,
@@ -404,9 +429,7 @@ where
                 // SAFETY: The above compare-exchange ensures that we transitioned the
                 // authoritative state from "published" to "retired" and prevents other
                 // threads from attempting the same transition.
-                unsafe {
-                    plugin::Plugin::retire(self.plugin(), i as u32, Lifecycle::new())
-                };
+                unsafe { plugin::Plugin::retire(self.plugin(), i as u32, Lifecycle::new()) };
                 guard.retire(i as u32);
                 Ok(())
             }
@@ -426,8 +449,8 @@ where
     /// Periodically, the freelist is checked to see if another thread has found an available
     /// slot for us.
     fn scan_acquire(&self) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
-        // This is potentially quite slow - but stop if we've scanned the entire range
-        // without finding anything.
+        // This is potentially quite slow, so scan approximately `1 / RETRY_LIMIT` of the
+        // writable range. The outer retry loop provides broader coverage.
         let mut remaining = self.unfrozen.value().div_ceil(RETRY_LIMIT);
         let mut chunks_since_freelist_check = 0;
         let mut acquired: Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> = None;
@@ -476,6 +499,10 @@ where
         None
     }
 
+    /// Attempt to directly acquire a [`Slot`] to id `i`.
+    ///
+    /// Returns `None` if `i` is not within [`Self::id_limit`] or if the slot is not currently
+    /// acquirable.
     pub(crate) fn slot(&self, i: u32) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
         let tag = &self.tags.get(i.into_usize())?;
 
@@ -501,12 +528,15 @@ where
         match tag.compare_exchange(
             Tag::AVAILABLE,
             Tag::OWNED,
-            Ordering::Relaxed,
+            Ordering::Acquire,
             Ordering::Relaxed,
         ) {
             Ok(_) => {
                 // SAFETY: The above compare-exchange ensures that this slot was previously
-                // "available" and prevents other threads from trying acquire this slot.
+                // "available" and prevents other threads from trying to acquire this slot.
+                // The acquire ordering synchronizes with the release transition to
+                // "available", making plugin reclamation or abort work visible before
+                // `Plugin::acquire`.
                 //
                 // The `Slot` data structure ensures that exactly one of the terminal methods
                 // for `plugin::Slot` is called.
@@ -525,10 +555,10 @@ where
 
     /// Return whether or not it is probably okay to read from the slot `i`.
     ///
-    /// This check is approximate and non-synchronizing. To fully check, [`Reader::can_read`]
-    /// must be used.
+    /// This check is approximate and non-synchronizing. A full check requires the
+    /// plugin-specific reader.
     ///
-    /// Returns `None` is index `i` is out-of-bounds.
+    /// Returns `None` if `i` is not within [`Self::id_limit`].
     pub(crate) fn can_read_approximate(&self, i: usize) -> Option<bool> {
         self.tags
             .get(i)
@@ -567,6 +597,10 @@ impl StoreError {
     {
         Self(StoreErrorInner::PluginError(ANNError::new(err)))
     }
+
+    fn invalid_construction(got: IdLimit, expected: IdLimit) -> Self {
+        Self(StoreErrorInner::InvalidConstruction { got, expected })
+    }
 }
 
 impl From<BufferError> for StoreError {
@@ -599,6 +633,8 @@ enum StoreErrorInner {
     NeighborsError(#[from] NeighborsError),
     #[error("error creating plugin")]
     PluginError(ANNError),
+    #[error("requested {} but the plugin returned {}", expected, got)]
+    InvalidConstruction { got: IdLimit, expected: IdLimit },
 }
 
 /// Error conditions for [`Store::retire`].
@@ -622,7 +658,7 @@ diskann::convert_error!(RetireError);
 
 /// A writable buffer into the data managed by a [`Store`], obtained from [`Store::acquire`].
 ///
-/// This is the only safe way to interace with a [`plugin::Slot`] since this ensure that one
+/// This is the only safe way to interact with a [`plugin::Slot`] since this ensures that one
 /// of the terminal methods is called. Dropping a [`Slot`] without calling [`Slot::publish`]
 /// or [`Slot::freeze`] automatically invokes [`plugin::Slot::abort`].
 #[derive(Debug)]
@@ -706,13 +742,27 @@ where
 
 /// These tests are basic functionality tests for the store.
 ///
-/// Longer running conurrency tests are in the integration test suite.
+/// Longer running concurrency tests are in the integration test suite.
 #[cfg(test)]
 mod tests {
     use super::{checked::Checked, *};
 
     use std::assert_matches;
 
+    /// A faulty config for [`Checked`] that doesn't respect the [`IdLimit`].
+    #[derive(Debug)]
+    struct FaultyConfig;
+
+    impl plugin::PluginConfig for FaultyConfig {
+        type Plugin = Checked;
+        type Error = diskann::error::Infallible;
+
+        fn build(self, id_limit: IdLimit) -> Result<Checked, diskann::error::Infallible> {
+            let faulty = id_limit.value().checked_sub(1).unwrap_or(1);
+            Ok(Checked::new(IdLimit::new(faulty)))
+        }
+    }
+
     // Build a store with `entries` writable slots of `entry_bytes` each, backed by `frozen`
     // zeroed frozen points. The frozen points occupy the highest slot indices.
     fn store(entries: usize, frozen: u32) -> Result<Store<Checked>, StoreError> {
@@ -721,7 +771,6 @@ mod tests {
 
         let layout = Layout::new(Capacity::new(entries), MaxDegree::new(0), frozen);
         let store = Store::new(layout, config, Checked::config())?;
-
         assert_eq!(store.frozen().len(), frozen.into_usize());
 
         for (i, id) in store.frozen().enumerate() {
@@ -750,7 +799,7 @@ mod tests {
             Checked::config(),
         )
         .unwrap_err();
-        assert!(matches!(err.0, StoreErrorInner::TooManyEntries { .. }));
+        assert_matches!(err.0, StoreErrorInner::TooManyEntries { .. });
     }
 
     #[test]
@@ -765,7 +814,21 @@ mod tests {
             Checked::config(),
         )
         .unwrap_err();
-        assert!(matches!(err.0, StoreErrorInner::TooManyNeighbors { .. }));
+        assert_matches!(err.0, StoreErrorInner::TooManyNeighbors { .. });
+    }
+
+    #[test]
+    fn new_rejects_faulty_plugin() {
+        let err = Store::new(
+            Layout::new(
+                Capacity::new(4),
+                MaxDegree::new(10),
+                0,
+            ),
+            Config::default(),
+            FaultyConfig,
+        ).unwrap_err();
+        assert_matches!(err.0, StoreErrorInner::InvalidConstruction { .. });
     }
 
     //--------//

From 1b1fda099e466fee5bee89617464f70709af2cb7 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Fri, 21 Aug 2026 15:07:26 -0700
Subject: [PATCH 21/34] Checkpoint.

---
 diskann-inmem/integration/store/checked.rs    |  10 +-
 diskann-inmem/integration/store/invasive.rs   |   8 +-
 diskann-inmem/integration/store/mod.rs        |  60 +++++++++--
 .../src/integration/store/checked.rs          |   1 +
 diskann-inmem/src/integration/store/mod.rs    |  14 ++-
 diskann-inmem/src/layers/full.rs              | 101 +++++++++++-------
 6 files changed, 136 insertions(+), 58 deletions(-)

diff --git a/diskann-inmem/integration/store/checked.rs b/diskann-inmem/integration/store/checked.rs
index 2078ac81c8..9a139349ad 100644
--- a/diskann-inmem/integration/store/checked.rs
+++ b/diskann-inmem/integration/store/checked.rs
@@ -13,7 +13,7 @@ pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::Registry
     registry.register("store-stress-test-checked", Stress)
 }
 
-/// Configuration for a [`StoreStress`] run.
+/// Configuration for a [`Stress`] run.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct Input {
     /// Shared stress test setup.
@@ -64,7 +64,7 @@ struct Stress;
 
 impl dbr::Benchmark for Stress {
     type Input = Input;
-    type Output = super::StoreStressStats;
+    type Output = super::Stats;
 
     fn try_match(
         &self,
@@ -126,11 +126,11 @@ impl super::Testable for checked::Store {
     }
 
     fn readable_slots(&self) -> usize {
-        <checked::Store>::slots(self)
+        <checked::Store>::readable_slots(self)
     }
 
     fn writable_slots(&self) -> usize {
-        <checked::Store>::writable(self)
+        <checked::Store>::writable_slots(self)
     }
 }
 
@@ -202,7 +202,7 @@ impl super::Reader for Reader<'_> {
                 }
             }
             // Readable -> unreadable: an allowed, terminal transition.
-            (Some(previous), None) => {
+            (Some(_), None) => {
                 self.shared.transitions.fetch_add(1, Relaxed);
             }
         }
diff --git a/diskann-inmem/integration/store/invasive.rs b/diskann-inmem/integration/store/invasive.rs
index 8f826381c1..ede89d8111 100644
--- a/diskann-inmem/integration/store/invasive.rs
+++ b/diskann-inmem/integration/store/invasive.rs
@@ -13,7 +13,7 @@ pub(super) fn register(registry: &mut dbr::Registry) -> Result<(), dbr::Registry
     registry.register("invasive-store-stress-test", Stress)
 }
 
-/// Configuration for a [`StoreStress`] run.
+/// Configuration for a [`Stress`] run.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct Input {
     /// Shared stress test setup.
@@ -77,7 +77,7 @@ struct Stress;
 
 impl dbr::Benchmark for Stress {
     type Input = Input;
-    type Output = super::StoreStressStats;
+    type Output = super::Stats;
 
     fn try_match(
         &self,
@@ -174,11 +174,11 @@ impl super::Testable for invasive::Store {
     }
 
     fn readable_slots(&self) -> usize {
-        <invasive::Store>::slots(self)
+        <invasive::Store>::readable_slots(self)
     }
 
     fn writable_slots(&self) -> usize {
-        <invasive::Store>::writable(self)
+        <invasive::Store>::writable_slots(self)
     }
 }
 
diff --git a/diskann-inmem/integration/store/mod.rs b/diskann-inmem/integration/store/mod.rs
index 04b2cf6884..799e106056 100644
--- a/diskann-inmem/integration/store/mod.rs
+++ b/diskann-inmem/integration/store/mod.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-//! Concurrency stress test for the in-memory [`Store`](diskann_inmem::integration::store::Store).
+//! Concurrency stress test for the in-memory [stores](diskann_inmem::integration::store).
 //!
 //! Reader, writer, and retirer threads hammer the epoch-based store concurrently while a
 //! per-guard invariant checker verifies the store's safety guarantees:
@@ -11,6 +11,8 @@
 //! 1. Reads are never torn.
 //! 2. A readable value is stable for the lifetime of a single reader guard.
 //! 3. A slot never resurrects (`readable -> unreadable -> readable`) within one guard.
+//!
+//! This module exposes shared functionality that is instantiated by different plugins.
 
 #![expect(
     clippy::unwrap_used,
@@ -143,45 +145,91 @@ impl std::fmt::Display for Setup {
     }
 }
 
+/// A testable store.
+///
+/// Readers are split into a [`ReaderState`], which is used to produce a shorter lived [`Reader`].
+///
+/// The reason for this is two fold:
+///
+/// 1. We rely on [`Reader`]s being dropped to allow epochs to advance.
+/// 2. A separate [`ReaderState`] allows correctness checking data structures (e.g. hash maps)
+///    to be allocated once and used for the duration of the test.
+///
+///    Since the goal is to hammer the underlying store as hard as possible, amortizing
+///    allocations makes a non-negligible difference.
 trait Testable: std::fmt::Debug + Sized + Sync {
+    /// The writer.
     type Writer<'a>: Writer
     where
         Self: 'a;
 
+    /// Shared reader-state to enable allocation amortization.
     type ReaderState<'a>: ReaderState
     where
         Self: 'a;
 
+    /// Construct a [`Writer`] into a slot. Returns `None` is an available slot could not
+    /// be found.
     fn writer(&self) -> Option<Self::Writer<'_>>;
 
+    /// Create the [`ReaderState`]. This will be called once per reader thread.
     fn reader_state<'a>(
         &'a self,
         capacity_hint: usize,
         shared: &'a Shared<Self>,
     ) -> Self::ReaderState<'a>;
 
+    /// Attempt to retire slot `i`. Return `true` on success, otherwise return `false`.
     fn retire(&self, i: usize) -> bool;
 
+    /// Attempt to advance the epoch and reclaim retired slots.
+    ///
+    /// Return `None` if we failed to advance the epoch. Otherwise, return the number of
+    /// slots reclaimed.
     fn reclaim(&self) -> Option<usize>;
 
+    /// Return the number of readable slots. Assume indices `[0..readable_slots)` are valid
+    /// for reading.
     fn readable_slots(&self) -> usize;
+
+    /// Return the number of writable slots. Assume indices `[0..writable_slots)` are valid
+    /// for writing.
     fn writable_slots(&self) -> usize;
 }
 
+/// A writable slot.
 trait Writer: std::fmt::Debug {
+    /// Perform any writes, using `stamp` as a unique tag.
     fn write(self, stamp: u64);
 }
 
+/// Amortized shared state for reader tasks.
 trait ReaderState: std::fmt::Debug {
+    /// The type of the [`Reader`].
     type Reader<'a>: Reader;
 
+    /// Attempt to run the closure `f` on a [`Reader`] into the state's parent store.
+    ///
+    /// Return `true` if a reader was obtained and `f` was called. Otherwise return `false`.
+    ///
+    /// The callback style mechanism is used to allow implementations to stack-allocate some
+    /// variables. This allows validation implementations to hold onto references into the
+    /// parent store by doing the following:
+    ///
+    /// 1. Stack allocate the internal "reader" to the store. The reader will have an epoch
+    ///    guard.
+    ///
+    /// 2. Borrow from that reader to construct [`Self::Reader`], allowing [`Self::Reader`]
+    ///    to borrow items directly from the internal reader.
     #[must_use]
     fn try_with_reader<F>(&mut self, f: F) -> bool
     where
         F: FnOnce(Self::Reader<'_>);
 }
 
+/// A read validator for a [`Testable`].
 trait Reader: std::fmt::Debug {
+    /// Observe the state of `i`, panicking if an invalid transition has been observed.
     fn observe(&mut self, i: usize);
 }
 
@@ -189,9 +237,9 @@ trait Reader: std::fmt::Debug {
 // Output //
 ////////////
 
-/// Summary statistics produced by a [`StoreStress`] run.
+/// Summary statistics produced by a [`Shared`] run.
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct StoreStressStats {
+pub struct Stats {
     elapsed_secs: f64,
     reads: u64,
     acquires_ok: u64,
@@ -205,7 +253,7 @@ pub struct StoreStressStats {
     peak_live: usize,
 }
 
-impl std::fmt::Display for StoreStressStats {
+impl std::fmt::Display for Stats {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let mut kv = KeyValue::new();
         kv.push("elapsed_secs", &self.elapsed_secs);
@@ -272,7 +320,7 @@ impl Drop for LocalMax<'_> {
     }
 }
 
-fn run_benchmark<T>(store: T, setup: &Setup) -> anyhow::Result<StoreStressStats>
+fn run_benchmark<T>(store: T, setup: &Setup) -> anyhow::Result<Stats>
 where
     T: Testable,
 {
@@ -334,7 +382,7 @@ where
     }
 
     let elapsed = start.elapsed();
-    let stats = StoreStressStats {
+    let stats = Stats {
         elapsed_secs: elapsed.as_secs_f64(),
         reads: shared.reads.load(Relaxed),
         acquires_ok: shared.acquires_ok.load(Relaxed),
diff --git a/diskann-inmem/src/integration/store/checked.rs b/diskann-inmem/src/integration/store/checked.rs
index fb84a13ad9..1aaa7e6fc2 100644
--- a/diskann-inmem/src/integration/store/checked.rs
+++ b/diskann-inmem/src/integration/store/checked.rs
@@ -5,6 +5,7 @@
 
 #![expect(
     clippy::expect_used,
+    clippy::unwrap_used,
     reason = "integration test tools are not production code"
 )]
 
diff --git a/diskann-inmem/src/integration/store/mod.rs b/diskann-inmem/src/integration/store/mod.rs
index 1d39b97cd2..2254e897ee 100644
--- a/diskann-inmem/src/integration/store/mod.rs
+++ b/diskann-inmem/src/integration/store/mod.rs
@@ -3,6 +3,12 @@
  * Licensed under the MIT license.
  */
 
+//! This module exposes "public" integration-test wrappers for the various internal store
+//! mechanisms to drive larger concurrency tests.
+//!
+//! These implementationa have a similar structure. A [`boilerplate`] macro is used to ensure
+//! the capabilities exposed are mostly the same.
+
 pub mod checked;
 pub mod invasive;
 
@@ -12,6 +18,7 @@ macro_rules! boilerplate {
         for<$read_lt:lifetime> $read:ty => $reader:ident,
         for<$slot_lt:lifetime> $slot:ty => $writer:ident,
     ) => {
+        /// A test store wraper.
         #[derive(Debug)]
         pub struct $store {
             store: $crate::store::Store<$plugin>,
@@ -19,12 +26,12 @@ macro_rules! boilerplate {
 
         impl $store {
             /// Return the total number of slots, including the frozen point.
-            pub fn slots(&self) -> usize {
+            pub fn readable_slots(&self) -> usize {
                 self.store.frozen().end as usize
             }
 
             /// Return the range of writable (non-frozen) slot indices.
-            pub fn writable(&self) -> usize {
+            pub fn writable_slots(&self) -> usize {
                 self.store.frozen().start as usize
             }
 
@@ -56,6 +63,7 @@ macro_rules! boilerplate {
             }
         }
 
+        /// A reader for the test store.
         #[derive(Debug)]
         pub struct $reader<$read_lt> {
             reader: $read,
@@ -67,6 +75,7 @@ macro_rules! boilerplate {
             }
         }
 
+        /// A writer for the test store.
         #[derive(Debug)]
         pub struct $writer<$slot_lt> {
             slot: $crate::store::Slot<$slot_lt, $slot>,
@@ -77,6 +86,7 @@ macro_rules! boilerplate {
                 Self { slot }
             }
 
+            /// Publish the slot - making it accessible to readers.
             pub fn publish(self) {
                 self.slot.publish();
             }
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index bbb3604ce6..8664fa5320 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -109,39 +109,6 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
 }
 
-impl<T> FullPrecision for T
-where
-    T: FullPrecisionImpl,
-{
-    fn __search_accessor<'a>(
-        layer: &'a Full<Self>,
-        query: &'a [Self],
-        provider: &'a (dyn std::any::Any + Send + Sync),
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        let expand_beam = T::make_expand_beam(layer, query)?;
-        Ok(crate::provider::SearchAccessor::new(
-            layer.store.neighbors(),
-            expand_beam,
-            provider,
-            layer.store.frozen(),
-            counters,
-        ))
-    }
-
-    fn __prune_accessor<'a>(
-        layer: &'a Full<Self>,
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
-        let prune = T::make_prune(layer)?;
-        Ok(crate::provider::PruneAccessor::new(
-            prune,
-            layer.store.neighbors(),
-            counters,
-        ))
-    }
-}
-
 /// Full-precision data layer.
 #[derive(Debug)]
 pub struct Full<T>
@@ -381,7 +348,8 @@ struct Prune<'a, T, D> {
 
 impl<'a, T, D> Prune<'a, T, D> {
     fn new(reader: store::invasive::Reader<'a>) -> Self {
-        assert!(
+        // This should be ensured at construction time
+        debug_assert!(
             reader
                 .bytes()
                 .value()
@@ -468,7 +436,7 @@ impl<T> std::ops::Deref for Calf<'_, T> {
 /// allow `f16` queries to be pre-converted to `f32`, saving on-the-fly conversion that
 /// would otherwise be needed.
 #[derive(Debug)]
-struct QueryDistance<'a, const PREFETCH: usize, T, U, D> {
+struct QueryDistance<'a, const PREFETCH_DIM: usize, T, U, D> {
     // The original query.
     query: Calf<'a, T>,
     // A reader into a layer's store.
@@ -479,9 +447,8 @@ struct QueryDistance<'a, const PREFETCH: usize, T, U, D> {
     _distance: PhantomData<D>,
 }
 
-impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
+impl<'a, const PREFETCH_DIM: usize, T, U, D> QueryDistance<'a, PREFETCH_DIM, T, U, D> {
     fn new(query: Calf<'a, T>, reader: store::invasive::Reader<'a>) -> Self {
-        // TODO: Check PREFETCH and `query` with the reader's size.
         Self {
             query,
             reader,
@@ -512,11 +479,19 @@ impl<'a, const PREFETCH: usize, T, U, D> QueryDistance<'a, PREFETCH, T, U, D> {
         if x.len() != self.bytes() {
             self.error(x.len())
         } else {
-            // SAFETY: We've validated that `x` has the correct length.
-            let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
-            Ok(D::run(ARCH, (*self.query).into(), x))
+            Ok(unsafe { self.run_unchecked(x) })
         }
     }
+
+    #[inline(always)]
+    unsafe fn run_unchecked(&self, x: &[u8]) -> f32
+    where
+        D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>,
+    {
+        // SAFETY: We've validated that `x` has the correct length.
+        let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
+        D::run(ARCH, (*self.query).into(), x)
+    }
 }
 
 // TEMPORARY DEFINITIONS
@@ -597,8 +572,10 @@ where
 
             // SAFETY: Caller asserts that `i` is in-bounds.
             if let Some(data) = unsafe { self.reader.read_in_bounds(i.into_usize()) } {
+                let distance = unsafe { self.run_unchecked(data) };
+
                 // SAFETY: Inherited from caller.
-                *unsafe { buffer.get_unchecked_mut(processed) } = (i, self.run(data)?);
+                *unsafe { buffer.get_unchecked_mut(processed) } = (i, distance);
                 processed += 1;
             }
         }
@@ -798,6 +775,48 @@ impl FullPrecisionImpl for i8 {
     }
 }
 
+/// We use a macro to stamp out implementations of [`FullPrecision`] instead of using a
+/// blanket implementation from [`FullPrecisionImpl`] to make implementations more
+/// discoverable through the generated rust-doc.
+macro_rules! impl_full_precision {
+    ($T:ty) => {
+        impl FullPrecision for $T {
+            fn __search_accessor<'a>(
+                layer: &'a Full<Self>,
+                query: &'a [Self],
+                provider: &'a (dyn std::any::Any + Send + Sync),
+                counters: LocalCounters<'a>,
+            ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
+                let expand_beam = <$T>::make_expand_beam(layer, query)?;
+                Ok(crate::provider::SearchAccessor::new(
+                    layer.store.neighbors(),
+                    expand_beam,
+                    provider,
+                    layer.store.frozen(),
+                    counters,
+                ))
+            }
+
+            fn __prune_accessor<'a>(
+                layer: &'a Full<Self>,
+                counters: LocalCounters<'a>,
+            ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
+                let prune = <$T>::make_prune(layer)?;
+                Ok(crate::provider::PruneAccessor::new(
+                    prune,
+                    layer.store.neighbors(),
+                    counters,
+                ))
+            }
+        }
+    };
+    ($($Ts:ty),* $(,)?) => {
+        $(impl_full_precision!($Ts);)*
+    }
+}
+
+impl_full_precision!(f32, f16, u8, i8);
+
 ///////////
 // Tests //
 ///////////

From ff98f58be6dfe874ac983a011bffccdb2a72570f Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Sat, 22 Aug 2026 10:52:52 -0700
Subject: [PATCH 22/34] Insanity.

---
 diskann-inmem/src/arch.rs           | 136 +++++++++++++++++++++++++++-
 diskann-inmem/src/layers/full.rs    |  78 +++++++++++-----
 diskann-inmem/src/lib.rs            |   2 +-
 diskann-inmem/src/num.rs            |  27 ++++--
 diskann-inmem/src/store/invasive.rs |  17 +++-
 5 files changed, 224 insertions(+), 36 deletions(-)

diff --git a/diskann-inmem/src/arch.rs b/diskann-inmem/src/arch.rs
index a40f9c0f00..dd7711a6bd 100644
--- a/diskann-inmem/src/arch.rs
+++ b/diskann-inmem/src/arch.rs
@@ -3,7 +3,78 @@
  * Licensed under the MIT license.
  */
 
-use crate::num::Bytes;
+use std::num::NonZeroUsize;
+
+use crate::num::{Bytes};
+
+pub(crate) unsafe trait Prefetch: std::fmt::Debug + Send + Sync + 'static + Copy {
+    fn bytes(self) -> Bytes;
+    unsafe fn prefetch(self, ptr: *const u8);
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct Loop(Bytes);
+
+impl Loop {
+    pub(crate) const fn new(bytes: Bytes) -> Self {
+        Self(bytes)
+    }
+}
+
+unsafe impl Prefetch for Loop {
+    fn bytes(self) -> Bytes {
+        self.0
+    }
+
+    #[inline(always)]
+    unsafe fn prefetch(self, ptr: *const u8) {
+        unsafe { prefetch(ptr, self.bytes().value()) }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct Unrolled<const BYTES: usize>;
+
+impl<const BYTES: usize> Unrolled<BYTES> {
+    pub(crate) const fn new() -> Self {
+        Self
+    }
+}
+
+unsafe impl<const BYTES: usize> Prefetch for Unrolled<BYTES> {
+    fn bytes(self) -> Bytes {
+        Bytes::new(BYTES)
+    }
+
+    #[inline(always)]
+    unsafe fn prefetch(self, ptr: *const u8) {
+        unsafe { prefetch(ptr, self.bytes().value()) }
+    }
+}
+
+// #[derive(Debug, Clone, Copy)]
+// pub(crate) struct JumpTable {
+//     bytes: Bytes,
+// }
+//
+// impl JumpTable {
+//     pub(crate) const fn new(bytes: Bytes) -> Self {
+//         Self {
+//             bytes,
+//         }
+//     }
+// }
+//
+// unsafe impl Prefetch for JumpTable {
+//     fn bytes(self) -> Bytes {
+//         self.bytes
+//     }
+//
+//     #[inline(always)]
+//     unsafe fn prefetch(self, ptr: *const u8) {
+//         unsafe { prefetch_up_to_8(ptr, self.bytes().value().div_ceil(Bytes::CACHELINE.value())) }
+//     }
+// }
 
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
@@ -27,7 +98,7 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
 
     // SAFETY: Inherited from caller.
     unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
-    for i in 0..(lines - 1) {
+    for i in 0..(lines - 1).min(8) {
         // SAFETY: Inherited from caller.
         unsafe {
             _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
@@ -35,6 +106,48 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
     }
 }
 
+// #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
+// #[inline(always)]
+// pub unsafe fn prefetch_up_to_8(ptr: *const u8, lines: usize) {
+//     use std::arch::x86_64::*;
+//
+//     const STRIDE: usize = Bytes::CACHELINE.value();
+//     const PREFETCH_INSTRUCTION_BYTES: usize = 7;
+//
+//     let lines = if lines > 8 {
+//         unsafe { _mm_prefetch(ptr.cast::<i8>().add(STRIDE * (lines - 1)), _MM_HINT_T0); }
+//         8
+//     } else {
+//         lines
+//     };
+//
+//     let back = PREFETCH_INSTRUCTION_BYTES * lines;
+//     let ptr = ptr.wrapping_sub(128);
+//
+//     unsafe {
+//         std::arch::asm! {
+//             // Obtain the address of the label - the base of our prefetch table.
+//             "lea {tmp}, [rip + 3f]",
+//             "sub {tmp}, {back}",
+//             "notrack jmp {tmp}",
+//             "2:",
+//             "prefetcht0 byte ptr [{base} + 576]",
+//             "prefetcht0 byte ptr [{base} + 512]",
+//             "prefetcht0 byte ptr [{base} + 448]",
+//             "prefetcht0 byte ptr [{base} + 384]",
+//             "prefetcht0 byte ptr [{base} + 320]",
+//             "prefetcht0 byte ptr [{base} + 256]",
+//             "prefetcht0 byte ptr [{base} + 192]",
+//             "prefetcht0 byte ptr [{base} + 128]",
+//             "3:",
+//             back = in(reg) back,
+//             base = in(reg_abcd) ptr,
+//             tmp = out(reg) _,
+//             options(readonly, nostack, preserves_flags),
+//         }
+//     }
+// }
+
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
 /// The last cache line prefetched first, followed by the rest in ascending order.
@@ -44,3 +157,22 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
 /// The memory range `[ptr, ptr.add(len))` must be valid.
 #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
 pub(crate) unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
+
+///////////
+// Tests //
+///////////
+
+// #
+// [cfg(test)]
+// mod test {
+//     use super::*;
+//
+//     #[test]
+//     fn test_prefetch_up_to_8() {
+//         let v = vec![0u8; 600];
+//         for lines in 0..10 {
+//             unsafe { prefetch_up_to_8(v.as_ptr(), lines) };
+//         }
+//     }
+// }
+
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 8664fa5320..7523774460 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -20,6 +20,7 @@ use half::f16;
 use thiserror::Error;
 
 use crate::{
+    arch::Prefetch,
     counters::LocalCounters,
     layers,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
@@ -115,9 +116,8 @@ pub struct Full<T>
 where
     T: 'static,
 {
-    dim: usize,
-    metric: Metric,
     store: Store<Invasive>,
+    metric: Metric,
     _type: PhantomData<T>,
 }
 
@@ -166,21 +166,20 @@ where
         }
 
         Ok(Self {
-            dim: start_points.ncols(),
-            metric,
             store,
+            metric,
             _type: PhantomData,
         })
     }
 
     /// Return the logical dimension of the data handled by this [`layers::Layer`].
     pub fn dim(&self) -> usize {
-        self.dim
+        self.bytes().value() / std::mem::size_of::<T>()
     }
 
     /// Return the number of bytes of the data handles by this [`layers::Layer`].
     pub fn bytes(&self) -> Bytes {
-        Bytes::new(self.dim() * std::mem::size_of::<T>())
+        self.store.plugin().bytes()
     }
 
     fn check_dim(&self, dim: usize) -> Result<(), QueryDistanceError> {
@@ -271,6 +270,7 @@ where
     }
 }
 
+/// A [`layers::Guard`] for [`Full`].
 #[derive(Debug)]
 pub struct Guard<'a> {
     slot: store::Slot<'a, invasive::Slot<'a>>,
@@ -436,22 +436,38 @@ impl<T> std::ops::Deref for Calf<'_, T> {
 /// allow `f16` queries to be pre-converted to `f32`, saving on-the-fly conversion that
 /// would otherwise be needed.
 #[derive(Debug)]
-struct QueryDistance<'a, const PREFETCH_DIM: usize, T, U, D> {
+struct QueryDistance<'a, P, T, U, D> {
     // The original query.
     query: Calf<'a, T>,
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
+    // The type of the data prefetcher.
+    prefetch: P,
     // The type of the data in the original dataset.
     _data: PhantomData<U>,
     // The type of the `PureDistanceFunction` used for the implementation.
     _distance: PhantomData<D>,
 }
 
-impl<'a, const PREFETCH_DIM: usize, T, U, D> QueryDistance<'a, PREFETCH_DIM, T, U, D> {
-    fn new(query: Calf<'a, T>, reader: store::invasive::Reader<'a>) -> Self {
+impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
+    fn new(
+        query: Calf<'a, T>,
+        reader: store::invasive::Reader<'a>,
+        prefetch: P,
+    ) -> Self
+    where
+        P: Prefetch,
+    {
+        assert_eq!(
+            prefetch.bytes(),
+            reader.bytes_plus_tag(),
+            "invalid prefetcher"
+        );
+
         Self {
             query,
             reader,
+            prefetch,
             _data: PhantomData,
             _distance: PhantomData,
         }
@@ -470,7 +486,6 @@ impl<'a, const PREFETCH_DIM: usize, T, U, D> QueryDistance<'a, PREFETCH_DIM, T,
         Err(ANNError::new(error))
     }
 
-    // TODO: Since we control the reader - we can avoid the length check.
     #[inline(always)]
     fn run(&self, x: &[u8]) -> ANNResult<f32>
     where
@@ -488,6 +503,8 @@ impl<'a, const PREFETCH_DIM: usize, T, U, D> QueryDistance<'a, PREFETCH_DIM, T,
     where
         D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>,
     {
+        debug_assert_eq!(x.len(), self.bytes());
+
         // SAFETY: We've validated that `x` has the correct length.
         let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
         D::run(ARCH, (*self.query).into(), x)
@@ -496,11 +513,11 @@ impl<'a, const PREFETCH_DIM: usize, T, U, D> QueryDistance<'a, PREFETCH_DIM, T,
 
 // TEMPORARY DEFINITIONS
 const LOOKAHEAD: usize = 8;
-const BYTES: usize = 0;
 
-unsafe impl<const PREFETCH: usize, T, U, D> layers::ExpandBeam
-    for QueryDistance<'_, PREFETCH, T, U, D>
+unsafe impl<P, T, U, D> layers::ExpandBeam
+    for QueryDistance<'_, P, T, U, D>
 where
+    P: Prefetch,
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
     D: for<'a> FTarget2<Current, f32, UnalignedSlice<'a, T>, UnalignedSlice<'a, U>>
@@ -513,7 +530,7 @@ where
             Err(ANNError::new(OutOfBounds(i)))
         } else {
             match unsafe { self.reader.read_in_bounds(i.into_usize()) } {
-                Some(data) => Ok(Some(self.run(data)?)),
+                Some(data) => Ok(Some(unsafe { self.run_unchecked(data) })),
                 None => Ok(None),
             }
         }
@@ -527,11 +544,11 @@ where
         let len = list.len();
         let lookahead = LOOKAHEAD.min(len);
 
-        let bytes = if PREFETCH == 0 {
-            self.reader.bytes().value()
-        } else {
-            PREFETCH * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
-        };
+        // let bytes = if PREFETCH == 0 {
+        //     self.reader.bytes().value()
+        // } else {
+        //     PREFETCH * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
+        // };
 
         for j in 0..lookahead {
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
@@ -539,12 +556,11 @@ where
             //
             // We do not materialize the `RawSlice` as a reference.
             unsafe {
-                crate::arch::prefetch(
+                self.prefetch.prefetch(
                     self.reader
                         .read_raw_unchecked(list.get_unchecked(j).into_usize())
                         .as_ptr()
                         .cast(),
-                    bytes,
                 )
             }
         }
@@ -559,12 +575,11 @@ where
                 //
                 // We do not materialize the `RawSlice` as a reference.
                 unsafe {
-                    crate::arch::prefetch(
+                    self.prefetch.prefetch(
                         self.reader
                             .read_raw_unchecked(list.get_unchecked(j).into_usize())
                             .as_ptr()
                             .cast(),
-                        bytes,
                     )
                 }
                 j += 1;
@@ -603,18 +618,31 @@ struct QueryDistanceError {
 
 diskann::convert_error!(QueryDistanceError);
 
+const fn compute_bytes<T>(dim: usize) -> usize {
+    dim * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
+}
+
 macro_rules! mint {
     ($query:ident, $reader:ident, $T:ty => { $N:literal, $f:ident }) => {{
         mint!($query, $reader, { $T, $T } => { $N, $f })
     }};
     ($query:ident, $reader:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
-        Box::new(QueryDistance::<$N, $T, $U, Specialize<$N, $f>>::new($query, $reader))
+        Box::new(QueryDistance::<_, $T, $U, Specialize<$N, $f>>::new(
+            $query,
+            $reader,
+            $crate::arch::Unrolled::<{ compute_bytes::<$U>($N) }>::new(),
+        ))
     }};
     ($query:ident, $reader:ident, $T:ty => $f:ident) => {{
         mint!($query, $reader, { $T, $T } => $f)
     }};
     ($query:ident, $reader:ident, { $T:ty, $U:ty } => $f:ident) => {{
-        Box::new(QueryDistance::<0, $T, $U, $f>::new($query, $reader))
+        let bytes = $reader.bytes_plus_tag();
+        Box::new(QueryDistance::<_, $T, $U, $f>::new(
+            $query,
+            $reader,
+            $crate::arch::Loop::new(bytes),
+        ))
     }};
 }
 
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 1c63e613cb..1c9351ea7e 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -9,7 +9,7 @@
 
 pub mod num;
 
-mod arch;
+pub mod arch;
 mod buffer;
 mod counters;
 mod epoch;
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 4a5ab74f9a..cefe1acb25 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -69,6 +69,12 @@ impl Bytes {
         Self::new(std::mem::size_of::<T>())
     }
 
+    /// Return the number of bytes occupied by the slice.
+    #[inline]
+    pub const fn of_slice<T>(x: &[T]) -> Self {
+        Self(std::mem::size_of_val::<[T]>(x))
+    }
+
     /// Return `true` if `self` is zero.
     pub const fn is_zero(self) -> bool {
         self.0 == 0
@@ -160,17 +166,18 @@ impl std::fmt::Display for Align {
 //-------------------------//
 
 macro_rules! typed_int {
-    ($(#[$doc:meta])* $name:ident, $T:ty $(,)?) => {
+    ($(#[$doc:meta])* $vis:vis $name:ident, $T:ty $(,)?) => {
         $(#[$doc])*
         #[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
-        pub struct $name($T);
+        #[repr(transparent)]
+        $vis struct $name($T);
 
         impl $name {
-            pub const fn new(value: $T) -> Self {
+            $vis const fn new(value: $T) -> Self {
                 Self(value)
             }
 
-            pub const fn value(self) -> $T {
+            $vis const fn value(self) -> $T {
                 self.0
             }
         }
@@ -189,13 +196,13 @@ typed_int!(
     /// The number of distinct slots a [`crate::Provider`] or [`crate::Layer`] has capacity
     /// for. This is logically distinct from [`MaximumId`], which may be greater due too
     /// immutable points within a storage container.
-    Capacity,
+    pub Capacity,
     usize,
 );
 
 typed_int!(
     /// The maximum degree of an adjacency list.
-    MaxDegree,
+    pub MaxDegree,
     usize
 );
 
@@ -207,7 +214,7 @@ typed_int!(
     ///
     /// [`Capacity`] is related, but the [`IdLimit`] for a collection may be larger due to
     /// immutable points.
-    IdLimit,
+    pub IdLimit,
     u32
 );
 
@@ -230,6 +237,12 @@ impl IdLimit {
     }
 }
 
+// typed_int!(
+//     /// Temp
+//     pub(crate) CacheLines,
+//     NonZeroUsize,
+// )
+
 ///////////
 // Tests //
 ///////////
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index a48fbcb64e..d6e728f63d 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -125,6 +125,16 @@ impl Invasive {
         IdLimit::new(self.buffer.len() as u32)
     }
 
+    /// Return the number of bytes for each entry.
+    pub(crate) fn bytes(&self) -> Bytes {
+        self.bytes_plus_tag().unchecked_sub(AtomicTag::SIZE)
+    }
+
+    /// Return the number of bytes plus the atomic tag.
+    pub(crate) fn bytes_plus_tag(&self) -> Bytes {
+        self.unpadded
+    }
+
     /// Return a [`Reader`] over [`Self`] inside `store`.
     pub(crate) fn reader(store: &Store<Self>) -> Result<Reader<'_>, epoch::Unavailable> {
         store.guard(|this, guard: epoch::Guard<'_>| Reader {
@@ -362,7 +372,12 @@ impl<'a> Reader<'a> {
 
     /// Return the number of bytes for each entry.
     pub(crate) fn bytes(&self) -> Bytes {
-        self.unpadded.unchecked_sub(AtomicTag::SIZE)
+        self.bytes_plus_tag().unchecked_sub(AtomicTag::SIZE)
+    }
+
+    /// Return the number of bytes plus the atomic tag.
+    pub(crate) fn bytes_plus_tag(&self) -> Bytes {
+        self.unpadded
     }
 }
 

From 4e6d2594cb5f6678d6636460f659c758a371f264 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Sat, 22 Aug 2026 14:11:14 -0700
Subject: [PATCH 23/34] Checkpoint.

---
 diskann-inmem/src/arch.rs        | 146 ++++++++--------
 diskann-inmem/src/layers/full.rs | 274 +++++++++++++++++++------------
 diskann-inmem/src/store/mod.rs   |   9 +-
 3 files changed, 245 insertions(+), 184 deletions(-)

diff --git a/diskann-inmem/src/arch.rs b/diskann-inmem/src/arch.rs
index dd7711a6bd..5deb96100e 100644
--- a/diskann-inmem/src/arch.rs
+++ b/diskann-inmem/src/arch.rs
@@ -3,11 +3,11 @@
  * Licensed under the MIT license.
  */
 
-use std::num::NonZeroUsize;
+use crate::num::Bytes;
 
-use crate::num::{Bytes};
-
-pub(crate) unsafe trait Prefetch: std::fmt::Debug + Send + Sync + 'static + Copy {
+pub(crate) unsafe trait Prefetch:
+    std::fmt::Debug + Send + Sync + 'static + Copy
+{
     fn bytes(self) -> Bytes;
     unsafe fn prefetch(self, ptr: *const u8);
 }
@@ -52,29 +52,43 @@ unsafe impl<const BYTES: usize> Prefetch for Unrolled<BYTES> {
     }
 }
 
-// #[derive(Debug, Clone, Copy)]
-// pub(crate) struct JumpTable {
-//     bytes: Bytes,
-// }
-//
-// impl JumpTable {
-//     pub(crate) const fn new(bytes: Bytes) -> Self {
-//         Self {
-//             bytes,
-//         }
-//     }
-// }
-//
-// unsafe impl Prefetch for JumpTable {
-//     fn bytes(self) -> Bytes {
-//         self.bytes
-//     }
-//
-//     #[inline(always)]
-//     unsafe fn prefetch(self, ptr: *const u8) {
-//         unsafe { prefetch_up_to_8(ptr, self.bytes().value().div_ceil(Bytes::CACHELINE.value())) }
-//     }
-// }
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct JumpTable {
+    bytes: Bytes,
+    back: usize,
+    last: usize,
+}
+
+impl JumpTable {
+    pub(crate) fn new(bytes: Bytes) -> Self {
+        let stride = Bytes::CACHELINE.value();
+        let lines = bytes.value().div_ceil(stride);
+
+        let back = 7 * lines.min(8);
+        let last = if lines > 8 {
+            stride * (lines - 1)
+        } else {
+            0
+        };
+
+        Self {
+            bytes,
+            back,
+            last,
+        }
+    }
+}
+
+unsafe impl Prefetch for JumpTable {
+    fn bytes(self) -> Bytes {
+        self.bytes
+    }
+
+    #[inline(always)]
+    unsafe fn prefetch(self, ptr: *const u8) {
+        unsafe { prefetch_up_to_8(ptr, self.back, self.last) }
+    }
+}
 
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
@@ -106,47 +120,42 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
     }
 }
 
-// #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
-// #[inline(always)]
-// pub unsafe fn prefetch_up_to_8(ptr: *const u8, lines: usize) {
-//     use std::arch::x86_64::*;
-//
-//     const STRIDE: usize = Bytes::CACHELINE.value();
-//     const PREFETCH_INSTRUCTION_BYTES: usize = 7;
-//
-//     let lines = if lines > 8 {
-//         unsafe { _mm_prefetch(ptr.cast::<i8>().add(STRIDE * (lines - 1)), _MM_HINT_T0); }
-//         8
-//     } else {
-//         lines
-//     };
-//
-//     let back = PREFETCH_INSTRUCTION_BYTES * lines;
-//     let ptr = ptr.wrapping_sub(128);
-//
-//     unsafe {
-//         std::arch::asm! {
-//             // Obtain the address of the label - the base of our prefetch table.
-//             "lea {tmp}, [rip + 3f]",
-//             "sub {tmp}, {back}",
-//             "notrack jmp {tmp}",
-//             "2:",
-//             "prefetcht0 byte ptr [{base} + 576]",
-//             "prefetcht0 byte ptr [{base} + 512]",
-//             "prefetcht0 byte ptr [{base} + 448]",
-//             "prefetcht0 byte ptr [{base} + 384]",
-//             "prefetcht0 byte ptr [{base} + 320]",
-//             "prefetcht0 byte ptr [{base} + 256]",
-//             "prefetcht0 byte ptr [{base} + 192]",
-//             "prefetcht0 byte ptr [{base} + 128]",
-//             "3:",
-//             back = in(reg) back,
-//             base = in(reg_abcd) ptr,
-//             tmp = out(reg) _,
-//             options(readonly, nostack, preserves_flags),
-//         }
-//     }
-// }
+#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
+#[inline(always)]
+pub unsafe fn prefetch_up_to_8(ptr: *const u8, back: usize, last: usize) {
+    use std::arch::x86_64::*;
+
+    // const STRIDE: usize = Bytes::CACHELINE.value();
+    // const PREFETCH_INSTRUCTION_BYTES: usize = 7;
+
+    if last != 0 {
+        unsafe { _mm_prefetch(ptr.cast::<i8>().add(last), _MM_HINT_T0); }
+    }
+
+    let ptr = ptr.wrapping_sub(128);
+    unsafe {
+        std::arch::asm! {
+            // Obtain the address of the label - the base of our prefetch table.
+            "lea {tmp}, [rip + 3f]",
+            "sub {tmp}, {back}",
+            "notrack jmp {tmp}",
+            "2:",
+            "prefetcht0 byte ptr [{base} + 576]",
+            "prefetcht0 byte ptr [{base} + 512]",
+            "prefetcht0 byte ptr [{base} + 448]",
+            "prefetcht0 byte ptr [{base} + 384]",
+            "prefetcht0 byte ptr [{base} + 320]",
+            "prefetcht0 byte ptr [{base} + 256]",
+            "prefetcht0 byte ptr [{base} + 192]",
+            "prefetcht0 byte ptr [{base} + 128]",
+            "3:",
+            back = in(reg) back,
+            base = in(reg_abcd) ptr,
+            tmp = out(reg) _,
+            options(readonly, nostack, preserves_flags),
+        }
+    }
+}
 
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
@@ -175,4 +184,3 @@ pub(crate) unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
 //         }
 //     }
 // }
-
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 7523774460..0814140202 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-use std::{fmt::Debug, marker::PhantomData};
+use std::{fmt::Debug, num::NonZeroUsize, marker::PhantomData};
 
 use diskann::{ANNError, ANNResult, utils::IntoUsize};
 use diskann_utils::views::Matrix;
@@ -22,7 +22,7 @@ use thiserror::Error;
 use crate::{
     arch::Prefetch,
     counters::LocalCounters,
-    layers,
+    epoch, layers,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
     store::{
         self, Store,
@@ -37,8 +37,11 @@ pub struct Config<T> {
     metric: Metric,
     start_points: Matrix<T>,
     store: store::Config,
+    lookahead: Option<NonZeroUsize>,
 }
 
+const DEFAULT_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(8).unwrap();
+
 impl<T> Config<T> {
     pub fn new(
         capacity: Capacity,
@@ -55,6 +58,7 @@ impl<T> Config<T> {
             metric,
             start_points,
             store: store::Config::default(),
+            lookahead: Some(DEFAULT_LOOKAHEAD),
         }
     }
 
@@ -63,6 +67,11 @@ impl<T> Config<T> {
         self
     }
 
+    pub fn prefetch(mut self, lookahead: Option<NonZeroUsize>) -> Self {
+        self.lookahead = lookahead;
+        self
+    }
+
     /// Return the vector dimension of this configuration and the resulting [`Full`].
     pub fn dim(&self) -> usize {
         self.start_points.ncols()
@@ -118,6 +127,7 @@ where
 {
     store: Store<Invasive>,
     metric: Metric,
+    lookahead: Option<NonZeroUsize>,
     _type: PhantomData<T>,
 }
 
@@ -149,6 +159,7 @@ where
             metric,
             start_points,
             store,
+            lookahead,
         } = config;
 
         let bytes = Bytes::new(start_points.ncols() * std::mem::size_of::<T>());
@@ -168,6 +179,7 @@ where
         Ok(Self {
             store,
             metric,
+            lookahead,
             _type: PhantomData,
         })
     }
@@ -193,7 +205,7 @@ where
         }
     }
 
-    fn reader(&self) -> ANNResult<invasive::Reader<'_>> {
+    fn reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
         Ok(Invasive::reader(&self.store)?)
     }
 }
@@ -429,6 +441,89 @@ impl<T> std::ops::Deref for Calf<'_, T> {
     }
 }
 
+/// A temporary precursor for [`QueryDistance`] to simplify macros.
+#[derive(Debug)]
+struct IntoQueryDistance<'a, T, U> {
+    query: Calf<'a, T>,
+    reader: store::invasive::Reader<'a>,
+    lookahead: Option<NonZeroUsize>,
+    _data: PhantomData<U>,
+}
+
+impl<'a, T, U> IntoQueryDistance<'a, T, U> {
+    /// Construct a new [`IntoQueryDistance`] - verifying that
+    fn new(full: &'a Full<U>, query: Calf<'a, T>) -> ANNResult<Self> {
+        full.check_dim(query.len())?;
+        let reader = full.reader()?;
+        let lookahead = full.lookahead;
+        Ok(Self {
+            query,
+            reader,
+            lookahead,
+            _data: PhantomData,
+        })
+    }
+
+    fn bytes_plus_tag(&self) -> Bytes {
+        self.reader.bytes_plus_tag()
+    }
+}
+
+trait Distance<T, U>: std::fmt::Debug + Send + Sync + 'static {
+    fn eval(&self, x: UnalignedSlice<'_, T>, u: UnalignedSlice<'_, U>) -> f32;
+}
+
+#[derive(Debug)]
+struct Pure<D>(PhantomData<D>);
+
+impl<D> Pure<D> {
+    const fn new() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<T, U, D> Distance<T, U> for Pure<D>
+where
+    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>
+        + std::fmt::Debug + Send + Sync + 'static,
+{
+    #[inline(always)]
+    fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
+        D::run(ARCH, x, y)
+    }
+}
+
+#[derive(Debug)]
+struct PureNoInline<D>(PhantomData<D>);
+
+impl<D> PureNoInline<D> {
+    const fn new() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<T, U, D> Distance<T, U> for PureNoInline<D>
+where
+    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>
+    + std::fmt::Debug + Send + Sync + 'static,
+{
+    #[inline(never)]
+    fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
+        D::run(ARCH, x, y)
+    }
+}
+
+impl<T, U> Distance<T, U> for diskann_vector::distance::Distance<T, U>
+where
+    T: std::fmt::Debug + 'static,
+    U: std::fmt::Debug + 'static,
+{
+    #[inline(always)]
+    fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
+        self.call_unaligned(x, y)
+    }
+}
+
 /// A fused query distance based on [`diskann_vector::PureDistanceFunction`] to enable
 /// inlining of the final distance function (`D`).
 ///
@@ -441,23 +536,28 @@ struct QueryDistance<'a, P, T, U, D> {
     query: Calf<'a, T>,
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
+    // THe prefetch look-ahead.
+    lookahead: Option<NonZeroUsize>,
     // The type of the data prefetcher.
     prefetch: P,
+    // The type of the distance used for the arguments
+    distance: D,
     // The type of the data in the original dataset.
     _data: PhantomData<U>,
-    // The type of the `PureDistanceFunction` used for the implementation.
-    _distance: PhantomData<D>,
 }
 
 impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
-    fn new(
-        query: Calf<'a, T>,
-        reader: store::invasive::Reader<'a>,
-        prefetch: P,
-    ) -> Self
+    fn new(into: IntoQueryDistance<'a, T, U>, prefetch: P, distance: D) -> Self
     where
         P: Prefetch,
     {
+        let IntoQueryDistance {
+            query,
+            reader,
+            lookahead,
+            _data,
+        } = into;
+
         assert_eq!(
             prefetch.bytes(),
             reader.bytes_plus_tag(),
@@ -467,9 +567,10 @@ impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
         Self {
             query,
             reader,
+            lookahead,
             prefetch,
-            _data: PhantomData,
-            _distance: PhantomData,
+            distance,
+            _data,
         }
     }
 
@@ -477,53 +578,25 @@ impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
         std::mem::size_of::<U>() * self.query.len()
     }
 
-    fn error(&self, len: usize) -> ANNResult<f32> {
-        let error = QueryDistanceError {
-            expected: self.bytes(),
-            xlen: len,
-        };
-
-        Err(ANNError::new(error))
-    }
-
-    #[inline(always)]
-    fn run(&self, x: &[u8]) -> ANNResult<f32>
-    where
-        D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>,
-    {
-        if x.len() != self.bytes() {
-            self.error(x.len())
-        } else {
-            Ok(unsafe { self.run_unchecked(x) })
-        }
-    }
-
     #[inline(always)]
     unsafe fn run_unchecked(&self, x: &[u8]) -> f32
     where
-        D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>,
+        D: Distance<T, U>,
     {
         debug_assert_eq!(x.len(), self.bytes());
 
         // SAFETY: We've validated that `x` has the correct length.
         let x = unsafe { UnalignedSlice::new(x.as_ptr().cast::<U>(), self.query.len()) };
-        D::run(ARCH, (*self.query).into(), x)
+        self.distance.eval((*self.query).into(), x)
     }
 }
 
-// TEMPORARY DEFINITIONS
-const LOOKAHEAD: usize = 8;
-
-unsafe impl<P, T, U, D> layers::ExpandBeam
-    for QueryDistance<'_, P, T, U, D>
+unsafe impl<P, T, U, D> layers::ExpandBeam for QueryDistance<'_, P, T, U, D>
 where
     P: Prefetch,
     T: Send + Sync + 'static + Debug,
     U: Send + Sync + 'static + Debug,
-    D: for<'a> FTarget2<Current, f32, UnalignedSlice<'a, T>, UnalignedSlice<'a, U>>
-        + Send
-        + Sync
-        + Debug,
+    D: Distance<T, U>,
 {
     fn evaluate(&self, i: u32) -> ANNResult<Option<f32>> {
         if !self.reader.is_in_bounds(i.into_usize()) {
@@ -542,13 +615,8 @@ where
 
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize> {
         let len = list.len();
-        let lookahead = LOOKAHEAD.min(len);
-
-        // let bytes = if PREFETCH == 0 {
-        //     self.reader.bytes().value()
-        // } else {
-        //     PREFETCH * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
-        // };
+        // let lookahead = self.lookahead.map(|l| l.get()).unwrap_or(0).min(len);
+        let lookahead = 8.min(len);
 
         for j in 0..lookahead {
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
@@ -623,25 +691,25 @@ const fn compute_bytes<T>(dim: usize) -> usize {
 }
 
 macro_rules! mint {
-    ($query:ident, $reader:ident, $T:ty => { $N:literal, $f:ident }) => {{
-        mint!($query, $reader, { $T, $T } => { $N, $f })
+    ($into:ident, $T:ty => { $N:literal, $f:ident }) => {{
+        mint!($into, { $T, $T } => { $N, $f })
     }};
-    ($query:ident, $reader:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
-        Box::new(QueryDistance::<_, $T, $U, Specialize<$N, $f>>::new(
-            $query,
-            $reader,
+    ($into:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
+        Box::new(QueryDistance::<_, $T, $U, _>::new(
+            $into,
             $crate::arch::Unrolled::<{ compute_bytes::<$U>($N) }>::new(),
+            Pure::<Specialize<$N, $f>>::new(),
         ))
     }};
-    ($query:ident, $reader:ident, $T:ty => $f:ident) => {{
-        mint!($query, $reader, { $T, $T } => $f)
+    ($into:ident, $T:ty => $f:ident) => {{
+        mint!($into, { $T, $T } => $f)
     }};
-    ($query:ident, $reader:ident, { $T:ty, $U:ty } => $f:ident) => {{
-        let bytes = $reader.bytes_plus_tag();
-        Box::new(QueryDistance::<_, $T, $U, $f>::new(
-            $query,
-            $reader,
+    ($into:ident, { $T:ty, $U:ty } => $f:ident) => {{
+        let bytes = $into.bytes_plus_tag();
+        Box::new(QueryDistance::<_, $T, $U, _>::new(
+            $into,
             $crate::arch::Loop::new(bytes),
+            Pure::<$f>::new(),
         ))
     }};
 }
@@ -651,23 +719,18 @@ impl FullPrecisionImpl for f32 {
         full: &'a Full<f32>,
         query: &'a [f32],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        full.check_dim(query.len())?;
-        let reader = full.reader()?;
-        let query = Calf::Borrowed(query);
-
+        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 100 {
-                    mint!(query, reader, f32 => { 100, SquaredL2 })
-                } else {
-                    mint!(query, reader, f32 => SquaredL2)
-                }
+                // if full.dim() == 100 {
+                //     mint!(into, f32 => { 100, SquaredL2 })
+                // } else {
+                    mint!(into, f32 => SquaredL2)
+                // }
             }
-            Metric::InnerProduct => {
-                mint!(query, reader, f32 => InnerProduct)
-            }
-            Metric::Cosine => mint!(query, reader, f32 => Cosine),
-            Metric::CosineNormalized => mint!(query, reader, f32 => CosineNormalized),
+            Metric::InnerProduct => mint!(into, f32 => InnerProduct),
+            Metric::Cosine => mint!(into, f32 => Cosine),
+            Metric::CosineNormalized => mint!(into, f32 => CosineNormalized),
         };
 
         Ok(output)
@@ -692,24 +755,23 @@ impl FullPrecisionImpl for f16 {
         full: &'a Full<f16>,
         query: &'a [f16],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        full.check_dim(query.len())?;
-        let reader = full.reader()?;
-
         let mut as_f32: Box<[f32]> = std::iter::repeat_n(0.0, full.dim()).collect();
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
         let query = Calf::Owned(as_f32);
 
+        let into = IntoQueryDistance::new(full, query)?;
+
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 100 {
-                    mint!(query, reader, { f32, f16 } => { 100, SquaredL2 })
-                } else {
-                    mint!(query, reader, { f32, f16 } => SquaredL2)
-                }
+                // if full.dim() == 100 {
+                //     mint!(into, { f32, f16 } => { 100, SquaredL2 })
+                // } else {
+                    mint!(into, { f32, f16 } => SquaredL2)
+                // }
             }
-            Metric::InnerProduct => mint!(query, reader, { f32, f16 } => InnerProduct),
-            Metric::Cosine => mint!(query, reader, { f32, f16 } => Cosine),
-            Metric::CosineNormalized => mint!(query, reader, { f32, f16 } => CosineNormalized),
+            Metric::InnerProduct => mint!(into, { f32, f16 } => InnerProduct),
+            Metric::Cosine => mint!(into, { f32, f16 } => Cosine),
+            Metric::CosineNormalized => mint!(into, { f32, f16 } => CosineNormalized),
         };
 
         Ok(output)
@@ -734,22 +796,19 @@ impl FullPrecisionImpl for u8 {
         full: &'a Full<u8>,
         query: &'a [u8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        full.check_dim(query.len())?;
-        let reader = full.reader()?;
-
-        let query = Calf::Borrowed(query);
+        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
 
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                if full.dim() == 128 {
-                    mint!(query, reader, u8 => { 128, SquaredL2 })
-                } else {
-                    mint!(query, reader, u8 => SquaredL2)
-                }
+                // if full.dim() == 128 {
+                //     mint!(into, u8 => { 128, SquaredL2 })
+                // } else {
+                    mint!(into, u8 => SquaredL2)
+                // }
             }
-            Metric::InnerProduct => mint!(query, reader, u8 => InnerProduct),
-            Metric::Cosine => mint!(query, reader, u8 => Cosine),
-            Metric::CosineNormalized => mint!(query, reader, u8 => Cosine),
+            Metric::InnerProduct => mint!(into, u8 => InnerProduct),
+            Metric::Cosine => mint!(into, u8 => Cosine),
+            Metric::CosineNormalized => mint!(into, u8 => Cosine),
         };
 
         Ok(output)
@@ -774,16 +833,13 @@ impl FullPrecisionImpl for i8 {
         full: &'a Full<i8>,
         query: &'a [i8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        full.check_dim(query.len())?;
-        let reader = full.reader()?;
-
-        let query = Calf::Borrowed(query);
+        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
 
         let output: Box<dyn layers::ExpandBeam + 'a> = match full.metric {
-            Metric::L2 => mint!(query, reader, i8 => SquaredL2),
-            Metric::InnerProduct => mint!(query, reader, i8 => InnerProduct),
-            Metric::Cosine => mint!(query, reader, i8 => Cosine),
-            Metric::CosineNormalized => mint!(query, reader, i8 => Cosine),
+            Metric::L2 => mint!(into, i8 => SquaredL2),
+            Metric::InnerProduct => mint!(into, i8 => InnerProduct),
+            Metric::Cosine => mint!(into, i8 => Cosine),
+            Metric::CosineNormalized => mint!(into, i8 => Cosine),
         };
 
         Ok(output)
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 14d62b9775..5634b8e887 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -820,14 +820,11 @@ mod tests {
     #[test]
     fn new_rejects_faulty_plugin() {
         let err = Store::new(
-            Layout::new(
-                Capacity::new(4),
-                MaxDegree::new(10),
-                0,
-            ),
+            Layout::new(Capacity::new(4), MaxDegree::new(10), 0),
             Config::default(),
             FaultyConfig,
-        ).unwrap_err();
+        )
+        .unwrap_err();
         assert_matches!(err.0, StoreErrorInner::InvalidConstruction { .. });
     }
 

From 2ba93d5fddf7eeb14a74fe42606cb9c86dff17f8 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Sat, 22 Aug 2026 19:00:59 -0700
Subject: [PATCH 24/34] Checkpoint.

---
 diskann-inmem/src/layers/full.rs           | 257 +++++++++++----------
 diskann-inmem/src/lib.rs                   |   2 +-
 diskann-inmem/src/num.rs                   |   5 +
 diskann-inmem/src/{arch.rs => prefetch.rs} | 137 ++++++++---
 4 files changed, 239 insertions(+), 162 deletions(-)
 rename diskann-inmem/src/{arch.rs => prefetch.rs} (56%)

diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 0814140202..630932d7c6 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -3,14 +3,14 @@
  * Licensed under the MIT license.
  */
 
-use std::{fmt::Debug, num::NonZeroUsize, marker::PhantomData};
+use std::{fmt::Debug, marker::PhantomData, num::NonZeroUsize};
 
 use diskann::{ANNError, ANNResult, utils::IntoUsize};
 use diskann_utils::views::Matrix;
 use diskann_vector::{
     UnalignedSlice,
     conversion::SliceCast,
-    distance::{Cosine, CosineNormalized, InnerProduct, Metric, Specialize, SquaredL2},
+    distance::{Cosine, CosineNormalized, InnerProduct, Metric, Specialize, SquaredL2, DistanceProvider},
 };
 use diskann_wide::{
     ARCH,
@@ -20,7 +20,7 @@ use half::f16;
 use thiserror::Error;
 
 use crate::{
-    arch::Prefetch,
+    prefetch::{self, Prefetch},
     counters::LocalCounters,
     epoch, layers,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
@@ -40,7 +40,7 @@ pub struct Config<T> {
     lookahead: Option<NonZeroUsize>,
 }
 
-const DEFAULT_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(8).unwrap();
+const DEFAULT_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(12).unwrap();
 
 impl<T> Config<T> {
     pub fn new(
@@ -194,9 +194,17 @@ where
         self.store.plugin().bytes()
     }
 
-    fn check_dim(&self, dim: usize) -> Result<(), QueryDistanceError> {
+    pub fn bytes_plus_tag(&self) -> Bytes {
+        self.store.plugin().bytes_plus_tag()
+    }
+
+    pub fn metric(&self) -> Metric {
+        self.metric
+    }
+
+    fn check_dim(&self, dim: usize) -> Result<(), ExpandBeamError> {
         if self.dim() != dim {
-            Err(QueryDistanceError {
+            Err(ExpandBeamError {
                 expected: self.dim(),
                 xlen: dim,
             })
@@ -355,11 +363,14 @@ struct Prune<'a, T, D> {
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
     // Type type of the `PureDistanceFunction` used for the implementation.
-    _distance: PhantomData<D>,
+    distance: D,
 }
 
 impl<'a, T, D> Prune<'a, T, D> {
-    fn new(reader: store::invasive::Reader<'a>) -> Self {
+    fn new(
+        reader: store::invasive::Reader<'a>,
+        distance: D,
+    ) -> Self {
         // This should be ensured at construction time
         debug_assert!(
             reader
@@ -372,18 +383,20 @@ impl<'a, T, D> Prune<'a, T, D> {
         Self {
             buffer: Vec::new(),
             reader,
-            _distance: PhantomData,
+            distance,
         }
     }
+
+    fn boxed(self) -> Box<Self> {
+        Box::new(self)
+    }
+
 }
 
 impl<T, D> layers::Prune for Prune<'_, T, D>
 where
-    T: Send + Sync + 'static + Debug,
-    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, T>>
-        + Send
-        + Sync
-        + Debug,
+    T: Debug + Send + Sync + 'static,
+    D: Distance<T, T>,
 {
     fn prepare(
         &mut self,
@@ -416,13 +429,13 @@ where
     }
 
     fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
-        D::run(ARCH, self.buffer[a.index()], self.buffer[b.index()])
+        self.distance.eval(self.buffer[a.index()], self.buffer[b.index()])
     }
 }
 
-///////////////////
-// QueryDistance //
-///////////////////
+////////////////
+// ExpandBeam //
+////////////////
 
 // A baby [`std::borrow::Cow`].
 #[derive(Debug)]
@@ -441,17 +454,17 @@ impl<T> std::ops::Deref for Calf<'_, T> {
     }
 }
 
-/// A temporary precursor for [`QueryDistance`] to simplify macros.
+/// A temporary precursor for [`ExpandBeam`] to simplify macros.
 #[derive(Debug)]
-struct IntoQueryDistance<'a, T, U> {
+struct IntoExpandBeam<'a, T, U> {
     query: Calf<'a, T>,
     reader: store::invasive::Reader<'a>,
     lookahead: Option<NonZeroUsize>,
     _data: PhantomData<U>,
 }
 
-impl<'a, T, U> IntoQueryDistance<'a, T, U> {
-    /// Construct a new [`IntoQueryDistance`] - verifying that
+impl<'a, T, U> IntoExpandBeam<'a, T, U> {
+    /// Construct a new [`IntoExpandBeam`] - verifying that
     fn new(full: &'a Full<U>, query: Calf<'a, T>) -> ANNResult<Self> {
         full.check_dim(query.len())?;
         let reader = full.reader()?;
@@ -485,7 +498,10 @@ impl<D> Pure<D> {
 impl<T, U, D> Distance<T, U> for Pure<D>
 where
     D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>
-        + std::fmt::Debug + Send + Sync + 'static,
+        + std::fmt::Debug
+        + Send
+        + Sync
+        + 'static,
 {
     #[inline(always)]
     fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
@@ -493,26 +509,6 @@ where
     }
 }
 
-#[derive(Debug)]
-struct PureNoInline<D>(PhantomData<D>);
-
-impl<D> PureNoInline<D> {
-    const fn new() -> Self {
-        Self(PhantomData)
-    }
-}
-
-impl<T, U, D> Distance<T, U> for PureNoInline<D>
-where
-    D: for<'any> FTarget2<Current, f32, UnalignedSlice<'any, T>, UnalignedSlice<'any, U>>
-    + std::fmt::Debug + Send + Sync + 'static,
-{
-    #[inline(never)]
-    fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32 {
-        D::run(ARCH, x, y)
-    }
-}
-
 impl<T, U> Distance<T, U> for diskann_vector::distance::Distance<T, U>
 where
     T: std::fmt::Debug + 'static,
@@ -531,7 +527,7 @@ where
 /// allow `f16` queries to be pre-converted to `f32`, saving on-the-fly conversion that
 /// would otherwise be needed.
 #[derive(Debug)]
-struct QueryDistance<'a, P, T, U, D> {
+struct ExpandBeam<'a, P, T, U, D> {
     // The original query.
     query: Calf<'a, T>,
     // A reader into a layer's store.
@@ -546,24 +542,19 @@ struct QueryDistance<'a, P, T, U, D> {
     _data: PhantomData<U>,
 }
 
-impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
-    fn new(into: IntoQueryDistance<'a, T, U>, prefetch: P, distance: D) -> Self
+impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
+    fn new(into: IntoExpandBeam<'a, T, U>, prefetch: P, distance: D) -> Self
     where
         P: Prefetch,
     {
-        let IntoQueryDistance {
+        let IntoExpandBeam {
             query,
             reader,
             lookahead,
             _data,
         } = into;
 
-        assert_eq!(
-            prefetch.bytes(),
-            reader.bytes_plus_tag(),
-            "invalid prefetcher"
-        );
-
+        prefetch.check(reader.bytes_plus_tag());
         Self {
             query,
             reader,
@@ -578,6 +569,10 @@ impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
         std::mem::size_of::<U>() * self.query.len()
     }
 
+    fn boxed(self) -> Box<Self> {
+        Box::new(self)
+    }
+
     #[inline(always)]
     unsafe fn run_unchecked(&self, x: &[u8]) -> f32
     where
@@ -591,7 +586,7 @@ impl<'a, P, T, U, D> QueryDistance<'a, P, T, U, D> {
     }
 }
 
-unsafe impl<P, T, U, D> layers::ExpandBeam for QueryDistance<'_, P, T, U, D>
+unsafe impl<P, T, U, D> layers::ExpandBeam for ExpandBeam<'_, P, T, U, D>
 where
     P: Prefetch,
     T: Send + Sync + 'static + Debug,
@@ -615,8 +610,8 @@ where
 
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize> {
         let len = list.len();
-        // let lookahead = self.lookahead.map(|l| l.get()).unwrap_or(0).min(len);
-        let lookahead = 8.min(len);
+        let lookahead = self.lookahead.map(|l| l.get()).unwrap_or(0).min(len);
+        // let lookahead = 8.min(len);
 
         for j in 0..lookahead {
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
@@ -679,58 +674,62 @@ diskann::convert_error!(OutOfBounds);
     self.expected,
     self.xlen,
 )]
-struct QueryDistanceError {
+struct ExpandBeamError {
     expected: usize,
     xlen: usize,
 }
 
-diskann::convert_error!(QueryDistanceError);
+diskann::convert_error!(ExpandBeamError);
 
 const fn compute_bytes<T>(dim: usize) -> usize {
     dim * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
 }
 
-macro_rules! mint {
-    ($into:ident, $T:ty => { $N:literal, $f:ident }) => {{
-        mint!($into, { $T, $T } => { $N, $f })
-    }};
-    ($into:ident, { $T:ty, $U:ty } => { $N:literal, $f:ident }) => {{
-        Box::new(QueryDistance::<_, $T, $U, _>::new(
+macro_rules! expand_beam {
+    ($into:ident, { $T:ty, $N:literal, $f:ident }) => {{
+        Box::new(ExpandBeam::<_, _, $T, _>::new(
             $into,
-            $crate::arch::Unrolled::<{ compute_bytes::<$U>($N) }>::new(),
+            prefetch::Unrolled::<{ compute_bytes::<$T>($N) }>::new(),
             Pure::<Specialize<$N, $f>>::new(),
         ))
     }};
-    ($into:ident, $T:ty => $f:ident) => {{
-        mint!($into, { $T, $T } => $f)
-    }};
-    ($into:ident, { $T:ty, $U:ty } => $f:ident) => {{
+    ($into:ident, $f:ident) => {{
         let bytes = $into.bytes_plus_tag();
-        Box::new(QueryDistance::<_, $T, $U, _>::new(
+        Box::new(ExpandBeam::new(
             $into,
-            $crate::arch::Loop::new(bytes),
+            prefetch::Loop::new(bytes),
             Pure::<$f>::new(),
         ))
     }};
 }
 
+macro_rules! prune {
+    ($self:ty, $reader:ident, $f:ident) => {{
+        Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed()
+    }};
+    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{
+        Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed()
+    }};
+}
+
 impl FullPrecisionImpl for f32 {
     fn make_expand_beam<'a>(
         full: &'a Full<f32>,
         query: &'a [f32],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
+        let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
+
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 100 {
-                //     mint!(into, f32 => { 100, SquaredL2 })
-                // } else {
-                    mint!(into, f32 => SquaredL2)
-                // }
+                if full.dim() == 100 {
+                    expand_beam!(into, { f32, 100, SquaredL2 })
+                } else {
+                    expand_beam!(into, SquaredL2)
+                }
             }
-            Metric::InnerProduct => mint!(into, f32 => InnerProduct),
-            Metric::Cosine => mint!(into, f32 => Cosine),
-            Metric::CosineNormalized => mint!(into, f32 => CosineNormalized),
+            Metric::InnerProduct => expand_beam!(into, InnerProduct),
+            Metric::Cosine => expand_beam!(into, Cosine),
+            Metric::CosineNormalized => expand_beam!(into, CosineNormalized),
         };
 
         Ok(output)
@@ -740,10 +739,10 @@ impl FullPrecisionImpl for f32 {
         let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
-            Metric::L2 => Box::new(Prune::<f32, SquaredL2>::new(reader)),
-            Metric::InnerProduct => Box::new(Prune::<f32, InnerProduct>::new(reader)),
-            Metric::Cosine => Box::new(Prune::<f32, Cosine>::new(reader)),
-            Metric::CosineNormalized => Box::new(Prune::<f32, CosineNormalized>::new(reader)),
+            Metric::L2 => prune!(Self, reader, SquaredL2),
+            Metric::InnerProduct => prune!(Self, reader, InnerProduct),
+            Metric::Cosine => prune!(Self, reader, Cosine),
+            Metric::CosineNormalized => prune!(Self, reader, CosineNormalized),
         };
 
         Ok(output)
@@ -759,19 +758,19 @@ impl FullPrecisionImpl for f16 {
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
         let query = Calf::Owned(as_f32);
 
-        let into = IntoQueryDistance::new(full, query)?;
+        let into = IntoExpandBeam::new(full, query)?;
 
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 100 {
-                //     mint!(into, { f32, f16 } => { 100, SquaredL2 })
-                // } else {
-                    mint!(into, { f32, f16 } => SquaredL2)
-                // }
+                if full.dim() == 100 {
+                    expand_beam!(into, { f16, 100, SquaredL2 })
+                } else {
+                    expand_beam!(into, SquaredL2)
+                }
             }
-            Metric::InnerProduct => mint!(into, { f32, f16 } => InnerProduct),
-            Metric::Cosine => mint!(into, { f32, f16 } => Cosine),
-            Metric::CosineNormalized => mint!(into, { f32, f16 } => CosineNormalized),
+            Metric::InnerProduct => expand_beam!(into, InnerProduct),
+            Metric::Cosine => expand_beam!(into, Cosine),
+            Metric::CosineNormalized => expand_beam!(into, CosineNormalized),
         };
 
         Ok(output)
@@ -781,10 +780,10 @@ impl FullPrecisionImpl for f16 {
         let reader = full.reader()?;
 
         let output: Box<dyn layers::Prune> = match full.metric {
-            Metric::L2 => Box::new(Prune::<f16, SquaredL2>::new(reader)),
-            Metric::InnerProduct => Box::new(Prune::<f16, InnerProduct>::new(reader)),
-            Metric::Cosine => Box::new(Prune::<f16, Cosine>::new(reader)),
-            Metric::CosineNormalized => Box::new(Prune::<f16, CosineNormalized>::new(reader)),
+            Metric::L2 => prune!(Self, reader, SquaredL2),
+            Metric::InnerProduct => prune!(Self, reader, InnerProduct),
+            Metric::Cosine => prune!(Self, reader, Cosine),
+            Metric::CosineNormalized => prune!(Self, reader, CosineNormalized),
         };
 
         Ok(output)
@@ -796,19 +795,18 @@ impl FullPrecisionImpl for u8 {
         full: &'a Full<u8>,
         query: &'a [u8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
+        let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
         let output: Box<dyn layers::ExpandBeam> = match full.metric {
             Metric::L2 => {
-                // if full.dim() == 128 {
-                //     mint!(into, u8 => { 128, SquaredL2 })
-                // } else {
-                    mint!(into, u8 => SquaredL2)
-                // }
+                if full.dim() == 128 {
+                    expand_beam!(into, { u8, 128, SquaredL2 })
+                } else {
+                    expand_beam!(into, SquaredL2)
+                }
             }
-            Metric::InnerProduct => mint!(into, u8 => InnerProduct),
-            Metric::Cosine => mint!(into, u8 => Cosine),
-            Metric::CosineNormalized => mint!(into, u8 => Cosine),
+            Metric::InnerProduct => expand_beam!(into, InnerProduct),
+            Metric::Cosine | Metric::CosineNormalized => expand_beam!(into, Cosine),
         };
 
         Ok(output)
@@ -816,12 +814,13 @@ impl FullPrecisionImpl for u8 {
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = full.reader()?;
+        let dim = full.dim();
 
         let output: Box<dyn layers::Prune> = match full.metric {
-            Metric::L2 => Box::new(Prune::<u8, SquaredL2>::new(reader)),
-            Metric::InnerProduct => Box::new(Prune::<u8, InnerProduct>::new(reader)),
-            Metric::Cosine => Box::new(Prune::<u8, Cosine>::new(reader)),
-            Metric::CosineNormalized => Box::new(Prune::<u8, CosineNormalized>::new(reader)),
+            Metric::L2 => prune!(Self, reader, SquaredL2),
+            Metric::InnerProduct => prune!(Self, reader, InnerProduct),
+            Metric::Cosine => prune!(Self, reader, Cosine),
+            Metric::CosineNormalized => prune!(Self, reader, CosineNormalized),
         };
 
         Ok(output)
@@ -833,28 +832,32 @@ impl FullPrecisionImpl for i8 {
         full: &'a Full<i8>,
         query: &'a [i8],
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
-        let into = IntoQueryDistance::new(full, Calf::Borrowed(query))?;
+        let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
-        let output: Box<dyn layers::ExpandBeam + 'a> = match full.metric {
-            Metric::L2 => mint!(into, i8 => SquaredL2),
-            Metric::InnerProduct => mint!(into, i8 => InnerProduct),
-            Metric::Cosine => mint!(into, i8 => Cosine),
-            Metric::CosineNormalized => mint!(into, i8 => Cosine),
-        };
+        let distance = <Self as DistanceProvider<Self>>::distance_comparer(
+            full.metric(),
+            Some(full.dim()),
+        );
+
+        let output: Box<dyn layers::ExpandBeam + 'a> = ExpandBeam::new(
+            into,
+            prefetch::Loop::new(full.bytes_plus_tag()),
+            distance,
+        ).boxed();
 
         Ok(output)
     }
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = full.reader()?;
+        let dim = full.dim();
 
-        let output: Box<dyn layers::Prune> = match full.metric {
-            Metric::L2 => Box::new(Prune::<i8, SquaredL2>::new(reader)),
-            Metric::InnerProduct => Box::new(Prune::<i8, InnerProduct>::new(reader)),
-            Metric::Cosine => Box::new(Prune::<i8, Cosine>::new(reader)),
-            Metric::CosineNormalized => Box::new(Prune::<i8, CosineNormalized>::new(reader)),
-        };
+        let distance = <Self as DistanceProvider<Self>>::distance_comparer(
+            full.metric(),
+            Some(full.dim()),
+        );
 
+        let output: Box<dyn layers::Prune> = Prune::<Self, _>::new(reader, distance).boxed();
         Ok(output)
     }
 }
@@ -913,7 +916,7 @@ impl_full_precision!(f32, f16, u8, i8);
 //     use rand::{Rng, SeedableRng, rngs::StdRng};
 //
 //     use super::*;
-//     // Bring the inherent-call traits into method scope. The `Distance` / `QueryDistance`
+//     // Bring the inherent-call traits into method scope. The `Distance` / `ExpandBeam`
 //     // traits are not imported: their methods are reached through `&dyn _` trait objects,
 //     // which does not require the trait to be in scope.
 //     use crate::layers::{AsDistance as _, QueryVisitor, Search as _, Set as _};
@@ -951,16 +954,16 @@ impl_full_precision!(f32, f16, u8, i8);
 //         (0..dim).map(|_| T::sample(rng)).collect()
 //     }
 //
-//     /// A [`QueryVisitor`] that simply boxes the minted kernel so the test can probe it
+//     /// A [`QueryVisitor`] that simply boxes the query kernel so the test can probe it
 //     /// directly. Exercises both `visit` (dynamic) and `visit_sized` (specialized) paths.
 //     struct Collect;
 //
 //     impl<'a> QueryVisitor<'a> for Collect {
-//         type Output = Box<dyn layers::QueryDistance + 'a>;
+//         type Output = Box<dyn layers::ExpandBeam + 'a>;
 //
 //         fn visit<Q>(self, distance: Q) -> Self::Output
 //         where
-//             Q: layers::QueryDistance + 'a,
+//             Q: layers::ExpandBeam + 'a,
 //         {
 //             Box::new(distance)
 //         }
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 1c9351ea7e..9ac33674ce 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -9,13 +9,13 @@
 
 pub mod num;
 
-pub mod arch;
 mod buffer;
 mod counters;
 mod epoch;
 mod freelist;
 mod ids;
 mod neighbors;
+mod prefetch;
 mod tag;
 
 mod store;
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index cefe1acb25..449b83a85d 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -46,6 +46,11 @@ impl Bytes {
         }
     }
 
+    #[inline]
+    const fn unchecked_mul(self, other: usize) -> Bytes {
+        Bytes::new(self.value() * other)
+    }
+
     /// Subtract `other` from `self` without checking for underflow.
     #[inline]
     pub(crate) const fn unchecked_sub(self, other: Bytes) -> Bytes {
diff --git a/diskann-inmem/src/arch.rs b/diskann-inmem/src/prefetch.rs
similarity index 56%
rename from diskann-inmem/src/arch.rs
rename to diskann-inmem/src/prefetch.rs
index 5deb96100e..e49c035eb1 100644
--- a/diskann-inmem/src/arch.rs
+++ b/diskann-inmem/src/prefetch.rs
@@ -8,10 +8,21 @@ use crate::num::Bytes;
 pub(crate) unsafe trait Prefetch:
     std::fmt::Debug + Send + Sync + 'static + Copy
 {
-    fn bytes(self) -> Bytes;
+    /// Check that slices of length `bytes` are compatible with this prefetcher.
+    fn check(self, bytes: Bytes);
     unsafe fn prefetch(self, ptr: *const u8);
 }
 
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct NoPrefetch;
+
+unsafe impl Prefetch for NoPrefetch {
+    fn check(self, bytes: Bytes) {}
+
+    #[inline(always)]
+    unsafe fn prefetch(self, ptr: *const u8) {}
+}
+
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct Loop(Bytes);
 
@@ -22,13 +33,13 @@ impl Loop {
 }
 
 unsafe impl Prefetch for Loop {
-    fn bytes(self) -> Bytes {
-        self.0
+    fn check(self, bytes: Bytes) {
+        assert!(bytes == self.0);
     }
 
     #[inline(always)]
     unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch(ptr, self.bytes().value()) }
+        unsafe { prefetch(ptr, self.0.value()) }
     }
 }
 
@@ -42,54 +53,110 @@ impl<const BYTES: usize> Unrolled<BYTES> {
 }
 
 unsafe impl<const BYTES: usize> Prefetch for Unrolled<BYTES> {
-    fn bytes(self) -> Bytes {
-        Bytes::new(BYTES)
+    fn check(self, bytes: Bytes) {
+        assert_eq!(bytes, Bytes::new(BYTES));
     }
 
     #[inline(always)]
     unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch(ptr, self.bytes().value()) }
+        unsafe { prefetch(ptr, BYTES) }
     }
 }
 
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct JumpTable {
-    bytes: Bytes,
-    back: usize,
-    last: usize,
+pub(crate) struct Binned<const LINES: usize>;
+
+impl<const LINES: usize> Binned<LINES> {
+    pub(crate) const fn new() -> Self {
+        Self
+    }
 }
 
-impl JumpTable {
-    pub(crate) fn new(bytes: Bytes) -> Self {
-        let stride = Bytes::CACHELINE.value();
-        let lines = bytes.value().div_ceil(stride);
-
-        let back = 7 * lines.min(8);
-        let last = if lines > 8 {
-            stride * (lines - 1)
-        } else {
-            0
-        };
-
-        Self {
-            bytes,
-            back,
-            last,
-        }
+unsafe impl<const LINES: usize> Prefetch for Binned<LINES> {
+    fn check(self, bytes: Bytes) {
+        let lower = Bytes::CACHELINE
+            .checked_mul(LINES.saturating_sub(1))
+            .unwrap();
+        let upper = Bytes::CACHELINE.checked_mul(LINES).unwrap();
+
+        assert!(bytes > lower);
+        assert!(bytes <= upper);
+    }
+
+    #[inline(always)]
+    unsafe fn prefetch(self, ptr: *const u8) {
+        unsafe { prefetch(ptr, Bytes::CACHELINE.value() * LINES) }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct BinnedPlus<const LINES: usize>(Bytes);
+
+impl<const LINES: usize> BinnedPlus<LINES> {
+    pub(crate) const fn new(bytes: Bytes) -> Self {
+        assert!(bytes.value() > Bytes::CACHELINE.value() * LINES);
+        Self(bytes)
     }
 }
 
-unsafe impl Prefetch for JumpTable {
-    fn bytes(self) -> Bytes {
-        self.bytes
+unsafe impl<const LINES: usize> Prefetch for BinnedPlus<LINES> {
+    fn check(self, bytes: Bytes) {
+        assert_eq!(bytes, self.0)
     }
 
     #[inline(always)]
     unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch_up_to_8(ptr, self.back, self.last) }
+        use std::arch::x86_64::*;
+
+        let ptr = ptr.cast::<i8>();
+
+        unsafe { _mm_prefetch(ptr.add(self.0.value()), _MM_HINT_T0) };
+
+        let stride = Bytes::CACHELINE.value();
+        for i in 0..LINES {
+            unsafe { _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0) };
+        }
     }
 }
 
+// #[derive(Debug, Clone, Copy)]
+// pub(crate) struct JumpTable {
+//     bytes: Bytes,
+//     back: usize,
+//     last: usize,
+// }
+//
+// impl JumpTable {
+//     pub(crate) fn new(bytes: Bytes) -> Self {
+//         let stride = Bytes::CACHELINE.value();
+//         let lines = bytes.value().div_ceil(stride);
+//
+//         let back = 7 * lines.min(8);
+//         let last = if lines > 8 {
+//             stride * (lines - 1)
+//         } else {
+//             0
+//         };
+//
+//         Self {
+//             bytes,
+//             back,
+//             last,
+//         }
+//     }
+// }
+//
+// unsafe impl Prefetch for JumpTable {
+//     fn bytes(self) -> Bytes {
+//         self.bytes
+//     }
+//
+//     #[inline(always)]
+//     unsafe fn prefetch(self, ptr: *const u8) {
+//         unsafe { prefetch_up_to_8(ptr, self.back, self.last) }
+//     }
+// }
+
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
 /// The last cache line prefetched first, followed by the rest in ascending order.
@@ -112,7 +179,7 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
 
     // SAFETY: Inherited from caller.
     unsafe { _mm_prefetch(ptr.add(stride * (lines - 1)), _MM_HINT_T0) };
-    for i in 0..(lines - 1).min(8) {
+    for i in 0..(lines - 1) {
         // SAFETY: Inherited from caller.
         unsafe {
             _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0);
@@ -129,7 +196,9 @@ pub unsafe fn prefetch_up_to_8(ptr: *const u8, back: usize, last: usize) {
     // const PREFETCH_INSTRUCTION_BYTES: usize = 7;
 
     if last != 0 {
-        unsafe { _mm_prefetch(ptr.cast::<i8>().add(last), _MM_HINT_T0); }
+        unsafe {
+            _mm_prefetch(ptr.cast::<i8>().add(last), _MM_HINT_T0);
+        }
     }
 
     let ptr = ptr.wrapping_sub(128);

From c31f617d558c0008769b615dcf929791805dc76e Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 24 Aug 2026 14:16:56 -0700
Subject: [PATCH 25/34] Tests and polish.

---
 diskann-inmem/src/layers/full.rs | 799 +++++++++++++++++++++----------
 diskann-inmem/src/layers/mod.rs  | 234 +++++----
 diskann-inmem/src/num.rs         |   5 -
 diskann-inmem/src/prefetch.rs    | 336 ++++++-------
 diskann-inmem/src/provider.rs    |  32 +-
 5 files changed, 887 insertions(+), 519 deletions(-)

diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 630932d7c6..f074526741 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -10,7 +10,9 @@ use diskann_utils::views::Matrix;
 use diskann_vector::{
     UnalignedSlice,
     conversion::SliceCast,
-    distance::{Cosine, CosineNormalized, InnerProduct, Metric, Specialize, SquaredL2, DistanceProvider},
+    distance::{
+        Cosine, CosineNormalized, DistanceProvider, InnerProduct, Metric, Specialize, SquaredL2,
+    },
 };
 use diskann_wide::{
     ARCH,
@@ -20,10 +22,10 @@ use half::f16;
 use thiserror::Error;
 
 use crate::{
-    prefetch::{self, Prefetch},
     counters::LocalCounters,
     epoch, layers,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
+    prefetch::{self, Prefetch},
     store::{
         self, Store,
         invasive::{self, Invasive},
@@ -232,7 +234,6 @@ where
         };
 
         let mut buf: Box<[_]> = std::iter::repeat_n(T::zeroed(), self.dim()).collect();
-
         bytemuck::must_cast_slice_mut::<T, u8>(&mut buf).copy_from_slice(data);
         Ok(buf)
     }
@@ -367,10 +368,7 @@ struct Prune<'a, T, D> {
 }
 
 impl<'a, T, D> Prune<'a, T, D> {
-    fn new(
-        reader: store::invasive::Reader<'a>,
-        distance: D,
-    ) -> Self {
+    fn new(reader: store::invasive::Reader<'a>, distance: D) -> Self {
         // This should be ensured at construction time
         debug_assert!(
             reader
@@ -390,7 +388,6 @@ impl<'a, T, D> Prune<'a, T, D> {
     fn boxed(self) -> Box<Self> {
         Box::new(self)
     }
-
 }
 
 impl<T, D> layers::Prune for Prune<'_, T, D>
@@ -401,7 +398,7 @@ where
     fn prepare(
         &mut self,
         items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
-    ) -> ANNResult<layers::PruneKey> {
+    ) -> ANNResult<usize> {
         let mut counter = layers::PruneKey::counter();
         self.buffer.clear();
         self.buffer.reserve(items.len());
@@ -421,15 +418,16 @@ where
                 // someone will provide a prune list exceeding `u16::MAX`.
                 //
                 // In addition, `diskann` limits this bound as well.
-                counter = counter.inc()?;
+                counter = counter.increment()?;
             }
         }
 
-        Ok(counter)
+        Ok(counter.index())
     }
 
     fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
-        self.distance.eval(self.buffer[a.index()], self.buffer[b.index()])
+        self.distance
+            .eval(self.buffer[a.index()], self.buffer[b.index()])
     }
 }
 
@@ -476,14 +474,10 @@ impl<'a, T, U> IntoExpandBeam<'a, T, U> {
             _data: PhantomData,
         })
     }
-
-    fn bytes_plus_tag(&self) -> Bytes {
-        self.reader.bytes_plus_tag()
-    }
 }
 
 trait Distance<T, U>: std::fmt::Debug + Send + Sync + 'static {
-    fn eval(&self, x: UnalignedSlice<'_, T>, u: UnalignedSlice<'_, U>) -> f32;
+    fn eval(&self, x: UnalignedSlice<'_, T>, y: UnalignedSlice<'_, U>) -> f32;
 }
 
 #[derive(Debug)]
@@ -535,7 +529,7 @@ struct ExpandBeam<'a, P, T, U, D> {
     // THe prefetch look-ahead.
     lookahead: Option<NonZeroUsize>,
     // The type of the data prefetcher.
-    prefetch: P,
+    prefetch: prefetch::Checked<P>,
     // The type of the distance used for the arguments
     distance: D,
     // The type of the data in the original dataset.
@@ -554,7 +548,10 @@ impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
             _data,
         } = into;
 
-        prefetch.check(reader.bytes_plus_tag());
+        // TAG: PREFETCH-CHECK
+        let prefetch = prefetch::Checked::new(prefetch, reader.bytes_plus_tag())
+            .expect("internal APIs should only provide valid prefetcher");
+
         Self {
             query,
             reader,
@@ -609,22 +606,19 @@ where
     }
 
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize> {
+        debug_assert!(buffer.len() >= list.len());
+
         let len = list.len();
         let lookahead = self.lookahead.map(|l| l.get()).unwrap_or(0).min(len);
-        // let lookahead = 8.min(len);
 
-        for j in 0..lookahead {
+        for j in list.iter().take(lookahead) {
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
             // as the validity of the prefetch bounds.
             //
             // We do not materialize the `RawSlice` as a reference.
             unsafe {
-                self.prefetch.prefetch(
-                    self.reader
-                        .read_raw_unchecked(list.get_unchecked(j).into_usize())
-                        .as_ptr()
-                        .cast(),
-                )
+                let raw = self.reader.read_raw_unchecked(j.into_usize());
+                self.prefetch.prefetch(raw.as_ptr(), raw.len());
             }
         }
 
@@ -638,12 +632,10 @@ where
                 //
                 // We do not materialize the `RawSlice` as a reference.
                 unsafe {
-                    self.prefetch.prefetch(
-                        self.reader
-                            .read_raw_unchecked(list.get_unchecked(j).into_usize())
-                            .as_ptr()
-                            .cast(),
-                    )
+                    let raw = self
+                        .reader
+                        .read_raw_unchecked(list.get_unchecked(j).into_usize());
+                    self.prefetch.prefetch(raw.as_ptr(), raw.len());
                 }
                 j += 1;
             }
@@ -694,22 +686,17 @@ macro_rules! expand_beam {
         ))
     }};
     ($into:ident, $f:ident) => {{
-        let bytes = $into.bytes_plus_tag();
         Box::new(ExpandBeam::new(
             $into,
-            prefetch::Loop::new(bytes),
+            prefetch::Loop::new(),
             Pure::<$f>::new(),
         ))
     }};
 }
 
 macro_rules! prune {
-    ($self:ty, $reader:ident, $f:ident) => {{
-        Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed()
-    }};
-    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{
-        Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed()
-    }};
+    ($self:ty, $reader:ident, $f:ident) => {{ Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed() }};
+    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{ Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed() }};
 }
 
 impl FullPrecisionImpl for f32 {
@@ -834,16 +821,11 @@ impl FullPrecisionImpl for i8 {
     ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
         let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
-        let distance = <Self as DistanceProvider<Self>>::distance_comparer(
-            full.metric(),
-            Some(full.dim()),
-        );
+        let distance =
+            <Self as DistanceProvider<Self>>::distance_comparer(full.metric(), Some(full.dim()));
 
-        let output: Box<dyn layers::ExpandBeam + 'a> = ExpandBeam::new(
-            into,
-            prefetch::Loop::new(full.bytes_plus_tag()),
-            distance,
-        ).boxed();
+        let output: Box<dyn layers::ExpandBeam + 'a> =
+            ExpandBeam::new(into, prefetch::Loop::new(), distance).boxed();
 
         Ok(output)
     }
@@ -852,10 +834,8 @@ impl FullPrecisionImpl for i8 {
         let reader = full.reader()?;
         let dim = full.dim();
 
-        let distance = <Self as DistanceProvider<Self>>::distance_comparer(
-            full.metric(),
-            Some(full.dim()),
-        );
+        let distance =
+            <Self as DistanceProvider<Self>>::distance_comparer(full.metric(), Some(full.dim()));
 
         let output: Box<dyn layers::Prune> = Prune::<Self, _>::new(reader, distance).boxed();
         Ok(output)
@@ -908,197 +888,520 @@ impl_full_precision!(f32, f16, u8, i8);
 // Tests //
 ///////////
 
-// #[cfg(test)]
-// #[cfg(not(miri))]
-// mod tests {
-//     use std::fmt::Display;
-//
-//     use rand::{Rng, SeedableRng, rngs::StdRng};
-//
-//     use super::*;
-//     // Bring the inherent-call traits into method scope. The `Distance` / `ExpandBeam`
-//     // traits are not imported: their methods are reached through `&dyn _` trait objects,
-//     // which does not require the trait to be in scope.
-//     use crate::layers::{AsDistance as _, QueryVisitor, Search as _, Set as _};
-//
-//     /// Generate random elements of a layer's data type from a seeded RNG.
-//     trait Sample: bytemuck::Pod {
-//         fn sample<R: Rng>(rng: &mut R) -> Self;
-//     }
-//
-//     impl Sample for f32 {
-//         fn sample<R: Rng>(rng: &mut R) -> Self {
-//             rng.random_range(-1.0f32..1.0f32)
-//         }
-//     }
-//
-//     impl Sample for f16 {
-//         fn sample<R: Rng>(rng: &mut R) -> Self {
-//             f16::from_f32(rng.random_range(-1.0f32..1.0f32))
-//         }
-//     }
-//
-//     impl Sample for u8 {
-//         fn sample<R: Rng>(rng: &mut R) -> Self {
-//             rng.random()
-//         }
-//     }
-//
-//     impl Sample for i8 {
-//         fn sample<R: Rng>(rng: &mut R) -> Self {
-//             rng.random()
-//         }
-//     }
-//
-//     fn gen_vec<T: Sample, R: Rng>(rng: &mut R, dim: usize) -> Vec<T> {
-//         (0..dim).map(|_| T::sample(rng)).collect()
-//     }
-//
-//     /// A [`QueryVisitor`] that simply boxes the query kernel so the test can probe it
-//     /// directly. Exercises both `visit` (dynamic) and `visit_sized` (specialized) paths.
-//     struct Collect;
-//
-//     impl<'a> QueryVisitor<'a> for Collect {
-//         type Output = Box<dyn layers::ExpandBeam + 'a>;
-//
-//         fn visit<Q>(self, distance: Q) -> Self::Output
-//         where
-//             Q: layers::ExpandBeam + 'a,
-//         {
-//             Box::new(distance)
-//         }
-//     }
-//
-//     /// Compare two distances allowing for floating-point reassociation between the
-//     /// specialized / converted kernels and the dynamic reference.
-//     fn approx_eq(got: f32, want: f32) -> bool {
-//         (got - want).abs() <= 1e-3 + 1e-4 * want.abs()
-//     }
-//
-//     /// Exercise every `Full<T>` API across dimensions `1..=max_dim`.
-//     ///
-//     /// For each dimension we check that `bytes`/`set` agree, that `distance` and
-//     /// `query_distance` are consistent with `DistanceProvider`, and that all of these
-//     /// reject byte slices that are too long or too short.
-//     fn test_impl<T>(max_dim: usize, ctx: &dyn Display)
-//     where
-//         T: FullPrecision + Sample + DistanceProvider<T>,
-//     {
-//         let mut rng = StdRng::seed_from_u64(0x0D15_0ACE ^ max_dim as u64);
-//         let metrics = [
-//             Metric::L2,
-//             Metric::InnerProduct,
-//             Metric::Cosine,
-//             Metric::CosineNormalized,
-//         ];
-//
-//         for dim in 1..=max_dim {
-//             let a = gen_vec::<T, _>(&mut rng, dim);
-//             let b = gen_vec::<T, _>(&mut rng, dim);
-//
-//             // `bytes` and `set` agree: the encoded buffer equals the raw cast bytes.
-//             let layer = Full::<T>::new(dim, Metric::L2);
-//             assert_eq!(
-//                 layer.bytes().value(),
-//                 dim * std::mem::size_of::<T>(),
-//                 "{ctx}: dim {dim}: unexpected byte length",
-//             );
-//
-//             let mut a_bytes = vec![0u8; layer.bytes().value()];
-//             layer.set(&a, &mut a_bytes).unwrap();
-//             assert_eq!(
-//                 a_bytes.as_slice(),
-//                 bytemuck::cast_slice::<T, u8>(&a),
-//                 "{ctx}: dim {dim}: set mismatch",
-//             );
-//
-//             let mut b_bytes = vec![0u8; layer.bytes().value()];
-//             layer.set(&b, &mut b_bytes).unwrap();
-//
-//             for metric in metrics {
-//                 let full = Full::<T>::new(dim, metric);
-//
-//                 // Reference value straight from `DistanceProvider`.
-//                 let reference =
-//                     <T as DistanceProvider<T>>::distance_comparer(metric, Some(dim)).call(&a, &b);
-//
-//                 // `distance` is built from the same comparer, so it must match exactly.
-//                 let distance = full.as_distance();
-//                 let via_distance = distance.evaluate(&a_bytes, &b_bytes).unwrap();
-//                 assert_eq!(
-//                     via_distance, reference,
-//                     "{ctx}: dim {dim}, metric {metric:?}: distance != DistanceProvider",
-//                 );
-//
-//                 // `query_distance` computes the same geometry. Specialized and f16-converted
-//                 // kernels may reassociate the summation, so compare approximately.
-//                 let query = full.query_distance(a.as_slice(), Collect).unwrap();
-//                 let via_query = query.evaluate(&b_bytes).unwrap();
-//                 assert!(
-//                     approx_eq(via_query, via_distance),
-//                     "{ctx}: dim {dim}, metric {metric:?}: query {via_query} != distance {via_distance}",
-//                 );
-//
-//                 // Every distance API rejects byte slices that are too long or too short.
-//                 let short = &a_bytes[..a_bytes.len() - 1];
-//                 let mut long = a_bytes.clone();
-//                 long.push(0);
-//
-//                 assert!(distance.evaluate(short, &b_bytes).is_err());
-//                 assert!(distance.evaluate(&long, &b_bytes).is_err());
-//                 assert!(distance.evaluate(&a_bytes, short).is_err());
-//                 assert!(distance.evaluate(&a_bytes, &long).is_err());
-//
-//                 assert!(query.evaluate(short).is_err());
-//                 assert!(query.evaluate(&long).is_err());
-//             }
-//
-//             // `set` rejects mis-sized element and buffer slices.
-//             let mut buf = vec![0u8; layer.bytes().value()];
-//             let too_many = gen_vec::<T, _>(&mut rng, dim + 1);
-//             assert!(
-//                 layer.set(&too_many, &mut buf).is_err(),
-//                 "{ctx}: dim {dim}: set accepted an over-long element slice",
-//             );
-//
-//             assert!(
-//                 layer.query_distance(&too_many, Collect).is_err(),
-//                 "{ctx}: dim {dim}: incorrect query lengths should be rejected"
-//             );
-//
-//             let mut short_buf = vec![0u8; layer.bytes().value().saturating_sub(1)];
-//             assert!(
-//                 layer.set(&a, &mut short_buf).is_err(),
-//                 "{ctx}: dim {dim}: set accepted an under-sized buffer",
-//             );
-//
-//             let too_few = gen_vec::<T, _>(&mut rng, dim - 1);
-//             assert!(
-//                 layer.query_distance(&too_few, Collect).is_err(),
-//                 "{ctx}: dim {dim}: incorrect query lengths should be rejected"
-//             );
-//         }
-//     }
-//
-//     // `max_dim` must exceed the largest specialized dimension for each type so the
-//     // const-generic (`visit_sized`) paths are covered alongside the dynamic ones.
-//     #[test]
-//     fn full_f32() {
-//         test_impl::<f32>(256, &"f32");
-//     }
-//
-//     #[test]
-//     fn full_f16() {
-//         test_impl::<f16>(256, &"f16");
-//     }
-//
-//     #[test]
-//     fn full_u8() {
-//         test_impl::<u8>(160, &"u8");
-//     }
-//
-//     #[test]
-//     fn full_i8() {
-//         test_impl::<i8>(160, &"i8");
-//     }
-// }
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use std::fmt::Display;
+
+    use diskann_utils::lazy_format;
+    use hashbrown::{HashMap, HashSet};
+    use rand::{Rng, SeedableRng, rngs::StdRng};
+
+    /// Generate random elements of a layer's data type from a seeded RNG.
+    trait Sample: bytemuck::Pod {
+        fn sample<R: Rng>(rng: &mut R) -> Self;
+    }
+
+    impl Sample for f32 {
+        fn sample<R: Rng>(rng: &mut R) -> Self {
+            rng.random_range(-1.0f32..1.0f32)
+        }
+    }
+
+    impl Sample for f16 {
+        fn sample<R: Rng>(rng: &mut R) -> Self {
+            diskann_wide::cast_f32_to_f16(rng.random_range(-1.0f32..1.0f32))
+        }
+    }
+
+    impl Sample for u8 {
+        fn sample<R: Rng>(rng: &mut R) -> Self {
+            rng.random()
+        }
+    }
+
+    impl Sample for i8 {
+        fn sample<R: Rng>(rng: &mut R) -> Self {
+            rng.random()
+        }
+    }
+
+    fn gen_vec<T: Sample>(dim: usize, rng: &mut impl Rng) -> Vec<T> {
+        (0..dim).map(|_| T::sample(rng)).collect()
+    }
+
+    /// Compare two distances allowing for floating-point reassociation between the
+    /// specialized / converted kernels and the dynamic reference.
+    #[must_use]
+    fn approx_eq(got: f32, want: f32) -> bool {
+        (got - want).abs() <= 1e-3 + 1e-4 * want.abs()
+    }
+
+    /// A simple test `Full` containing 1-dimensional `f32` values.
+    ///
+    /// This is used in dedicated `ExpandBeam` and `Prune` tests in a miri-friendly way.
+    ///
+    /// Two start points are included, initializd to `capacity` and `capacity + 1`.
+    fn test_full(capacity: Capacity) -> (Full<f32>, HashMap<u32, f32>) {
+        let start_points = [capacity.value() as f32, (capacity.value() + 1) as f32];
+
+        let full = <_ as layers::LayerConfig>::build(Full::<f32>::config(
+            capacity,
+            MaxDegree::new(0),
+            Metric::L2,
+            Matrix::column_vector(Box::new(start_points)),
+        ))
+        .unwrap();
+
+        assert_eq!(full.dim(), 1, "start points only have one dimension");
+        assert_eq!(full.bytes(), Bytes::size_of::<f32>());
+        assert_eq!(
+            full.bytes_plus_tag(),
+            Bytes::size_of::<f32>()
+                .checked_add(Bytes::size_of::<AtomicTag>())
+                .unwrap()
+        );
+        assert_eq!(full.metric(), Metric::L2);
+        assert_eq!(
+            <_ as layers::Layer>::id_limit(&full),
+            IdLimit::new(capacity.value() as u32 + 2)
+        );
+        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
+
+        let points: HashMap<u32, f32> = {
+            let reader = full.reader().unwrap();
+            assert_eq!(
+                reader.read(capacity.value()).unwrap(),
+                bytemuck::bytes_of(&start_points[0])
+            );
+            assert_eq!(
+                reader.read(capacity.value() + 1).unwrap(),
+                bytemuck::bytes_of(&start_points[1])
+            );
+
+            [
+                (capacity.value() as u32, start_points[0]),
+                ((capacity.value() + 1) as u32, start_points[1]),
+            ]
+            .into_iter()
+            .collect()
+        };
+
+        (full, points)
+    }
+
+    #[derive(Debug)]
+    struct TestDistance;
+
+    impl Distance<f32, f32> for TestDistance {
+        fn eval(&self, x: UnalignedSlice<'_, f32>, y: UnalignedSlice<'_, f32>) -> f32 {
+            assert_eq!(x.len(), 1);
+            assert_eq!(y.len(), 1);
+
+            unsafe { x.as_ptr().read_unaligned() + y.as_ptr().read_unaligned() }
+        }
+    }
+
+    /// A Miri-friendly test for [`ExpandBeam`].
+    ///
+    /// This test covers the following:
+    ///
+    /// 1. Prefetches are in-bounds for all lookaheads.
+    /// 2. [`ExpandBeam`] doesn't lie about its [`IdLimit`].
+    /// 3. [`ExpandBeam`] various methods are internally consistent with eachother and
+    ///    consistent with the parent [`Full`] for item readability.
+    /// 4. [`ExpandBeam::expand_beam`] calls in order and visits all items in the input list.
+    #[test]
+    fn test_expand_beam() {
+        let capacity = Capacity::new(20);
+        let id_limit = IdLimit::new(22);
+
+        let (mut full, mut points) = test_full(capacity);
+
+        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
+        assert_eq!(<_ as layers::Layer>::id_limit(&full), id_limit);
+
+        let mut available: HashSet<u32> = (0..capacity.value()).map(|i| i as u32).collect();
+
+        // Insert the values 0 to 10.
+        for i in 0u32..10 {
+            let guard = <_ as layers::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
+
+            let id = <_ as layers::Guard>::id(&guard);
+
+            assert!(
+                available.remove(&id),
+                "insertion should return available slots",
+            );
+
+            assert!(
+                points.insert(id, i as f32).is_none(),
+                "insertion should not repeat",
+            );
+
+            <_ as layers::Guard>::publish(guard);
+        }
+
+        // Lookaheads to try.
+        let lookaheads: &[Option<NonZeroUsize>] = &[
+            None,
+            NonZeroUsize::new(1),
+            NonZeroUsize::new(2),
+            NonZeroUsize::new(5),
+            NonZeroUsize::new(10),
+            NonZeroUsize::new(100),
+        ];
+
+        // This is the main loop for testing `ExpandBeam`.
+        //
+        // We do several things.
+        //
+        // 1. We try to insert two additional IDs, but hold don't publish their guards.
+        //    This tests that we avoid items being readable until they are published.
+        //
+        // 2. We commit and insert two new points but immediately retire them.
+        //    This tests that we correctly make these points unreadable.
+        for lookahead in lookaheads {
+            full.lookahead = *lookahead;
+
+            let g0 = <_ as layers::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
+            let g1 = <_ as layers::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
+            let g2 = <_ as layers::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
+            let g3 = <_ as layers::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
+
+            {
+                let g0_id = <_ as layers::Guard>::id(&g0);
+                <_ as layers::Guard>::publish(g0);
+                <_ as layers::Layer>::retire(&full, g0_id).unwrap();
+            }
+
+            {
+                let g1_id = <_ as layers::Guard>::id(&g1);
+                <_ as layers::Guard>::publish(g1);
+                <_ as layers::Layer>::retire(&full, g1_id).unwrap();
+            }
+
+            let query = -1.0f32;
+
+            let into =
+                IntoExpandBeam::new(&full, Calf::Borrowed(std::slice::from_ref(&query))).unwrap();
+
+            let mut expand = ExpandBeam::new(into, prefetch::Loop::new(), TestDistance);
+
+            assert_eq!(<_ as layers::ExpandBeam>::id_limit(&expand), id_limit);
+
+            let mut buf = Vec::<(u32, f32)>::new();
+            let mut list = Vec::<u32>::new();
+
+            // Use triangular indexing from `0..id_limit` with `points` serving as the
+            // groundtruth.
+            //
+            // Note that we purposely make `list` extra long with redundant indices to help
+            // catch indexing bugs inside `ExpanBeam`.
+            for i in 0..=id_limit.value() {
+                list.clear();
+                list.extend((0..i).rev());
+                list.extend(0..i);
+
+                buf.resize(list.len(), Default::default());
+
+                let read =
+                    unsafe { <_ as layers::ExpandBeam>::expand_beam(&expand, &list, &mut buf) }
+                        .unwrap();
+
+                let expected: Vec<(u32, f32)> = list
+                    .iter()
+                    .copied()
+                    .filter_map(|id| match points.get(&id) {
+                        Some(point) => {
+                            let expected = point + query;
+
+                            assert!(
+                                <_ as layers::Layer>::is_readable(&full, id).unwrap(),
+                                "point should be readable"
+                            );
+
+                            assert_eq!(
+                                <_ as layers::ExpandBeam>::evaluate(&expand, id).unwrap(),
+                                Some(expected),
+                                "readable points should return valid distances",
+                            );
+
+                            Some((id, expected))
+                        }
+                        None => {
+                            assert!(
+                                !<_ as layers::Layer>::is_readable(&full, id).unwrap(),
+                                "points not yielded by ExpandBeam should be unreadable"
+                            );
+
+                            assert!(
+                                <_ as layers::ExpandBeam>::evaluate(&expand, id)
+                                    .unwrap()
+                                    .is_none(),
+                                "unreable points should return `None` for their distance",
+                            );
+
+                            None
+                        }
+                    })
+                    .collect();
+
+                assert_eq!(&buf[..read], &*expected);
+            }
+
+            assert!(
+                <_ as layers::ExpandBeam>::evaluate(&expand, id_limit.value()).is_err(),
+                "`ExpandBeam::evaluate` should catch out-of-bounds errors",
+            );
+
+            // Ensure we hold onto `g2` and `g3` for the duration of the above check.
+            drop(g2);
+            drop(g3);
+        }
+    }
+
+    fn test_prune_inner(
+        points: &HashMap<u32, f32>,
+        prune: &mut Prune<f32, TestDistance>,
+        ids: &[u32],
+    ) {
+        let mut items: HashMap<u32, Option<layers::PruneKey>> =
+            ids.iter().map(|id| (*id, None)).collect();
+
+        let processed = <_ as layers::Prune>::prepare(prune, items.iter_mut()).unwrap();
+        assert_eq!(processed, items.values().filter(|i| i.is_some()).count());
+
+        // Ensure that `prepare` agrees with `points`.
+        for (k, v) in items.iter() {
+            match v {
+                Some(_) => assert!(points.contains_key(k)),
+                None => assert!(!points.contains_key(k)),
+            }
+        }
+
+        fn filter((k, v): (&u32, &Option<layers::PruneKey>)) -> Option<(u32, layers::PruneKey)> {
+            v.map(|v| (*k, v))
+        }
+
+        // Ensure that distances agree.
+        for (k0, v0) in items.iter().filter_map(filter) {
+            for (k1, v1) in items.iter().filter_map(filter) {
+                // Manually implement `TestDistance`.
+                let expected = points[&k0] + points[&k1];
+                let got = <_ as layers::Prune>::evaluate(prune, v0, v1);
+                assert_eq!(expected, got);
+            }
+        }
+    }
+
+    /// A Miri-friendly test for `Prune`.
+    #[test]
+    fn test_prune() {
+        let capacity = Capacity::new(20);
+        let id_limit = IdLimit::new(22);
+
+        let (mut full, mut points) = test_full(capacity);
+
+        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
+        assert_eq!(<_ as layers::Layer>::id_limit(&full), id_limit);
+
+        let mut available: HashSet<u32> = (0..capacity.value()).map(|i| i as u32).collect();
+
+        // Insert the values 0 to 10.
+        for i in 0u32..10 {
+            let guard = <_ as layers::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
+
+            let id = <_ as layers::Guard>::id(&guard);
+
+            assert!(
+                available.remove(&id),
+                "insertion should return available slots",
+            );
+
+            assert!(
+                points.insert(id, i as f32).is_none(),
+                "insertion should not repeat",
+            );
+
+            <_ as layers::Guard>::publish(guard);
+        }
+
+        // We do several things.
+        //
+        // 1. We try to insert two additional IDs, but hold don't publish their guards.
+        //    This tests that we avoid items being readable until they are published.
+        //
+        // 2. We commit and insert two new points but immediately retire them.
+        //    This tests that we correctly make these points unreadable.
+        let g0 = <_ as layers::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
+        let g1 = <_ as layers::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
+        let g2 = <_ as layers::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
+        let g3 = <_ as layers::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
+
+        {
+            let g0_id = <_ as layers::Guard>::id(&g0);
+            <_ as layers::Guard>::publish(g0);
+            <_ as layers::Layer>::retire(&full, g0_id).unwrap();
+        }
+
+        {
+            let g1_id = <_ as layers::Guard>::id(&g1);
+            <_ as layers::Guard>::publish(g1);
+            <_ as layers::Layer>::retire(&full, g1_id).unwrap();
+        }
+
+        let mut prune = Prune::new(full.reader().unwrap(), TestDistance);
+
+        // Note that we emit reads above the `IdLimit`, which we expect to be silently
+        // rejected.
+        for i in 0..=(id_limit.value() + 5) {
+            let mut ids: Vec<u32> = (0..i).collect();
+            test_prune_inner(&points, &mut prune, &ids);
+
+            ids.reverse();
+            test_prune_inner(&points, &mut prune, &ids);
+        }
+
+        // Drop the guards - verifying that they are held in-limbo during the test.
+        drop(g2);
+        drop(g3);
+    }
+
+    //----------------------//
+    // Specialization Tests //
+    //----------------------//
+
+    // These test make sure that the mapping for metrics and specializations are routed
+    // correctly. They do not exhaustively test the `ExpandBeam` kernls as these are left
+    // to tests that are more Miri friendly.
+    fn test_dispatch<T>(dim: usize, metric: Metric, seed: u64, ctx: &dyn Display)
+    where
+        T: FullPrecision + FullPrecisionImpl + Sample + DistanceProvider<T>,
+    {
+        let mut rng = StdRng::seed_from_u64(seed);
+
+        let start_point = gen_vec::<T>(dim, &mut rng);
+        let query = gen_vec::<T>(dim, &mut rng);
+
+        let full = <_ as layers::LayerConfig>::build(Full::<T>::config(
+            Capacity::new(1),
+            MaxDegree::new(0),
+            metric,
+            Matrix::<T>::row_vector(start_point.clone().into()),
+        ))
+        .unwrap();
+
+        let start_id: u32 = 1;
+
+        let internal_query = {
+            let guard = <_ as layers::Set<&[T]>>::set(&full, &query).unwrap();
+            let id = <_ as layers::Guard>::id(&guard);
+            <_ as layers::Guard>::publish(guard);
+            id
+        };
+
+        let distance = <T as DistanceProvider<T>>::distance_comparer(metric, None);
+        let expected = distance.call(&start_point, &query);
+
+        // Expand Beam - both `evaluate` and `expand_beam` share the same distance computer,
+        // so we can just test `evaluate`.
+        {
+            let expand_beam = <T as FullPrecisionImpl>::make_expand_beam(&full, &query).unwrap();
+            let got = expand_beam.evaluate(start_id).unwrap().unwrap();
+            assert!(
+                approx_eq(expected, got),
+                "{ctx} - expected {expected}, got {got}"
+            );
+        }
+
+        // Prune
+        {
+            let mut prune = <T as FullPrecisionImpl>::make_prune(&full).unwrap();
+            let mut points: HashMap<u32, Option<layers::PruneKey>> =
+                [(internal_query, None), (start_id, None)]
+                    .into_iter()
+                    .collect();
+            prune.prepare(points.iter_mut()).unwrap();
+            let got = prune.evaluate(points[&internal_query].unwrap(), points[&start_id].unwrap());
+            assert!(
+                approx_eq(expected, got),
+                "{ctx} - expected {expected}, got {got}"
+            );
+        }
+    }
+
+    fn metrics() -> [Metric; 4] {
+        [
+            Metric::L2,
+            Metric::InnerProduct,
+            Metric::Cosine,
+            Metric::CosineNormalized,
+        ]
+    }
+
+    #[test]
+    fn test_f32_dynamic() {
+        let dim = 10;
+        for m in metrics() {
+            test_dispatch::<f32>(dim, m, 0x917a80fc68f66e04, &lazy_format!("dynamic-{m}-f32"));
+        }
+    }
+
+    // Test the specialized dispatches.
+    #[test]
+    fn test_f32_specialized() {
+        test_dispatch::<f32>(
+            100,
+            Metric::L2,
+            0x917a80fc68f66e04,
+            &lazy_format!("dynamic-l2-f32-100"),
+        );
+    }
+
+    #[test]
+    fn test_f16_dynamic() {
+        let dim = 10;
+        for m in metrics() {
+            test_dispatch::<f16>(dim, m, 0x917a80fc68f66e04, &lazy_format!("dynamic-{m}-f16"));
+        }
+    }
+
+    // Test the specialized dispatches.
+    #[test]
+    fn test_f16_specialized() {
+        test_dispatch::<f16>(
+            100,
+            Metric::L2,
+            0x917a80fc68f66e04,
+            &lazy_format!("dynamic-l2-f16-100"),
+        );
+    }
+
+    #[test]
+    fn test_u8_dynamic() {
+        let dim = 10;
+        for m in [Metric::L2, Metric::InnerProduct, Metric::Cosine] {
+            test_dispatch::<u8>(dim, m, 0x917a80fc68f66e04, &lazy_format!("dynamic-{m}-u8"));
+        }
+    }
+
+    #[test]
+    fn test_u8_specialized() {
+        test_dispatch::<u8>(
+            128,
+            Metric::L2,
+            0x917a80fc68f66e04,
+            &lazy_format!("dynamic-l2-u8-100"),
+        );
+    }
+
+    #[test]
+    fn test_i8_dynamic() {
+        let dim = 10;
+        for m in [Metric::L2, Metric::InnerProduct, Metric::Cosine] {
+            test_dispatch::<i8>(dim, m, 0x917a80fc68f66e04, &lazy_format!("dynamic-{m}-i8"));
+        }
+    }
+}
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/layers/mod.rs
index a6f9dd6edd..5169975f7f 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/layers/mod.rs
@@ -3,28 +3,9 @@
  * Licensed under the MIT license.
  */
 
-//! Distance layers indexing.
+//! # Layering
 //!
-//! An important assumption made by this module is that the data within each layer is
-//! uniformly sized: each entry occupies the same number of bytes. Furthermore, the data
-//! to be stored may not assume any particular alignment. Implementations will strive to
-//! achieve a reasonable alignment, but this may not be relied on.
-//!
-//! # Query Distance Specialization
-//!
-//! The design of this module allows aggressive optimization of graph search kernels via
-//! the [`Search`] and [`QueryVisitor`] pairs of traits.
-//!
-//! Implementations of [`Search`] can pass a [`QueryDistance`] kernel specialized to
-//! a specific geometry (dimensionality or metric type) which upstream [`QueryVisitor`]
-//! will fuse into larger kernels. While this allows for high performance graph kernels,
-//! some considerations should be taken into account:
-//!
-//! 1. For correctness purposes, upstream callers cannot do any kind of caching. As such,
-//!    the dispatch layer used to select the kernel passed to the [`QueryVisitor`] should
-//!    be relatively efficient.
-//!
-//! 2. Keep the number of specializations bounded for compile time reasons.
+//! A simplified interface for [`crate::Provider`] to use for building a graph index.
 
 use std::num::NonZeroU16;
 
@@ -39,48 +20,137 @@ use crate::{
 pub mod full;
 pub use full::{Full, FullPrecision};
 
+/// Deferred creation of [`Layer`]s.
+///
+/// This is used in APIs like [`crate::Provider::new`] to defer allocation of large
+/// in-memory data structures.
 pub trait LayerConfig {
+    /// The type of the resulting [`Layer`].
     type Layer: Layer;
 
+    /// Build the target [`Layer`].
     fn build(self) -> ANNResult<Self::Layer>;
 }
 
-/// Base layer for data representations.
+/// Configurable data layer for [`crate::Provider`].
+///
+/// Layers consist of the adjacency list and data for a concurrent in-memory graph index.
+/// These are expected to be indexed using `u32` IDs from `0..self.id_limit()`, with
+/// internal IDs in `0..self.capacity()` available for writing.
+///
+/// See also:
+///
+/// - [`Set`]: For assigning into the store.
+/// - [`Search`]: Search compatibility with [`crate::Provider`].
+/// - [`Insert`]: Insert compatibility with [`crate::Provider`].
 pub trait Layer: Send + Sync + 'static {
+    /// Return the [`MaxDegree`] of the internal graph.
     fn max_degree(&self) -> MaxDegree;
 
+    /// Return the [`IdLimit`] for the data store.
     fn id_limit(&self) -> IdLimit;
 
+    /// Return the functional [`Capacity`] for the data store.
     fn capacity(&self) -> Capacity;
 
+    /// Retire the internal ID `i`. Such IDs will eventually be recycled for reuse.
     fn retire(&self, i: u32) -> ANNResult<()>;
 
+    /// Return `true` if internal ID `i` is currently readable, returning `None` if `i`
+    /// is outside `0..self.id_limit()`.
     fn is_readable(&self, i: u32) -> Option<bool>;
 }
 
+/// Attempt to write data into a [`Layer`].
+///
+/// This will attempt to find an available internal ID to which `element` can be assigned,
+/// failing if no such ID can be found. The write is not eagerly committed. Instead, a
+/// [`Guard`] is returned, allowing writes to be aborted if necessary.
 pub trait Set<T>: Layer {
+    /// The type of the [`Guard`] used to defer commitment of the write.
     type Guard<'a>: Guard;
+
+    /// Attempt to write the data in `element` into the [`Layer`].
+    ///
+    /// Returns [`Self::Guard`] to retrieve the allocated internal ID for `element` and to
+    /// defer commitment of the write until external code is ready.
+    ///
+    /// Dropping an unpublished [`Guard`] must abort the write and leave its data unpublished.
     fn set(&self, element: T) -> ANNResult<Self::Guard<'_>>;
 }
 
+/// An insert guard for [`Set`] providing deferred commitment of pending writes.
+///
+/// Guards provide several services:
+///
+/// * [`Guard::id`] returns the associated internal ID.
+/// * [`Guard::publish`] commits the change, making the data at the slot publicly visible.
+///
+/// Dropping a guard without calling [`Guard::publish`] indicates a failed insert and
+/// implementations should abort the write and leave the data unpublished.
 pub trait Guard {
     fn id(&self) -> u32;
     fn publish(self);
 }
 
-/// Trait object based implementation of [`diskann::graph::glue::SearchAccessor::expand_beam`].
+/// Enable search over vectors defined by a [`Layer`].
+pub trait Search: Send + Sync + 'static {
+    /// The type of the query. This should be equivalent to the generic parameter in
+    /// [`Set`], but needs to be replicated here due to limitations in the current trait
+    /// design.
+    type Query<'a>;
+
+    /// Construct a [`crate::provider::SearchAccessor`] for the query.
+    #[doc(hidden)]
+    fn search_accessor<'a>(
+        &'a self,
+        query: Self::Query<'a>,
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
+}
+
+/// An insert-specific specialization of [`Search`].
+///
+/// Note that the bounds for this trait are unnecessarily complicated, but require changes
+/// to [`diskann`] to fully resolve.
+pub trait Insert: Search + for<'a> Set<Self::Query<'a>> {
+    #[doc(hidden)]
+    fn insert_search_accessor<'a>(
+        &'a self,
+        query: Self::Query<'a>,
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
+        self.search_accessor(query, provider, counters)
+    }
+
+    #[doc(hidden)]
+    fn prune_accessor<'a>(
+        &'a self,
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
+}
+
+//-----------------//
+// Internal Traits //
+//-----------------//
+
+/// Trait-object-based implementation of
+/// [`diskann::graph::glue::SearchAccessor::expand_beam`].
 ///
 /// Dynamic dispatch is used to enable aggressive specialization of this primitive without
-/// monomorphizing the entire search algorithm. Examples specializations include:
+/// monomorphizing the entire search algorithm. Example specializations include:
 ///
 /// * Optimizing for certain fixed dimensions.
-/// * Inlining metric specific distance functions.
+/// * Inlining metric-specific distance functions.
 /// * Tailoring prefetching to the dimension.
 ///
 /// # Safety
 ///
-/// This trait is `unsafe` because [`Self::id_limit`] **must** work for [`Self::expand_beam`]'s
-/// safety pre-conditions.
+/// Implementors must ensure that every ID accepted by [`Self::id_limit`] can be passed to
+/// [`Self::expand_beam`] without an out-of-bounds memory access, including accesses performed
+/// for prefetching.
 pub(crate) unsafe trait ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Evaluate a raw distance against index `i`.
     fn evaluate(&self, i: u32) -> ANNResult<Option<f32>>;
@@ -90,102 +160,84 @@ pub(crate) unsafe trait ExpandBeam: Send + Sync + std::fmt::Debug {
     /// Callers must be able to use this limit to satisfy the safety pre-conditions for
     /// [`Self::expand_beam`].
     ///
-    /// See also: [`IdLimit::is_in_bound`].
+    /// See also: [`IdLimit::is_in_bounds`].
     fn id_limit(&self) -> IdLimit;
 
     /// Compute the distance between the query and each neighbor in `list`.
     ///
+    /// Unreadable entries may be omitted. Return the number of entries in `buffer` that were
+    /// written so that `buffer[..returned]` contains the expansion IDs and distances.
+    ///
     /// # Safety
     ///
-    /// * All items in `list` must in-bounds with respect to [`Self::id_limit`].
+    /// * All items in `list` must be in bounds with respect to [`Self::id_limit`].
     /// * `buffer.len() >= list.len()`.
     unsafe fn expand_beam(&self, list: &[u32], buffer: &mut [(u32, f32)]) -> ANNResult<usize>;
 }
 
+/// Trait-object-based implementation for [`diskann::graph::glue::PruneAccessor`].
+///
+/// [`Self::prepare`] assigns a [`PruneKey`] to each retrieved entry, allowing implementations
+/// to buffer data in a representation suitable for pruning. A prune session consists of one
+/// call to [`Self::prepare`] followed by calls to [`Self::evaluate`] using the keys it
+/// produced.
+///
+/// [`Self::prepare`] can be called multiple times for a single [`Prune`] object.
+/// Implementations may assume that each call starts a new prune session and do not need to
+/// buffer entries from all calls.
+pub(crate) trait Prune: Send + Sync + std::fmt::Debug {
+    /// Prepare items for pruning.
+    ///
+    /// The `items` iterator contains all the internal IDs that need to be pruned as keys.
+    /// Each value in `items` is an output slot that should be populated with a [`PruneKey`]
+    /// if retrieval succeeds. Implementations are responsible for providing unique
+    /// [`PruneKey`]s and may assume these slots are initialized to `None`.
+    ///
+    /// Returns the total number of internal IDs that were successfully buffered.
+    fn prepare(
+        &mut self,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
+    ) -> ANNResult<usize>;
+
+    /// Compute the distance between the elements referenced by the two [`PruneKey`]s.
+    ///
+    /// Implementations may assume that `a` and `b` were produced by the most recent call to
+    /// [`Self::prepare`] and are allowed to panic if this is violated. This property may
+    /// **not** be relied on for `unsafe` code.
+    fn evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
+}
+
+/// A [`Prune`]-local miniature internal ID.
 #[derive(Debug, Clone, Copy)]
-pub struct PruneKey(NonZeroU16);
+pub(crate) struct PruneKey(NonZeroU16);
 
 impl PruneKey {
     const ONE: Self = Self(NonZeroU16::new(1).unwrap());
 
+    /// Return the initial [`PruneKey`] in a sequence.
     pub(crate) fn counter() -> Self {
         Self::ONE
     }
 
-    pub(crate) fn inc(self) -> Result<Self, Overflow> {
+    /// Increment `self` by 1.
+    ///
+    /// Overflows if called approximately 2^16 times.
+    pub(crate) fn increment(self) -> Result<Self, Overflow> {
         match self.0.checked_add(1) {
             Some(v) => Ok(Self(v)),
             None => Err(Overflow),
         }
     }
 
-    pub(crate) fn as_u64(self) -> u64 {
-        u64::from(self.0.get()) - 1
-    }
-
+    /// Return the zero-based index represented by this key.
     pub(crate) fn index(self) -> usize {
         usize::from(self.0.get()) - 1
     }
 }
 
-impl<'a> diskann_utils::Reborrow<'a> for PruneKey {
-    type Target = PruneKey;
-    fn reborrow(&'a self) -> Self::Target {
-        *self
-    }
-}
-
+/// Incrementing a [`PruneKey`] overflowed.
 #[derive(Debug, Error)]
 #[error("prune list exceeded u16::MAX")]
 pub(crate) struct Overflow;
 
 diskann::convert_error!(Overflow);
-
-/// Enable search over vectors defined by a [`Layer`].
-pub trait Search: Send + Sync + 'static {
-    /// The type of the query. This should be equivalent to the generic parameter in
-    /// [`Set`], but needs to be replicated here due to limitations in the current trait
-    /// design.
-    type Query<'a>;
-
-    #[doc(hidden)]
-    fn search_accessor<'a>(
-        &'a self,
-        query: Self::Query<'a>,
-        provider: &'a (dyn std::any::Any + Send + Sync),
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
-}
-
-// TODO: Try to hide?
-#[doc(hidden)]
-pub(crate) trait Prune: Send + Sync + std::fmt::Debug {
-    fn prepare(
-        &mut self,
-        items: hashbrown::hash_map::IterMut<'_, u32, Option<PruneKey>>,
-    ) -> ANNResult<PruneKey>;
-
-    fn evaluate(&self, a: PruneKey, b: PruneKey) -> f32;
-}
-
-/// A insert-specific specialization of [`Search`].
-///
-/// Note that the bounds for this trait are unnecessarily complicated, but rely on changes
-/// to `diskann` to full resolve.
-pub trait Insert: Search + for<'a> Set<Self::Query<'a>> {
-    #[doc(hidden)]
-    fn insert_search_accessor<'a>(
-        &'a self,
-        query: Self::Query<'a>,
-        provider: &'a (dyn std::any::Any + Send + Sync),
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-        self.search_accessor(query, provider, counters)
-    }
-
-    #[doc(hidden)]
-    fn prune_accessor<'a>(
-        &'a self,
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
-}
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 449b83a85d..cefe1acb25 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -46,11 +46,6 @@ impl Bytes {
         }
     }
 
-    #[inline]
-    const fn unchecked_mul(self, other: usize) -> Bytes {
-        Bytes::new(self.value() * other)
-    }
-
     /// Subtract `other` from `self` without checking for underflow.
     #[inline]
     pub(crate) const fn unchecked_sub(self, other: Bytes) -> Bytes {
diff --git a/diskann-inmem/src/prefetch.rs b/diskann-inmem/src/prefetch.rs
index e49c035eb1..9f2469db83 100644
--- a/diskann-inmem/src/prefetch.rs
+++ b/diskann-inmem/src/prefetch.rs
@@ -3,163 +3,168 @@
  * Licensed under the MIT license.
  */
 
-use crate::num::Bytes;
-
-pub(crate) unsafe trait Prefetch:
-    std::fmt::Debug + Send + Sync + 'static + Copy
-{
-    /// Check that slices of length `bytes` are compatible with this prefetcher.
-    fn check(self, bytes: Bytes);
-    unsafe fn prefetch(self, ptr: *const u8);
-}
-
-#[derive(Debug, Clone, Copy)]
-pub(crate) struct NoPrefetch;
+//! Utilities for prefetching.
 
-unsafe impl Prefetch for NoPrefetch {
-    fn check(self, bytes: Bytes) {}
-
-    #[inline(always)]
-    unsafe fn prefetch(self, ptr: *const u8) {}
-}
+use crate::num::Bytes;
 
+/// A validated [`Prefetch`].
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct Loop(Bytes);
+pub(crate) struct Checked<P>(P);
 
-impl Loop {
-    pub(crate) const fn new(bytes: Bytes) -> Self {
-        Self(bytes)
+impl<P> Checked<P>
+where
+    P: Prefetch,
+{
+    /// Construct a new [`Checked`] containing `prefetcher`, validating the prefetcher against
+    /// `len`.
+    pub(crate) fn new(prefetcher: P, len: Bytes) -> Result<Self, InvalidPrefetch> {
+        prefetcher.check(len)?;
+        Ok(Self(prefetcher))
     }
-}
 
-unsafe impl Prefetch for Loop {
-    fn check(self, bytes: Bytes) {
-        assert!(bytes == self.0);
+    /// Check if `self` can prefetch slices of length `len`.
+    pub(crate) fn check(self, len: Bytes) -> Result<(), InvalidPrefetch> {
+        self.0.check(len)
     }
 
-    #[inline(always)]
-    unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch(ptr, self.0.value()) }
+    /// Prefetch the slice defined by `[ptr, ptr.add(len.value()))`.
+    ///
+    /// # Safety
+    ///
+    /// * The slice must point to valid memory within a single allocation. There are no
+    ///   aliasing requirements.
+    ///
+    /// * `self` must be compatible with `len`, either through [`Self::new`] or [`Self::check`].
+    pub(crate) unsafe fn prefetch(self, ptr: *const u8, len: Bytes) {
+        debug_assert!(self.0.check(len).is_ok());
+        unsafe { self.0.prefetch(ptr, len) }
     }
-}
 
-#[derive(Debug, Clone, Copy)]
-pub(crate) struct Unrolled<const BYTES: usize>;
+    #[cfg(test)]
+    fn safe_prefetch(self, x: &[u8]) -> Result<(), InvalidPrefetch> {
+        let bytes = Bytes::new(x.len());
+        self.check(bytes)?;
 
-impl<const BYTES: usize> Unrolled<BYTES> {
-    pub(crate) const fn new() -> Self {
-        Self
+        // SAFETY: We've checked the length, and slices satisfy the memory and lifetime
+        // requirements.
+        unsafe { self.prefetch(x.as_ptr(), bytes) };
+        Ok(())
     }
 }
 
-unsafe impl<const BYTES: usize> Prefetch for Unrolled<BYTES> {
-    fn check(self, bytes: Bytes) {
-        assert_eq!(bytes, Bytes::new(BYTES));
+/// Prefetch contiguous chunks of memory.
+///
+/// Prefetchers are created ahead of time and reused. This allows specialized prefetchers
+/// (e.g. fully or partially unrolled) to be created.
+///
+/// # Safety
+///
+/// The function [`Self::check`] **must** be accurate. A successful return from
+/// [`Self::check`] must imply that [`Self::prefetch`] on a valid slice of that length is safe.
+pub(crate) unsafe trait Prefetch:
+    std::fmt::Debug + Send + Sync + 'static + Copy
+{
+    /// Check that slices of length `len` are compatible with this prefetcher.
+    ///
+    /// If this function returns `Ok(())`, calling [`Self::prefetch`] with a valid slice of
+    /// length `len` must be safe.
+    fn check(self, len: Bytes) -> Result<(), InvalidPrefetch>;
+
+    /// Prefetch the slice defined by `[ptr, ptr.add(len.value()))`.
+    ///
+    /// # Safety
+    ///
+    /// * The slice must point to valid memory within a single allocation. There are no
+    ///   aliasing requirements.
+    ///
+    /// * [`Self::check`] must return `Ok(())` for `len`.
+    unsafe fn prefetch(self, ptr: *const u8, len: Bytes);
+}
+
+/// A call to [`Prefetch::check`] failed.
+#[derive(Debug)]
+pub(crate) struct InvalidPrefetch(());
+
+impl InvalidPrefetch {
+    const fn new() -> Self {
+        Self(())
     }
+}
 
-    #[inline(always)]
-    unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch(ptr, BYTES) }
+impl std::fmt::Display for InvalidPrefetch {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "invalid prefetch")
     }
 }
 
+impl std::error::Error for InvalidPrefetch {}
+
+diskann::convert_error!(InvalidPrefetch);
+
+/// Prefetch data using a simple `for` loop.
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct Binned<const LINES: usize>;
+pub(crate) struct Loop(());
 
-impl<const LINES: usize> Binned<LINES> {
+impl Loop {
+    /// Construct a new `Loop`.
     pub(crate) const fn new() -> Self {
-        Self
+        Self(())
     }
 }
 
-unsafe impl<const LINES: usize> Prefetch for Binned<LINES> {
-    fn check(self, bytes: Bytes) {
-        let lower = Bytes::CACHELINE
-            .checked_mul(LINES.saturating_sub(1))
-            .unwrap();
-        let upper = Bytes::CACHELINE.checked_mul(LINES).unwrap();
-
-        assert!(bytes > lower);
-        assert!(bytes <= upper);
+// SAFETY: The `Loop` prefetcher is compatible with all lengths of slices as long as the
+// provided slice is valid.
+unsafe impl Prefetch for Loop {
+    fn check(self, _len: Bytes) -> Result<(), InvalidPrefetch> {
+        Ok(())
     }
 
     #[inline(always)]
-    unsafe fn prefetch(self, ptr: *const u8) {
-        unsafe { prefetch(ptr, Bytes::CACHELINE.value() * LINES) }
+    unsafe fn prefetch(self, ptr: *const u8, len: Bytes) {
+        // SAFETY: Inherited from caller.
+        unsafe { prefetch(ptr, len.value()) }
     }
 }
 
+/// A prefetcher for a fixed number of bytes.
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct BinnedPlus<const LINES: usize>(Bytes);
+pub(crate) struct Unrolled<const BYTES: usize>(());
 
-impl<const LINES: usize> BinnedPlus<LINES> {
-    pub(crate) const fn new(bytes: Bytes) -> Self {
-        assert!(bytes.value() > Bytes::CACHELINE.value() * LINES);
-        Self(bytes)
+impl<const BYTES: usize> Unrolled<BYTES> {
+    /// Construct a new `Unrolled`.
+    pub(crate) const fn new() -> Self {
+        Self(())
     }
 }
 
-unsafe impl<const LINES: usize> Prefetch for BinnedPlus<LINES> {
-    fn check(self, bytes: Bytes) {
-        assert_eq!(bytes, self.0)
+// SAFETY: This prefetcher is only valid for slices of length `BYTES`, which is correctly
+// reported in the implementation of `check`.
+unsafe impl<const BYTES: usize> Prefetch for Unrolled<BYTES> {
+    fn check(self, bytes: Bytes) -> Result<(), InvalidPrefetch> {
+        if bytes == Bytes::new(BYTES) {
+            Ok(())
+        } else {
+            Err(InvalidPrefetch::new())
+        }
     }
 
     #[inline(always)]
-    unsafe fn prefetch(self, ptr: *const u8) {
-        use std::arch::x86_64::*;
-
-        let ptr = ptr.cast::<i8>();
-
-        unsafe { _mm_prefetch(ptr.add(self.0.value()), _MM_HINT_T0) };
+    unsafe fn prefetch(self, ptr: *const u8, _len: Bytes) {
+        debug_assert!(self.check(_len).is_ok());
 
-        let stride = Bytes::CACHELINE.value();
-        for i in 0..LINES {
-            unsafe { _mm_prefetch(ptr.add(stride * i), _MM_HINT_T0) };
-        }
+        // SAFETY: Inherited from caller.
+        unsafe { prefetch(ptr, BYTES) }
     }
 }
 
-// #[derive(Debug, Clone, Copy)]
-// pub(crate) struct JumpTable {
-//     bytes: Bytes,
-//     back: usize,
-//     last: usize,
-// }
-//
-// impl JumpTable {
-//     pub(crate) fn new(bytes: Bytes) -> Self {
-//         let stride = Bytes::CACHELINE.value();
-//         let lines = bytes.value().div_ceil(stride);
-//
-//         let back = 7 * lines.min(8);
-//         let last = if lines > 8 {
-//             stride * (lines - 1)
-//         } else {
-//             0
-//         };
-//
-//         Self {
-//             bytes,
-//             back,
-//             last,
-//         }
-//     }
-// }
-//
-// unsafe impl Prefetch for JumpTable {
-//     fn bytes(self) -> Bytes {
-//         self.bytes
-//     }
-//
-//     #[inline(always)]
-//     unsafe fn prefetch(self, ptr: *const u8) {
-//         unsafe { prefetch_up_to_8(ptr, self.back, self.last) }
-//     }
-// }
+//------------------------//
+// Architecture Dependent //
+//------------------------//
 
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
-/// The last cache line prefetched first, followed by the rest in ascending order.
+/// Prefetch locations are spaced one cache-line width apart. The final location is
+/// prefetched first, followed by the rest in ascending order.
 ///
 /// # Safety
 ///
@@ -169,7 +174,7 @@ unsafe impl<const LINES: usize> Prefetch for BinnedPlus<LINES> {
 pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
     use std::arch::x86_64::*;
 
-    // Fetch the last cache line (the one with the tag) first.
+    // Fetch the final location (the one containing the tag) first.
     let stride = Bytes::CACHELINE.value();
     let ptr = ptr.cast::<i8>();
     let lines = len.div_ceil(stride);
@@ -187,48 +192,10 @@ pub(crate) unsafe fn prefetch(ptr: *const u8, len: usize) {
     }
 }
 
-#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
-#[inline(always)]
-pub unsafe fn prefetch_up_to_8(ptr: *const u8, back: usize, last: usize) {
-    use std::arch::x86_64::*;
-
-    // const STRIDE: usize = Bytes::CACHELINE.value();
-    // const PREFETCH_INSTRUCTION_BYTES: usize = 7;
-
-    if last != 0 {
-        unsafe {
-            _mm_prefetch(ptr.cast::<i8>().add(last), _MM_HINT_T0);
-        }
-    }
-
-    let ptr = ptr.wrapping_sub(128);
-    unsafe {
-        std::arch::asm! {
-            // Obtain the address of the label - the base of our prefetch table.
-            "lea {tmp}, [rip + 3f]",
-            "sub {tmp}, {back}",
-            "notrack jmp {tmp}",
-            "2:",
-            "prefetcht0 byte ptr [{base} + 576]",
-            "prefetcht0 byte ptr [{base} + 512]",
-            "prefetcht0 byte ptr [{base} + 448]",
-            "prefetcht0 byte ptr [{base} + 384]",
-            "prefetcht0 byte ptr [{base} + 320]",
-            "prefetcht0 byte ptr [{base} + 256]",
-            "prefetcht0 byte ptr [{base} + 192]",
-            "prefetcht0 byte ptr [{base} + 128]",
-            "3:",
-            back = in(reg) back,
-            base = in(reg_abcd) ptr,
-            tmp = out(reg) _,
-            options(readonly, nostack, preserves_flags),
-        }
-    }
-}
-
 /// Prefetch `len` bytes beginning at `ptr`.
 ///
-/// The last cache line prefetched first, followed by the rest in ascending order.
+/// Prefetch locations are spaced one cache-line width apart. The final location is
+/// prefetched first, followed by the rest in ascending order.
 ///
 /// # Safety
 ///
@@ -240,16 +207,55 @@ pub(crate) unsafe fn prefetch(_ptr: *const u8, _len: usize) {}
 // Tests //
 ///////////
 
-// #
-// [cfg(test)]
-// mod test {
-//     use super::*;
-//
-//     #[test]
-//     fn test_prefetch_up_to_8() {
-//         let v = vec![0u8; 600];
-//         for lines in 0..10 {
-//             unsafe { prefetch_up_to_8(v.as_ptr(), lines) };
-//         }
-//     }
-// }
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    // The safety of this test is mainly dependent on running under Miri.
+    #[test]
+    fn test_loop() {
+        let p = Loop::new();
+        for i in (0..=1024) {
+            let v = vec![0u8; i];
+
+            let checked = Checked::new(p, Bytes::new(v.len())).unwrap();
+            checked.safe_prefetch(&v).unwrap();
+        }
+    }
+
+    fn unrolled_prefetch<const BYTES: usize>() {
+        let p = Unrolled::<BYTES>::new();
+
+        // Happy Path
+        {
+            let c = Checked::new(p, Bytes::new(BYTES)).unwrap();
+            let v = vec![0u8; BYTES];
+            c.safe_prefetch(&v).unwrap();
+        }
+
+        if let Some(under) = BYTES.checked_sub(1) {
+            assert!(Checked::new(p, Bytes::new(under)).is_err());
+            let v = vec![0u8; under];
+            let c = Checked::new(p, Bytes::new(BYTES)).unwrap();
+            assert!(c.safe_prefetch(&v).is_err());
+        }
+
+        if let Some(over) = BYTES.checked_add(1) {
+            assert!(Checked::new(p, Bytes::new(over)).is_err());
+            let v = vec![0u8; over];
+            let c = Checked::new(p, Bytes::new(BYTES)).unwrap();
+            assert!(c.safe_prefetch(&v).is_err());
+        }
+    }
+
+    #[test]
+    fn test_unrolled() {
+        unrolled_prefetch::<0>();
+        unrolled_prefetch::<63>();
+        unrolled_prefetch::<64>();
+        unrolled_prefetch::<65>();
+        unrolled_prefetch::<127>();
+        unrolled_prefetch::<128>();
+        unrolled_prefetch::<129>();
+    }
+}
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index dcd03d7264..e20bc6463f 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -449,11 +449,23 @@ impl<'a> Distance<'a> {
     }
 }
 
-impl diskann_vector::DistanceFunction<layers::PruneKey, layers::PruneKey, f32> for Distance<'_> {
+/// An opaque element-ref for [`PruneAccessor`].
+#[derive(Debug, Clone, Copy)]
+#[repr(transparent)]
+pub struct ElementRef(layers::PruneKey);
+
+impl<'a> diskann_utils::Reborrow<'a> for ElementRef {
+    type Target = ElementRef;
+    fn reborrow(&'a self) -> Self {
+        *self
+    }
+}
+
+impl diskann_vector::DistanceFunction<ElementRef, ElementRef, f32> for Distance<'_> {
     #[inline]
-    fn evaluate_similarity(&self, x: layers::PruneKey, y: layers::PruneKey) -> f32 {
+    fn evaluate_similarity(&self, x: ElementRef, y: ElementRef) -> f32 {
         self.counters.distance_ref(1);
-        self.prune.evaluate(x, y)
+        self.prune.evaluate(x.0, y.0)
     }
 }
 
@@ -467,7 +479,7 @@ impl glue::PruneAccessor for PruneAccessor<'_> {
     where
         Self: 'a;
 
-    type ElementRef<'a> = layers::PruneKey;
+    type ElementRef<'a> = ElementRef;
 
     type View<'a>
         = &'a Self
@@ -492,8 +504,8 @@ impl glue::PruneAccessor for PruneAccessor<'_> {
     {
         self.keys.clear();
         self.keys.extend(itr.map(|i| (i, None)));
-        let count = self.prune.prepare(self.keys.iter_mut())?;
-        self.counters.get_vector(count.as_u64());
+        let count: usize = self.prune.prepare(self.keys.iter_mut())?;
+        self.counters.get_vector(count as u64);
 
         Ok((self, Distance::new(&*self.prune, self.counters.fork())))
     }
@@ -557,14 +569,14 @@ impl provider::NeighborAccessorMut for PruneAccessor<'_> {
 }
 
 impl workingset::View<u32> for &PruneAccessor<'_> {
-    type ElementRef<'a> = layers::PruneKey;
+    type ElementRef<'a> = ElementRef;
     type Element<'a>
-        = layers::PruneKey
+        = ElementRef
     where
         Self: 'a;
 
-    fn get(&self, id: u32) -> Option<layers::PruneKey> {
-        *self.keys.get(&id)?
+    fn get(&self, id: u32) -> Option<ElementRef> {
+        self.keys.get(&id)?.map(|v| ElementRef(v))
     }
 }
 

From 03ea4acb405125f17c2f90c25a3c1140eaa6d550 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 24 Aug 2026 15:18:32 -0700
Subject: [PATCH 26/34] Prepare last cleanups.

---
 diskann-inmem/integration/index/runner.rs |   8 +-
 diskann-inmem/src/layers/full.rs          | 442 ++++++++++++++--------
 diskann-inmem/src/prefetch.rs             |  17 +-
 diskann-inmem/src/provider.rs             |   5 +-
 4 files changed, 311 insertions(+), 161 deletions(-)

diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index 74994f3cc6..3b2057ec66 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -381,7 +381,7 @@ impl Test {
                             MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
-                        ))?,
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::F16(v) => finish(
@@ -390,7 +390,7 @@ impl Test {
                             MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
-                        ))?,
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::U8(v) => finish(
@@ -399,7 +399,7 @@ impl Test {
                             MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
-                        ))?,
+                        )?)?,
                         index_config,
                     ),
                     DatasetView::I8(v) => finish(
@@ -408,7 +408,7 @@ impl Test {
                             MaxDegree::new(max_degree),
                             metric,
                             v.to_owned(),
-                        ))?,
+                        )?)?,
                         index_config,
                     ),
                 };
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index f074526741..f1c57c3f88 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -3,6 +3,38 @@
  * Licensed under the MIT license.
  */
 
+//! # Full-Precision
+//!
+//! A concurrent data store for [`crate::Provider`] enabling full-precision searches and
+//! inserts for collections consisting of `f32`, `f16`, `u8`, or `i8` data types.
+//!
+//! The [`FullPrecision`] generic bound can be used to constrain these data types.
+
+mod internal_docs {
+    //! Internally, the [`super::layers::Search`] and [`super::layers::Insert`] traits
+    //! are implemented via [`super::FullPrecisionImpl`], which creates:
+    //!
+    //! * [`super::ExpandBeam`]: For index search.
+    //! * [`super::Prune`]: For index construction.
+    //!
+    //! These two structs are modular with respect to their exact distance function and
+    //! prefetcher. Since [`super::layers::ExpandBeam`] and [`super::layers::Prune`] are
+    //! used as trait objects, this allows the implementation structs in this module to be
+    //! highly specialized, including:
+    //!
+    //! * Inlining of distance functions.
+    //! * Specializing distance functions on dimension.
+    //! * Specializing prefetches on dimension.
+    //! * Dispatching to different micro-architecture levels.
+    //! * Specialized query preprocessing.
+    //!
+    //! Picking the best combination of all of these requires extensive experimentation.
+    //! The choices made here are mainly heuristic defaults, meant to try to balance
+    //! performance with compile time.
+    //!
+    //! Feel free to experiment and create optimized implementations for workloads that need it.
+}
+
 use std::{fmt::Debug, marker::PhantomData, num::NonZeroUsize};
 
 use diskann::{ANNError, ANNResult, utils::IntoUsize};
@@ -33,6 +65,27 @@ use crate::{
     tag::AtomicTag,
 };
 
+/// A useful trait bound for types compatible with [`Full`].
+///
+/// This encompasses *everything* required for `Full: layers::Insert` and can be used as
+/// a single bound.
+pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
+    #[doc(hidden)]
+    fn __search_accessor<'a>(
+        layer: &'a Full<Self>,
+        query: &'a [Self],
+        provider: &'a (dyn std::any::Any + Send + Sync),
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
+
+    #[doc(hidden)]
+    fn __prune_accessor<'a>(
+        layer: &'a Full<Self>,
+        counters: LocalCounters<'a>,
+    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
+}
+
+/// A configuration struct for [`Full`].
 #[derive(Debug, Clone)]
 pub struct Config<T> {
     layout: store::Layout,
@@ -45,30 +98,62 @@ pub struct Config<T> {
 const DEFAULT_LOOKAHEAD: NonZeroUsize = NonZeroUsize::new(12).unwrap();
 
 impl<T> Config<T> {
+    /// Create a new [`Config`] for a [`Full`].
+    ///
+    /// The resulting store will hold `capacity` writable items and `start_points.nrows()`
+    /// frozen points at internal IDs `[capacity, capacity + start_points.nrows())`. The
+    /// dimensionality of the full-precision data will be inferred from
+    /// `start_points.ncols()`.
+    ///
+    /// The associated graph will be bounded with `max_degree` and `metric` will be used to
+    /// compute distances among the stored points.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the number of start points exceeds `u32::MAX` or the number of
+    /// bytes required for each point exceeds `usize::MAX`.
     pub fn new(
         capacity: Capacity,
         max_degree: MaxDegree,
         metric: Metric,
         start_points: Matrix<T>,
-    ) -> Self {
-        Self {
-            layout: store::Layout::new(
-                capacity,
-                max_degree,
-                start_points.nrows().try_into().unwrap(),
-            ),
+    ) -> Result<Self, ConfigError> {
+        let num_start_points: u32 = match start_points.nrows().try_into() {
+            Ok(points) => points,
+            Err(_) => return Err(ConfigError::TooManyStartPoints(start_points.nrows())),
+        };
+
+        // Check that we won't overflow when computing the number of bytes required for each
+        // data point. This can happen if `start_points` has 0 rows but a large number of
+        // columns.
+        if start_points
+            .ncols()
+            .checked_mul(std::mem::size_of::<T>())
+            .is_none()
+        {
+            return Err(ConfigError::DimTooLarge(start_points.ncols()));
+        }
+
+        Ok(Self {
+            layout: store::Layout::new(capacity, max_degree, num_start_points),
             metric,
             start_points,
             store: store::Config::default(),
             lookahead: Some(DEFAULT_LOOKAHEAD),
-        }
+        })
     }
 
+    /// Override the [`store::Config`] for tailoring concurrency details.
     pub fn store(mut self, config: store::Config) -> Self {
         self.store = config;
         self
     }
 
+    /// Set the prefetch lookahead.
+    ///
+    /// This controls how many iterations ahead in
+    /// [`diskann::graph::glue::SearchAccessor::expand_beam`] data is prefetched into the CPU
+    /// cache. Passing `None` disables prefetching.
     pub fn prefetch(mut self, lookahead: Option<NonZeroUsize>) -> Self {
         self.lookahead = lookahead;
         self
@@ -80,6 +165,21 @@ impl<T> Config<T> {
     }
 }
 
+/// Errors that can arise when constructing [`Config`].
+#[derive(Debug, Error)]
+#[non_exhaustive]
+pub enum ConfigError {
+    #[error("{} start points exceed `u32::MAX`", 0)]
+    TooManyStartPoints(usize),
+    #[error(
+        "the number of bytes to hold {}-dimensional data exceeds `usize::MAX`",
+        0
+    )]
+    DimTooLarge(usize),
+}
+
+diskann::convert_error!(ConfigError);
+
 impl<T> layers::LayerConfig for Config<T>
 where
     T: FullPrecision,
@@ -91,6 +191,7 @@ where
     }
 }
 
+/// Internal helper for implementing [`FullPrecision`].
 trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_expand_beam<'a>(
         full: &'a Full<Self>,
@@ -101,26 +202,6 @@ trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>>;
 }
 
-/// A useful trait bound for types compatible with [`Full`].
-///
-/// This encompasses *everything* required for `Full: layers::Insert` and can be used as
-/// a single bound.
-pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
-    #[doc(hidden)]
-    fn __search_accessor<'a>(
-        layer: &'a Full<Self>,
-        query: &'a [Self],
-        provider: &'a (dyn std::any::Any + Send + Sync),
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::SearchAccessor<'a>>;
-
-    #[doc(hidden)]
-    fn __prune_accessor<'a>(
-        layer: &'a Full<Self>,
-        counters: LocalCounters<'a>,
-    ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
-}
-
 /// Full-precision data layer.
 #[derive(Debug)]
 pub struct Full<T>
@@ -140,16 +221,20 @@ where
     /// Initialize a [`Config`] for this layer.
     ///
     /// See also: [`Config::new`].
+    ///
+    /// # Errors
+    ///
+    /// Returns the errors described by [`Config::new`].
     pub fn config(
         capacity: Capacity,
         max_degree: MaxDegree,
         metric: Metric,
         start_points: Matrix<T>,
-    ) -> Config<T> {
+    ) -> Result<Config<T>, ConfigError> {
         Config::new(capacity, max_degree, metric, start_points)
     }
 
-    /// Create a new full-precision layer for data with the given `dim` and `metric`.
+    /// Create a new full-precision layer from `config`.
     ///
     /// See: [`Config::build`].
     fn new(config: Config<T>) -> ANNResult<Self>
@@ -170,7 +255,13 @@ where
 
         // Initialize start points.
         for (i, row) in std::iter::zip(store.frozen(), start_points.row_iter()) {
-            let mut slot = store.slot(i).unwrap();
+            #[expect(
+                clippy::expect_used,
+                reason = "failing this is an internal, unrecoverable bug"
+            )]
+            let mut slot = store
+                .slot(i)
+                .expect("internal store should leave frozen-points available for writing");
             slot.data()
                 .as_mut_slice()
                 .copy_from_slice(bytemuck::must_cast_slice::<T, u8>(row));
@@ -191,15 +282,17 @@ where
         self.bytes().value() / std::mem::size_of::<T>()
     }
 
-    /// Return the number of bytes of the data handles by this [`layers::Layer`].
+    /// Return the number of payload bytes in each stored vector.
     pub fn bytes(&self) -> Bytes {
         self.store.plugin().bytes()
     }
 
-    pub fn bytes_plus_tag(&self) -> Bytes {
+    #[cfg(test)]
+    fn bytes_plus_tag(&self) -> Bytes {
         self.store.plugin().bytes_plus_tag()
     }
 
+    /// Return the [`Metric`] for this layer.
     pub fn metric(&self) -> Metric {
         self.metric
     }
@@ -216,7 +309,7 @@ where
     }
 
     fn reader(&self) -> Result<invasive::Reader<'_>, epoch::Unavailable> {
-        Ok(Invasive::reader(&self.store)?)
+        Invasive::reader(&self.store)
     }
 }
 
@@ -353,87 +446,9 @@ where
     }
 }
 
-///////////
-// Prune //
-///////////
-
-#[derive(Debug)]
-struct Prune<'a, T, D> {
-    // Buffered data to prune over.
-    buffer: Vec<UnalignedSlice<'a, T>>,
-    // A reader into a layer's store.
-    reader: store::invasive::Reader<'a>,
-    // Type type of the `PureDistanceFunction` used for the implementation.
-    distance: D,
-}
-
-impl<'a, T, D> Prune<'a, T, D> {
-    fn new(reader: store::invasive::Reader<'a>, distance: D) -> Self {
-        // This should be ensured at construction time
-        debug_assert!(
-            reader
-                .bytes()
-                .value()
-                .is_multiple_of(std::mem::size_of::<T>()),
-            "internal inveriant violated",
-        );
-
-        Self {
-            buffer: Vec::new(),
-            reader,
-            distance,
-        }
-    }
-
-    fn boxed(self) -> Box<Self> {
-        Box::new(self)
-    }
-}
-
-impl<T, D> layers::Prune for Prune<'_, T, D>
-where
-    T: Debug + Send + Sync + 'static,
-    D: Distance<T, T>,
-{
-    fn prepare(
-        &mut self,
-        items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
-    ) -> ANNResult<usize> {
-        let mut counter = layers::PruneKey::counter();
-        self.buffer.clear();
-        self.buffer.reserve(items.len());
-
-        for (id, key) in items {
-            if let Some(v) = self.reader.read(id.into_usize()) {
-                self.buffer.push(unsafe {
-                    UnalignedSlice::new(
-                        v.as_ptr().cast::<T>(),
-                        self.reader.bytes().value() / std::mem::size_of::<T>(),
-                    )
-                });
-
-                *key = Some(counter);
-
-                // Potential overflow issue - but it's exceedingly unlikely that
-                // someone will provide a prune list exceeding `u16::MAX`.
-                //
-                // In addition, `diskann` limits this bound as well.
-                counter = counter.increment()?;
-            }
-        }
-
-        Ok(counter.index())
-    }
-
-    fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
-        self.distance
-            .eval(self.buffer[a.index()], self.buffer[b.index()])
-    }
-}
-
-////////////////
-// ExpandBeam //
-////////////////
+//----------------------//
+// Expand Beam (Search) //
+//----------------------//
 
 // A baby [`std::borrow::Cow`].
 #[derive(Debug)]
@@ -462,7 +477,8 @@ struct IntoExpandBeam<'a, T, U> {
 }
 
 impl<'a, T, U> IntoExpandBeam<'a, T, U> {
-    /// Construct a new [`IntoExpandBeam`] - verifying that
+    /// Construct a new [`IntoExpandBeam`], validating the query dimension and acquiring a
+    /// reader for `full`.
     fn new(full: &'a Full<U>, query: Calf<'a, T>) -> ANNResult<Self> {
         full.check_dim(query.len())?;
         let reader = full.reader()?;
@@ -526,7 +542,7 @@ struct ExpandBeam<'a, P, T, U, D> {
     query: Calf<'a, T>,
     // A reader into a layer's store.
     reader: store::invasive::Reader<'a>,
-    // THe prefetch look-ahead.
+    // The prefetch lookahead.
     lookahead: Option<NonZeroUsize>,
     // The type of the data prefetcher.
     prefetch: prefetch::Checked<P>,
@@ -549,8 +565,12 @@ impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
         } = into;
 
         // TAG: PREFETCH-CHECK
+        #[expect(
+            clippy::expect_used,
+            reason = "internal APIs should only provide valid prefetchers"
+        )]
         let prefetch = prefetch::Checked::new(prefetch, reader.bytes_plus_tag())
-            .expect("internal APIs should only provide valid prefetcher");
+            .expect("internal APIs should only provide valid prefetchers");
 
         Self {
             query,
@@ -570,6 +590,12 @@ impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
         Box::new(self)
     }
 
+    /// Compute the distance between the embedded query and `x`.
+    ///
+    /// # Safety
+    ///
+    /// `x.len()` must be exactly `self.bytes()` bytes long and contain
+    /// `self.query.len()` valid values of `U`.
     #[inline(always)]
     unsafe fn run_unchecked(&self, x: &[u8]) -> f32
     where
@@ -583,6 +609,9 @@ impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
     }
 }
 
+// SAFETY: Our implementation of `layers::ExpandBeam::id_limit` is consistent with our
+// `layers::ExpandBeam::expand_beam` implementation. They are both dependent on
+// `invasive::Reader`'s internal bounds.
 unsafe impl<P, T, U, D> layers::ExpandBeam for ExpandBeam<'_, P, T, U, D>
 where
     P: Prefetch,
@@ -594,8 +623,14 @@ where
         if !self.reader.is_in_bounds(i.into_usize()) {
             Err(ANNError::new(OutOfBounds(i)))
         } else {
+            // SAFETY: We have checked that `i` is in-bounds.
             match unsafe { self.reader.read_in_bounds(i.into_usize()) } {
-                Some(data) => Ok(Some(unsafe { self.run_unchecked(data) })),
+                Some(data) => {
+                    // SAFETY: Since we just read `data` from `self.reader`, we know it's
+                    // exactly `self.bytes()` long.
+                    let distance = unsafe { self.run_unchecked(data) };
+                    Ok(Some(distance))
+                }
                 None => Ok(None),
             }
         }
@@ -615,6 +650,8 @@ where
             // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as well
             // as the validity of the prefetch bounds.
             //
+            // We validated `self.prefetch` with `self.reader.bytes_with_tag()` upon construction.
+            //
             // We do not materialize the `RawSlice` as a reference.
             unsafe {
                 let raw = self.reader.read_raw_unchecked(j.into_usize());
@@ -630,6 +667,9 @@ where
                 // SAFETY: The in-bounds constraint is assured by the caller, both for `j` as
                 // well as the validity of the prefetch bounds.
                 //
+                // We validated `self.prefetch` with `self.reader.bytes_with_tag()` upon
+                // construction.
+                //
                 // We do not materialize the `RawSlice` as a reference.
                 unsafe {
                     let raw = self
@@ -642,6 +682,8 @@ where
 
             // SAFETY: Caller asserts that `i` is in-bounds.
             if let Some(data) = unsafe { self.reader.read_in_bounds(i.into_usize()) } {
+                // SAFETY: We just read `data` from `self.reader`, so it has a length of
+                // exactly `self.bytes()`.
                 let distance = unsafe { self.run_unchecked(data) };
 
                 // SAFETY: Inherited from caller.
@@ -673,6 +715,95 @@ struct ExpandBeamError {
 
 diskann::convert_error!(ExpandBeamError);
 
+//-------//
+// Prune //
+//-------//
+
+#[derive(Debug)]
+struct Prune<'a, T, D> {
+    // Buffered data to prune over.
+    buffer: Vec<UnalignedSlice<'a, T>>,
+    // A reader into a layer's store.
+    reader: store::invasive::Reader<'a>,
+    // The distance implementation used for pruning.
+    distance: D,
+}
+
+impl<'a, T, D> Prune<'a, T, D> {
+    fn new(reader: store::invasive::Reader<'a>, distance: D) -> Self {
+        // This should be ensured at construction time
+        debug_assert!(
+            reader
+                .bytes()
+                .value()
+                .is_multiple_of(std::mem::size_of::<T>()),
+            "internal invariant violated",
+        );
+
+        Self {
+            buffer: Vec::new(),
+            reader,
+            distance,
+        }
+    }
+
+    fn boxed(self) -> Box<Self> {
+        Box::new(self)
+    }
+}
+
+impl<T, D> layers::Prune for Prune<'_, T, D>
+where
+    T: Debug + Send + Sync + 'static,
+    D: Distance<T, T>,
+{
+    fn prepare(
+        &mut self,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
+    ) -> ANNResult<usize> {
+        let mut counter = layers::PruneKey::counter();
+        self.buffer.clear();
+        self.buffer.reserve(items.len());
+
+        for (id, key) in items {
+            if let Some(v) = self.reader.read(id.into_usize()) {
+                // SAFETY: We have checked that it is safe to read this data vector and
+                // `self.reader` is preventing any mutation for `self`'s lifetime.
+                //
+                // Further, we know the raw slice has a length exactly `self.reader.bytes()`,
+                // so the formed `UnalignedSlice` is within a single allocated object.
+                let unaligned = unsafe {
+                    UnalignedSlice::new(
+                        v.as_ptr().cast::<T>(),
+                        self.reader.bytes().value() / std::mem::size_of::<T>(),
+                    )
+                };
+
+                self.buffer.push(unaligned);
+
+                *key = Some(counter);
+
+                // Potential overflow issue - but it's exceedingly unlikely that
+                // someone will provide a prune list exceeding `u16::MAX`.
+                //
+                // In addition, `diskann` limits this bound as well.
+                counter = counter.increment()?;
+            }
+        }
+
+        Ok(counter.index())
+    }
+
+    fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
+        self.distance
+            .eval(self.buffer[a.index()], self.buffer[b.index()])
+    }
+}
+
+/////////////////
+// Dispatching //
+/////////////////
+
 const fn compute_bytes<T>(dim: usize) -> usize {
     dim * std::mem::size_of::<T>() + (AtomicTag::SIZE).value()
 }
@@ -695,8 +826,12 @@ macro_rules! expand_beam {
 }
 
 macro_rules! prune {
-    ($self:ty, $reader:ident, $f:ident) => {{ Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed() }};
-    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{ Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed() }};
+    ($self:ty, $reader:ident, $f:ident) => {{
+        Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed()
+    }};
+    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{
+        Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed()
+    }};
 }
 
 impl FullPrecisionImpl for f32 {
@@ -801,7 +936,6 @@ impl FullPrecisionImpl for u8 {
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = full.reader()?;
-        let dim = full.dim();
 
         let output: Box<dyn layers::Prune> = match full.metric {
             Metric::L2 => prune!(Self, reader, SquaredL2),
@@ -832,7 +966,6 @@ impl FullPrecisionImpl for i8 {
 
     fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
         let reader = full.reader()?;
-        let dim = full.dim();
 
         let distance =
             <Self as DistanceProvider<Self>>::distance_comparer(full.metric(), Some(full.dim()));
@@ -942,16 +1075,19 @@ mod tests {
     ///
     /// This is used in dedicated `ExpandBeam` and `Prune` tests in a miri-friendly way.
     ///
-    /// Two start points are included, initializd to `capacity` and `capacity + 1`.
+    /// Two start points are included, initialized to `capacity` and `capacity + 1`.
     fn test_full(capacity: Capacity) -> (Full<f32>, HashMap<u32, f32>) {
         let start_points = [capacity.value() as f32, (capacity.value() + 1) as f32];
 
-        let full = <_ as layers::LayerConfig>::build(Full::<f32>::config(
-            capacity,
-            MaxDegree::new(0),
-            Metric::L2,
-            Matrix::column_vector(Box::new(start_points)),
-        ))
+        let full = <_ as layers::LayerConfig>::build(
+            Full::<f32>::config(
+                capacity,
+                MaxDegree::new(0),
+                Metric::L2,
+                Matrix::column_vector(Box::new(start_points)),
+            )
+            .unwrap(),
+        )
         .unwrap();
 
         assert_eq!(full.dim(), 1, "start points only have one dimension");
@@ -999,6 +1135,8 @@ mod tests {
             assert_eq!(x.len(), 1);
             assert_eq!(y.len(), 1);
 
+            // SAFETY: `UnalignedSlice`s must point to valid data, and we've checked that
+            // the length of each slice is exactly 1. Therefore, the pointer read is safe.
             unsafe { x.as_ptr().read_unaligned() + y.as_ptr().read_unaligned() }
         }
     }
@@ -1009,9 +1147,10 @@ mod tests {
     ///
     /// 1. Prefetches are in-bounds for all lookaheads.
     /// 2. [`ExpandBeam`] doesn't lie about its [`IdLimit`].
-    /// 3. [`ExpandBeam`] various methods are internally consistent with eachother and
+    /// 3. [`ExpandBeam`]'s methods are internally consistent with each other and
     ///    consistent with the parent [`Full`] for item readability.
-    /// 4. [`ExpandBeam::expand_beam`] calls in order and visits all items in the input list.
+    /// 4. [`ExpandBeam::expand_beam`] preserves input order and visits every item in the
+    ///    input list.
     #[test]
     fn test_expand_beam() {
         let capacity = Capacity::new(20);
@@ -1057,10 +1196,10 @@ mod tests {
         //
         // We do several things.
         //
-        // 1. We try to insert two additional IDs, but hold don't publish their guards.
-        //    This tests that we avoid items being readable until they are published.
+        // 1. We insert two additional IDs but hold their guards without publishing.
+        //    This tests that items remain unreadable until they are published.
         //
-        // 2. We commit and insert two new points but immediately retire them.
+        // 2. We publish two new points and immediately retire them.
         //    This tests that we correctly make these points unreadable.
         for lookahead in lookaheads {
             full.lookahead = *lookahead;
@@ -1087,7 +1226,7 @@ mod tests {
             let into =
                 IntoExpandBeam::new(&full, Calf::Borrowed(std::slice::from_ref(&query))).unwrap();
 
-            let mut expand = ExpandBeam::new(into, prefetch::Loop::new(), TestDistance);
+            let expand = ExpandBeam::new(into, prefetch::Loop::new(), TestDistance);
 
             assert_eq!(<_ as layers::ExpandBeam>::id_limit(&expand), id_limit);
 
@@ -1098,7 +1237,7 @@ mod tests {
             // groundtruth.
             //
             // Note that we purposely make `list` extra long with redundant indices to help
-            // catch indexing bugs inside `ExpanBeam`.
+            // catch indexing bugs inside `ExpandBeam`.
             for i in 0..=id_limit.value() {
                 list.clear();
                 list.extend((0..i).rev());
@@ -1106,6 +1245,10 @@ mod tests {
 
                 buf.resize(list.len(), Default::default());
 
+                // SAFETY: By construction, all entries in `list` are within `id_limit`
+                // (verified against this `ExpandBeam` instance.
+                //
+                // Also by construction `buf` is at least as long as `list`.
                 let read =
                     unsafe { <_ as layers::ExpandBeam>::expand_beam(&expand, &list, &mut buf) }
                         .unwrap();
@@ -1140,7 +1283,7 @@ mod tests {
                                 <_ as layers::ExpandBeam>::evaluate(&expand, id)
                                     .unwrap()
                                     .is_none(),
-                                "unreable points should return `None` for their distance",
+                                "unreadable points should return `None` for their distance",
                             );
 
                             None
@@ -1202,7 +1345,7 @@ mod tests {
         let capacity = Capacity::new(20);
         let id_limit = IdLimit::new(22);
 
-        let (mut full, mut points) = test_full(capacity);
+        let (full, mut points) = test_full(capacity);
 
         assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
         assert_eq!(<_ as layers::Layer>::id_limit(&full), id_limit);
@@ -1230,10 +1373,10 @@ mod tests {
 
         // We do several things.
         //
-        // 1. We try to insert two additional IDs, but hold don't publish their guards.
-        //    This tests that we avoid items being readable until they are published.
+        // 1. We insert two additional IDs but hold their guards without publishing.
+        //    This tests that items remain unreadable until they are published.
         //
-        // 2. We commit and insert two new points but immediately retire them.
+        // 2. We publish two new points and immediately retire them.
         //    This tests that we correctly make these points unreadable.
         let g0 = <_ as layers::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
         let g1 = <_ as layers::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
@@ -1285,12 +1428,15 @@ mod tests {
         let start_point = gen_vec::<T>(dim, &mut rng);
         let query = gen_vec::<T>(dim, &mut rng);
 
-        let full = <_ as layers::LayerConfig>::build(Full::<T>::config(
-            Capacity::new(1),
-            MaxDegree::new(0),
-            metric,
-            Matrix::<T>::row_vector(start_point.clone().into()),
-        ))
+        let full = <_ as layers::LayerConfig>::build(
+            Full::<T>::config(
+                Capacity::new(1),
+                MaxDegree::new(0),
+                metric,
+                Matrix::<T>::row_vector(start_point.clone().into()),
+            )
+            .unwrap(),
+        )
         .unwrap();
 
         let start_id: u32 = 1;
diff --git a/diskann-inmem/src/prefetch.rs b/diskann-inmem/src/prefetch.rs
index 9f2469db83..8e0a8ba040 100644
--- a/diskann-inmem/src/prefetch.rs
+++ b/diskann-inmem/src/prefetch.rs
@@ -22,11 +22,6 @@ where
         Ok(Self(prefetcher))
     }
 
-    /// Check if `self` can prefetch slices of length `len`.
-    pub(crate) fn check(self, len: Bytes) -> Result<(), InvalidPrefetch> {
-        self.0.check(len)
-    }
-
     /// Prefetch the slice defined by `[ptr, ptr.add(len.value()))`.
     ///
     /// # Safety
@@ -36,10 +31,18 @@ where
     ///
     /// * `self` must be compatible with `len`, either through [`Self::new`] or [`Self::check`].
     pub(crate) unsafe fn prefetch(self, ptr: *const u8, len: Bytes) {
-        debug_assert!(self.0.check(len).is_ok());
+        debug_assert!(self.check(len).is_ok());
+
+        // SAFETY: Inherited from caller.
         unsafe { self.0.prefetch(ptr, len) }
     }
 
+    /// Check if `self` can prefetch slices of length `len`.
+    #[cfg(debug_assertions)]
+    pub(crate) fn check(self, len: Bytes) -> Result<(), InvalidPrefetch> {
+        self.0.check(len)
+    }
+
     #[cfg(test)]
     fn safe_prefetch(self, x: &[u8]) -> Result<(), InvalidPrefetch> {
         let bytes = Bytes::new(x.len());
@@ -215,7 +218,7 @@ mod test {
     #[test]
     fn test_loop() {
         let p = Loop::new();
-        for i in (0..=1024) {
+        for i in 0..=1024 {
             let v = vec![0u8; i];
 
             let checked = Checked::new(p, Bytes::new(v.len())).unwrap();
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index e20bc6463f..0eba46526d 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -576,7 +576,7 @@ impl workingset::View<u32> for &PruneAccessor<'_> {
         Self: 'a;
 
     fn get(&self, id: u32) -> Option<ElementRef> {
-        self.keys.get(&id)?.map(|v| ElementRef(v))
+        self.keys.get(&id)?.map(ElementRef)
     }
 }
 
@@ -800,7 +800,8 @@ mod tests {
             MaxDegree::new(degree),
             Metric::L2,
             Matrix::row_vector(start.into()),
-        );
+        )
+        .unwrap();
 
         // let full = Full::<f32>::new(grid.dim().into(), Metric::L2);
 

From dcd70c5cd828a292125d123fecf2517c2a992aab Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Mon, 24 Aug 2026 15:24:22 -0700
Subject: [PATCH 27/34] Cleanups.

---
 diskann-benchmark/src/index/inmem2.rs | 5 ++---
 diskann-vector/src/unaligned.rs       | 3 ---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs
index f8013ec767..fca63991c2 100644
--- a/diskann-benchmark/src/index/inmem2.rs
+++ b/diskann-benchmark/src/index/inmem2.rs
@@ -480,7 +480,7 @@ where
             MaxDegree::new(input.build.config.max_degree().get()),
             input.data.distance,
             start,
-        );
+        )?;
 
         let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(
@@ -843,7 +843,6 @@ where
         let queries: Arc<Matrix<T>> = Arc::new(datafiles::load_dataset(datafiles::BinFile(
             &input.search.queries,
         ))?);
-        let dim = dataset.ncols();
 
         // Compute the medoid of the dataset as the single start point.
         let start = StartPointStrategy::Medoid.compute(dataset.as_view())?;
@@ -853,7 +852,7 @@ where
             MaxDegree::new(index_config.max_degree().get()),
             input.data.distance,
             start,
-        );
+        )?;
 
         let provider = Provider::<_, u32>::new(config)?;
         let index = Arc::new(DiskANNIndex::new(index_config, provider, None));
diff --git a/diskann-vector/src/unaligned.rs b/diskann-vector/src/unaligned.rs
index fb47c8bcbc..9a65211870 100644
--- a/diskann-vector/src/unaligned.rs
+++ b/diskann-vector/src/unaligned.rs
@@ -89,9 +89,6 @@ impl<'a, T, const N: usize> From<&'a [T; N]> for UnalignedSlice<'a, T> {
     }
 }
 
-unsafe impl<T> Send for UnalignedSlice<'_, T> where T: Sync {}
-unsafe impl<T> Sync for UnalignedSlice<'_, T> where T: Sync {}
-
 /// View `self` as an [`UnalignedSlice`].
 pub trait AsUnaligned {
     /// The element type of the slice.

From bf7903a21e3342a838db0461bb8c1ff2e14f570d Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 25 Aug 2026 12:44:13 -0700
Subject: [PATCH 28/34] Prepare docs.

---
 .github/workflows/ci.yml                   |   6 +-
 diskann-inmem/integration/main.rs          |   2 +
 diskann-inmem/src/integration/counters.rs  |   2 +-
 diskann-inmem/src/integration/store/mod.rs |   5 +-
 diskann-inmem/src/layers/full.rs           |  18 ++--
 diskann-inmem/src/lib.rs                   |   3 +-
 diskann-inmem/src/num.rs                   |  10 +-
 diskann-inmem/src/store/invasive.rs        |   8 +-
 diskann-inmem/src/store/mod.rs             | 106 ++++++++++++---------
 diskann-inmem/src/store/plugin.rs          |   3 +-
 diskann-utils/src/views.rs                 |  13 ---
 diskann/src/error/ann_error.rs             |   1 -
 12 files changed, 92 insertions(+), 85 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c874e1b361..85e5808431 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -231,7 +231,11 @@ jobs:
         run: rustup show
       - uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
       - name: "doc --workspace --no-deps"
-        run: cargo doc --locked --workspace --no-deps --features linalg,flatbuffers,experimental_diversity_search
+        run: cargo doc --locked --workspace --no-deps  --features linalg,flatbuffers,experimental_diversity_search,integration-test
+        env:
+          RUSTDOCFLAGS: -D rustdoc::all
+      - name: "diskann-inmem private docs"
+        run: cargo doc --locked --workspace --no-deps --all-features --document-private-items
         env:
           RUSTDOCFLAGS: -D rustdoc::all
 
diff --git a/diskann-inmem/integration/main.rs b/diskann-inmem/integration/main.rs
index 9a9d0be453..af16d69f17 100644
--- a/diskann-inmem/integration/main.rs
+++ b/diskann-inmem/integration/main.rs
@@ -3,6 +3,8 @@
  * Licensed under the MIT license.
  */
 
+//! Integration test runner for [`diskann_inmem`].
+
 mod index;
 mod store;
 mod support;
diff --git a/diskann-inmem/src/integration/counters.rs b/diskann-inmem/src/integration/counters.rs
index fcb16d6cde..b6af15ef4b 100644
--- a/diskann-inmem/src/integration/counters.rs
+++ b/diskann-inmem/src/integration/counters.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-/// A snapshot of global [`Counters`](crate::counters::Counters).
+/// A snapshot of global counters.
 #[derive(Debug, Clone)]
 #[non_exhaustive]
 pub struct CounterSnapshot {
diff --git a/diskann-inmem/src/integration/store/mod.rs b/diskann-inmem/src/integration/store/mod.rs
index 2254e897ee..ae2ec9b1b8 100644
--- a/diskann-inmem/src/integration/store/mod.rs
+++ b/diskann-inmem/src/integration/store/mod.rs
@@ -5,13 +5,12 @@
 
 //! This module exposes "public" integration-test wrappers for the various internal store
 //! mechanisms to drive larger concurrency tests.
-//!
-//! These implementationa have a similar structure. A [`boilerplate`] macro is used to ensure
-//! the capabilities exposed are mostly the same.
 
 pub mod checked;
 pub mod invasive;
 
+/// These implementationa have a similar structure. A [`boilerplate`] macro is used to ensure
+/// the capabilities exposed are mostly the same.
 macro_rules! boilerplate {
     (
         $plugin:ty => $store:ident,
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index f1c57c3f88..129ad6ee77 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -163,6 +163,14 @@ impl<T> Config<T> {
     pub fn dim(&self) -> usize {
         self.start_points.ncols()
     }
+
+    /// Construct a [`Full`] from the [`Config`].
+    pub fn build(self) -> ANNResult<Full<T>>
+    where
+        T: FullPrecision,
+    {
+        Full::new(self)
+    }
 }
 
 /// Errors that can arise when constructing [`Config`].
@@ -187,7 +195,7 @@ where
     type Layer = Full<T>;
 
     fn build(self) -> ANNResult<Full<T>> {
-        Full::new(self)
+        <Config<T>>::build(self)
     }
 }
 
@@ -826,12 +834,8 @@ macro_rules! expand_beam {
 }
 
 macro_rules! prune {
-    ($self:ty, $reader:ident, $f:ident) => {{
-        Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed()
-    }};
-    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{
-        Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed()
-    }};
+    ($self:ty, $reader:ident, $f:ident) => {{ Prune::<$self, _>::new($reader, Pure::<$f>::new()).boxed() }};
+    ($self:ty, $reader:ident, { $N:literal, $f:ident }) => {{ Prune::<$self, _>::new($reader, Pure::<Specialize<$N, $f>>::new()).boxed() }};
 }
 
 impl FullPrecisionImpl for f32 {
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 9ac33674ce..4680dac23e 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -18,10 +18,9 @@ mod neighbors;
 mod prefetch;
 mod tag;
 
-mod store;
-
 pub mod layers;
 pub mod provider;
+pub mod store;
 
 pub use provider::{Context, Provider, Strategy};
 
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index cefe1acb25..ce1ffeeb5b 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -193,9 +193,9 @@ macro_rules! typed_int {
 // TODO: Provide a linkable reference for "immutable" points.
 
 typed_int!(
-    /// The number of distinct slots a [`crate::Provider`] or [`crate::Layer`] has capacity
-    /// for. This is logically distinct from [`MaximumId`], which may be greater due too
-    /// immutable points within a storage container.
+    /// The number of distinct slots a [`crate::Provider`] or [`crate::layers::Layer`] has
+    /// capacity for. This is logically distinct from [`IdLimit`], which may be greater due
+    /// too immutable points within a storage container.
     pub Capacity,
     usize,
 );
@@ -207,8 +207,8 @@ typed_int!(
 );
 
 typed_int!(
-    /// One larger than the maximum ID that a [`crate::Provider`], [`crate::Layer`], or
-    /// other such store in this crate can access in-bounds.
+    /// One larger than the maximum ID that a [`crate::Provider`], [`crate::layers::Layer`],
+    /// or other such store in this crate can access in-bounds.
     ///
     /// This implies that access to ids `[0..self)` are in-bounds.
     ///
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index d6e728f63d..7d3508641b 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -23,10 +23,10 @@
 //! prevents the slot from transitioning from "retiring" back to "available" and being
 //! reused while the guard remains active.
 //!
-//! The transitions [`Slot::publish`] and [`Slot::freeze`] use release stores. Since these
-//! are terminal slot operations, their release stores occur after all payload writes. The
-//! acquire load in [`Reader::read`] makes those writes visible before creating a shared
-//! slice.
+//! The transitions [`plugin::Slot::publish`] and [`plugin::Slot::freeze`] use release stores.
+//! Since these are terminal slot operations, their release stores occur after all payload
+//! writes. The acquire load in [`Reader::read`] makes those writes visible before creating
+//! a shared slice.
 //!
 //! ## Safety
 //!
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index 5634b8e887..b5536d1fd5 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -3,53 +3,65 @@
  * Licensed under the MIT license.
  */
 
-//! A concurrent in-memory data store for driving [`plugin::Plugin`]s.
-//!
-//! This supports concurrent data access, deletes, and inserts through a safe interface.
-//! Data is stored internally in slots indexed from `[0..N)` with `K` points reserved at the
-//! end at positions `[N..N+K)`.
-//!
-//! ## Reading
-//!
-//! A [`Store`] provides no direct way of reading data. Instead, the [`plugin::Plugin`] is
-//! responsible for exposing an appropriate reader (e.g., [`invasive::Invasive::reader`]) in
-//! accordance with its lifecycle implementation. [`Store::guard`] can be used for
-//! this purpose by acquiring an [`epoch::Guard`] for a [`Store`].
-//!
-//! ## Writing
-//!
-//! [`Store::acquire`] is used to find and claim an unused internal [`Slot`]. A [`Slot`]
-//! provides write access to its corresponding [`Slot::data`]. Either [`Slot::publish`] or
-//! [`Slot::freeze`] can be used to make data readable.
-//!
-//! If a [`Slot`] is dropped, its corresponding slot is returned to the [`Store`] without
-//! publishing its contents.
-//!
-//! The index of the slot chosen may be obtained via [`Slot::slot`].
-//!
-//! ## Deleting
-//!
-//! Data is deleted via [`Store::retire`]. This immediately marks the corresponding slot as
-//! unavailable for future readers. However, the retired slot will not be reused until the
-//! [`Store`] can guarantee that no readers that could be using the data are active.
-//!
-//! Slots are automatically reclaimed as part of slot acquisition in the "writing" phase.
-//!
-//! ## Neighbor Access
-//!
-//! The [`Store`] also contains a [`Neighbors`] instance to store adjacency lists. Since
-//! neighbors are generally accessed less frequently than data with a higher volume of write
-//! traffic, fine-grained locks are used for this data structure.
-//!
-//! # Details
-//!
-//! This uses an implementation of the epoch-based reclamation (EBR) provided by [`Registry`].
-//! Plugins follow the lifecycle process defined in the [plugin module docs](plugin).
-//!
-//! The EBR scheme allows readers to safely access data while only generating read traffic to
-//! the CPU caches. The cost is that there is a delay between when slots are retired and when
-//! they can be reused, with a long lived reader blocking this reclamation. As such, users of
-//! this data structure should ensure that readers are reasonably short lived.
+//! Concurrency configuration.
+
+mod internal_docs {
+    //! A concurrent in-memory data store for driving [`plugin::Plugin`]s.
+    //!
+    //! This supports concurrent data access, deletes, and inserts through a safe interface.
+    //! Data is stored internally in slots indexed from `[0..N)` with `K` points reserved at the
+    //! end at positions `[N..N+K)`.
+    //!
+    //! ## Reading
+    //!
+    //! A [`Store`] provides no direct way of reading data. Instead, the [`plugin::Plugin`] is
+    //! responsible for exposing an appropriate reader (e.g., [`invasive::Invasive::reader`]) in
+    //! accordance with its lifecycle implementation. [`Store::guard`] can be used for
+    //! this purpose by acquiring an [`epoch::Guard`] for a [`Store`].
+    //!
+    //! ## Writing
+    //!
+    //! [`Store::acquire`] is used to find and claim an unused internal [`Slot`]. A [`Slot`]
+    //! provides write access to its corresponding [`Slot::data`]. Either [`Slot::publish`]
+    //! or [`Slot::freeze`] can be used to make data readable.
+    //!
+    //! If a [`Slot`] is dropped, its corresponding slot is returned to the [`Store`] without
+    //! publishing its contents.
+    //!
+    //! The index of the slot chosen may be obtained via [`Slot::slot`].
+    //!
+    //! ## Deleting
+    //!
+    //! Data is deleted via [`Store::retire`]. This immediately marks the corresponding slot
+    //! as unavailable for future readers. However, the retired slot will not be reused until
+    //! the [`Store`] can guarantee that no readers that could be using the data are active.
+    //!
+    //! Slots are automatically reclaimed as part of slot acquisition in the "writing" phase.
+    //!
+    //! ## Neighbor Access
+    //!
+    //! The [`Store`] also contains a [`Neighbors`] instance to store adjacency lists. Since
+    //! neighbors are generally accessed less frequently than data with a higher volume of
+    //! write traffic, fine-grained locks are used for this data structure.
+    //!
+    //! # Details
+    //!
+    //! This uses an implementation of the epoch-based reclamation (EBR) provided by
+    //! [`Registry`]. Plugins follow the lifecycle process defined in the
+    //! [plugin module docs](plugin).
+    //!
+    //! The EBR scheme allows readers to safely access data while only generating read traffic
+    //! to the CPU caches. The cost is that there is a delay between when slots are retired
+    //! and when they can be reused, with a long lived reader blocking this reclamation. As
+    //! such, users of this data structure should ensure that readers are reasonably short
+    //! lived.
+
+    #[expect(
+        unused_imports,
+        reason = "this keeps cross-references nicer for internal docs"
+    )]
+    use super::*;
+}
 
 use std::{
     iter::repeat_n,
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/plugin.rs
index ac3977026f..1d49094f43 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/plugin.rs
@@ -50,7 +50,8 @@
 //!
 //! * [`Slot`]: Slots are a little spooky. Plugins can assume that a [`Slot`] for an index
 //!   `i` is exclusive for its duration. This means that [`Slot`] implementations can lend
-//!   out mutable references to its contents (for example, [`invasive::Slot::as_mut_slice`]).
+//!   out mutable references to its contents (for example,
+//!   [`super::invasive::Slot::as_mut_slice`]).
 //!
 //!   Code in [`super`] is very careful to maintain this invariant and all users of [`Slot`]
 //!   must carefully maintain this as well.
diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs
index 2d0cd49dbc..a9352918c9 100644
--- a/diskann-utils/src/views.rs
+++ b/diskann-utils/src/views.rs
@@ -637,19 +637,6 @@ where
             ncols: self.ncols,
         }
     }
-
-    pub fn as_bytes(&self) -> MatrixView<'_, u8>
-    where
-        T::Elem: bytemuck::Pod,
-    {
-        let data = bytemuck::must_cast_slice::<T::Elem, u8>(self.as_slice());
-
-        MatrixView {
-            data,
-            nrows: self.nrows(),
-            ncols: self.ncols() * std::mem::size_of::<T::Elem>(),
-        }
-    }
 }
 
 /// Represents an owning, 2-dimensional view of a contiguous block of memory,
diff --git a/diskann/src/error/ann_error.rs b/diskann/src/error/ann_error.rs
index be12669ef9..bfd211281d 100644
--- a/diskann/src/error/ann_error.rs
+++ b/diskann/src/error/ann_error.rs
@@ -242,7 +242,6 @@ macro_rules! convert_error {
     ($T:ty) => {
         impl From<$T> for $crate::ANNError {
             #[track_caller]
-            #[inline]
             fn from(e: $T) -> $crate::ANNError {
                 $crate::ANNError::new(e)
             }

From a89583b47972901c044318a5984bb5311cda7a13 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 25 Aug 2026 12:55:42 -0700
Subject: [PATCH 29/34] Just build private docs for diskann-inmem.

---
 .github/workflows/ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 85e5808431..932a3bc3c3 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -235,7 +235,7 @@ jobs:
         env:
           RUSTDOCFLAGS: -D rustdoc::all
       - name: "diskann-inmem private docs"
-        run: cargo doc --locked --workspace --no-deps --all-features --document-private-items
+        run: cargo doc --locked --package diskann-inmem --no-deps --all-features --document-private-items
         env:
           RUSTDOCFLAGS: -D rustdoc::all
 

From d35a982901cae8c0b411adcd952731d77ab206a3 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 25 Aug 2026 12:59:56 -0700
Subject: [PATCH 30/34] Shuffle deps.

---
 diskann-inmem/Cargo.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/diskann-inmem/Cargo.toml b/diskann-inmem/Cargo.toml
index 1277ff0e0c..72dec19bf5 100644
--- a/diskann-inmem/Cargo.toml
+++ b/diskann-inmem/Cargo.toml
@@ -16,8 +16,9 @@ diskann-utils = { workspace = true, default-features = false }
 diskann-vector = { workspace = true }
 diskann-wide = { workspace = true }
 parking_lot = "0.12.5"
-thiserror = { workspace = true }
 half = { workspace = true }
+hashbrown = { workspace = true }
+thiserror = { workspace = true }
 
 # Integration Test Dependencies
 diskann-benchmark-runner = { workspace = true, optional = true, features = ["ux-tools"] }
@@ -27,7 +28,6 @@ anyhow = { workspace = true, optional = true }
 rand = { workspace = true, optional = true }
 diskann-benchmark-core = { workspace = true, optional = true }
 tokio = { workspace = true, optional = true }
-hashbrown.workspace = true
 
 [lints.clippy]
 undocumented_unsafe_blocks = "warn"

From 5cba726993b32243498fded6906eb655801fd61a Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 25 Aug 2026 14:48:53 -0700
Subject: [PATCH 31/34] Fix `Checked::check`.

---
 diskann-inmem/src/prefetch.rs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/diskann-inmem/src/prefetch.rs b/diskann-inmem/src/prefetch.rs
index 8e0a8ba040..c6997bab0f 100644
--- a/diskann-inmem/src/prefetch.rs
+++ b/diskann-inmem/src/prefetch.rs
@@ -38,7 +38,6 @@ where
     }
 
     /// Check if `self` can prefetch slices of length `len`.
-    #[cfg(debug_assertions)]
     pub(crate) fn check(self, len: Bytes) -> Result<(), InvalidPrefetch> {
         self.0.check(len)
     }

From e6d840091e716608c9beeab66754dcc108e60b8c Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Tue, 25 Aug 2026 18:10:28 -0700
Subject: [PATCH 32/34] Cleanup docs.

---
 diskann-inmem/src/integration/store/checked.rs  |  6 +-----
 diskann-inmem/src/integration/store/invasive.rs |  8 ++------
 diskann-inmem/src/integration/store/mod.rs      |  2 +-
 diskann-inmem/src/num.rs                        | 14 +-------------
 diskann-inmem/src/provider.rs                   |  2 --
 5 files changed, 5 insertions(+), 27 deletions(-)

diff --git a/diskann-inmem/src/integration/store/checked.rs b/diskann-inmem/src/integration/store/checked.rs
index 1aaa7e6fc2..f536631f30 100644
--- a/diskann-inmem/src/integration/store/checked.rs
+++ b/diskann-inmem/src/integration/store/checked.rs
@@ -32,11 +32,7 @@ boilerplate!(
 );
 
 impl Store {
-    /// Construct a store with `config.capacity` writable slots.
-    ///
-    /// A single zeroed frozen point is created internally to satisfy the underlying
-    /// store's requirement of at least one frozen entry; it occupies the highest slot
-    /// index and is always readable.
+    /// Construct a store with `config.capacity` writable slots and no frozen points.
     ///
     /// # Panics
     ///
diff --git a/diskann-inmem/src/integration/store/invasive.rs b/diskann-inmem/src/integration/store/invasive.rs
index b21300dda9..e71703e473 100644
--- a/diskann-inmem/src/integration/store/invasive.rs
+++ b/diskann-inmem/src/integration/store/invasive.rs
@@ -33,11 +33,7 @@ boilerplate!(
 
 impl Store {
     /// Construct a store with `config.capacity` writable slots, each holding
-    /// `config.entry_bytes` bytes.
-    ///
-    /// A single zeroed frozen point is created internally to satisfy the underlying
-    /// store's requirement of at least one frozen entry; it occupies the highest slot
-    /// index and is always readable.
+    /// `config.entry_bytes` bytes. No frozen points are created.
     ///
     /// # Panics
     ///
@@ -45,7 +41,7 @@ impl Store {
     /// the frozen point exceeds `u32::MAX`) or if other configuration parameters such as
     /// the number of epoch guard slots are invalid (e.g. zero).
     pub fn new(config: Config) -> Self {
-        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 1);
+        let store_layout = store::Layout::new(Capacity::new(config.capacity), MaxDegree::new(0), 0);
 
         let store_config = store::Config::__exhaustive(
             NonZeroUsize::new(config.epoch_guard_slots)
diff --git a/diskann-inmem/src/integration/store/mod.rs b/diskann-inmem/src/integration/store/mod.rs
index ae2ec9b1b8..6d53dfeebf 100644
--- a/diskann-inmem/src/integration/store/mod.rs
+++ b/diskann-inmem/src/integration/store/mod.rs
@@ -9,7 +9,7 @@
 pub mod checked;
 pub mod invasive;
 
-/// These implementationa have a similar structure. A [`boilerplate`] macro is used to ensure
+/// These implementations have a similar structure. A [`boilerplate`] macro is used to ensure
 /// the capabilities exposed are mostly the same.
 macro_rules! boilerplate {
     (
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index ce1ffeeb5b..5d05649225 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -69,12 +69,6 @@ impl Bytes {
         Self::new(std::mem::size_of::<T>())
     }
 
-    /// Return the number of bytes occupied by the slice.
-    #[inline]
-    pub const fn of_slice<T>(x: &[T]) -> Self {
-        Self(std::mem::size_of_val::<[T]>(x))
-    }
-
     /// Return `true` if `self` is zero.
     pub const fn is_zero(self) -> bool {
         self.0 == 0
@@ -195,7 +189,7 @@ macro_rules! typed_int {
 typed_int!(
     /// The number of distinct slots a [`crate::Provider`] or [`crate::layers::Layer`] has
     /// capacity for. This is logically distinct from [`IdLimit`], which may be greater due
-    /// too immutable points within a storage container.
+    /// to immutable points within a storage container.
     pub Capacity,
     usize,
 );
@@ -237,12 +231,6 @@ impl IdLimit {
     }
 }
 
-// typed_int!(
-//     /// Temp
-//     pub(crate) CacheLines,
-//     NonZeroUsize,
-// )
-
 ///////////
 // Tests //
 ///////////
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index 0eba46526d..dacf3027ae 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -205,8 +205,6 @@ where
         _context: &Context,
         id: u32,
     ) -> ANNResult<diskann::provider::ElementStatus> {
-        // Not that this check is approximate. A full check requires materialization of
-        // a `reader`.
         match <L as layers::Layer>::is_readable(&self.layer, id) {
             Some(true) => Ok(diskann::provider::ElementStatus::Valid),
             Some(false) => Ok(diskann::provider::ElementStatus::Deleted),

From 4e76e9a68e6c00377894455c64a5b8b64d3c88e4 Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 27 Aug 2026 16:07:54 -0700
Subject: [PATCH 33/34] "Plugin" -> "Slots"

---
 diskann-inmem/integration/store/mod.rs        |   4 +-
 .../src/integration/store/checked.rs          |   4 +-
 .../src/integration/store/invasive.rs         |   4 +-
 diskann-inmem/src/integration/store/mod.rs    |   6 +-
 diskann-inmem/src/layers/full.rs              |   4 +-
 diskann-inmem/src/store/checked.rs            |  28 ++---
 diskann-inmem/src/store/invasive.rs           |  34 +++---
 diskann-inmem/src/store/mod.rs                | 112 +++++++++---------
 .../src/store/{plugin.rs => slots.rs}         |  44 +++----
 9 files changed, 118 insertions(+), 122 deletions(-)
 rename diskann-inmem/src/store/{plugin.rs => slots.rs} (78%)

diff --git a/diskann-inmem/integration/store/mod.rs b/diskann-inmem/integration/store/mod.rs
index 799e106056..8e6991f709 100644
--- a/diskann-inmem/integration/store/mod.rs
+++ b/diskann-inmem/integration/store/mod.rs
@@ -12,7 +12,7 @@
 //! 2. A readable value is stable for the lifetime of a single reader guard.
 //! 3. A slot never resurrects (`readable -> unreadable -> readable`) within one guard.
 //!
-//! This module exposes shared functionality that is instantiated by different plugins.
+//! This module exposes shared functionality that is instantiated by different implementations.
 
 #![expect(
     clippy::unwrap_used,
@@ -31,8 +31,6 @@ use diskann_benchmark_runner::{Registry, RegistryError, utils::fmt::KeyValue};
 use rand::{Rng, SeedableRng, distr::Uniform, rngs::StdRng};
 use serde::{Deserialize, Serialize};
 
-// use diskann_inmem::integration::store::invasive::{self, Config, Store};
-
 /// Number of slots a reader inspects per guard. Kept small so guards are short-lived,
 /// allowing the epoch to advance and reclamation to make progress.
 const READER_WINDOW: usize = 64;
diff --git a/diskann-inmem/src/integration/store/checked.rs b/diskann-inmem/src/integration/store/checked.rs
index f536631f30..8aee490b9b 100644
--- a/diskann-inmem/src/integration/store/checked.rs
+++ b/diskann-inmem/src/integration/store/checked.rs
@@ -54,8 +54,8 @@ impl Store {
             .expect("`freelist_recycle_capacity` must be non-zero"),
         );
 
-        let plugin_config = checked::Checked::config();
-        let store = store::Store::new(store_layout, store_config, plugin_config)
+        let slots_config = checked::Checked::config();
+        let store = store::Store::new(store_layout, store_config, slots_config)
             .expect("failed to construct store");
 
         Self { store }
diff --git a/diskann-inmem/src/integration/store/invasive.rs b/diskann-inmem/src/integration/store/invasive.rs
index e71703e473..274e80b77f 100644
--- a/diskann-inmem/src/integration/store/invasive.rs
+++ b/diskann-inmem/src/integration/store/invasive.rs
@@ -55,8 +55,8 @@ impl Store {
             .expect("`freelist_recycle_capacity` must be non-zero"),
         );
 
-        let plugin_config = store::invasive::Invasive::config(Bytes::new(config.entry_bytes));
-        let store = store::Store::new(store_layout, store_config, plugin_config)
+        let slots_config = store::invasive::Invasive::config(Bytes::new(config.entry_bytes));
+        let store = store::Store::new(store_layout, store_config, slots_config)
             .expect("failed to construct store");
 
         Self { store }
diff --git a/diskann-inmem/src/integration/store/mod.rs b/diskann-inmem/src/integration/store/mod.rs
index 6d53dfeebf..ef7b32cf52 100644
--- a/diskann-inmem/src/integration/store/mod.rs
+++ b/diskann-inmem/src/integration/store/mod.rs
@@ -13,14 +13,14 @@ pub mod invasive;
 /// the capabilities exposed are mostly the same.
 macro_rules! boilerplate {
     (
-        $plugin:ty => $store:ident,
+        $slots:ty => $store:ident,
         for<$read_lt:lifetime> $read:ty => $reader:ident,
         for<$slot_lt:lifetime> $slot:ty => $writer:ident,
     ) => {
         /// A test store wraper.
         #[derive(Debug)]
         pub struct $store {
-            store: $crate::store::Store<$plugin>,
+            store: $crate::store::Store<$slots>,
         }
 
         impl $store {
@@ -55,7 +55,7 @@ macro_rules! boilerplate {
             /// Attain a reader into the store. Returns `None` if all epoch guard slots
             /// are used.
             pub fn reader(&self) -> Option<$reader<'_>> {
-                match <$plugin>::reader(&self.store) {
+                match <$slots>::reader(&self.store) {
                     Ok(reader) => Some($reader::new(reader)),
                     Err($crate::epoch::Unavailable) => None,
                 }
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/layers/full.rs
index 129ad6ee77..91b3a3e440 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/layers/full.rs
@@ -292,12 +292,12 @@ where
 
     /// Return the number of payload bytes in each stored vector.
     pub fn bytes(&self) -> Bytes {
-        self.store.plugin().bytes()
+        self.store.slots().bytes()
     }
 
     #[cfg(test)]
     fn bytes_plus_tag(&self) -> Bytes {
-        self.store.plugin().bytes_plus_tag()
+        self.store.slots().bytes_plus_tag()
     }
 
     /// Return the [`Metric`] for this layer.
diff --git a/diskann-inmem/src/store/checked.rs b/diskann-inmem/src/store/checked.rs
index 607cbd669d..9f7eea836c 100644
--- a/diskann-inmem/src/store/checked.rs
+++ b/diskann-inmem/src/store/checked.rs
@@ -3,9 +3,9 @@
  * Licensed under the MIT license.
  */
 
-//! # A [`Store`] plugin for pedantically testing the EBR protocol
+//! # Pedantically testing the EBR protocol
 //!
-//! The goal here is to detect violations of the state machine outlined in [`plugin`].
+//! The goal here is to detect violations of the state machine outlined in [`slots`].
 //! This is accomplished by using a [`RwLock`] to protect internal [`State`] with guards
 //! **only** acquired via [`RwLock::try_read`] and [`RwLock::try_write`]. A correct EBR
 //! protocol should ensure that:
@@ -31,12 +31,12 @@
 //!
 //! ## Lifecycle Details
 //!
-//! The plugin lifecycle is carefully designed to allow readers of the plugin to avoid any
+//! Lifecycle transitions are carefully designed to allow readers of the slots to avoid any
 //! accesses to the authoritative [`Store`] for read-only operations. The [`Checked`] test
 //! code follows this pattern, but this does introduce a subtle detail that is worth
 //! highlighting. [`Reader::read`] needs to be able to check a slot for readability
 //! **without** trying to acquire a [`RwLockReadGuard`] for that slot. Doing so even briefly
-//! will cause a [`RwLock::try_write`] on an otherwise correct [`plugin::Plugin`] state
+//! will cause a [`RwLock::try_write`] on an otherwise correct [`slots::Slots`] state
 //! transition to fail.
 //!
 //! To circumvent this, an additional [`AtomicBool`] is bundled with [`Entry`] to broadcast
@@ -82,7 +82,7 @@ use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use crate::{epoch, num::IdLimit, store::Store};
 
-use super::{Lifecycle, plugin};
+use super::{Lifecycle, slots};
 
 /// The state of a slot.
 #[derive(Debug, Default)]
@@ -99,7 +99,7 @@ enum State {
 
 /// A slot entry. See the [module level docs](self) for a discussion on the contents of this
 /// struct. The table below describes how the combination of fields in this struct maps to
-/// the plugin lifecycle states.
+/// the lifecycle states.
 /// ```text
 /// +-----------------+----------+-----------+------------------------+
 /// | Lifecycle state | readable | State     | Lock Expectation       |
@@ -285,7 +285,7 @@ impl WriteEntry<'_> {
     }
 }
 
-/// A [`plugin::PluginConfig`] for [`Checked`].
+/// A [`slots::SlotsConfig`] for [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Config(());
 
@@ -295,8 +295,8 @@ impl Config {
     }
 }
 
-impl plugin::PluginConfig for Config {
-    type Plugin = Checked;
+impl slots::SlotsConfig for Config {
+    type Slots = Checked;
     type Error = diskann::error::Infallible;
 
     fn build(self, id_limit: IdLimit) -> Result<Checked, diskann::error::Infallible> {
@@ -304,7 +304,7 @@ impl plugin::PluginConfig for Config {
     }
 }
 
-/// A correctness checking [`plugin::Plugin`]. See the [module level docs](self) for details.
+/// A correctness checking [`slots::Slots`]. See the [module level docs](self) for details.
 #[derive(Debug)]
 pub(crate) struct Checked {
     entries: Vec<Entry>,
@@ -320,7 +320,7 @@ impl Checked {
         }
     }
 
-    /// Return the [`plugin::PluginConfig`] for [`Self`].
+    /// Return the [`slots::SlotsConfig`] for [`Self`].
     pub(crate) fn config() -> Config {
         Config::new()
     }
@@ -381,7 +381,7 @@ impl Reader<'_> {
     }
 }
 
-impl plugin::Plugin for Checked {
+impl slots::Slots for Checked {
     type Slot<'a> = Slot<'a>;
 
     fn id_limit(&self) -> IdLimit {
@@ -401,7 +401,7 @@ impl plugin::Plugin for Checked {
     }
 }
 
-/// A writable [`plugin::Slot`] for [`Checked`].
+/// A writable [`slots::Slot`] for [`Checked`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
     entry: WriteEntry<'a>,
@@ -419,7 +419,7 @@ impl<'a> Slot<'a> {
     }
 }
 
-impl plugin::Slot for Slot<'_> {
+impl slots::Slot for Slot<'_> {
     fn publish(self, _: Lifecycle) {
         let value = self.value.expect("`value` was not set");
         self.entry.publish(value);
diff --git a/diskann-inmem/src/store/invasive.rs b/diskann-inmem/src/store/invasive.rs
index 7d3508641b..03d6902a6f 100644
--- a/diskann-inmem/src/store/invasive.rs
+++ b/diskann-inmem/src/store/invasive.rs
@@ -3,15 +3,15 @@
  * Licensed under the MIT license.
  */
 
-//! A store [`plugin::Plugin`] that maintains an invasive slot state where the data in each
-//! slot is a contiguous slice of memory.
+//! A [`slots::Slots`] that maintains an invasive slot state where the data in each slot is
+//! a contiguous slice of memory.
 //!
 //! Slot state is stored as an [`AtomicTag`] immediately after the slot data.
 //!
 //! ## Lifecycle Details
 //!
-//! The plugin lifecycle details are relatively straightforward. The invasive [`AtomicTag`]
-//! mostly follows the transitions made by the [`Store`]. A [`Reader`] checks the tag for
+//! Lifecycle details are relatively straightforward. The invasive [`AtomicTag`] mostly
+//! follows the transitions made by the [`Store`]. A [`Reader`] checks the tag for
 //! readability before creating a shared reference to the data payload.
 //!
 //! The problematic transition from "published" to "retiring" is made safe because
@@ -23,7 +23,7 @@
 //! prevents the slot from transitioning from "retiring" back to "available" and being
 //! reused while the guard remains active.
 //!
-//! The transitions [`plugin::Slot::publish`] and [`plugin::Slot::freeze`] use release stores.
+//! The transitions [`slots::Slot::publish`] and [`slots::Slot::freeze`] use release stores.
 //! Since these are terminal slot operations, their release stores occur after all payload
 //! writes. The acquire load in [`Reader::read`] makes those writes visible before creating
 //! a shared slice.
@@ -31,8 +31,8 @@
 //! ## Safety
 //!
 //! The safety of this module depends on [`Invasive`] being embedded in a [`Store`] that
-//! observes the plugin lifecycle. Every lifecycle operation requires a [`Lifecycle`] token,
-//! which is constructible only by the parent store module. The unsafe [`plugin::Plugin`]
+//! observes the slot lifecycle. Every lifecycle operation requires a [`Lifecycle`] token,
+//! which is constructible only by the parent store module. The unsafe [`slots::Slots`]
 //! methods additionally rely on [`Store`] to satisfy their documented state and exclusivity
 //! preconditions.
 
@@ -45,11 +45,11 @@ use crate::{
     buffer::{Buffer, BufferError, RawSlice},
     epoch,
     num::{Align, Bytes, IdLimit},
-    store::{Lifecycle, Store, plugin},
+    store::{Lifecycle, Store, slots},
     tag::{AtomicTag, Tag},
 };
 
-/// A [`plugin::PluginConfig`] for [`Invasive`].
+/// A [`slots::SlotsConfig`] for [`Invasive`].
 #[derive(Debug, Clone)]
 pub(crate) struct Config {
     /// The number of bytes held in each slot.
@@ -69,8 +69,8 @@ impl Config {
     }
 }
 
-impl plugin::PluginConfig for Config {
-    type Plugin = Invasive;
+impl slots::SlotsConfig for Config {
+    type Slots = Invasive;
     type Error = InvasiveError;
     fn build(self, id_limit: IdLimit) -> Result<Invasive, InvasiveError> {
         <Config>::build(self, id_limit)
@@ -194,7 +194,7 @@ enum InvasiveErrorInner {
     BufferError(BufferError),
 }
 
-impl plugin::Plugin for Invasive {
+impl slots::Slots for Invasive {
     type Slot<'a> = Slot<'a>;
 
     fn id_limit(&self) -> IdLimit {
@@ -208,7 +208,7 @@ impl plugin::Plugin for Invasive {
         };
 
         // This is a pessimistic check to ensure that the caller is correctly using the
-        // `plugin` API.
+        // `slots` API.
         debug_assert_eq!(
             tag.load(Ordering::Relaxed),
             Tag::AVAILABLE,
@@ -381,7 +381,7 @@ impl<'a> Reader<'a> {
     }
 }
 
-/// A [`plugin::Slot`] for [`Invasive`].
+/// A [`slots::Slot`] for [`Invasive`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a> {
     // NOTE: `tag` and `data` must belong to the same slot.
@@ -395,15 +395,15 @@ impl<'a> Slot<'a> {
     /// The length of this slice is guaranteed to be the number of bytes passed to
     /// [`Invasive::new`] or [`Config::new`].
     pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
-        // SAFETY: Users of the `plugin::Slot` are obligated to ensure exclusivity.
+        // SAFETY: Users of the `slots::Slot` are obligated to ensure exclusivity.
         //
-        // Since `Reader` obeys the plugin life-cycle requirements, a concurrent reader
+        // Since `Reader` obeys the slots life-cycle requirements, a concurrent reader
         // of this data should not be possible.
         unsafe { self.data.as_mut_slice() }
     }
 }
 
-impl plugin::Slot for Slot<'_> {
+impl slots::Slot for Slot<'_> {
     fn publish(self, _: Lifecycle) {
         self.tag.store(Tag::PUBLISHED, Ordering::Release);
     }
diff --git a/diskann-inmem/src/store/mod.rs b/diskann-inmem/src/store/mod.rs
index b5536d1fd5..4d9825b598 100644
--- a/diskann-inmem/src/store/mod.rs
+++ b/diskann-inmem/src/store/mod.rs
@@ -6,7 +6,7 @@
 //! Concurrency configuration.
 
 mod internal_docs {
-    //! A concurrent in-memory data store for driving [`plugin::Plugin`]s.
+    //! A concurrent in-memory data store for driving [`slots::Slots`].
     //!
     //! This supports concurrent data access, deletes, and inserts through a safe interface.
     //! Data is stored internally in slots indexed from `[0..N)` with `K` points reserved at the
@@ -14,7 +14,7 @@ mod internal_docs {
     //!
     //! ## Reading
     //!
-    //! A [`Store`] provides no direct way of reading data. Instead, the [`plugin::Plugin`] is
+    //! A [`Store`] provides no direct way of reading data. Instead, the [`slots::Slots`] is
     //! responsible for exposing an appropriate reader (e.g., [`invasive::Invasive::reader`]) in
     //! accordance with its lifecycle implementation. [`Store::guard`] can be used for
     //! this purpose by acquiring an [`epoch::Guard`] for a [`Store`].
@@ -47,8 +47,7 @@ mod internal_docs {
     //! # Details
     //!
     //! This uses an implementation of the epoch-based reclamation (EBR) provided by
-    //! [`Registry`]. Plugins follow the lifecycle process defined in the
-    //! [plugin module docs](plugin).
+    //! [`Registry`]. Slots follow the lifecycle process defined in the [module docs](slots).
     //!
     //! The EBR scheme allows readers to safely access data while only generating read traffic
     //! to the CPU caches. The cost is that there is a delay between when slots are retired
@@ -83,12 +82,12 @@ use crate::{
 };
 
 pub(crate) mod invasive;
-pub(crate) mod plugin;
+pub(crate) mod slots;
 
 #[cfg(any(test, feature = "integration-test"))]
 pub(crate) mod checked;
 
-/// To make extra sure that [`plugin::Plugin`] life-cycle arguments are not callable outside
+/// To make extra sure that [`slots::Slots`] life-cycle arguments are not callable outside
 /// of this module (i.e., elsewhere in this crate), this [`Lifecycle`] marker type is used
 /// that is only constructible in this module.
 #[derive(Debug)]
@@ -98,7 +97,7 @@ impl Lifecycle {
     /// Construct a new [`Lifecycle`].
     ///
     /// DO NOT MAKE THIS `pub(anything)`. It helps prevent accidentally interacting with
-    /// plugins when all uses should be managed in this file instead.
+    /// slots when all uses should be managed in this file instead.
     const fn new() -> Self {
         Self(())
     }
@@ -172,7 +171,7 @@ impl Default for Config {
     }
 }
 
-/// Layout parameters for [`Store`] and the corresponding [`plugin::Plugin`].
+/// Layout parameters for [`Store`] and the corresponding [`slots::Slots`].
 #[derive(Debug, Clone)]
 pub(crate) struct Layout {
     /// The number of non-frozen slots to create space for.
@@ -206,9 +205,9 @@ impl Layout {
 
 /// A concurrent data and graph store.
 #[derive(Debug)]
-pub(crate) struct Store<P> {
-    // The [`plugin::Plugin`] managed by this [`Store`].
-    plugin: P,
+pub(crate) struct Store<T> {
+    // The [`slots::Slots`] managed by this [`Store`].
+    slots: T,
 
     // The number of unfrozen points.
     unfrozen: Capacity,
@@ -229,14 +228,14 @@ pub(crate) struct Store<P> {
 // TODO: This is a guess and probably needs tuning.
 const RETRY_LIMIT: usize = 20;
 
-impl<P> Store<P>
+impl<T> Store<T>
 where
-    P: plugin::Plugin,
+    T: slots::Slots,
 {
     /// Create a new [`Store`].
-    pub(crate) fn new<C>(layout: Layout, config: Config, plugin: C) -> Result<Self, StoreError>
+    pub(crate) fn new<C>(layout: Layout, config: Config, slots: C) -> Result<Self, StoreError>
     where
-        C: plugin::PluginConfig<Plugin = P>,
+        C: slots::SlotsConfig<Slots = T>,
     {
         let Layout {
             capacity,
@@ -266,15 +265,15 @@ where
             .try_into()
             .map_err(|_| StoreError::too_many_neighbors(max_degree))?;
 
-        let plugin = plugin::PluginConfig::build(plugin, id_limit).map_err(StoreError::plugin)?;
+        let slots = slots::SlotsConfig::build(slots, id_limit).map_err(StoreError::slots)?;
 
-        let plugin_id_limit = plugin.id_limit();
-        if plugin_id_limit != id_limit {
-            return Err(StoreError::invalid_construction(plugin_id_limit, id_limit));
+        let slots_id_limit = slots.id_limit();
+        if slots_id_limit != id_limit {
+            return Err(StoreError::invalid_construction(slots_id_limit, id_limit));
         }
 
         let me = Self {
-            plugin,
+            slots,
             unfrozen: capacity,
             tags: repeat_n(Tag::AVAILABLE, id_limit.as_usize())
                 .map(AtomicTag::new)
@@ -290,9 +289,9 @@ where
         Ok(me)
     }
 
-    /// Return the [`plugin::Plugin`] for this store.
-    pub(crate) fn plugin(&self) -> &P {
-        &self.plugin
+    /// Return the [`slots::Slots`] for this store.
+    pub(crate) fn slots(&self) -> &T {
+        &self.slots
     }
 
     /// Return the range of slots containing frozen items in `self`.
@@ -333,7 +332,7 @@ where
                 );
             };
 
-            // We release the plugin before the main tag. The other direction would
+            // We release the slot before the main tag. The other direction would
             // prematurely advertise availability.
             //
             // SAFETY: IDs only get added to the `epoch::Registry` upon retiring, and are
@@ -341,7 +340,7 @@ where
             // was retired have been dropped.
             //
             // Therefore, this slot has no accessors and is ready to be reclaimed.
-            unsafe { plugin::Plugin::reclaim(self.plugin(), i, Lifecycle::new()) };
+            unsafe { slots::Slots::reclaim(self.slots(), i, Lifecycle::new()) };
 
             // Use `Release` ordering to ensure that the store to the mirror cannot get moved
             // after the store to the authoritative list.
@@ -367,21 +366,21 @@ where
     /// Create an [`epoch::Guard`] for the [`epoch::Registry`] within `self` and invoke `f`
     /// with that guard, returning the result.
     ///
-    /// This can be used by [`plugin::Plugin`] readers to establish a verifiable chain of
-    /// custody for an [`epoch::Guard`] over the plugin.
+    /// This can be used by [`slots::Slots`] readers to establish a verifiable chain of
+    /// custody for an [`epoch::Guard`] over the slots.
     pub(crate) fn guard<'a, F, R>(&'a self, f: F) -> Result<R, epoch::Unavailable>
     where
-        F: FnOnce(&'a P, epoch::Guard<'a>) -> R,
+        F: FnOnce(&'a T, epoch::Guard<'a>) -> R,
     {
         let guard = self.registry.guard()?;
-        Ok(f(self.plugin(), guard))
+        Ok(f(self.slots(), guard))
     }
 
     /// Attempt to acquire a new [`Slot`] for writing.
     ///
     /// This method first consults the freelist and falls back to scanning the tags list
     /// if no ID is available from the fast path.
-    pub(crate) fn acquire(&self) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
+    pub(crate) fn acquire(&self) -> Option<Slot<'_, <T as slots::Slots>::Slot<'_>>> {
         for _ in 0..RETRY_LIMIT {
             match self.freelist.pop() {
                 freelist::Id::Found(id) => {
@@ -441,7 +440,7 @@ where
                 // SAFETY: The above compare-exchange ensures that we transitioned the
                 // authoritative state from "published" to "retired" and prevents other
                 // threads from attempting the same transition.
-                unsafe { plugin::Plugin::retire(self.plugin(), i as u32, Lifecycle::new()) };
+                unsafe { slots::Slots::retire(self.slots(), i as u32, Lifecycle::new()) };
                 guard.retire(i as u32);
                 Ok(())
             }
@@ -460,12 +459,12 @@ where
     ///
     /// Periodically, the freelist is checked to see if another thread has found an available
     /// slot for us.
-    fn scan_acquire(&self) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
+    fn scan_acquire(&self) -> Option<Slot<'_, <T as slots::Slots>::Slot<'_>>> {
         // This is potentially quite slow, so scan approximately `1 / RETRY_LIMIT` of the
         // writable range. The outer retry loop provides broader coverage.
         let mut remaining = self.unfrozen.value().div_ceil(RETRY_LIMIT);
         let mut chunks_since_freelist_check = 0;
-        let mut acquired: Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> = None;
+        let mut acquired: Option<Slot<'_, <T as slots::Slots>::Slot<'_>>> = None;
 
         while remaining != 0 {
             let chunk = self.freelist.scan();
@@ -515,7 +514,7 @@ where
     ///
     /// Returns `None` if `i` is not within [`Self::id_limit`] or if the slot is not currently
     /// acquirable.
-    pub(crate) fn slot(&self, i: u32) -> Option<Slot<'_, <P as plugin::Plugin>::Slot<'_>>> {
+    pub(crate) fn slot(&self, i: u32) -> Option<Slot<'_, <T as slots::Slots>::Slot<'_>>> {
         let tag = &self.tags.get(i.into_usize())?;
 
         // SAFETY: We've guaranteed that `tag` belongs to `slot`.
@@ -532,7 +531,7 @@ where
         &'a self,
         tag: &'a AtomicTag,
         slot: u32,
-    ) -> Option<Slot<'a, <P as plugin::Plugin>::Slot<'a>>> {
+    ) -> Option<Slot<'a, <T as slots::Slots>::Slot<'a>>> {
         if tag.load(Ordering::Relaxed) != Tag::AVAILABLE {
             return None;
         }
@@ -547,13 +546,12 @@ where
                 // SAFETY: The above compare-exchange ensures that this slot was previously
                 // "available" and prevents other threads from trying to acquire this slot.
                 // The acquire ordering synchronizes with the release transition to
-                // "available", making plugin reclamation or abort work visible before
-                // `Plugin::acquire`.
+                // "available", making slot reclamation or abort work visible before
+                // `Slots::acquire`.
                 //
                 // The `Slot` data structure ensures that exactly one of the terminal methods
-                // for `plugin::Slot` is called.
-                let data =
-                    unsafe { plugin::Plugin::acquire(self.plugin(), slot, Lifecycle::new()) };
+                // for `slot::Slot` is called.
+                let data = unsafe { slots::Slots::acquire(self.slots(), slot, Lifecycle::new()) };
 
                 Some(Slot {
                     tag,
@@ -568,7 +566,7 @@ where
     /// Return whether or not it is probably okay to read from the slot `i`.
     ///
     /// This check is approximate and non-synchronizing. A full check requires the
-    /// plugin-specific reader.
+    /// slot's specific reader.
     ///
     /// Returns `None` if `i` is not within [`Self::id_limit`].
     pub(crate) fn can_read_approximate(&self, i: usize) -> Option<bool> {
@@ -603,11 +601,11 @@ impl StoreError {
     }
 
     #[track_caller]
-    fn plugin<E>(err: E) -> Self
+    fn slots<E>(err: E) -> Self
     where
         E: std::error::Error + Send + Sync + 'static,
     {
-        Self(StoreErrorInner::PluginError(ANNError::new(err)))
+        Self(StoreErrorInner::SlotsError(ANNError::new(err)))
     }
 
     fn invalid_construction(got: IdLimit, expected: IdLimit) -> Self {
@@ -643,9 +641,9 @@ enum StoreErrorInner {
     BufferError(#[from] BufferError),
     #[error(transparent)]
     NeighborsError(#[from] NeighborsError),
-    #[error("error creating plugin")]
-    PluginError(ANNError),
-    #[error("requested {} but the plugin returned {}", expected, got)]
+    #[error("error creating slots")]
+    SlotsError(ANNError),
+    #[error("requested {} but the slots returned {}", expected, got)]
     InvalidConstruction { got: IdLimit, expected: IdLimit },
 }
 
@@ -670,13 +668,13 @@ diskann::convert_error!(RetireError);
 
 /// A writable buffer into the data managed by a [`Store`], obtained from [`Store::acquire`].
 ///
-/// This is the only safe way to interact with a [`plugin::Slot`] since this ensures that one
+/// This is the only safe way to interact with a [`slots::Slot`] since this ensures that one
 /// of the terminal methods is called. Dropping a [`Slot`] without calling [`Slot::publish`]
-/// or [`Slot::freeze`] automatically invokes [`plugin::Slot::abort`].
+/// or [`Slot::freeze`] automatically invokes [`slots::Slot::abort`].
 #[derive(Debug)]
 pub(crate) struct Slot<'a, S>
 where
-    S: plugin::Slot,
+    S: slots::Slot,
 {
     tag: &'a AtomicTag,
     data: ManuallyDrop<S>,
@@ -685,7 +683,7 @@ where
 
 impl<'a, S> Slot<'a, S>
 where
-    S: plugin::Slot,
+    S: slots::Slot,
 {
     /// View the raw inner slot.
     pub(crate) fn data(&mut self) -> &mut S {
@@ -702,7 +700,7 @@ where
         let mut me = ManuallyDrop::new(self);
 
         // Freeze the inner slot.
-        plugin::Slot::freeze(
+        slots::Slot::freeze(
             // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut me.data) },
             Lifecycle::new(),
@@ -722,7 +720,7 @@ where
         let mut me = ManuallyDrop::new(self);
 
         // Publish the inner slot.
-        plugin::Slot::publish(
+        slots::Slot::publish(
             // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut me.data) },
             Lifecycle::new(),
@@ -736,10 +734,10 @@ where
 
 impl<S> Drop for Slot<'_, S>
 where
-    S: plugin::Slot,
+    S: slots::Slot,
 {
     fn drop(&mut self) {
-        plugin::Slot::abort(
+        slots::Slot::abort(
             // SAFETY: The `ManuallyDrop` `data` is not used after this call.
             unsafe { ManuallyDrop::take(&mut self.data) },
             Lifecycle::new(),
@@ -765,8 +763,8 @@ mod tests {
     #[derive(Debug)]
     struct FaultyConfig;
 
-    impl plugin::PluginConfig for FaultyConfig {
-        type Plugin = Checked;
+    impl slots::SlotsConfig for FaultyConfig {
+        type Slots = Checked;
         type Error = diskann::error::Infallible;
 
         fn build(self, id_limit: IdLimit) -> Result<Checked, diskann::error::Infallible> {
@@ -830,7 +828,7 @@ mod tests {
     }
 
     #[test]
-    fn new_rejects_faulty_plugin() {
+    fn new_rejects_faulty_slots() {
         let err = Store::new(
             Layout::new(Capacity::new(4), MaxDegree::new(10), 0),
             Config::default(),
diff --git a/diskann-inmem/src/store/plugin.rs b/diskann-inmem/src/store/slots.rs
similarity index 78%
rename from diskann-inmem/src/store/plugin.rs
rename to diskann-inmem/src/store/slots.rs
index 1d49094f43..f011c8cec8 100644
--- a/diskann-inmem/src/store/plugin.rs
+++ b/diskann-inmem/src/store/slots.rs
@@ -8,7 +8,7 @@
 //! Please read this section carefully - the protocol is not difficult, but it *is* subtle.
 //!
 //! The transitions are a simplified version of the protocol described in [`crate::tag`]
-//! that storage plugins need to implement to be compatible. A state diagram is shown below:
+//! that storage slots need to implement to be compatible. A state diagram is shown below:
 //!
 //! ```text
 //!         +--------------- `reclaim` ----------------+
@@ -48,7 +48,7 @@
 //!
 //! ## Writable States
 //!
-//! * [`Slot`]: Slots are a little spooky. Plugins can assume that a [`Slot`] for an index
+//! * [`Slot`]: Slots are a little spooky. Slots can assume that a [`Slot`] for an index
 //!   `i` is exclusive for its duration. This means that [`Slot`] implementations can lend
 //!   out mutable references to its contents (for example,
 //!   [`super::invasive::Slot::as_mut_slice`]).
@@ -56,17 +56,17 @@
 //!   Code in [`super`] is very careful to maintain this invariant and all users of [`Slot`]
 //!   must carefully maintain this as well.
 //!
-//! * `reclaim`: On a call to [`Plugin::reclaim`], it can be assumed that the plugin has
-//!   exclusive access to the indicated slot for the duration of the function call.
+//! * `reclaim`: On a call to [`Slots::reclaim`], implementations may assume exclusive access
+//!   to the indicated slot for the duration of the function call.
 //!
 //! ## Contracts
 //!
-//! Users of [`Plugin`] must ensure that the lifecycle shown above is strictly observed.
+//! Users of [`Slots`] must ensure that the lifecycle shown above is strictly observed.
 //! Furthermore, for [`Slot`]s, exactly one of the terminal methods **must** be called.
 //!
 //! State transitions are driven by the authoritative [`super::Store`]. Before invoking a
-//! plugin transition, the store ensures the slot is not externally available in its previous
-//! state. Further, the store commits the destination state only after the plugin API call
+//! transition, the store ensures the slot is not externally available in its previous
+//! state. Further, the store commits the destination state only after the lifecycle API call
 //! completes.
 
 use std::fmt::Debug;
@@ -75,23 +75,23 @@ use crate::num::IdLimit;
 
 use super::Lifecycle;
 
-/// A configuration for a [`Plugin`].
-pub(crate) trait PluginConfig: Debug {
-    /// The type of the resulting [`Plugin`].
-    type Plugin: Plugin;
+/// A configuration for a [`Slots`].
+pub(crate) trait SlotsConfig: Debug {
+    /// The type of the resulting [`Slots`].
+    type Slots: Slots;
 
     /// Construction errors.
     type Error: std::error::Error + Send + Sync + 'static;
 
-    /// Build the associated [`Plugin`] from self with the [`IdLimit`].
-    fn build(self, id_limit: IdLimit) -> Result<Self::Plugin, Self::Error>;
+    /// Build the associated [`Slots`] from self with the [`IdLimit`].
+    fn build(self, id_limit: IdLimit) -> Result<Self::Slots, Self::Error>;
 }
 
 /// A lifecycle backend for [`super::Store`]'s EBR scheme.
 ///
 /// See the [module level documentation](self) for details.
-pub(crate) trait Plugin: Debug + 'static {
-    /// The writable [`Slot`] for this plugin.
+pub(crate) trait Slots: Debug + 'static {
+    /// The writable [`Slot`].
     type Slot<'a>: Slot;
 
     /// Return the exclusive upper bound for indices provided to this API.
@@ -108,7 +108,7 @@ pub(crate) trait Plugin: Debug + 'static {
     ///
     /// Callers must ensure **all** of the following:
     ///
-    /// 1. The plugin is in the implicit "available" state according to the [module docs](self).
+    /// 1. The slot is in the implicit "available" state according to the [module docs](self).
     ///
     /// 2. Access to slot `i` is exclusive before invoking this method and that exclusivity
     ///    is maintained until the returned [`Slot`] is consumed by a terminal method.
@@ -124,7 +124,7 @@ pub(crate) trait Plugin: Debug + 'static {
     ///
     /// # Safety
     ///
-    /// The plugin is in the implicit "published" state.
+    /// The slot is in the implicit "published" state.
     unsafe fn retire(&self, i: u32, _: Lifecycle);
 
     /// Transition slot `i` from the "retiring" state to the "available" state.
@@ -136,17 +136,17 @@ pub(crate) trait Plugin: Debug + 'static {
     ///
     /// Callers must ensure **all** of the following:
     ///
-    /// 1. The plugin is in the implicit "retiring" state.
+    /// 1. The slot is in the implicit "retiring" state.
     ///
-    /// 2. All [`crate::epoch::Guard`]s for this [`Plugin`] that could have obtained a
+    /// 2. All [`crate::epoch::Guard`]s for this [`Slots`] that could have obtained a
     ///    reference while this slot was in the "published" state have been dropped.
     unsafe fn reclaim(&self, i: u32, _: Lifecycle);
 }
 
-/// A writable slot for [`Plugin`].
+/// A writable slot for [`Slots`].
 ///
-/// [`Slot`]s may assume that they have exclusive ownership of their plugin slots for their
-/// duration in accordance with [`Plugin::acquire`].
+/// [`Slot`]s may assume that they have exclusive ownership of their slots for their duration
+/// in accordance with [`Slots::acquire`].
 pub(crate) trait Slot: Debug {
     /// Mark this slot as readable, transition it to the "published" state.
     fn publish(self, _: Lifecycle);

From 6269e7838d0450c6e3180be213efbcce52f054bb Mon Sep 17 00:00:00 2001
From: Mark Hildebrand <mhildebrand@microsoft.com>
Date: Thu, 27 Aug 2026 16:49:13 -0700
Subject: [PATCH 34/34] "Layer" -> "Repr".

---
 diskann-benchmark/src/index/inmem2.rs         |   4 +-
 diskann-inmem/integration/index/object.rs     |   6 +-
 diskann-inmem/integration/index/runner.rs     |  40 ++--
 .../jsons/integration-baseline.json           |  10 +-
 .../integration/jsons/integration.json        |   8 +-
 diskann-inmem/src/lib.rs                      |   2 +-
 diskann-inmem/src/num.rs                      |  10 +-
 diskann-inmem/src/provider.rs                 | 114 +++++-----
 diskann-inmem/src/{layers => repr}/full.rs    | 209 +++++++++---------
 diskann-inmem/src/{layers => repr}/mod.rs     |  32 +--
 10 files changed, 217 insertions(+), 218 deletions(-)
 rename diskann-inmem/src/{layers => repr}/full.rs (87%)
 rename diskann-inmem/src/{layers => repr}/mod.rs (89%)

diff --git a/diskann-benchmark/src/index/inmem2.rs b/diskann-benchmark/src/index/inmem2.rs
index fca63991c2..60b776d635 100644
--- a/diskann-benchmark/src/index/inmem2.rs
+++ b/diskann-benchmark/src/index/inmem2.rs
@@ -29,8 +29,8 @@ use diskann_benchmark_runner::{
     Benchmark, Checker, Checkpoint, Input, Registry,
 };
 use diskann_inmem::{
-    layers::{Full, FullPrecision},
     num::{Capacity, MaxDegree},
+    repr::{Full, FullPrecision},
     Provider, Strategy,
 };
 use diskann_utils::views::{Matrix, MatrixView};
@@ -423,7 +423,7 @@ impl<T> Build<T> {
 
 impl<T> Benchmark for Build<T>
 where
-    T: diskann_inmem::layers::FullPrecision + diskann::graph::SampleableForStart + AsDataType,
+    T: diskann_inmem::repr::FullPrecision + diskann::graph::SampleableForStart + AsDataType,
 {
     type Input = StaticBuild;
     type Output = ();
diff --git a/diskann-inmem/integration/index/object.rs b/diskann-inmem/integration/index/object.rs
index cc5f4dc6f9..32e357d69d 100644
--- a/diskann-inmem/integration/index/object.rs
+++ b/diskann-inmem/integration/index/object.rs
@@ -14,7 +14,7 @@ use diskann_benchmark_runner::utils::fmt::KeyValue;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-use diskann_inmem::{Context, Provider, Strategy, integration, layers};
+use diskann_inmem::{Context, Provider, Strategy, integration, repr};
 
 use crate::support::{
     check::{CheckMatch, Match, check_all_fields},
@@ -186,9 +186,9 @@ impl CheckMatch for Counters {
 // Impls //
 ///////////
 
-impl<T> Index for DiskANNIndex<Provider<layers::Full<T>, u64>>
+impl<T> Index for DiskANNIndex<Provider<repr::Full<T>, u64>>
 where
-    T: layers::FullPrecision + FromSlice + AsDataType,
+    T: repr::FullPrecision + FromSlice + AsDataType,
 {
     fn search<'a>(
         &'a self,
diff --git a/diskann-inmem/integration/index/runner.rs b/diskann-inmem/integration/index/runner.rs
index 3b2057ec66..97e711a5d1 100644
--- a/diskann-inmem/integration/index/runner.rs
+++ b/diskann-inmem/integration/index/runner.rs
@@ -19,8 +19,8 @@ use serde::{Deserialize, Serialize};
 
 use diskann_inmem::{
     Provider,
-    layers::Full,
     num::{Capacity, MaxDegree},
+    repr::Full,
 };
 
 use crate::{
@@ -109,7 +109,7 @@ mod dto {
     }
 
     #[derive(Debug, Serialize, Deserialize)]
-    pub(super) enum Layer {
+    pub(super) enum Representation {
         FullPrecision { data_type: DataType },
     }
 
@@ -137,7 +137,7 @@ mod dto {
     #[derive(Debug, Serialize, Deserialize)]
     pub(super) struct Test {
         pub(super) data: Data,
-        pub(super) layer: Layer,
+        pub(super) representation: Representation,
         pub(super) build: Build,
         pub(super) search: Search,
     }
@@ -230,20 +230,20 @@ struct Bundle {
 }
 
 #[derive(Debug)]
-enum Layer {
+enum Representation {
     FullPrecision { data_type: DataType },
 }
 
-impl Layer {
-    fn from_raw(raw: dto::Layer) -> Self {
+impl Representation {
+    fn from_raw(raw: dto::Representation) -> Self {
         match raw {
-            dto::Layer::FullPrecision { data_type } => Self::FullPrecision { data_type },
+            dto::Representation::FullPrecision { data_type } => Self::FullPrecision { data_type },
         }
     }
 
-    fn as_raw(&self) -> dto::Layer {
+    fn as_raw(&self) -> dto::Representation {
         match self {
-            Self::FullPrecision { data_type } => dto::Layer::FullPrecision {
+            Self::FullPrecision { data_type } => dto::Representation::FullPrecision {
                 data_type: *data_type,
             },
         }
@@ -326,7 +326,7 @@ impl Search {
 #[derive(Debug)]
 struct Test {
     data: Data,
-    layer: Layer,
+    representation: Representation,
     build: Build,
     search: Search,
 }
@@ -334,13 +334,13 @@ struct Test {
 impl Test {
     fn from_raw(raw: dto::Test, checker: Option<&mut Checker>) -> anyhow::Result<Self> {
         let data = Data::from_raw(raw.data, checker)?;
-        let layer = Layer::from_raw(raw.layer);
+        let representation = Representation::from_raw(raw.representation);
         let build = Build::from_raw(raw.build, data.metric)?;
         let search = Search::from_raw(raw.search)?;
 
         Ok(Self {
             data,
-            layer,
+            representation,
             build,
             search,
         })
@@ -349,7 +349,7 @@ impl Test {
     fn as_raw(&self) -> anyhow::Result<dto::Test> {
         Ok(dto::Test {
             data: self.data.as_raw()?,
-            layer: self.layer.as_raw(),
+            representation: self.representation.as_raw(),
             build: self.build.as_raw(),
             search: self.search.as_raw(),
         })
@@ -360,8 +360,8 @@ impl Test {
         capacity: usize,
         start_points: DatasetView<'_>,
     ) -> anyhow::Result<Arc<dyn Index>> {
-        match self.layer {
-            Layer::FullPrecision { data_type } => {
+        match self.representation {
+            Representation::FullPrecision { data_type } => {
                 if start_points.data_type() != data_type {
                     anyhow::bail!(
                         "mismatched data types for start point - expected {}, got {}",
@@ -457,7 +457,7 @@ impl diskann_benchmark_runner::Input for Test {
                 data_type: DataType::F32,
                 preprocess: vec![],
             },
-            layer: dto::Layer::FullPrecision {
+            representation: dto::Representation::FullPrecision {
                 data_type: DataType::F32,
             },
             build: dto::Build {
@@ -501,8 +501,8 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision {
     type Output = BuildAndSearch;
 
     fn try_match(&self, input: &Test, context: &MatchContext) -> Score {
-        // Future-proof against additional enums in `input.layer`.
-        let Layer::FullPrecision { .. } = input.layer;
+        // Future-proof against additional enums in `input.representation`.
+        let Representation::FullPrecision { .. } = input.representation;
         context.success(0)
     }
 
@@ -516,8 +516,8 @@ impl diskann_benchmark_runner::Benchmark for FullPrecision {
         _checkpoint: Checkpoint<'_>,
         mut output: &mut dyn Output,
     ) -> anyhow::Result<Self::Output> {
-        // Future-proof against additional enums in `input.layer`.
-        let Layer::FullPrecision { data_type } = input.layer;
+        // Future-proof against additional enums in `input.representation`.
+        let Representation::FullPrecision { data_type } = input.representation;
 
         // Load the data and perform any necessary data conversions.
         let Bundle {
diff --git a/diskann-inmem/integration/jsons/integration-baseline.json b/diskann-inmem/integration/jsons/integration-baseline.json
index 941a5cb752..6433eee856 100644
--- a/diskann-inmem/integration/jsons/integration-baseline.json
+++ b/diskann-inmem/integration/jsons/integration-baseline.json
@@ -16,7 +16,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f32"
           }
@@ -137,7 +137,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f16"
           }
@@ -258,7 +258,7 @@
           "preprocess": [],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "u8"
           }
@@ -382,7 +382,7 @@
           ],
           "queries": "/yfcc/yfcc_query_100.fbin"
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "i8"
           }
@@ -486,4 +486,4 @@
       ]
     }
   }
-]
\ No newline at end of file
+]
diff --git a/diskann-inmem/integration/jsons/integration.json b/diskann-inmem/integration/jsons/integration.json
index 57f2a338ce..c00dbae441 100644
--- a/diskann-inmem/integration/jsons/integration.json
+++ b/diskann-inmem/integration/jsons/integration.json
@@ -21,7 +21,7 @@
           "queries": "yfcc_query_100.fbin",
           "preprocess": []
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f32"
           }
@@ -64,7 +64,7 @@
           "queries": "yfcc_query_100.fbin",
           "preprocess": []
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "f16"
           }
@@ -107,7 +107,7 @@
           "queries": "yfcc_query_100.fbin",
           "preprocess": []
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "u8"
           }
@@ -153,7 +153,7 @@
             "floor"
           ]
         },
-        "layer": {
+        "representation": {
           "FullPrecision": {
             "data_type": "i8"
           }
diff --git a/diskann-inmem/src/lib.rs b/diskann-inmem/src/lib.rs
index 4680dac23e..eb250bf518 100644
--- a/diskann-inmem/src/lib.rs
+++ b/diskann-inmem/src/lib.rs
@@ -18,8 +18,8 @@ mod neighbors;
 mod prefetch;
 mod tag;
 
-pub mod layers;
 pub mod provider;
+pub mod repr;
 pub mod store;
 
 pub use provider::{Context, Provider, Strategy};
diff --git a/diskann-inmem/src/num.rs b/diskann-inmem/src/num.rs
index 5d05649225..2d734c11ec 100644
--- a/diskann-inmem/src/num.rs
+++ b/diskann-inmem/src/num.rs
@@ -187,9 +187,9 @@ macro_rules! typed_int {
 // TODO: Provide a linkable reference for "immutable" points.
 
 typed_int!(
-    /// The number of distinct slots a [`crate::Provider`] or [`crate::layers::Layer`] has
-    /// capacity for. This is logically distinct from [`IdLimit`], which may be greater due
-    /// to immutable points within a storage container.
+    /// The number of distinct slots a [`crate::Provider`] or [`crate::repr::Representation`]
+    /// has capacity for. This is logically distinct from [`IdLimit`], which may be greater
+    /// due to immutable points within a storage container.
     pub Capacity,
     usize,
 );
@@ -201,8 +201,8 @@ typed_int!(
 );
 
 typed_int!(
-    /// One larger than the maximum ID that a [`crate::Provider`], [`crate::layers::Layer`],
-    /// or other such store in this crate can access in-bounds.
+    /// One larger than the maximum ID that a [`crate::Provider`],
+    /// [`crate::repr::Representation`], or other such store in this crate can access in-bounds.
     ///
     /// This implies that access to ids `[0..self)` are in-bounds.
     ///
diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs
index dacf3027ae..d2d2c8f82e 100644
--- a/diskann-inmem/src/provider.rs
+++ b/diskann-inmem/src/provider.rs
@@ -46,9 +46,9 @@ use diskann::{
 use crate::{
     counters::{Counters, LocalCounters},
     ids::IdMap,
-    layers,
     neighbors::Neighbors,
     num::{IdLimit, MaxDegree},
+    repr,
 };
 
 /// Aggregate trait for the external ID type of [`Provider`].
@@ -58,16 +58,16 @@ impl<T> Id for T where T: Send + Sync + Hash + Eq + Clone + 'static {}
 
 /// An in-memory data-provider for DiskANN's graph indexing algorithms.
 ///
-/// The first type parameter `L` is a [`layers::Layer`] for describing the kind of data
+/// The first type parameter `R` is a [`repr::Representation`] for describing the kind of data
 /// stored within the provider. The second parameter `M` is the associated data for items
 /// inserted into the provider.
 #[derive(Debug)]
-pub struct Provider<L, M = u32>
+pub struct Provider<R, M = u32>
 where
     M: Id,
 {
     // Data representation and storage.
-    layer: L,
+    representation: R,
     // ID translation.
     mapping: IdMap<M>,
     // `Counters` is only non-trivial under the `integration-test` feature flag. Otherwise,
@@ -75,7 +75,7 @@ where
     counters: Counters,
 }
 
-impl<L, M> Provider<L, M>
+impl<R, M> Provider<R, M>
 where
     M: Id,
 {
@@ -91,23 +91,21 @@ where
     }
 }
 
-impl<L, M> Provider<L, M>
+impl<R, M> Provider<R, M>
 where
-    L: layers::Layer,
+    R: repr::Representation,
     M: Id,
 {
     /// Construct a new [`Provider`].
-    ///
-    /// The list of `start_points` must be must be compatible with `layer`.
     pub fn new<C>(config: C) -> ANNResult<Self>
     where
-        C: layers::LayerConfig<Layer = L>,
+        C: repr::RepresentationConfig<Representation = R>,
     {
-        let layer = <_ as layers::LayerConfig>::build(config)?;
-        let mapping = IdMap::new(layer.capacity());
+        let representation = <_ as repr::RepresentationConfig>::build(config)?;
+        let mapping = IdMap::new(representation.capacity());
 
         Ok(Self {
-            layer,
+            representation,
             mapping,
             counters: Counters::new(),
         })
@@ -115,7 +113,7 @@ where
 
     /// Return the maximum number of neighbors that can be stored in the provider's graph.
     pub fn max_degree(&self) -> MaxDegree {
-        self.layer.max_degree()
+        self.representation.max_degree()
     }
 }
 
@@ -169,9 +167,9 @@ where
 //
 // `diskann` has plans to move deletion checks behind an accessor trait, which will help
 // with this situation.
-impl<L, M> diskann::provider::Delete for Provider<L, M>
+impl<R, M> diskann::provider::Delete for Provider<R, M>
 where
-    L: layers::Layer,
+    R: repr::Representation,
     M: Id,
 {
     async fn delete(&self, _context: &Context, gid: &M) -> ANNResult<()> {
@@ -188,7 +186,7 @@ where
 
         // An early return here will cause `entry` to be dropped, which will *not* cause
         // the delete to commit.
-        <L as layers::Layer>::retire(&self.layer, entry.internal())?;
+        <R as repr::Representation>::retire(&self.representation, entry.internal())?;
 
         // Successfully retired the internal slot. We can safely release the ID mapping.
         entry.delete();
@@ -205,7 +203,7 @@ where
         _context: &Context,
         id: u32,
     ) -> ANNResult<diskann::provider::ElementStatus> {
-        match <L as layers::Layer>::is_readable(&self.layer, id) {
+        match <R as repr::Representation>::is_readable(&self.representation, id) {
             Some(true) => Ok(diskann::provider::ElementStatus::Valid),
             Some(false) => Ok(diskann::provider::ElementStatus::Deleted),
             None => Err(ANNError::message("accessed invalid internal ID")),
@@ -232,9 +230,9 @@ where
     std::future::ready(f())
 }
 
-impl<T, L, M> diskann::provider::SetElement<T> for Provider<L, M>
+impl<T, R, M> diskann::provider::SetElement<T> for Provider<R, M>
 where
-    L: layers::Set<T>,
+    R: repr::Set<T>,
     M: Id,
 {
     type SetError = ANNError;
@@ -251,13 +249,13 @@ where
             //
             // The internal `Guard` is sufficient for local rollback, but not after
             // `set_element` returns.
-            let guard = <L as layers::Set<T>>::set(&self.layer, element)?;
-            let internal = <_ as layers::Guard>::id(&guard);
+            let guard = <R as repr::Set<T>>::set(&self.representation, element)?;
+            let internal = <_ as repr::Guard>::id(&guard);
             self.mapping.insert(id.clone(), internal)?;
 
             // Now that insert has succeeded - publish the slot. This method cannot fail, so
             // we do not need to worry about potentially unwinding the ID mapping.
-            <_ as layers::Guard>::publish(guard);
+            <_ as repr::Guard>::publish(guard);
 
             // This is a rather expensive update.
             //
@@ -285,7 +283,7 @@ where
 pub struct SearchAccessor<'a> {
     neighbors: &'a Neighbors,
     ids: AdjacencyList<u32>,
-    expand_beam: Box<dyn layers::ExpandBeam + 'a>,
+    expand_beam: Box<dyn repr::ExpandBeam + 'a>,
     id_limit: IdLimit,
     buffer: Vec<(u32, f32)>,
 
@@ -298,7 +296,7 @@ pub struct SearchAccessor<'a> {
 impl<'a> SearchAccessor<'a> {
     pub(crate) fn new(
         neighbors: &'a Neighbors,
-        expand_beam: Box<dyn layers::ExpandBeam + 'a>,
+        expand_beam: Box<dyn repr::ExpandBeam + 'a>,
         provider: &'a (dyn std::any::Any + Send + Sync),
         start_points: std::ops::Range<u32>,
         counters: LocalCounters<'a>,
@@ -413,15 +411,15 @@ impl glue::SearchAccessor for SearchAccessor<'_> {
 /// This type implements zero-copy access to the data within its parent provider during prunes.
 #[derive(Debug)]
 pub struct PruneAccessor<'a> {
-    prune: Box<dyn layers::Prune + 'a>,
-    keys: hashbrown::HashMap<u32, Option<layers::PruneKey>>,
+    prune: Box<dyn repr::Prune + 'a>,
+    keys: hashbrown::HashMap<u32, Option<repr::PruneKey>>,
     neighbors: &'a Neighbors,
     counters: LocalCounters<'a>,
 }
 
 impl<'a> PruneAccessor<'a> {
     pub(crate) fn new(
-        prune: Box<dyn layers::Prune + 'a>,
+        prune: Box<dyn repr::Prune + 'a>,
         neighbors: &'a Neighbors,
         counters: LocalCounters<'a>,
     ) -> Self {
@@ -437,12 +435,12 @@ impl<'a> PruneAccessor<'a> {
 /// The distance computer for [`PruneAccessor`].
 #[derive(Debug)]
 pub struct Distance<'a> {
-    prune: &'a dyn layers::Prune,
+    prune: &'a dyn repr::Prune,
     counters: LocalCounters<'a>,
 }
 
 impl<'a> Distance<'a> {
-    fn new(prune: &'a dyn layers::Prune, counters: LocalCounters<'a>) -> Self {
+    fn new(prune: &'a dyn repr::Prune, counters: LocalCounters<'a>) -> Self {
         Self { prune, counters }
     }
 }
@@ -450,7 +448,7 @@ impl<'a> Distance<'a> {
 /// An opaque element-ref for [`PruneAccessor`].
 #[derive(Debug, Clone, Copy)]
 #[repr(transparent)]
-pub struct ElementRef(layers::PruneKey);
+pub struct ElementRef(repr::PruneKey);
 
 impl<'a> diskann_utils::Reborrow<'a> for ElementRef {
     type Target = ElementRef;
@@ -585,9 +583,9 @@ impl workingset::View<u32> for &PruneAccessor<'_> {
 #[derive(Debug, Clone, Copy)]
 pub struct Strategy;
 
-impl<'a, L, M> glue::SearchStrategy<'a, Provider<L, M>, L::Query<'a>> for Strategy
+impl<'a, R, M> glue::SearchStrategy<'a, Provider<R, M>, R::Query<'a>> for Strategy
 where
-    L: layers::Search,
+    R: repr::Search,
     M: Id,
 {
     type SearchAccessor = SearchAccessor<'a>;
@@ -595,12 +593,12 @@ where
 
     fn search_accessor(
         &'a self,
-        provider: &'a Provider<L, M>,
+        provider: &'a Provider<R, M>,
         _context: &'a Context,
-        query: L::Query<'a>,
+        query: R::Query<'a>,
     ) -> ANNResult<SearchAccessor<'a>> {
-        <L as layers::Search>::search_accessor(
-            &provider.layer,
+        <R as repr::Search>::search_accessor(
+            &provider.representation,
             query,
             provider,
             provider.local_counters(),
@@ -611,7 +609,7 @@ where
 // This is a utility for helping inspect the generated code for `ExpandBeam`.
 //
 pub fn test_function<'a>(
-    x: &'a Provider<layers::Full<u8>>,
+    x: &'a Provider<repr::Full<u8>>,
     strategy: &'a Strategy,
     context: &'a Context,
     query: &'a [u8],
@@ -621,17 +619,17 @@ pub fn test_function<'a>(
 
 /// Perform ID translation during post-processing.
 #[derive(Debug, Clone, Copy)]
-pub struct Translate<L, M>(std::marker::PhantomData<(L, M)>);
+pub struct Translate<R, M>(std::marker::PhantomData<(R, M)>);
 
-impl<L, M> Default for Translate<L, M> {
+impl<R, M> Default for Translate<R, M> {
     fn default() -> Self {
         Self(std::marker::PhantomData)
     }
 }
 
-impl<'a, L, M> glue::SearchPostProcess<SearchAccessor<'a>, L::Query<'a>, M> for Translate<L, M>
+impl<'a, R, M> glue::SearchPostProcess<SearchAccessor<'a>, R::Query<'a>, M> for Translate<R, M>
 where
-    L: layers::Search,
+    R: repr::Search,
     M: Id,
 {
     type Error = ANNError;
@@ -639,7 +637,7 @@ where
     fn post_process<I, B>(
         &self,
         accessor: &mut SearchAccessor<'_>,
-        _query: L::Query<'a>,
+        _query: R::Query<'a>,
         candidates: I,
         output: &mut B,
     ) -> impl std::future::Future<Output = Result<usize, Self::Error>> + Send
@@ -649,7 +647,7 @@ where
     {
         let work = move || {
             // By construction - the downcast should succeed. Otherwise, this is a program bug.
-            let provider = match accessor.provider.downcast_ref::<Provider<L, M>>() {
+            let provider = match accessor.provider.downcast_ref::<Provider<R, M>>() {
                 Some(provider) => provider,
                 None => return Err(ANNError::message("bad any cast")),
             };
@@ -674,17 +672,17 @@ where
     }
 }
 
-impl<'a, L, M> glue::DefaultPostProcessor<'a, Provider<L, M>, L::Query<'a>, M> for Strategy
+impl<'a, R, M> glue::DefaultPostProcessor<'a, Provider<R, M>, R::Query<'a>, M> for Strategy
 where
-    L: layers::Search,
+    R: repr::Search,
     M: Id,
 {
-    diskann::default_post_processor!(Translate<L, M>);
+    diskann::default_post_processor!(Translate<R, M>);
 }
 
-impl<L, M> glue::PruneStrategy<Provider<L, M>> for Strategy
+impl<R, M> glue::PruneStrategy<Provider<R, M>> for Strategy
 where
-    L: layers::Insert,
+    R: repr::Insert,
     M: Id,
 {
     type PruneAccessor<'a> = PruneAccessor<'a>;
@@ -692,17 +690,17 @@ where
 
     fn prune_accessor<'a>(
         &self,
-        provider: &'a Provider<L, M>,
+        provider: &'a Provider<R, M>,
         _context: &'a Context,
         _capacity: usize,
     ) -> ANNResult<PruneAccessor<'a>> {
-        <L as layers::Insert>::prune_accessor(&provider.layer, provider.local_counters())
+        <R as repr::Insert>::prune_accessor(&provider.representation, provider.local_counters())
     }
 }
 
-impl<'a, L, M> glue::InsertStrategy<'a, Provider<L, M>, L::Query<'a>> for Strategy
+impl<'a, R, M> glue::InsertStrategy<'a, Provider<R, M>, R::Query<'a>> for Strategy
 where
-    L: layers::Insert,
+    R: repr::Insert,
     M: Id,
 {
     type PruneStrategy = Self;
@@ -711,10 +709,10 @@ where
     }
 }
 
-impl<T, M> glue::InplaceDeleteStrategy<Provider<layers::Full<T>, M>> for Strategy
+impl<T, M> glue::InplaceDeleteStrategy<Provider<repr::Full<T>, M>> for Strategy
 where
     M: Id,
-    T: layers::FullPrecision,
+    T: repr::FullPrecision,
 {
     type DeleteElement<'a> = &'a [T];
     type DeleteElementGuard = Box<[T]>;
@@ -739,12 +737,12 @@ where
 
     fn get_delete_element<'a>(
         &'a self,
-        provider: &'a Provider<layers::Full<T>, M>,
+        provider: &'a Provider<repr::Full<T>, M>,
         _context: &'a Context,
         id: u32,
     ) -> impl Future<Output = Result<Self::DeleteElementGuard, Self::DeleteElementError>> + Send
     {
-        let work = move || provider.layer.get(id);
+        let work = move || provider.representation.get(id);
         ready(work)
     }
 }
@@ -793,7 +791,7 @@ mod tests {
         let start = grid.start_point(size);
         let degree = 6;
 
-        let config = layers::full::Config::new(
+        let config = repr::full::Config::new(
             Capacity::new(grid.num_points(size)),
             MaxDegree::new(degree),
             Metric::L2,
diff --git a/diskann-inmem/src/layers/full.rs b/diskann-inmem/src/repr/full.rs
similarity index 87%
rename from diskann-inmem/src/layers/full.rs
rename to diskann-inmem/src/repr/full.rs
index 91b3a3e440..743300e020 100644
--- a/diskann-inmem/src/layers/full.rs
+++ b/diskann-inmem/src/repr/full.rs
@@ -11,14 +11,14 @@
 //! The [`FullPrecision`] generic bound can be used to constrain these data types.
 
 mod internal_docs {
-    //! Internally, the [`super::layers::Search`] and [`super::layers::Insert`] traits
+    //! Internally, the [`super::repr::Search`] and [`super::repr::Insert`] traits
     //! are implemented via [`super::FullPrecisionImpl`], which creates:
     //!
     //! * [`super::ExpandBeam`]: For index search.
     //! * [`super::Prune`]: For index construction.
     //!
     //! These two structs are modular with respect to their exact distance function and
-    //! prefetcher. Since [`super::layers::ExpandBeam`] and [`super::layers::Prune`] are
+    //! prefetcher. Since [`super::repr::ExpandBeam`] and [`super::repr::Prune`] are
     //! used as trait objects, this allows the implementation structs in this module to be
     //! highly specialized, including:
     //!
@@ -55,9 +55,10 @@ use thiserror::Error;
 
 use crate::{
     counters::LocalCounters,
-    epoch, layers,
+    epoch,
     num::{Bytes, Capacity, IdLimit, MaxDegree},
     prefetch::{self, Prefetch},
+    repr,
     store::{
         self, Store,
         invasive::{self, Invasive},
@@ -67,12 +68,12 @@ use crate::{
 
 /// A useful trait bound for types compatible with [`Full`].
 ///
-/// This encompasses *everything* required for `Full: layers::Insert` and can be used as
+/// This encompasses *everything* required for `Full: repr::Insert` and can be used as
 /// a single bound.
 pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     #[doc(hidden)]
     fn __search_accessor<'a>(
-        layer: &'a Full<Self>,
+        representation: &'a Full<Self>,
         query: &'a [Self],
         provider: &'a (dyn std::any::Any + Send + Sync),
         counters: LocalCounters<'a>,
@@ -80,7 +81,7 @@ pub trait FullPrecision: bytemuck::Pod + std::fmt::Debug + Send + Sync {
 
     #[doc(hidden)]
     fn __prune_accessor<'a>(
-        layer: &'a Full<Self>,
+        representation: &'a Full<Self>,
         counters: LocalCounters<'a>,
     ) -> ANNResult<crate::provider::PruneAccessor<'a>>;
 }
@@ -188,11 +189,11 @@ pub enum ConfigError {
 
 diskann::convert_error!(ConfigError);
 
-impl<T> layers::LayerConfig for Config<T>
+impl<T> repr::RepresentationConfig for Config<T>
 where
     T: FullPrecision,
 {
-    type Layer = Full<T>;
+    type Representation = Full<T>;
 
     fn build(self) -> ANNResult<Full<T>> {
         <Config<T>>::build(self)
@@ -204,13 +205,13 @@ trait FullPrecisionImpl: bytemuck::Pod + std::fmt::Debug + Send + Sync {
     fn make_expand_beam<'a>(
         full: &'a Full<Self>,
         query: &'a [Self],
-    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>>;
+    ) -> ANNResult<Box<dyn repr::ExpandBeam + 'a>>;
 
     #[doc(hidden)]
-    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>>;
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn repr::Prune + 'a>>;
 }
 
-/// Full-precision data layer.
+/// Full-precision data representation.
 #[derive(Debug)]
 pub struct Full<T>
 where
@@ -226,7 +227,7 @@ impl<T> Full<T>
 where
     T: 'static,
 {
-    /// Initialize a [`Config`] for this layer.
+    /// Initialize a [`Config`] for this representation.
     ///
     /// See also: [`Config::new`].
     ///
@@ -242,7 +243,7 @@ where
         Config::new(capacity, max_degree, metric, start_points)
     }
 
-    /// Create a new full-precision layer from `config`.
+    /// Create a new full-precision representation from `config`.
     ///
     /// See: [`Config::build`].
     fn new(config: Config<T>) -> ANNResult<Self>
@@ -285,7 +286,7 @@ where
         })
     }
 
-    /// Return the logical dimension of the data handled by this [`layers::Layer`].
+    /// Return the logical dimension of the data handled by this [`repr::Representation`].
     pub fn dim(&self) -> usize {
         self.bytes().value() / std::mem::size_of::<T>()
     }
@@ -300,7 +301,7 @@ where
         self.store.slots().bytes_plus_tag()
     }
 
-    /// Return the [`Metric`] for this layer.
+    /// Return the [`Metric`] for this representation.
     pub fn metric(&self) -> Metric {
         self.metric
     }
@@ -340,7 +341,7 @@ where
     }
 }
 
-impl<T> layers::Layer for Full<T>
+impl<T> repr::Representation for Full<T>
 where
     T: FullPrecision,
 {
@@ -365,7 +366,7 @@ where
     }
 }
 
-impl<T> layers::Set<&[T]> for Full<T>
+impl<T> repr::Set<&[T]> for Full<T>
 where
     T: FullPrecision,
 {
@@ -392,7 +393,7 @@ where
     }
 }
 
-/// A [`layers::Guard`] for [`Full`].
+/// A [`repr::Guard`] for [`Full`].
 #[derive(Debug)]
 pub struct Guard<'a> {
     slot: store::Slot<'a, invasive::Slot<'a>>,
@@ -404,7 +405,7 @@ impl<'a> Guard<'a> {
     }
 }
 
-impl layers::Guard for Guard<'_> {
+impl repr::Guard for Guard<'_> {
     fn publish(self) {
         self.slot.publish();
     }
@@ -415,7 +416,7 @@ impl layers::Guard for Guard<'_> {
 
 #[derive(Debug, Error)]
 #[error(
-    "data of dimension {} does not match full precision layer's dimension {}",
+    "data of dimension {} does not match full precision representation's dimension {}",
     self.got,
     self.expected
 )]
@@ -426,7 +427,7 @@ struct SetError {
 
 diskann::convert_error!(SetError);
 
-impl<T> layers::Search for Full<T>
+impl<T> repr::Search for Full<T>
 where
     T: FullPrecision,
 {
@@ -442,7 +443,7 @@ where
     }
 }
 
-impl<T> layers::Insert for Full<T>
+impl<T> repr::Insert for Full<T>
 where
     T: FullPrecision,
 {
@@ -548,7 +549,7 @@ where
 struct ExpandBeam<'a, P, T, U, D> {
     // The original query.
     query: Calf<'a, T>,
-    // A reader into a layer's store.
+    // A reader into a representation's store.
     reader: store::invasive::Reader<'a>,
     // The prefetch lookahead.
     lookahead: Option<NonZeroUsize>,
@@ -617,10 +618,10 @@ impl<'a, P, T, U, D> ExpandBeam<'a, P, T, U, D> {
     }
 }
 
-// SAFETY: Our implementation of `layers::ExpandBeam::id_limit` is consistent with our
-// `layers::ExpandBeam::expand_beam` implementation. They are both dependent on
+// SAFETY: Our implementation of `repr::ExpandBeam::id_limit` is consistent with our
+// `repr::ExpandBeam::expand_beam` implementation. They are both dependent on
 // `invasive::Reader`'s internal bounds.
-unsafe impl<P, T, U, D> layers::ExpandBeam for ExpandBeam<'_, P, T, U, D>
+unsafe impl<P, T, U, D> repr::ExpandBeam for ExpandBeam<'_, P, T, U, D>
 where
     P: Prefetch,
     T: Send + Sync + 'static + Debug,
@@ -731,7 +732,7 @@ diskann::convert_error!(ExpandBeamError);
 struct Prune<'a, T, D> {
     // Buffered data to prune over.
     buffer: Vec<UnalignedSlice<'a, T>>,
-    // A reader into a layer's store.
+    // A reader into a representation's store.
     reader: store::invasive::Reader<'a>,
     // The distance implementation used for pruning.
     distance: D,
@@ -760,16 +761,16 @@ impl<'a, T, D> Prune<'a, T, D> {
     }
 }
 
-impl<T, D> layers::Prune for Prune<'_, T, D>
+impl<T, D> repr::Prune for Prune<'_, T, D>
 where
     T: Debug + Send + Sync + 'static,
     D: Distance<T, T>,
 {
     fn prepare(
         &mut self,
-        items: hashbrown::hash_map::IterMut<'_, u32, Option<layers::PruneKey>>,
+        items: hashbrown::hash_map::IterMut<'_, u32, Option<repr::PruneKey>>,
     ) -> ANNResult<usize> {
-        let mut counter = layers::PruneKey::counter();
+        let mut counter = repr::PruneKey::counter();
         self.buffer.clear();
         self.buffer.reserve(items.len());
 
@@ -802,7 +803,7 @@ where
         Ok(counter.index())
     }
 
-    fn evaluate(&self, a: layers::PruneKey, b: layers::PruneKey) -> f32 {
+    fn evaluate(&self, a: repr::PruneKey, b: repr::PruneKey) -> f32 {
         self.distance
             .eval(self.buffer[a.index()], self.buffer[b.index()])
     }
@@ -842,10 +843,10 @@ impl FullPrecisionImpl for f32 {
     fn make_expand_beam<'a>(
         full: &'a Full<f32>,
         query: &'a [f32],
-    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn repr::ExpandBeam + 'a>> {
         let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
-        let output: Box<dyn layers::ExpandBeam> = match full.metric {
+        let output: Box<dyn repr::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 100 {
                     expand_beam!(into, { f32, 100, SquaredL2 })
@@ -861,10 +862,10 @@ impl FullPrecisionImpl for f32 {
         Ok(output)
     }
 
-    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn repr::Prune + 'a>> {
         let reader = full.reader()?;
 
-        let output: Box<dyn layers::Prune> = match full.metric {
+        let output: Box<dyn repr::Prune> = match full.metric {
             Metric::L2 => prune!(Self, reader, SquaredL2),
             Metric::InnerProduct => prune!(Self, reader, InnerProduct),
             Metric::Cosine => prune!(Self, reader, Cosine),
@@ -879,14 +880,14 @@ impl FullPrecisionImpl for f16 {
     fn make_expand_beam<'a>(
         full: &'a Full<f16>,
         query: &'a [f16],
-    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn repr::ExpandBeam + 'a>> {
         let mut as_f32: Box<[f32]> = std::iter::repeat_n(0.0, full.dim()).collect();
         diskann_wide::arch::dispatch2(SliceCast::new(), &mut *as_f32, query);
         let query = Calf::Owned(as_f32);
 
         let into = IntoExpandBeam::new(full, query)?;
 
-        let output: Box<dyn layers::ExpandBeam> = match full.metric {
+        let output: Box<dyn repr::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 100 {
                     expand_beam!(into, { f16, 100, SquaredL2 })
@@ -902,10 +903,10 @@ impl FullPrecisionImpl for f16 {
         Ok(output)
     }
 
-    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn repr::Prune + 'a>> {
         let reader = full.reader()?;
 
-        let output: Box<dyn layers::Prune> = match full.metric {
+        let output: Box<dyn repr::Prune> = match full.metric {
             Metric::L2 => prune!(Self, reader, SquaredL2),
             Metric::InnerProduct => prune!(Self, reader, InnerProduct),
             Metric::Cosine => prune!(Self, reader, Cosine),
@@ -920,10 +921,10 @@ impl FullPrecisionImpl for u8 {
     fn make_expand_beam<'a>(
         full: &'a Full<u8>,
         query: &'a [u8],
-    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn repr::ExpandBeam + 'a>> {
         let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
-        let output: Box<dyn layers::ExpandBeam> = match full.metric {
+        let output: Box<dyn repr::ExpandBeam> = match full.metric {
             Metric::L2 => {
                 if full.dim() == 128 {
                     expand_beam!(into, { u8, 128, SquaredL2 })
@@ -938,10 +939,10 @@ impl FullPrecisionImpl for u8 {
         Ok(output)
     }
 
-    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn repr::Prune + 'a>> {
         let reader = full.reader()?;
 
-        let output: Box<dyn layers::Prune> = match full.metric {
+        let output: Box<dyn repr::Prune> = match full.metric {
             Metric::L2 => prune!(Self, reader, SquaredL2),
             Metric::InnerProduct => prune!(Self, reader, InnerProduct),
             Metric::Cosine => prune!(Self, reader, Cosine),
@@ -956,25 +957,25 @@ impl FullPrecisionImpl for i8 {
     fn make_expand_beam<'a>(
         full: &'a Full<i8>,
         query: &'a [i8],
-    ) -> ANNResult<Box<dyn layers::ExpandBeam + 'a>> {
+    ) -> ANNResult<Box<dyn repr::ExpandBeam + 'a>> {
         let into = IntoExpandBeam::new(full, Calf::Borrowed(query))?;
 
         let distance =
             <Self as DistanceProvider<Self>>::distance_comparer(full.metric(), Some(full.dim()));
 
-        let output: Box<dyn layers::ExpandBeam + 'a> =
+        let output: Box<dyn repr::ExpandBeam + 'a> =
             ExpandBeam::new(into, prefetch::Loop::new(), distance).boxed();
 
         Ok(output)
     }
 
-    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn layers::Prune + 'a>> {
+    fn make_prune<'a>(full: &'a Full<Self>) -> ANNResult<Box<dyn repr::Prune + 'a>> {
         let reader = full.reader()?;
 
         let distance =
             <Self as DistanceProvider<Self>>::distance_comparer(full.metric(), Some(full.dim()));
 
-        let output: Box<dyn layers::Prune> = Prune::<Self, _>::new(reader, distance).boxed();
+        let output: Box<dyn repr::Prune> = Prune::<Self, _>::new(reader, distance).boxed();
         Ok(output)
     }
 }
@@ -986,29 +987,29 @@ macro_rules! impl_full_precision {
     ($T:ty) => {
         impl FullPrecision for $T {
             fn __search_accessor<'a>(
-                layer: &'a Full<Self>,
+                representation: &'a Full<Self>,
                 query: &'a [Self],
                 provider: &'a (dyn std::any::Any + Send + Sync),
                 counters: LocalCounters<'a>,
             ) -> ANNResult<crate::provider::SearchAccessor<'a>> {
-                let expand_beam = <$T>::make_expand_beam(layer, query)?;
+                let expand_beam = <$T>::make_expand_beam(representation, query)?;
                 Ok(crate::provider::SearchAccessor::new(
-                    layer.store.neighbors(),
+                    representation.store.neighbors(),
                     expand_beam,
                     provider,
-                    layer.store.frozen(),
+                    representation.store.frozen(),
                     counters,
                 ))
             }
 
             fn __prune_accessor<'a>(
-                layer: &'a Full<Self>,
+                representation: &'a Full<Self>,
                 counters: LocalCounters<'a>,
             ) -> ANNResult<crate::provider::PruneAccessor<'a>> {
-                let prune = <$T>::make_prune(layer)?;
+                let prune = <$T>::make_prune(representation)?;
                 Ok(crate::provider::PruneAccessor::new(
                     prune,
-                    layer.store.neighbors(),
+                    representation.store.neighbors(),
                     counters,
                 ))
             }
@@ -1035,7 +1036,7 @@ mod tests {
     use hashbrown::{HashMap, HashSet};
     use rand::{Rng, SeedableRng, rngs::StdRng};
 
-    /// Generate random elements of a layer's data type from a seeded RNG.
+    /// Generate random elements of a representation's data type from a seeded RNG.
     trait Sample: bytemuck::Pod {
         fn sample<R: Rng>(rng: &mut R) -> Self;
     }
@@ -1083,7 +1084,7 @@ mod tests {
     fn test_full(capacity: Capacity) -> (Full<f32>, HashMap<u32, f32>) {
         let start_points = [capacity.value() as f32, (capacity.value() + 1) as f32];
 
-        let full = <_ as layers::LayerConfig>::build(
+        let full = <_ as repr::RepresentationConfig>::build(
             Full::<f32>::config(
                 capacity,
                 MaxDegree::new(0),
@@ -1104,10 +1105,10 @@ mod tests {
         );
         assert_eq!(full.metric(), Metric::L2);
         assert_eq!(
-            <_ as layers::Layer>::id_limit(&full),
+            <_ as repr::Representation>::id_limit(&full),
             IdLimit::new(capacity.value() as u32 + 2)
         );
-        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
+        assert_eq!(<_ as repr::Representation>::capacity(&full), capacity);
 
         let points: HashMap<u32, f32> = {
             let reader = full.reader().unwrap();
@@ -1162,16 +1163,16 @@ mod tests {
 
         let (mut full, mut points) = test_full(capacity);
 
-        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
-        assert_eq!(<_ as layers::Layer>::id_limit(&full), id_limit);
+        assert_eq!(<_ as repr::Representation>::capacity(&full), capacity);
+        assert_eq!(<_ as repr::Representation>::id_limit(&full), id_limit);
 
         let mut available: HashSet<u32> = (0..capacity.value()).map(|i| i as u32).collect();
 
         // Insert the values 0 to 10.
         for i in 0u32..10 {
-            let guard = <_ as layers::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
+            let guard = <_ as repr::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
 
-            let id = <_ as layers::Guard>::id(&guard);
+            let id = <_ as repr::Guard>::id(&guard);
 
             assert!(
                 available.remove(&id),
@@ -1183,7 +1184,7 @@ mod tests {
                 "insertion should not repeat",
             );
 
-            <_ as layers::Guard>::publish(guard);
+            <_ as repr::Guard>::publish(guard);
         }
 
         // Lookaheads to try.
@@ -1208,21 +1209,21 @@ mod tests {
         for lookahead in lookaheads {
             full.lookahead = *lookahead;
 
-            let g0 = <_ as layers::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
-            let g1 = <_ as layers::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
-            let g2 = <_ as layers::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
-            let g3 = <_ as layers::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
+            let g0 = <_ as repr::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
+            let g1 = <_ as repr::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
+            let g2 = <_ as repr::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
+            let g3 = <_ as repr::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
 
             {
-                let g0_id = <_ as layers::Guard>::id(&g0);
-                <_ as layers::Guard>::publish(g0);
-                <_ as layers::Layer>::retire(&full, g0_id).unwrap();
+                let g0_id = <_ as repr::Guard>::id(&g0);
+                <_ as repr::Guard>::publish(g0);
+                <_ as repr::Representation>::retire(&full, g0_id).unwrap();
             }
 
             {
-                let g1_id = <_ as layers::Guard>::id(&g1);
-                <_ as layers::Guard>::publish(g1);
-                <_ as layers::Layer>::retire(&full, g1_id).unwrap();
+                let g1_id = <_ as repr::Guard>::id(&g1);
+                <_ as repr::Guard>::publish(g1);
+                <_ as repr::Representation>::retire(&full, g1_id).unwrap();
             }
 
             let query = -1.0f32;
@@ -1232,7 +1233,7 @@ mod tests {
 
             let expand = ExpandBeam::new(into, prefetch::Loop::new(), TestDistance);
 
-            assert_eq!(<_ as layers::ExpandBeam>::id_limit(&expand), id_limit);
+            assert_eq!(<_ as repr::ExpandBeam>::id_limit(&expand), id_limit);
 
             let mut buf = Vec::<(u32, f32)>::new();
             let mut list = Vec::<u32>::new();
@@ -1254,7 +1255,7 @@ mod tests {
                 //
                 // Also by construction `buf` is at least as long as `list`.
                 let read =
-                    unsafe { <_ as layers::ExpandBeam>::expand_beam(&expand, &list, &mut buf) }
+                    unsafe { <_ as repr::ExpandBeam>::expand_beam(&expand, &list, &mut buf) }
                         .unwrap();
 
                 let expected: Vec<(u32, f32)> = list
@@ -1265,12 +1266,12 @@ mod tests {
                             let expected = point + query;
 
                             assert!(
-                                <_ as layers::Layer>::is_readable(&full, id).unwrap(),
+                                <_ as repr::Representation>::is_readable(&full, id).unwrap(),
                                 "point should be readable"
                             );
 
                             assert_eq!(
-                                <_ as layers::ExpandBeam>::evaluate(&expand, id).unwrap(),
+                                <_ as repr::ExpandBeam>::evaluate(&expand, id).unwrap(),
                                 Some(expected),
                                 "readable points should return valid distances",
                             );
@@ -1279,12 +1280,12 @@ mod tests {
                         }
                         None => {
                             assert!(
-                                !<_ as layers::Layer>::is_readable(&full, id).unwrap(),
+                                !<_ as repr::Representation>::is_readable(&full, id).unwrap(),
                                 "points not yielded by ExpandBeam should be unreadable"
                             );
 
                             assert!(
-                                <_ as layers::ExpandBeam>::evaluate(&expand, id)
+                                <_ as repr::ExpandBeam>::evaluate(&expand, id)
                                     .unwrap()
                                     .is_none(),
                                 "unreadable points should return `None` for their distance",
@@ -1299,7 +1300,7 @@ mod tests {
             }
 
             assert!(
-                <_ as layers::ExpandBeam>::evaluate(&expand, id_limit.value()).is_err(),
+                <_ as repr::ExpandBeam>::evaluate(&expand, id_limit.value()).is_err(),
                 "`ExpandBeam::evaluate` should catch out-of-bounds errors",
             );
 
@@ -1314,10 +1315,10 @@ mod tests {
         prune: &mut Prune<f32, TestDistance>,
         ids: &[u32],
     ) {
-        let mut items: HashMap<u32, Option<layers::PruneKey>> =
+        let mut items: HashMap<u32, Option<repr::PruneKey>> =
             ids.iter().map(|id| (*id, None)).collect();
 
-        let processed = <_ as layers::Prune>::prepare(prune, items.iter_mut()).unwrap();
+        let processed = <_ as repr::Prune>::prepare(prune, items.iter_mut()).unwrap();
         assert_eq!(processed, items.values().filter(|i| i.is_some()).count());
 
         // Ensure that `prepare` agrees with `points`.
@@ -1328,7 +1329,7 @@ mod tests {
             }
         }
 
-        fn filter((k, v): (&u32, &Option<layers::PruneKey>)) -> Option<(u32, layers::PruneKey)> {
+        fn filter((k, v): (&u32, &Option<repr::PruneKey>)) -> Option<(u32, repr::PruneKey)> {
             v.map(|v| (*k, v))
         }
 
@@ -1337,7 +1338,7 @@ mod tests {
             for (k1, v1) in items.iter().filter_map(filter) {
                 // Manually implement `TestDistance`.
                 let expected = points[&k0] + points[&k1];
-                let got = <_ as layers::Prune>::evaluate(prune, v0, v1);
+                let got = <_ as repr::Prune>::evaluate(prune, v0, v1);
                 assert_eq!(expected, got);
             }
         }
@@ -1351,16 +1352,16 @@ mod tests {
 
         let (full, mut points) = test_full(capacity);
 
-        assert_eq!(<_ as layers::Layer>::capacity(&full), capacity);
-        assert_eq!(<_ as layers::Layer>::id_limit(&full), id_limit);
+        assert_eq!(<_ as repr::Representation>::capacity(&full), capacity);
+        assert_eq!(<_ as repr::Representation>::id_limit(&full), id_limit);
 
         let mut available: HashSet<u32> = (0..capacity.value()).map(|i| i as u32).collect();
 
         // Insert the values 0 to 10.
         for i in 0u32..10 {
-            let guard = <_ as layers::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
+            let guard = <_ as repr::Set<&[f32]>>::set(&full, &[i as f32]).unwrap();
 
-            let id = <_ as layers::Guard>::id(&guard);
+            let id = <_ as repr::Guard>::id(&guard);
 
             assert!(
                 available.remove(&id),
@@ -1372,7 +1373,7 @@ mod tests {
                 "insertion should not repeat",
             );
 
-            <_ as layers::Guard>::publish(guard);
+            <_ as repr::Guard>::publish(guard);
         }
 
         // We do several things.
@@ -1382,21 +1383,21 @@ mod tests {
         //
         // 2. We publish two new points and immediately retire them.
         //    This tests that we correctly make these points unreadable.
-        let g0 = <_ as layers::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
-        let g1 = <_ as layers::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
-        let g2 = <_ as layers::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
-        let g3 = <_ as layers::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
+        let g0 = <_ as repr::Set<&[f32]>>::set(&full, &[1000.0]).unwrap();
+        let g1 = <_ as repr::Set<&[f32]>>::set(&full, &[2000.0]).unwrap();
+        let g2 = <_ as repr::Set<&[f32]>>::set(&full, &[3000.0]).unwrap();
+        let g3 = <_ as repr::Set<&[f32]>>::set(&full, &[4000.0]).unwrap();
 
         {
-            let g0_id = <_ as layers::Guard>::id(&g0);
-            <_ as layers::Guard>::publish(g0);
-            <_ as layers::Layer>::retire(&full, g0_id).unwrap();
+            let g0_id = <_ as repr::Guard>::id(&g0);
+            <_ as repr::Guard>::publish(g0);
+            <_ as repr::Representation>::retire(&full, g0_id).unwrap();
         }
 
         {
-            let g1_id = <_ as layers::Guard>::id(&g1);
-            <_ as layers::Guard>::publish(g1);
-            <_ as layers::Layer>::retire(&full, g1_id).unwrap();
+            let g1_id = <_ as repr::Guard>::id(&g1);
+            <_ as repr::Guard>::publish(g1);
+            <_ as repr::Representation>::retire(&full, g1_id).unwrap();
         }
 
         let mut prune = Prune::new(full.reader().unwrap(), TestDistance);
@@ -1432,7 +1433,7 @@ mod tests {
         let start_point = gen_vec::<T>(dim, &mut rng);
         let query = gen_vec::<T>(dim, &mut rng);
 
-        let full = <_ as layers::LayerConfig>::build(
+        let full = <_ as repr::RepresentationConfig>::build(
             Full::<T>::config(
                 Capacity::new(1),
                 MaxDegree::new(0),
@@ -1446,9 +1447,9 @@ mod tests {
         let start_id: u32 = 1;
 
         let internal_query = {
-            let guard = <_ as layers::Set<&[T]>>::set(&full, &query).unwrap();
-            let id = <_ as layers::Guard>::id(&guard);
-            <_ as layers::Guard>::publish(guard);
+            let guard = <_ as repr::Set<&[T]>>::set(&full, &query).unwrap();
+            let id = <_ as repr::Guard>::id(&guard);
+            <_ as repr::Guard>::publish(guard);
             id
         };
 
@@ -1469,7 +1470,7 @@ mod tests {
         // Prune
         {
             let mut prune = <T as FullPrecisionImpl>::make_prune(&full).unwrap();
-            let mut points: HashMap<u32, Option<layers::PruneKey>> =
+            let mut points: HashMap<u32, Option<repr::PruneKey>> =
                 [(internal_query, None), (start_id, None)]
                     .into_iter()
                     .collect();
diff --git a/diskann-inmem/src/layers/mod.rs b/diskann-inmem/src/repr/mod.rs
similarity index 89%
rename from diskann-inmem/src/layers/mod.rs
rename to diskann-inmem/src/repr/mod.rs
index 5169975f7f..c06cfe51d9 100644
--- a/diskann-inmem/src/layers/mod.rs
+++ b/diskann-inmem/src/repr/mod.rs
@@ -3,7 +3,7 @@
  * Licensed under the MIT license.
  */
 
-//! # Layering
+//! # Data Representation
 //!
 //! A simplified interface for [`crate::Provider`] to use for building a graph index.
 
@@ -20,30 +20,30 @@ use crate::{
 pub mod full;
 pub use full::{Full, FullPrecision};
 
-/// Deferred creation of [`Layer`]s.
+/// Deferred creation of [`Representation`]s.
 ///
 /// This is used in APIs like [`crate::Provider::new`] to defer allocation of large
 /// in-memory data structures.
-pub trait LayerConfig {
-    /// The type of the resulting [`Layer`].
-    type Layer: Layer;
+pub trait RepresentationConfig {
+    /// The type of the resulting [`Representation`].
+    type Representation: Representation;
 
-    /// Build the target [`Layer`].
-    fn build(self) -> ANNResult<Self::Layer>;
+    /// Build the target [`Representation`].
+    fn build(self) -> ANNResult<Self::Representation>;
 }
 
-/// Configurable data layer for [`crate::Provider`].
+/// Configurable data representation for [`crate::Provider`].
 ///
-/// Layers consist of the adjacency list and data for a concurrent in-memory graph index.
-/// These are expected to be indexed using `u32` IDs from `0..self.id_limit()`, with
-/// internal IDs in `0..self.capacity()` available for writing.
+/// A representation owns the adjacency lists and vector data for a concurrent in-memory graph
+/// index. Representations are expected to be indexed using `u32` IDs from
+/// `0..self.id_limit()`, with internal IDs in `0..self.capacity()` available for writing.
 ///
 /// See also:
 ///
 /// - [`Set`]: For assigning into the store.
 /// - [`Search`]: Search compatibility with [`crate::Provider`].
 /// - [`Insert`]: Insert compatibility with [`crate::Provider`].
-pub trait Layer: Send + Sync + 'static {
+pub trait Representation: Send + Sync + 'static {
     /// Return the [`MaxDegree`] of the internal graph.
     fn max_degree(&self) -> MaxDegree;
 
@@ -61,16 +61,16 @@ pub trait Layer: Send + Sync + 'static {
     fn is_readable(&self, i: u32) -> Option<bool>;
 }
 
-/// Attempt to write data into a [`Layer`].
+/// Attempt to write data into a [`Representation`].
 ///
 /// This will attempt to find an available internal ID to which `element` can be assigned,
 /// failing if no such ID can be found. The write is not eagerly committed. Instead, a
 /// [`Guard`] is returned, allowing writes to be aborted if necessary.
-pub trait Set<T>: Layer {
+pub trait Set<T>: Representation {
     /// The type of the [`Guard`] used to defer commitment of the write.
     type Guard<'a>: Guard;
 
-    /// Attempt to write the data in `element` into the [`Layer`].
+    /// Attempt to write the data in `element` into the [`Representation`].
     ///
     /// Returns [`Self::Guard`] to retrieve the allocated internal ID for `element` and to
     /// defer commitment of the write until external code is ready.
@@ -93,7 +93,7 @@ pub trait Guard {
     fn publish(self);
 }
 
-/// Enable search over vectors defined by a [`Layer`].
+/// Enable search over vectors defined by a [`Representation`].
 pub trait Search: Send + Sync + 'static {
     /// The type of the query. This should be equivalent to the generic parameter in
     /// [`Set`], but needs to be replicated here due to limitations in the current trait