stringtheory/interning/

map.rs

#[cfg(not(feature = "loom"))]
use std::sync::{
    atomic::{
        AtomicUsize,
        Ordering::{AcqRel, Acquire},
    },
    Arc, Mutex,
};
use std::{collections::HashMap, num::NonZeroUsize, ptr::NonNull};

#[cfg(feature = "loom")]
use loom::sync::{
    atomic::{
        AtomicUsize,
        Ordering::{AcqRel, Acquire},
    },
    Arc, Mutex,
};

use super::{
    helpers::{layout_for_data, PackedLengthCapacity},
    InternedString, Interner,
};
use crate::interning::helpers::{aligned_string, ReclaimedEntries, ReclaimedEntry};

const HEADER_LEN: usize = std::mem::size_of::<EntryHeader>();
const HEADER_ALIGN: usize = std::mem::align_of::<EntryHeader>();

/// The minimum possible length of an entry.
///
/// For any entry in the interner, there is already an `EntryHeader` followed by the string data itself. In order to
/// ensure that entries can be written contiguously, we additionally ensure that the number of bytes we utilize for the
/// string data is aligned at least as much as `EntryHeader` itself.
///
/// This means that the minimum possible length of an entry, or the minimum number of bytes a valid entry could consume,
/// is the length of the header plus the alignment of the header.
const MINIMUM_ENTRY_LEN: usize = HEADER_LEN + HEADER_ALIGN;

#[derive(Debug)]
pub(crate) struct StringState {
    interner: Arc<Mutex<InternerState>>,
    header: NonNull<EntryHeader>,
}

impl StringState {
    #[inline]
    pub const fn as_str(&self) -> &str {
        // SAFETY: We ensure `self.header` is well-aligned and points to an initialized `EntryHeader` value when creating `StringState`.
        unsafe { get_entry_string(self.header) }
    }
}

impl PartialEq for StringState {
    fn eq(&self, other: &Self) -> bool {
        self.header == other.header
    }
}

impl Clone for StringState {
    fn clone(&self) -> Self {
        // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
        // and points to an initialized `EntryHeader` value.
        let header = unsafe { self.header.as_ref() };
        header.increment_active_refs();

        Self {
            interner: self.interner.clone(),
            header: self.header,
        }
    }
}

impl Drop for StringState {
    fn drop(&mut self) {
        // SAFETY: The caller that creates `StringState` is responsible for ensuring that `self.header` is well-aligned
        // and points to an initialized `EntryHeader` value.
        let header = unsafe { self.header.as_ref() };
        if header.decrement_active_refs() {
            // We decremented the reference count to zero, so try to mark this entry for reclamation.
            let mut interner = self.interner.lock().unwrap();
            interner.mark_for_reclamation(self.header);
        }
    }
}

/// Metadata about an interner entry.
///
/// `EntryHeader` represents the smallest amount of information about an interned entry that's needed to support both
/// lookup of existing interned strings, as well as the ability to reclaim space in the interner when an entry is no
/// longer in use.
struct EntryHeader {
    /// The number of active references to this entry.
    ///
    /// Only incremented by the interner itself, and decremented by `InternedString` when it's dropped.
    refs: AtomicUsize,

    /// Combined length/capacity of the entry, in terms of the string itself.
    ///
    /// Notably, this does _not_ include the length of the header itself. For example, an entry holding the string
    /// "hello, world!" has a string length of 13 bytes, but since we've to pad out to meet our alignment requirements
    /// for `EntryHeader`, we would end up with a capacity of 16 bytes. As such, `EntryHeader::len` would report `13`,
    /// while `EntryHeader::capacity` would report `16`. Likewise, `EntryHeader::entry_len` would report `40`,
    /// accounting for the string capacity (16) as well as the header length itself (24).
    ///
    /// As explained in the description of `PackedLengthCapacity`, this does mean strings can't be larger than ~4 GB on
    /// 64-bit platforms, which isn't a problem we've.
    len_cap: PackedLengthCapacity,
}

impl EntryHeader {
    /// Creates a tombstone entry with the given capacity.
    ///
    /// This is to allow for updating a region in the data buffer, which has been reclaimed, such that it's
    /// identifiable as being unused.
    fn tombstone(entry: ReclaimedEntry) -> Self {
        // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
        // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
        // string portion itself.
        let str_cap = EntryHeader::usable_from_reclaimed(entry);

        Self {
            refs: AtomicUsize::new(0),
            len_cap: PackedLengthCapacity::new(str_cap, 0),
        }
    }

    /// Creates a new entry for the given string.
    fn from_string(s: &str) -> Self {
        // We're dictating the necessary capacity here, which is the length of the string rounded to the nearest
        // multiple of the alignment of `EntryHeader`, which ensures that any subsequent entry will be properly aligned.
        let str_cap = aligned_string::<Self>(s);

        Self {
            refs: AtomicUsize::new(1),
            len_cap: PackedLengthCapacity::new(str_cap, s.len()),
        }
    }

    /// Creates a new entry for the given string, based on the given reclaimed entry.
    ///
    /// This maps the entry header to the underlying capacity of the given reclaimed entry, which is done in cases where
    /// a reclaimed entry is being used when interning a new string, and the reclaimed entry is larger than the string
    /// being interned, but not large enough that we could split the excess capacity into a new reclaimed entry.
    fn from_reclaimed_entry(mut entry: ReclaimedEntry, s: &str) -> (Self, Option<ReclaimedEntry>) {
        // The usable capacity for a reclaimed entry is the full capacity minus the size of `EntryHeader` itself, as
        // reclaimed entries represent the _entire_ region in the data buffer, but `EntryHeader` only cares about the
        // string portion itself.
        let entry_str_cap = EntryHeader::usable_from_reclaimed(entry);
        let required_str_cap = aligned_string::<Self>(s);

        // If the reclaimed entry has enough additional space beyond what we need for the string, we'll split it off and
        // return it for the caller to keep around in the reclaimed entries list.
        let remainder = entry_str_cap - required_str_cap;
        let (adjusted_str_cap, maybe_split_entry) = if remainder >= MINIMUM_ENTRY_LEN {
            let entry_len = EntryHeader::len_for(s);
            let split_entry = entry.split_off(entry_len);

            (entry_len - HEADER_LEN, Some(split_entry))
        } else {
            (entry_str_cap, None)
        };

        let header = Self {
            refs: AtomicUsize::new(1),
            len_cap: PackedLengthCapacity::new(adjusted_str_cap, s.len()),
        };

        (header, maybe_split_entry)
    }

