kevy_map/

map.rs

//! The KevyMap implementation: struct, allocation, probing, and the live
//! lookup / insert / remove API. Helpers (`h2`, `prefetch_t0`, metadata
//! constants) and the private `ProbeOutcome` enum are all map-scoped.
//!
//! Layout (single allocation):
//!
//! ```text
//! +------------+------------+-----+---------------+---------+--------+
//! | slot[0]    | slot[1]    | ... | slot[cap-1]   | padding | meta   |
//! +------------+------------+-----+---------------+---------+--------+
//! ^                                                          ^
//! slots_ptr                                                  metadata_ptr
//! ```
//!
//! Both pointers are precomputed at `alloc_table` time and never re-derived
//! in the hot path. The single allocation cuts one alloc/dealloc pair vs
//! the previous two-`Box<[…]>` layout, and keeps metadata + slots in
//! adjacent pages (warmer TLB, contiguous OS-prefetch).

use core::alloc::Layout;
use core::fmt;
use core::marker::PhantomData;
use core::mem::MaybeUninit;
use core::ptr::{self, NonNull};

use kevy_hash::KevyHash;

use crate::iter::{Iter, IterMut, Keys, Values};

/// SIMD group width (16 metadata bytes loaded per probe iteration).
pub(crate) const GROUP_WIDTH: usize = 16;

/// Metadata byte for an empty slot (top bit set, value bits 1's — distinct
/// from DELETED so the probe loop can stop at EMPTY but skip DELETED).
pub(crate) const EMPTY: u8 = 0xFF;
/// Metadata byte for a tombstone (top bit set, value bits 0).
pub(crate) const DELETED: u8 = 0x80;
/// Minimum table size. ≥ 16 (one SSE2 group) so the future SIMD path can run
/// a full group scan unconditionally.
pub(crate) const MIN_CAP: usize = 16;

/// Top-7 bits of the hash, used as the per-slot metadata byte for occupied
/// slots. The top bit is always 0 (so occupancy = `meta & 0x80 == 0`).
#[inline]
pub(crate) fn h2(hash: u64) -> u8 {
    ((hash >> 57) & 0x7F) as u8
}

/// Issue a hint to fetch the cache line containing `ptr` into L1 ("T0" =
/// "all levels"). Stable on x86_64 / aarch64; no-op elsewhere AND under
/// `cfg(miri)` (miri cannot model inline asm / arch intrinsics, so the hint
/// degrades to a no-op for unsafe-correctness testing — the semantic
/// contract of `prefetch_t0` is "may do nothing", so this is sound).
///
/// `inline(always)` is the point: prefetch only helps when the load-latency
/// window the call hides is bigger than the call site itself. A non-inlined
/// call_site / ret pair (~10 ns out-of-order) wipes the benefit.
#[allow(clippy::inline_always)]
#[inline(always)]
fn prefetch_t0(ptr: *const u8) {
    #[cfg(all(target_arch = "x86_64", not(miri)))]
    {
        // SAFETY: _mm_prefetch reads no memory; any aligned/unaligned/
        // out-of-bounds pointer is permitted by the ISA.
        unsafe {
            core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0);
        }
    }
    #[cfg(all(target_arch = "aarch64", not(miri)))]
    {
        // SAFETY: prfm reads no memory; any pointer permitted.
        unsafe {
            core::arch::asm!(
                "prfm pldl1keep, [{p}]",
                p = in(reg) ptr,
                options(nostack, preserves_flags, readonly),
            );
        }
    }
    #[cfg(any(miri, not(any(target_arch = "x86_64", target_arch = "aarch64"))))]
    {
        let _ = ptr;
    }
}

