baryl 0.0.2

//! A fixed-capacity map, allocated out of a pool region.

use std::{alloc::Allocator, hash::Hash};

use zerocopy::FromZeros;

use crate::abi::{BVec, SandboxSafe, hash64};

/// A slot nothing has been written to. Zero, so zeroed bytes are an empty map.
pub const EMPTY: u32 = 0;
/// A slot holding a live key and value.
pub const VALID: u32 = 1;

const _: () = assert!(
    EMPTY == 0,
    "BMap relies on EMPTY=0 for zero-init empty-table semantics"
);

/// One bucket of a [`BMap`]: whether it is filled, and what is in it.
///
/// `#[repr(C)]`, so a separately built binary can walk the slot array directly.
/// Going through `insert` and `get` never needs one.
#[repr(C)]
#[derive(FromZeros)]
pub struct Slot<K, V> {
    /// [`EMPTY`] or [`VALID`].
    pub state: u32,
    // FIXME: free 32 bits -- a fingerprint of the hash here would cut probe misses to a u32 compare.
    _pad: u32,
    pub key: K,
    pub value: V,
}

unsafe impl<K: SandboxSafe, V: SandboxSafe> SandboxSafe for Slot<K, V> {}

/// An open-addressed map of fixed capacity, allocated once out of `A`.
///
/// Insert and look up; there is no remove, no resize and no overwrite. The
/// capacity is chosen at construction and a full map simply refuses further
/// keys. `#[repr(C)]` throughout, so a checkpoint can hold one and a separately
/// built binary can read it.
///
/// All-zeroes is a valid empty map, so freshly claimed pool bytes need no
/// construction pass.
///
/// # Examples
///
/// ```ignore
/// let heap: SbxAlloc = ctl.subs.alloc.ck_heap();
/// let Some(mut first_seen): Option<BMap<u64, u64, SbxAlloc>> =
///     BMap::try_with_capacity_in(heap, 1000)
/// else {
///     return;  // heap full
/// };
///
/// assert!(first_seen.insert(0xffff_ffff_8100_0000, 41));
/// // Second write of the same key is refused; the first value stands.
/// assert!(!first_seen.insert(0xffff_ffff_8100_0000, 99));
/// assert_eq!(first_seen.get(&0xffff_ffff_8100_0000), Some(&41));
/// assert_eq!(first_seen.len(), 1);
/// ```
#[repr(C)]
pub struct BMap<K, V, A: Allocator> {
    slots: BVec<Slot<K, V>, A>,
    len: u64,
}

// SAFETY: the slot array is a `BVec` over `A`, so every pointer inside targets
// region-managed memory.
unsafe impl<K: SandboxSafe, V: SandboxSafe, A: Allocator + SandboxSafe> SandboxSafe
    for BMap<K, V, A>
{
}

impl<K, V, A: Allocator> BMap<K, V, A> {
    /// How many keys are held. Never falls, because nothing is ever removed.
    pub fn len(&self) -> usize {
        self.len as usize
    }

    /// True until the first successful `insert`.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

impl<K: Eq + Hash + FromZeros, V: FromZeros, A: Allocator> BMap<K, V, A> {
    /// A map sized to hold `want` keys without crowding: room for `want / 0.7`,
    /// rounded up to a power of two, and never fewer than 8 slots.
    ///
    /// `None` when the allocator has no room. The map never grows past this, so
    /// ask for what the run will actually reach.
    pub fn try_with_capacity_in(alloc: A, want: usize) -> Option<Self> {
        let cap = (want.saturating_mul(10) / 7)
            .max(8)
            .checked_next_power_of_two()?;
        Some(Self {
            slots: BVec::try_new_in(alloc, cap)?,
            len: 0,
        })
    }

    /// Store `value` under `key` if the key is new. `true` when it was stored.
    ///
    /// First write wins: inserting a key that is already there leaves the
    /// original value alone and answers `false`. A full map answers `false`
    /// too, and the two are not told apart — watch `len()` across the call if
    /// you need to know which it was.
    pub fn insert(&mut self, key: K, value: V) -> bool {
        let cap = self.slots.len();
        if cap == 0 {
            return false;
        }
        let (mask, start) = (cap - 1, hash64(&key) as usize);
        for probe in 0..cap {
            let slot = &mut self.slots[(start + probe) & mask];
            match slot.state {
                EMPTY => {
                    *slot = Slot {
                        state: VALID,
                        _pad: 0,
                        key: key,
                        value: value,
                    };
                    self.len += 1;
                    return true;
                },
                _ if slot.key == key => return false,
                _ => {},
            }
        }
        false
    }

    /// The value stored under `key`, or `None`.
    pub fn get(&self, key: &K) -> Option<&V> {
        let cap = self.slots.len();
        if cap == 0 {
            return None;
        }
        let (mask, start) = (cap - 1, hash64(key) as usize);
        for probe in 0..cap {
            let slot = &self.slots[(start + probe) & mask];
            match slot.state {
                EMPTY => return None,
                _ if &slot.key == key => return Some(&slot.value),
                _ => {},
            }
        }
        None
    }
}