    /// Returns the computed length of a complete entry, in bytes, for the given string.
    ///
    /// This includes the size of the entry header itself and the string data, when padded for alignment, and represents
    /// the number of bytes that would be consumed in the data buffer.
    const fn len_for(s: &str) -> usize {
        HEADER_LEN + aligned_string::<Self>(s)
    }

    /// Returns the usable capacity of a reclaimed entry, in bytes.
    ///
    /// Usable refers to the number of bytes in a reclaimed entry that could be used for string data, after accounting
    /// for the size of `EntryHeader` itself.
    const fn usable_from_reclaimed(entry: ReclaimedEntry) -> usize {
        entry.capacity() - HEADER_LEN
    }

    /// Returns the size of the string, in bytes, that this entry can hold.
    const fn capacity(&self) -> usize {
        self.len_cap.capacity()
    }

    /// Returns the size of the string, in bytes, that this entry _actually_ holds.
    const fn len(&self) -> usize {
        self.len_cap.len()
    }

    /// Returns the total length of the entry, in bytes.
    ///
    /// This includes the length of the header in addition to the string data.
    const fn entry_len(&self) -> usize {
        HEADER_LEN + self.capacity()
    }

    /// Returns `true` if this entry is currently referenced.
    fn is_active(&self) -> bool {
        self.refs.load(Acquire) != 0
    }

    /// Increments the active reference count by one.
    fn increment_active_refs(&self) {
        self.refs.fetch_add(1, AcqRel);
    }

    /// Decrements the active reference count by one.
    ///
    /// Returns `true` if the active reference count is zero _after_ calling this method.
    fn decrement_active_refs(&self) -> bool {
        self.refs.fetch_sub(1, AcqRel) == 1
    }
}

// SAFETY: We don't take references to the entry header pointer that outlast `StringState`, and the only modification we
// do to the entry header is through atomic operations, so it's safe to both send and share `StringState` between
// threads.
unsafe impl Send for StringState {}
unsafe impl Sync for StringState {}

#[derive(Debug)]
struct InternerStorage {
    // Direct pieces of our buffer allocation.
    ptr: NonNull<u8>,
    offset: usize,
    capacity: NonZeroUsize,

    // Length of all active entries, in bytes.
    //
    // This is equivalent to `self.offset` minus the total size of all reclaimed entries.
    len: usize,

    // Markers for entries that can be reused.
    reclaimed: ReclaimedEntries,
}

impl InternerStorage {
    fn with_capacity(capacity: NonZeroUsize) -> Self {
        assert!(
            capacity.get() <= isize::MAX as usize,
            "capacity would overflow isize::MAX, which violates layout constraints"
        );

        // Allocate our data buffer. This is the main backing allocation for all interned strings, and is well-aligned
        // for `EntryHeader`.
        //
        // SAFETY: `layout_for_data` ensures the layout is non-zero.
        let data_layout = layout_for_data::<EntryHeader>(capacity);
        let data_ptr = unsafe { std::alloc::alloc(data_layout) };
        let ptr = match NonNull::new(data_ptr) {
            Some(ptr) => ptr,
            None => std::alloc::handle_alloc_error(data_layout),
        };

        Self {
            ptr,
            offset: 0,
            capacity,
            len: 0,
            reclaimed: ReclaimedEntries::new(),
        }
    }

    #[cfg(test)]
    /// Returns the total number of unused bytes that are available for interning.
    fn available(&self) -> usize {
        self.capacity.get() - self.len
    }

    /// Returns the total number of bytes of contiguous, unoccupied space at the end of the data buffer.
    fn available_unoccupied(&self) -> usize {
        self.capacity.get() - self.offset
    }

    fn get_entry_ptr(&self, offset: usize) -> NonNull<EntryHeader> {
        debug_assert!(
            offset + MINIMUM_ENTRY_LEN <= self.capacity.get(),
            "offset would point to entry that cannot possibly avoid extending past end of data buffer"
        );

        // SAFETY: The caller is responsible for ensuring that `offset` is within the bounds of the data buffer, and
        // that `offset` is well-aligned for `EntryHeader`.
        let entry_ptr = unsafe { self.ptr.as_ptr().add(offset).cast::<EntryHeader>() };
        debug_assert!(entry_ptr.is_aligned(), "entry header pointer must be well-aligned");

        // SAFETY: `entry_ptr` is derived from `self.ptr`, which itself is `NonNull<u8>`, and the caller is responsible
        // for ensuring that `offset` is within the bounds of the data buffer, so we know `entry_ptr` is non-null.
        unsafe { NonNull::new_unchecked(entry_ptr) }
    }

    fn write_entry(&mut self, offset: usize, entry_header: EntryHeader, s: &str) -> NonNull<EntryHeader> {
        debug_assert_eq!(
            entry_header.len(),
            s.len(),
            "entry header length must match string length"
        );

        let entry_ptr = self.get_entry_ptr(offset);
        let entry_len = entry_header.entry_len();

        // Write the entry header.
        unsafe { entry_ptr.as_ptr().write(entry_header) };

        let s_buf = s.as_bytes();

        // Write the string.
        let entry_s_buf = unsafe {
            // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
            let entry_s_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
            std::slice::from_raw_parts_mut(entry_s_ptr, s_buf.len())
        };
        entry_s_buf.copy_from_slice(s_buf);

        // Update our internal statistics.
        self.len += entry_len;

        entry_ptr
    }