/// Compute the single-buffer layout for a table of `cap` slots: returns the
/// combined `Layout` and the byte offset to the metadata array. Panics on
/// arithmetic overflow (only reachable for cap ≈ usize::MAX which would OOM
/// anyway).
///
/// Metadata size is `cap + GROUP_WIDTH` (hashbrown 0.15 layout): the first
/// `cap` bytes are the real per-slot metadata, the trailing `GROUP_WIDTH`
/// bytes mirror the leading ones so the branchless `set_meta` formula
/// `index2 = ((i - GW) & mask) + GW` always lands inside the buffer (for
/// `i = GROUP_WIDTH - 1` the formula evaluates to `cap + GROUP_WIDTH - 1`,
/// the very last byte). That last byte is written by `set_meta` but never
/// read by `Group::load` — SIMD loads from `group_start ∈ [0, cap)` reach
/// at most `cap + GROUP_WIDTH - 2`.
#[inline]
pub(crate) fn table_layout<KV>(cap: usize) -> (Layout, usize) {
    let slots = Layout::array::<MaybeUninit<KV>>(cap).expect("slots layout overflow");
    let meta = Layout::array::<u8>(cap + GROUP_WIDTH).expect("metadata layout overflow");
    let (combined, meta_offset) = slots.extend(meta).expect("layout extend overflow");
    (combined.pad_to_align(), meta_offset)
}

/// An open-addressing Swiss-style hashtable keyed by [`KevyHash`].
///
/// Power-of-two capacity (`mask = cap - 1`); 7/8 load factor; linear probing
/// over the metadata array; full slots' (K, V) live AoS in a parallel slot
/// array of `MaybeUninit<(K, V)>` co-allocated with the metadata.
///
/// When `cap == 0` both pointers are dangling and no allocation is held.
/// # Examples
///
/// Keys are byte strings or integers — whatever implements `KevyHash`.
/// There is deliberately no `&str` impl: the keyspace deals in bytes, and a
/// key that exists only as UTF-8 would need validating on every lookup.
///
/// ```
/// let mut m: kevy_map::KevyMap<Vec<u8>, u32> = kevy_map::KevyMap::new();
/// assert_eq!(m.insert(b"k".to_vec(), 1), None, "insert returns what it displaced");
/// assert_eq!(m.insert(b"k".to_vec(), 2), Some(1));
/// assert_eq!(m.get(b"k".as_slice()), Some(&2));
/// assert_eq!(m.remove(b"k".as_slice()), Some(2));
/// assert_eq!(m.get(b"k".as_slice()), None);
/// assert!(m.is_empty());
/// ```
pub struct KevyMap<K, V> {
    /// Slot array. `cap` initialised iff the corresponding metadata byte is
    /// in `0x00..=0x7F`. Dangling when `cap == 0`.
    pub(crate) slots_ptr: NonNull<MaybeUninit<(K, V)>>,
    /// Metadata array (`cap + GROUP_WIDTH` bytes; trailing
    /// `GROUP_WIDTH - 1` bytes mirror the leading ones for SIMD-safe
    /// wraparound loads — the hashbrown layout). Dangling when `cap == 0`.
    pub(crate) metadata_ptr: NonNull<u8>,
    /// Allocated slot count. `0` when no allocation is held.
    pub(crate) cap: usize,
    /// `cap - 1` when `cap > 0`; `0` when `cap == 0`.
    pub(crate) mask: usize,
    /// Live entries.
    pub(crate) occupied: usize,
    /// Tombstones (not yet reclaimed).
    pub(crate) deleted: usize,
    /// `true` when the table buffer was obtained from
    /// [`kevy_madvise::mmap_anon_aligned_2mb`] (large tables that wanted
    /// THP-aligned storage); `false` when it came from the global allocator.
    /// Drives the dispatch in [`Drop`] between `munmap_2mb` and `dealloc`.
    pub(crate) mmap_backed: bool,
    /// Marker so dropck and variance treat us as owning `(K, V)` like a
    /// `Box<[MaybeUninit<(K, V)>]>` would.
    pub(crate) _marker: PhantomData<(K, V)>,
}

// SAFETY: KevyMap owns its `(K, V)` entries (via the slot allocation). The
// `NonNull<...>` fields are conceptually `Box<[…]>` and inherit the same
// Send/Sync bounds: send-K + send-V ⇒ KevyMap is Send. Same for Sync.
unsafe impl<K: Send, V: Send> Send for KevyMap<K, V> {}
unsafe impl<K: Sync, V: Sync> Sync for KevyMap<K, V> {}