    fn write_to_unoccupied(&mut self, s: &str) -> (usize, NonNull<EntryHeader>) {
        let entry_header = EntryHeader::from_string(s);

        // Write the entry to the end of the data buffer.
        let entry_offset = self.offset;
        self.offset += entry_header.entry_len();

        (entry_offset, self.write_entry(entry_offset, entry_header, s))
    }

    fn write_to_reclaimed_entry(&mut self, entry: ReclaimedEntry, s: &str) -> (usize, NonNull<EntryHeader>) {
        let entry_offset = entry.offset();
        let (entry_header, maybe_split_entry) = EntryHeader::from_reclaimed_entry(entry, s);

        // If we had enough capacity in the reclaimed entry to hold this string _and_ potentially hold another entry, we
        // split it off and store that remainder entry.
        if let Some(split_entry) = maybe_split_entry {
            self.add_reclaimed(split_entry);
        }

        // Write the entry in place of the reclaimed entry.
        (entry_offset, self.write_entry(entry_offset, entry_header, s))
    }

    fn add_reclaimed(&mut self, entry: ReclaimedEntry) {
        // Reclamation is a two-step process: first, we have to actually keep track of the reclaimed entry, which
        // potentially involves merging adjacent reclaimed entries, and then once all of that has happened, we tombstone
        // the entry (whether merged or not).
        //
        // However, if the merged reclaimed entry immediately precedes any available capacity, we can skip tombstoning
        // it, since we can just wind back `offset` to reclaim the space.
        let merged_entry = self.reclaimed.insert(entry);
        if merged_entry.offset() + merged_entry.capacity() == self.offset {
            self.offset -= merged_entry.capacity();
            self.reclaimed.remove(&merged_entry);
            return;
        }

        self.clear_reclaimed_entry(merged_entry);
    }

    fn clear_reclaimed_entry(&mut self, entry: ReclaimedEntry) {
        let entry_ptr = self.get_entry_ptr(entry.offset);

        // Write the entry tombstone itself, which clears out the hash and sets the reference count to zero.
        //
        // SAFETY: We know that `entry_ptr` is valid for writes (reclaimed entries are, by definition, inactive regions
        // in the data buffer) and is well-aligned for `EntryHeader`.
        let tombstone = EntryHeader::tombstone(entry);
        let str_cap = tombstone.capacity();

        unsafe {
            entry_ptr.as_ptr().write(EntryHeader::tombstone(entry));
        }

        // Write a magic value to the entire string capacity for the entry. This ensures that there's a known repeating
        // value which, in the case of debugging issues, can be a signal that offsets/reclaimed entries are incorrect
        // and overlapping with active entries.
        //
        // SAFETY: Like above, the caller is responsible for ensuring that `offset` is within the bounds of the data
        // buffer, and that `offset + capacity` does not extend past the bounds of the data buffer.
        unsafe {
            // Take the entry pointer and add 1, which sets our pointer to right _after_ the header.
            let str_ptr = entry_ptr.as_ptr().add(1).cast::<u8>();
            let str_buf = std::slice::from_raw_parts_mut(str_ptr, str_cap);
            str_buf.fill(0x21);
        }
    }

    /// Returns the byte offset of the given entry header within the data buffer.
    fn offset_of(&self, header_ptr: NonNull<EntryHeader>) -> usize {
        // SAFETY: The caller is responsible for ensuring the entry header reference belongs to this interner. If that
        // is upheld, then we know that entry header belongs to our data buffer, and that the pointer to the entry
        // header is not less than the base pointer of the data buffer, ensuring the offset is non-negative.
        let entry_offset = unsafe {
            header_ptr
                .cast::<u8>()
                .as_ptr()
                .offset_from(self.ptr.as_ptr().cast_const())
        };
        debug_assert!(entry_offset >= 0, "entry offset must be non-negative");

        entry_offset as usize
    }

    fn mark_for_reclamation(&mut self, header_ptr: NonNull<EntryHeader>) {
        let entry_offset = self.offset_of(header_ptr);

        let header = unsafe { header_ptr.as_ref() };
        let entry_len = header.entry_len();

        let entry = ReclaimedEntry::new(entry_offset, entry_len);
        self.len -= entry.capacity();
        self.add_reclaimed(entry);
    }
}

impl Drop for InternerStorage {
    fn drop(&mut self) {
        // SAFETY: We allocated this buffer with the global allocator, and we're generating the same layout that was
        // used to allocate it in the first place.
        unsafe {
            std::alloc::dealloc(self.ptr.as_ptr(), layout_for_data::<EntryHeader>(self.capacity));
        }
    }
}

#[derive(Debug)]
struct InternerState {
    // Backing storage for the interned strings.
    storage: InternerStorage,

    // Active entries in the interner.
    entries: HashMap<&'static str, usize>,
}

impl InternerState {
    /// Creates a new `InternerState` with a pre-allocated buffer that has the given capacity.
    pub fn with_capacity(capacity: NonZeroUsize) -> Self {
        Self {
            storage: InternerStorage::with_capacity(capacity),
            entries: HashMap::new(),
        }
    }

    fn mark_for_reclamation(&mut self, header_ptr: NonNull<EntryHeader>) {
        // Reclamation must happen exactly once per entry, even though more than one thread can legitimately observe the
        // reference count drop to zero for the same slot.
        //
        // `StringState::drop` decrements the reference count _before_ taking this lock, so within that window the entry
        // can still be found -- and resurrected (`refs` 0 -> 1) -- by a concurrent `try_intern`. If that resurrecting
        // handle is then also dropped (`refs` 1 -> 0), two different droppers each saw a 1 -> 0 transition and will
        // each reach this method observing `refs == 0`. A bare `is_active()` check is therefore not enough on its own
        // to decide who reclaims: reclaiming on both would double-subtract `storage.len` and insert an overlapping
        // reclaimed entry, corrupting the data buffer.
        //
        // We guard against that with two checks, both performed under this lock:
        //
        //   1. `is_active()`: if a resurrecting handle is still live (`refs > 0`), we must NOT reclaim. The reference
        //      count cannot rise while we hold the lock -- cloning needs a live handle, and `try_intern` needs this
        //      lock -- so observing zero here means it stays zero through the reclamation below (no use-after-free).
        //   2. Map identity: reclaim only if the entries map still maps this entry's string to _this_ entry's offset.
        //      The first dropper to reclaim removes that mapping, so any later dropper of the same slot sees the string
        //      gone (or remapped to a newer entry) and skips -- making reclamation exactly-once.
        //
        // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
        // `EntryHeader` value that was acquired from this interner.
        let header = unsafe { header_ptr.as_ref() };
        if header.is_active() {
            return;
        }

        let entry_str = unsafe { get_entry_string(header_ptr) };
        let entry_offset = self.storage.offset_of(header_ptr);
        if self.entries.get(entry_str) == Some(&entry_offset) {
            // Remove the entry from the entries map first before we reclaim the entry, since doing so overwrites the
            // entry data in the data buffer.
            self.entries.remove(entry_str);
            self.storage.mark_for_reclamation(header_ptr);
        }
    }

    fn try_intern(&mut self, s: &str) -> Option<NonNull<EntryHeader>> {
        // We can only intern strings with a size that fits within a packed length/capacity value, so if `s` is larger
        // than that, we can't intern it, and there's existing entry we could have for it either.
        if s.len() > PackedLengthCapacity::maximum_value() {
            return None;
        }

        // Try and find an existing entry for this string.
        if let Some(entry_offset) = self.entries.get(s) {
            let entry_ptr = self.storage.get_entry_ptr(*entry_offset);

            let header = unsafe { entry_ptr.as_ref() };
            header.increment_active_refs();

            return Some(entry_ptr);
        }

        let required_cap = EntryHeader::len_for(s);

        // We didn't find an existing entry, so we're going to intern it.
        //
        // First, try and see if we have a reclaimed entry that can fit this string. If nothing suitable is found, or we
        // have no reclaimed entries, then we'll just try to fit it in the remaining capacity of our data buffer.
        let maybe_reclaimed_entry = self.storage.reclaimed.take_if(|entry| entry.capacity() >= required_cap);
        let (entry_offset, entry_ptr) = if let Some(reclaimed_entry) = maybe_reclaimed_entry {
            saluki_antithesis::sometimes!(true, "reclaimed interner slot reused for a new string");
            self.storage.write_to_reclaimed_entry(reclaimed_entry, s)
        } else if required_cap <= self.storage.available_unoccupied() {
            self.storage.write_to_unoccupied(s)
        } else {
            // We don't have enough space to intern this string at all, so we'll just return `None`.
            saluki_antithesis::sometimes!(true, "interner full — no capacity to intern");
            return None;
        };

        // SAFETY: Callers of `get_entry_string` are responsible for ensuring that the chosen lifetime of the string is
        // valid with respect to ensuring that the underlying entry lives long enough, and that the string data is valid
        // UTF-8 from as long as the reference is live.
        //
        // We're creating a `'static` reference here, which is generally frowned upon when the data is in fact not
        // `'static`, but this is safe because we never leak this lifetime outside of the interner, and we ensure that
        // we don't actually keep this reference around longer than the entry itself, as we only need a `'static`
        // reference to use it as the key in the entries map.
        let entry_str = unsafe { get_entry_string(entry_ptr) };
        self.entries.insert(entry_str, entry_offset);

        Some(entry_ptr)
    }
}

// SAFETY: We don't take references to the data buffer pointer that outlast `InternerState`, and all access to
// `InternerState` itself is mediated through a mutex, so we're safe to send it around and share it between threads.
unsafe impl Send for InternerState {}
unsafe impl Sync for InternerState {}

/// A string interner based on a single, fixed-size backing buffer with support for reclamation.
///
/// ## Overview
///
/// This interner uses a single, fixed-size backing buffer where interned strings are stored contiguously. This provides
/// bounded memory usage, and the interner won't allocate additional memory for new strings once the buffer is full.
/// Since interned strings aren't likely to need to live for the life of the program, the interner supports
/// reclamation. Once all references to an interned string have been dropped, the storage for that string is reclaimed
/// and can be used to hold new strings.
///
/// ## Storage layout
///
/// The backing buffer stores strings contiguously, with an entry "header" prepended to each string. The header contains
/// relevant data -- hash of the string, reference count, and length of the string -- needed to work with the entry
/// either when searching for existing entries or when using the entry itself.
///
/// The layout of an entry is as follows:
///
/// ```text
/// ┌───────────────────────── entry #1 ──────────────────────────┐ ┌─ entry #2 ─┐ ┌─ entry .. ─┐
/// ▼                                                             ▼ ▼            ▼ ▼            ▼
/// ┏━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━┯━━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓ ┏━━━━━━━━━━━━┓
/// ┃ str hash  │  ref cnt  │  str len  │ str data  │   padding   ┃ ┃   header   ┃ ┃   header   ┃
/// ┃ (8 bytes) │ (8 bytes) │ (8 bytes) │ (N bytes) │ (1-7 bytes) ┃ ┃  & string  ┃ ┃  & string  ┃
/// ┗━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━┷━━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛ ┗━━━━━━━━━━━━┛
/// ▲                                   ▲                           ▲
/// └────────── `EntryHeader` ──────────┘                           └── aligned for `EntryHeader`
///          (8 byte alignment)                                         via trailing padding
/// ```
///
/// The backing buffer is always aligned properly for `EntryHeader`, so that the first entry can be referenced
/// correctly. However, when appending additional entries to the buffer, we need to ensure that those entries also have
/// an aligned start for accessing the header. This is complicated due to the variable number of bytes for the string
/// data.
///
/// Alignment padding is added to the end of the entry to ensure that when appending the next entry, the start of the
/// entry is properly aligned for `EntryHeader`. In the worst case, up to 7 bytes could be added (and thus wasted) on
/// this alignment padding.
///
/// ## `InternedString`
///
/// The `InternedString` type is a handle to the entry header, and thus the string data, for an interned string. It's
/// designed to be small -- 8 bytes. -- and cheap to clone, as it contains an atomic reference to the entry header and a
/// reference to the interner that owns the string. It dereferences to the underlying string with relatively low
/// overhead: two pointer indirections.
///
/// When an `InternedString` is dropped, it decrements the reference count for the entry it points to. If the reference
/// count drops to zero, it will attempt to mark the entry for reclamation.
///
/// ## Reclamation
///
/// As we want to bound the memory used by the interner, but also not allow it to be filled up with strings that
/// eventually end up going entirely unused, we need a way to remove those unused strings so their underlying storage
/// can be used for new strings. This is where reclamation comes in.
///
/// When a string is interned, the entry header tracks how many active references there are to it. When that reference
/// count drops to zero, the last reference to the string attempts to mark the entry for reclamation. Assuming no other
/// reference has been taken out on the entry in the meantime, the entry gets added to a list of "reclaimed" entries.
///
/// Reclaimed entries are simple markers -- start and end position in the data buffer -- which are stored in a freelist.
/// When attempting to intern a new string, this freelist is searched to see if there's an entry large enough to fit the
/// new string, and if so, it's used.
///
/// Additionally, when entries are reclaimed, adjacent entries are merged together where possible. This helps to avoid
/// unnecessary fragmentation over time, although not as effectively as reconstructing the data buffer to re-pack
/// entries.
#[derive(Clone, Debug)]
pub struct GenericMapInterner {
    state: Arc<Mutex<InternerState>>,
}