/// `(metadata, slots)` parallel-slice pair returned by [`KevyMap::as_slices`].
/// Aliased so the long `(&[u8], &[MaybeUninit<(K, V)>])` signature doesn't
/// trip clippy's `type_complexity` lint on a member-by-member basis.
type SlotSlices<'a, K, V> = (&'a [u8], &'a [MaybeUninit<(K, V)>]);

/// Mutable-slot variant of [`SlotSlices`], returned by
/// [`KevyMap::as_mut_slices`] for [`KevyMap::iter_mut`].
type SlotSlicesMut<'a, K, V> = (&'a [u8], &'a mut [MaybeUninit<(K, V)>]);

pub(crate) enum ProbeOutcome {
    Found(usize),
    NotFound { insert_at: usize, via_tombstone: bool },
}

impl<K, V> KevyMap<K, V> {
    /// Construct an empty map without allocating.
    /// # Examples
    ///
    /// ```
    /// // The key must implement `KevyHash`, which is deliberately a small
    /// // set: byte strings and integers, the things a KV engine keys on.
    /// let m: kevy_map::KevyMap<Vec<u8>, u32> = kevy_map::KevyMap::new();
    /// assert!(m.is_empty());
    /// assert_eq!(m.len(), 0);
    /// ```
    pub fn new() -> Self {
        Self {
            slots_ptr: NonNull::dangling(),
            metadata_ptr: NonNull::dangling(),
            cap: 0,
            mask: 0,
            occupied: 0,
            deleted: 0,
            mmap_backed: false,
            _marker: PhantomData,
        }
    }

    /// Construct a map sized to hold `cap_hint` entries without growing
    /// (accounting for the 7/8 load factor).
    /// # Examples
    ///
    /// ```
    /// // A hint, not a promise: capacity is rounded to the table's own
    /// // shape, and `len` still starts at zero.
    /// let m: kevy_map::KevyMap<u32, u32> = kevy_map::KevyMap::with_capacity(100);
    /// assert_eq!(m.len(), 0);
    /// ```
    pub fn with_capacity(cap_hint: usize) -> Self {
        if cap_hint == 0 {
            return Self::new();
        }
        // ceil(cap_hint * 8 / 7) → smallest table where cap_hint fits below 7/8.
        let needed = cap_hint.saturating_mul(8).div_ceil(7);
        let cap = needed.next_power_of_two().max(MIN_CAP);
        Self::alloc_table(cap)
    }

    // alloc_table + Drop live in `crate::alloc` so this file stays under
    // the 500-LOC house rule.

    /// Write `v` into metadata slot `i`, also updating the mirror byte
    /// at `cap + i` when `i < GROUP_WIDTH`. Every metadata mutation goes
    /// through this helper so the mirror stays consistent with the real
    /// metadata.
    ///
    /// Branchless: the formula `index2 = ((i - GW) & mask) + GW`
    /// (hashbrown 0.15's `set_ctrl`) yields the real mirror position
    /// `cap + i` when `i < GW`, and yields `i` itself when `i >= GW`.
    /// The second write is therefore either to the mirror byte or a
    /// duplicate write to the same real byte (a no-op). No branch.
    #[inline]
    pub(crate) fn set_meta(&mut self, i: usize, v: u8) {
        debug_assert!(i < self.cap);
        // SAFETY: i ∈ [0, cap); i2 ∈ [GROUP_WIDTH, cap + GROUP_WIDTH);
        // both in-bounds since metadata buffer length is cap + GROUP_WIDTH.
        let i2 = (i.wrapping_sub(GROUP_WIDTH) & self.mask) + GROUP_WIDTH;
        unsafe {
            *self.metadata_ptr.as_ptr().add(i) = v;
            *self.metadata_ptr.as_ptr().add(i2) = v;
        }
    }

    /// Live entry count.
    #[inline]
    /// # Examples
    ///
    /// ```
    /// let mut m = kevy_map::KevyMap::new();
    /// m.insert(1u32, "a");
    /// m.insert(1u32, "b");
    /// assert_eq!(m.len(), 1, "the second insert REPLACED the first");
    /// ```
    pub fn len(&self) -> usize {
        self.occupied
    }