impl GenericMapInterner {
    /// Creates a new `GenericMapInterner` with the given capacity.
    ///
    /// The given capacity will potentially be rounded up by a small number of bytes (up to 7) in order to ensure the
    /// backing buffer is properly aligned.
    pub fn new(capacity: NonZeroUsize) -> Self {
        Self {
            state: Arc::new(Mutex::new(InternerState::with_capacity(capacity))),
        }
    }
}

impl Interner for GenericMapInterner {
    fn is_empty(&self) -> bool {
        self.state.lock().unwrap().entries.is_empty()
    }

    fn len(&self) -> usize {
        self.state.lock().unwrap().entries.len()
    }

    fn len_bytes(&self) -> usize {
        self.state.lock().unwrap().storage.len
    }

    fn capacity_bytes(&self) -> usize {
        self.state.lock().unwrap().storage.capacity.get()
    }

    fn try_intern(&self, s: &str) -> Option<InternedString> {
        let header = {
            let mut state = self.state.lock().unwrap();
            state.try_intern(s)?
        };

        Some(InternedString::from(StringState {
            interner: Arc::clone(&self.state),
            header,
        }))
    }
}

#[inline]
const unsafe fn get_entry_string_parts(header_ptr: NonNull<EntryHeader>) -> (NonNull<u8>, usize) {
    // SAFETY: The caller is responsible for ensuring that `header_ptr` is well-aligned and points to an initialized
    // `EntryHeader` value.
    let header = header_ptr.as_ref();

    // Advance past the header and get the pointer to the string.
    //
    // SAFETY: We know that we're simply skipping over the header by advancing the pointer by one when it's still typed
    // as `*mut EntryHeader`.
    let s_ptr = header_ptr.add(1).cast::<u8>();
    (s_ptr, header.len())
}

#[inline]
const unsafe fn get_entry_string<'a>(header_ptr: NonNull<EntryHeader>) -> &'a str {
    let (s_ptr, s_len) = get_entry_string_parts(header_ptr);

    // SAFETY: We depend on `get_entry_string_parts` to give us a valid pointer and length for the string.
    std::str::from_utf8_unchecked(std::slice::from_raw_parts(s_ptr.as_ptr() as *const _, s_len))
}

#[cfg(test)]
mod tests {
    use std::ops::Deref as _;

    use prop::sample::Index;
    use proptest::{collection::vec as arb_vec, prelude::*};

    use super::*;
    use crate::interning::{test_support, InternedStringState};

    pub(super) fn create_interner(capacity: usize) -> GenericMapInterner {
        assert!(capacity > 0, "capacity must be greater than zero");
        GenericMapInterner::new(NonZeroUsize::new(capacity).unwrap())
    }

    fn available_len(interner: &GenericMapInterner) -> usize {
        interner.state.lock().unwrap().storage.available()
    }

    fn reclaimed_len(interner: &GenericMapInterner) -> usize {
        interner.state.lock().unwrap().storage.reclaimed.len()
    }

    fn first_reclaimed_entry(interner: &GenericMapInterner) -> ReclaimedEntry {
        interner.state.lock().unwrap().storage.reclaimed.first().unwrap()
    }

    const fn entry_len(s: &str) -> usize {
        EntryHeader::len_for(s)
    }

    pub(super) fn get_reclaimed_entry_for_string(s: &InternedString) -> ReclaimedEntry {
        let state = match &s.state {
            InternedStringState::GenericMap(state) => state,
            _ => panic!("unexpected string state"),
        };

        let ptr = state.interner.lock().unwrap().storage.ptr.as_ptr();
        let header = unsafe { state.header.as_ref() };
        let offset = unsafe { state.header.as_ptr().cast::<u8>().offset_from(ptr) as usize };
        ReclaimedEntry::new(offset, header.entry_len())
    }

    #[test]
    fn basic() {
        let interner = create_interner(1024);

        let s1 = interner.try_intern("hello").unwrap();
        let s2 = interner.try_intern("world").unwrap();
        let s3 = interner.try_intern("hello").unwrap();

        assert_eq!(s1.deref(), "hello");
        assert_eq!(s2.deref(), "world");
        assert_eq!(s3.deref(), "hello");

        // The pointers from the interned strings should be the same, but not between the interned string and a pointer
        // to an equivalent (but not interned) string:
        assert!(std::ptr::eq(s1.deref() as *const _, s3.deref() as *const _));

        let local_hello = "hello";
        assert!(!std::ptr::eq(s1.deref() as *const _, local_hello as *const _));
    }

    #[test]
    fn try_intern_without_capacity() {
        // Big enough to fit a single "hello world!" string, but not big enough to fit two.
        let interner = create_interner(48);

        let s1 = interner.try_intern("hello world!");
        assert!(s1.is_some());

        let s2 = interner.try_intern("hello, world");
        assert!(s2.is_none());
    }

    #[test]
    fn reclaim_after_dropped() {
        let interner = create_interner(1024);

        let s1 = interner.try_intern("hello world!").expect("should not fail to intern");
        let s1_entry_len = entry_len(&s1);

        assert_eq!(interner.len(), 1);
        assert_eq!(available_len(&interner), 1024 - s1_entry_len);
        assert_eq!(reclaimed_len(&interner), 0);

        // Drop the interned string, which should decrement the reference count to zero and then reclaim the entry.
        drop(s1);

        assert_eq!(interner.len(), 0);
        assert_eq!(available_len(&interner), 1024);
        assert_eq!(reclaimed_len(&interner), 0);
    }

    #[test]
    fn interns_to_reclaimed_entry_with_leftover() {
        // We want to intern a string initially, which takes up almost all of the capacity, and then drop it so it gets
        // reclaimed. After that, we'll intern a much smaller string which should lead to utilizing that reclaimed
        // entry, but only a part of it. Finally, we'll intern another new string.
        //
        // The point is to demonstrate that our reclamation logic is sound in terms of allowing reclaimed entries to be
        // split while the search/insertion logic is operating.
        let capacity = 256;
        let interner = create_interner(capacity);

        // We craft four strings such that the first two (`s_large` and `s_medium1`) will take up enough capacity that
        // `s_small` can't possibly be interned in the available capacity. We'll also craft `s_medium2` so it can fit
        // within the reclaimed entry for `s_large` but takes enough capacity that `s_small` cannot fit in the leftover
        // reclaimed entry that we split off.
        let s_large = "99 bottles of beer on the wall, 99 bottles of beer! take one down, pass it around, 98 bottles of beer on the wall!";
        let s_medium1 = "no act of kindness, no matter how small, is ever wasted";
        let s_medium2 = "if you want to go fast, go alone; if you want to go far, go together";
        let s_small = "are you there god? it's me, margaret";

        let phase1_available_capacity = capacity - entry_len(s_large) - entry_len(s_medium1);
        assert!(phase1_available_capacity < entry_len(s_small));
        assert!((entry_len(s_large) - entry_len(s_medium2)) < entry_len(s_small));
        assert!(entry_len(s_medium2) < entry_len(s_large));

        // Phase 1: intern our two larger strings.
        let s1 = interner.try_intern(s_large).expect("should not fail to intern");
        let s2 = interner.try_intern(s_medium1).expect("should not fail to intern");

        assert_eq!(interner.len(), 2);
        assert_eq!(available_len(&interner), phase1_available_capacity);
        assert_eq!(reclaimed_len(&interner), 0);

        // Phase 2: drop `s_large` so it gets reclaimed.
        drop(s1);
        assert_eq!(interner.len(), 1);
        assert_eq!(reclaimed_len(&interner), 1);

        // Phase 3: intern `s_medium2`, which should fit in the reclaimed entry for `s_large`. This should leave a
        // small, split off reclaimed entry.
        let s3 = interner.try_intern(s_medium2).expect("should not fail to intern");

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 1);