    /// Whether the map has zero live entries.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.occupied == 0
    }

    /// Allocated slot count (NOT live entries).
    #[inline]
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Drop every live entry and reset the metadata. Keeps the allocation.
    /// # Examples
    ///
    /// ```
    /// let mut m = kevy_map::KevyMap::new();
    /// for i in 0..8u32 { m.insert(i, i); }
    /// m.clear();
    /// assert!(m.is_empty());
    /// assert_eq!(m.get(&0), None);
    /// ```
    pub fn clear(&mut self) {
        if self.cap == 0 {
            return;
        }
        if core::mem::needs_drop::<(K, V)>() {
            for i in 0..self.cap {
                // SAFETY: i < cap ⇒ metadata pointer in-bounds.
                let meta = unsafe { *self.metadata_ptr.as_ptr().add(i) };
                if meta & 0x80 == 0 {
                    // SAFETY: full slot ⇒ initialised.
                    unsafe {
                        ptr::drop_in_place(self.slots_ptr.as_ptr().add(i).cast::<(K, V)>());
                    }
                }
            }
        }
        // Reset entire metadata buffer (real range + mirror tail) in one memset.
        // SAFETY: metadata buffer is exactly cap + GROUP_WIDTH bytes wide.
        unsafe {
            ptr::write_bytes(self.metadata_ptr.as_ptr(), EMPTY, self.cap + GROUP_WIDTH);
        }
        self.occupied = 0;
        self.deleted = 0;
    }

    /// `(&K, &V)` over all live entries; order is unspecified.
    /// # Examples
    ///
    /// ```
    /// let mut m = kevy_map::KevyMap::new();
    /// for i in 0..4u32 { m.insert(i, i * 10); }
    /// // Unordered, like any hash table: sort before comparing.
    /// let mut got: Vec<_> = m.iter().map(|(k, v)| (*k, *v)).collect();
    /// got.sort_unstable();
    /// assert_eq!(got, vec![(0, 0), (1, 10), (2, 20), (3, 30)]);
    /// ```
    pub fn iter(&self) -> Iter<'_, K, V> {
        let (metadata, slots) = self.as_slices();
        Iter::new(metadata, slots)
    }

    /// `iter` that begins at bucket `start` (clamped to `capacity()`) and
    /// walks to the end. To sweep the full ring beginning at a random offset
    /// — the pattern the kevy-store eviction sampler uses — chain it with a
    /// second `iter_from_bucket(0)` and `take(start)`.
    pub fn iter_from_bucket(&self, start: usize) -> Iter<'_, K, V> {
        let (metadata, slots) = self.as_slices();
        Iter::with_start(metadata, slots, start)
    }

    /// `(&K, &mut V)` over all live entries; order is unspecified. Keys stay
    /// shared (mutating a key in place would corrupt its bucket).
    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        let (metadata, slots) = self.as_mut_slices();
        IterMut::new(metadata, slots)
    }

    /// `&K` over all live entries.
    /// # Examples
    ///
    /// ```
    /// let mut m = kevy_map::KevyMap::new();
    /// m.insert(b"a".to_vec(), 1u32);
    /// m.insert(b"b".to_vec(), 2);
    /// let mut ks: Vec<_> = m.keys().cloned().collect();
    /// ks.sort_unstable();
    /// assert_eq!(ks, vec![b"a".to_vec(), b"b".to_vec()]);
    /// ```
    pub fn keys(&self) -> Keys<'_, K, V> {
        Keys::new(self.iter())
    }

    /// `&V` over all live entries.
    pub fn values(&self) -> Values<'_, K, V> {
        Values::new(self.iter())
    }

    /// Borrow the metadata and slots as parallel slices of length `cap`.
    /// Used by [`KevyMap::iter`] (which only needs the real slot range,
    /// not the mirror tail). When `cap == 0` returns two empty slices —
    /// the dangling pointer is never dereferenced.
    #[inline]
    fn as_slices(&self) -> SlotSlices<'_, K, V> {
        if self.cap == 0 {
            return (&[], &[]);
        }
        // SAFETY: cap > 0 ⇒ both pointers are valid for `cap` reads; we hand
        // out shared borrows tied to `&self`'s lifetime, so the allocation
        // outlives the returned slices.
        unsafe {
            (
                core::slice::from_raw_parts(self.metadata_ptr.as_ptr(), self.cap),
                core::slice::from_raw_parts(self.slots_ptr.as_ptr(), self.cap),
            )
        }
    }

    /// [`KevyMap::as_slices`] with mutable slots (metadata stays shared —
    /// [`KevyMap::iter_mut`] only reads it). The disjoint borrows are sound:
    /// metadata and slots are separate allocations.
    #[inline]
    fn as_mut_slices(&mut self) -> SlotSlicesMut<'_, K, V> {
        if self.cap == 0 {
            return (&[], &mut []);
        }
        // SAFETY: cap > 0 ⇒ both pointers are valid; `&mut self` guarantees
        // exclusive access for the lifetime of the returned slices.
        unsafe {
            (
                core::slice::from_raw_parts(self.metadata_ptr.as_ptr(), self.cap),
                core::slice::from_raw_parts_mut(self.slots_ptr.as_ptr(), self.cap),
            )
        }
    }

    /// Hint the CPU to fetch the bucket cache line that a probe at `hash`
    /// would start at. The prefetch lever against the bucket-probe DRAM
    /// miss: the command-batch driver calls this for command N+1 while
    /// finishing command N, so by the time N+1 actually probes the
    /// metadata, the line is in L1.
    ///
    /// No-op when the table is empty. Cheap when not empty (a single
    /// `prefetcht0` on x86_64 / `prfm pldl1keep` on aarch64; a regular
    /// volatile load on other arches via [`std::intrinsics`] — but we
    /// only use stable intrinsics here, so non-x86/aarch64 architectures
    /// degrade to a no-op rather than a fake hint).
    ///
    /// `inline(always)` is the point — see the `prefetch_t0` rationale.
    #[allow(clippy::inline_always)]
    #[inline(always)]
    pub fn prefetch_for_hash(&self, hash: u64) {
        if self.cap == 0 {
            return;
        }
        let idx = (hash as usize) & self.mask;
        // SAFETY: idx < cap ≤ metadata length ⇒ pointer in-bounds; prefetch
        // reads never trap and never observe values.
        let ptr = unsafe { self.metadata_ptr.as_ptr().add(idx) };
        prefetch_t0(ptr);
    }

    /// 7/8 of the capacity — the inclusive max for `occupied + deleted`.
    #[inline]
    pub(crate) fn threshold(&self) -> usize {
        self.cap - (self.cap / 8)
    }
}

impl<K, V> Default for KevyMap<K, V> {
    fn default() -> Self {
        Self::new()
    }
}

/// `m[&q]` panics on missing key (matches `std::HashMap::Index` semantics).
impl<K, Q, V> core::ops::Index<&Q> for KevyMap<K, V>
where
    K: core::borrow::Borrow<Q>,
    Q: KevyHash + Eq + ?Sized,
{
    type Output = V;
    fn index(&self, key: &Q) -> &V {
        self.get(key).expect("no entry found for key")
    }
}

impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for KevyMap<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

impl<K: KevyHash + Eq, V> FromIterator<(K, V)> for KevyMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let iter = iter.into_iter();
        let mut m = match iter.size_hint() {
            (lo, Some(hi)) if hi <= lo.saturating_mul(2) => Self::with_capacity(hi),
            (lo, _) => Self::with_capacity(lo),
        };
        for (k, v) in iter {
            m.insert(k, v);
        }
        m
    }
}

impl<K: KevyHash + Eq, V> Extend<(K, V)> for KevyMap<K, V> {
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        for (k, v) in iter {
            self.insert(k, v);
        }
    }
}

#[cfg(test)]
#[path = "map_tests.rs"]
mod tests;