        // Phase 4: intern `s_small`, which should not fit in the leftover reclaimed entry from `s_large` _or_ the
        // available capacity.
        let s4 = interner.try_intern(s_small);
        assert_eq!(s4, None);

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 1);

        // And make sure we can still dereference the interned strings we _do_ have left:
        assert_eq!(s2.deref(), s_medium1);
        assert_eq!(s3.deref(), s_medium2);
    }

    #[test]
    fn has_reclaimed_entries_string_fits_exactly() {
        // The interner is large enough to fit two of the identically-sized strings, but not all three. We show that
        // when we drop an entry and it is reclaimed, a string of identical size should always be able to reuse that
        // reclaimed entry.
        const S1_VALUE: &str = "hello world!";
        const S2_VALUE: &str = "hello, world";
        const S3_VALUE: &str = "hello--world";

        let interner = create_interner(80);

        // Intern the first two strings, which should fit without issue.
        let s1 = interner.try_intern(S1_VALUE).expect("should not fail to intern");
        let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
        let _s2 = interner.try_intern(S2_VALUE).expect("should not fail to intern");

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 0);

        // Try to intern a third string, which should fail as we don't have the space.
        let s3 = interner.try_intern(S3_VALUE);
        assert_eq!(s3, None);

        // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
        drop(s1);

        assert_eq!(interner.len(), 1);
        assert_eq!(reclaimed_len(&interner), 1);

        let s1_reclaimed = first_reclaimed_entry(&interner);
        assert_eq!(s1_reclaimed_expected, s1_reclaimed);

        // Try again to intern the third string, which should now succeed and take over the reclaimed entry entirely
        // as the strings are identical in length.
        let _s3 = interner.try_intern(S3_VALUE).expect("should not fail to intern");

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 0);
    }

    #[test]
    fn reclaimed_entry_reuse_split_too_small() {
        // This situation is slightly contrived, but: we want to test that when there's a reclaimed entry of a certain
        // size, reusing that reclaimed entry won't lead to it being split if the resulting split entry would be too
        // "small": unable to hold another minimum-sized entry.
        //
        // We have to intern three strings to do this because we only track reclaimed entries when they're followed by
        // in-use entries, and the string we drop to create a reclaimed entry has to be big enough, but not _too_ big,
        // to hold the string we want to intern after it.
        let interner = create_interner(128);

        // Declare our strings to intern and just check some preconditions by hand.
        let s_one = "a horse, a horse, my kingdom for a horse!";
        let s_one_entry_len = entry_len(s_one);
        let s_two = "why hello there, beautiful";
        let s_two_entry_len = entry_len(s_two);
        let s_three = "real gs move in silence like lasagna";
        let s_three_entry_len = entry_len(s_three);

        assert!(s_one_entry_len <= interner.capacity_bytes());
        assert!(s_two_entry_len <= interner.capacity_bytes());
        assert!(s_one_entry_len + s_two_entry_len + s_three_entry_len > interner.capacity_bytes());
        assert!(s_one_entry_len > s_two_entry_len);
        assert!(s_three_entry_len > s_two_entry_len);
        assert!((s_one_entry_len - s_three_entry_len) < MINIMUM_ENTRY_LEN);

        // Intern the first two strings, which should fit without issue.
        let s1 = interner.try_intern(s_one).expect("should not fail to intern");
        let s1_reclaimed_expected = get_reclaimed_entry_for_string(&s1);
        let _s2 = interner.try_intern(s_two).expect("should not fail to intern");

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 0);

        // Try to intern the third string, which should fail as we don't have the space.
        let s3 = interner.try_intern(s_three);
        assert_eq!(s3, None);

        // Drop the first string, which should decrement the reference count to zero and then reclaim the entry.
        drop(s1);

        assert_eq!(interner.len(), 1);
        assert_eq!(reclaimed_len(&interner), 1);

        let s1_reclaimed = first_reclaimed_entry(&interner);
        assert_eq!(s1_reclaimed_expected, s1_reclaimed);

        // Try again to intern the third string, which should now succeed and take over the reclaimed entry, but since
        // the remainder of the reclaimed entry after taking the necessary capacity for `s_three` is not large enough
        // (`MINIMUM_ENTRY_LEN`), we shouldn't end up splitting the reclaimed entry, and instead, `s3` should consume
        // the entire reclaimed entry.
        let s3 = interner.try_intern(s_three).expect("should not fail to intern");
        let s3_reclaimed_expected = get_reclaimed_entry_for_string(&s3);

        assert_eq!(interner.len(), 2);
        assert_eq!(reclaimed_len(&interner), 0);
        assert_eq!(s1_reclaimed_expected, s3_reclaimed_expected);
    }

    #[test]
    fn reclaimed_entry_adjacent_to_spare_capacity() {
        let interner = create_interner(128);

        // Intern two smallish strings that fit without issue.
        let s1 = interner.try_intern("hello, world!").expect("should not fail to intern");
        let s2 = interner
            .try_intern("cheeeeeehooooo!")
            .expect("should not fail to intern");
        let s1_entry_len = entry_len(&s1);
        let s2_entry_len = entry_len(&s2);

        assert_eq!(reclaimed_len(&interner), 0);
        assert_eq!(available_len(&interner), 128 - s1_entry_len - s2_entry_len);

        drop(s2);
        assert_eq!(reclaimed_len(&interner), 0);
        assert_eq!(available_len(&interner), 128 - s1_entry_len);

        drop(s1);
        assert_eq!(reclaimed_len(&interner), 0);
        assert_eq!(available_len(&interner), 128);
    }

    #[test]
    fn len_bytes_reps_active_interned_entries() {
        let interner = create_interner(256);
        let mut active_interned_entries_len = 0;
        let string_1 = "hello world!";
        let string_2 = "hello again, world";
        let string_3 = "this is another string";

        // Intern a string.
        let s1 = interner.try_intern(string_1).expect("should not fail to intern");
        active_interned_entries_len += entry_len(string_1);
        assert_eq!(interner.len(), 1);
        assert_eq!(interner.len_bytes(), active_interned_entries_len);

        // Intern a second string.
        let _s2 = interner.try_intern(string_2).expect("should not fail to intern");
        active_interned_entries_len += entry_len(string_2);
        assert_eq!(interner.len(), 2);
        assert_eq!(interner.len_bytes(), active_interned_entries_len);

        // Drop the first string.
        drop(s1);
        active_interned_entries_len -= entry_len(string_1);
        assert_eq!(interner.len(), 1);
        assert_eq!(interner.len_bytes(), active_interned_entries_len);

        // Intern a new string.
        let _s3 = interner.try_intern(string_3).expect("should not fail to intern");
        active_interned_entries_len += entry_len(string_3);
        assert_eq!(interner.len(), 2);
        assert_eq!(interner.len_bytes(), active_interned_entries_len);
    }

    proptest! {
        #[test]
        #[cfg_attr(miri, ignore)]
        fn property_test_entry_count_accurate(
            strs in test_support::arb_alphanum_strings(1..=128, 16..=512),
            indices in arb_vec(any::<Index>(), 1..=1000),
        ) {
            // Create an interner with enough capacity to hold all of the strings we've generated. This is the maximum
            // string size multiplied by the number of strings we've generated... plus a little constant factor per
            // string to account for the entry header. The header size differs between interner implementations, which
            // is exactly why the interner construction (and thus this `property_test_`-named entry point) stays
            // per-implementation while the invariant check is shared.
            const ENTRY_SIZE: usize = 128 + HEADER_LEN;
            let interner = create_interner(ENTRY_SIZE * indices.len());

            test_support::assert_entry_count_matches_unique_strings(&interner, &strs, &indices)?;
        }
    }
}

#[cfg(all(test, feature = "loom"))]
mod loom_tests {
    use std::ops::Deref;

    use super::{
        tests::{create_interner, get_reclaimed_entry_for_string},
        *,
    };

    #[test]
    fn concurrent_drop_and_intern() {
        fn reclaimed_entries(interner: &GenericMapInterner) -> Vec<ReclaimedEntry> {
            interner
                .state
                .lock()
                .unwrap()
                .storage
                .reclaimed
                .iter()
                .copied()
                .collect()
        }

        fn do_reclaimed_entries_overlap(a: ReclaimedEntry, b: ReclaimedEntry) -> bool {
            let a_start = a.offset;
            let a_end = a.offset + a.capacity - 1;

            let b_start = b.offset;
            let b_end = b.offset + b.capacity - 1;

            (a_start <= b_start && b_start <= a_end) || (a_start <= b_end && b_end <= a_end)
        }

        const STRING_TO_INTERN: &str = "hello, world!";

        // This test is meant to explore the thread orderings when one thread is trying to drop (and thus reclaim) the
        // last active reference to an interned string, and another thread is trying to intern that very same string.
        //
        // We accept, as a caveat, that a possible outcome is that we intern the "new" string again, even though an
        // existing entry to that string may have existed in an alternative ordering.
        loom::model(|| {
            let interner = create_interner(1024);
            let t2_interner = interner.clone();

            // Intern the string from thread T1.
            let t1_interned_s = interner
                .try_intern(STRING_TO_INTERN)
                .expect("should not fail to intern");
            assert_eq!(t1_interned_s.deref(), STRING_TO_INTERN);
            let t1_reclaimed_entry = get_reclaimed_entry_for_string(&t1_interned_s);

            // Spawn thread T2, which tries to intern the same string and returns the handle to it.
            let t2_result = loom::thread::spawn(move || {
                let interned_s = t2_interner
                    .try_intern(STRING_TO_INTERN)
                    .expect("should not fail to intern");
                let reclaimed_entry = get_reclaimed_entry_for_string(&interned_s);

                (interned_s, reclaimed_entry)
            });

            drop(t1_interned_s);

            let (t2_interned_s, t2_reclaimed_entry) = t2_result.join().expect("should not fail to join T2");
            assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);

            // What we're checking for here is that either:
            // - there's no reclaimed entries (T2 found the existing entry for the string before T1 dropped it)
            // - there's a reclaimed entry (T2 didn't find the existing entry for the string before T1 marked it as
            //   inactive) but the reclaimed entry does _not_ overlap with the interned string from T2, meaning we
            //   didn't get confused and allow T2 to use an existing entry that T1 then later marked as reclaimed
            let reclaimed_entries = reclaimed_entries(&interner);
            assert!(reclaimed_entries.len() <= 1, "should have at most one reclaimed entry");

            if !reclaimed_entries.is_empty() {
                // If we do have a reclaimed entry, it needs to match exactly with only one of the interned strings.
                let is_t1_entry = reclaimed_entries.first().unwrap() == &t1_reclaimed_entry;
                let is_t2_entry = reclaimed_entries.first().unwrap() == &t2_reclaimed_entry;

                assert!(
                    (is_t1_entry || is_t2_entry) && !(is_t1_entry && is_t2_entry),
                    "should only match one interned string"
                );

                // Additionally, we ensure that the reclaimed entry does not overlap with the other interned string.
                assert!(
                    !do_reclaimed_entries_overlap(t1_reclaimed_entry, t2_reclaimed_entry),
                    "reclaimed entry should not overlap with remaining interned string"
                );
            }
        });
    }

    #[test]
    fn concurrent_resurrect_and_double_drop() {
        // Regression test for a double-reclamation race.
        //
        // Unlike `concurrent_drop_and_intern` -- where the second thread keeps its handle alive -- here the second
        // thread interns the string _and_ drops it within the race window. This models the dangerous interleaving:
        //
        //   1. T1 drops the last reference: `refs` 1 -> 0 (but T1 hasn't taken the interner lock yet).
        //   2. T2 calls `try_intern`, finds the still-mapped entry, and resurrects it: `refs` 0 -> 1.
        //   3. T2 drops its handle: `refs` 1 -> 0.
        //   4. T1 and T2 each take the lock and observe `refs == 0`, so BOTH mark the same slot for reclamation.
        //
        // A double reclamation either underflows `storage.len` (a panic under overflow-checks, which the loom test
        // profile enables) or inserts a spurious/overlapping reclaimed entry. After every handle to the single interned
        // string has been dropped, the only correct end state is a fully empty interner.
        const STRING_TO_INTERN: &str = "hello, world!";

        loom::model(|| {
            let interner = create_interner(1024);
            let t2_interner = interner.clone();

            // T1 interns the string and holds the only reference.
            let t1_interned_s = interner
                .try_intern(STRING_TO_INTERN)
                .expect("should not fail to intern");

            // T2 interns the same string and immediately drops it. The interesting orderings are those where this runs
            // while T1's drop has decremented the refcount to zero but not yet taken the interner lock.
            let t2 = loom::thread::spawn(move || {
                let t2_interned_s = t2_interner
                    .try_intern(STRING_TO_INTERN)
                    .expect("should not fail to intern");
                assert_eq!(t2_interned_s.deref(), STRING_TO_INTERN);
                drop(t2_interned_s);
            });

            // T1 drops its (originally last) reference.
            drop(t1_interned_s);

            t2.join().expect("should not fail to join T2");

            // Every handle has been dropped, so the slot must have been reclaimed exactly once: the interner is fully
            // empty, with consistent accounting and no leftover reclaimed entries. A double reclaim breaks at least one
            // of these (and typically panics first via the `storage.len` underflow).
            let state = interner.state.lock().unwrap();
            assert!(state.entries.is_empty(), "entries map should be empty");
            assert_eq!(
                state.storage.len, 0,
                "storage.len should be zero (it is double-subtracted on a double reclaim)"
            );
            assert_eq!(state.storage.offset, 0, "offset should wind back to zero");
            assert!(
                state.storage.reclaimed.is_empty(),
                "no reclaimed entries should remain (a double reclaim inserts a spurious/overlapping entry)"
            );
        });
    }
}