ijson 0.1.7

//! The heap interned-string representation (tag `String`).
//!
//! Strings too long to store inline are interned in a global, reference-counted
//! cache so that equal strings share one allocation and compare by pointer.
//! Interning uses `DashSet`, a concurrent hash-set, so many strings can be
//! interned at once without contention. The header is 8-aligned so the tag bits
//! stay free.

use std::alloc::{Layout, LayoutError};
use std::borrow::Borrow;
use std::cmp::Ordering;
use std::fmt::{self, Formatter};
use std::hash::Hash;
use std::ops::Deref;
use std::ptr::{copy_nonoverlapping, NonNull};
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};

use dashmap::{DashSet, SharedValue};
use lazy_static::lazy_static;

use crate::alloc::{alloc_infallible, dealloc_infallible};
use crate::string::IString;
use crate::thin::{ThinMut, ThinMutExt, ThinRef, ThinRefExt};
use crate::value::{
    string_cmp, string_debug, Destructured, DestructuredMut, DestructuredRef, IValue, ValueRepr,
    ValueType,
};

#[repr(C)]
#[repr(align(8))]
struct Header {
    rc: AtomicUsize,
    // We use 48 bits for the length and 16 bits for the shard index.
    len_lower: u32,
    len_upper: u16,
    shard_index: u16,
}

trait HeaderRef<'a>: ThinRefExt<'a, Header> {
    fn len(&self) -> usize {
        (u64::from(self.len_lower) | (u64::from(self.len_upper) << 32)) as usize
    }
    fn shard_index(&self) -> usize {
        self.shard_index as usize
    }
    fn str_ptr(&self) -> *const u8 {
        // Safety: pointers to the end of structs are allowed
        unsafe { self.ptr().add(1).cast() }
    }
    fn bytes(&self) -> &'a [u8] {
        // Safety: Header `len` must be accurate
        unsafe { std::slice::from_raw_parts(self.str_ptr(), self.len()) }
    }
    fn str(&self) -> &'a str {
        // Safety: UTF-8 enforced on construction
        unsafe { std::str::from_utf8_unchecked(self.bytes()) }
    }
}

trait HeaderMut<'a>: ThinMutExt<'a, Header> {
    fn str_ptr_mut(mut self) -> *mut u8 {
        // Safety: pointers to the end of structs are allowed
        unsafe { self.ptr_mut().add(1).cast() }
    }
}

impl<'a, T: ThinRefExt<'a, Header>> HeaderRef<'a> for T {}
impl<'a, T: ThinMutExt<'a, Header>> HeaderMut<'a> for T {}

lazy_static! {
    static ref STRING_CACHE: DashSet<WeakIString> = DashSet::new();
}

// Eagerly initialize the string cache during tests or when the
// `ctor` feature is enabled.
#[cfg(any(test, feature = "ctor"))]
#[ctor::ctor]
fn ctor_init_cache() {
    lazy_static::initialize(&STRING_CACHE);
}

pub(crate) fn init_cache() {
    lazy_static::initialize(&STRING_CACHE);
}

struct WeakIString {
    ptr: NonNull<Header>,
}

unsafe impl Send for WeakIString {}
unsafe impl Sync for WeakIString {}
impl PartialEq for WeakIString {
    fn eq(&self, other: &Self) -> bool {
        **self == **other
    }
}
impl Eq for WeakIString {}
impl Hash for WeakIString {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}

impl Deref for WeakIString {
    type Target = str;
    fn deref(&self) -> &str {
        self.borrow()
    }
}

impl Borrow<str> for WeakIString {
    fn borrow(&self) -> &str {
        self.header().str()
    }
}
impl WeakIString {
    fn header<'a>(&'a self) -> ThinRef<'a, Header> {
        // Safety: pointer is always valid
        unsafe { ThinRef::new(self.ptr) }
    }
    // Bumps the reference count and returns the (aligned) header pointer.
    fn upgrade(&self) -> NonNull<u8> {
        // Safety: `self.ptr` is a live header for as long as this `WeakIString` is in
        // the cache.
        unsafe {
            self.ptr.as_ref().rc.fetch_add(1, AtomicOrdering::Relaxed);
        }
        self.ptr.cast::<u8>()
    }
}

/// The heap interned-string representation of a JSON string.
///
/// Every operation on the interned allocation is an associated function here: the
/// `Header`-plus-bytes layout, allocation, interning, and reference-count management.
/// The only external entry point is [`InternedRepr::intern`] (called by
/// `IValue::new_string`); the rest are reached through this type's `ValueRepr` impl.
/// The global cache itself (`STRING_CACHE`, `WeakIString`) is a
/// module concern, not per-value state, so it stays free-standing above.
pub(crate) struct InternedRepr;

impl InternedRepr {
    /// The heap layout for an interned string of `len` UTF-8 bytes: the `Header`
    /// followed by the bytes, padded to the header's alignment.
    fn layout(len: usize) -> Result<Layout, LayoutError> {
        Ok(Layout::new::<Header>()
            .extend(Layout::array::<u8>(len)?)?
            .0
            .pad_to_align())
    }

    /// Allocates and initialises a header-plus-bytes block for `s` (reference count
    /// zero), returning the aligned header pointer.
    fn alloc(s: &str, shard_index: usize) -> NonNull<Header> {
        assert!((s.len() as u64) < (1 << 48));
        assert!(shard_index < (1 << 16));
        // Safety: freshly allocated to `Self::layout(s.len())` and written before any
        // read; `str_ptr_mut` points at the trailing byte array we just sized.
        unsafe {
            let ptr = alloc_infallible(Self::layout(s.len()).unwrap()).cast::<Header>();
            ptr.write(Header {
                len_lower: s.len() as u32,
                len_upper: ((s.len() as u64) >> 32) as u16,
                shard_index: shard_index as u16,
                rc: AtomicUsize::new(0),
            });
            let hd = ThinMut::new(ptr);
            copy_nonoverlapping(s.as_ptr(), hd.str_ptr_mut(), s.len());
            ptr
        }
    }

    /// Frees an interned allocation. Safety: `ptr` must be a live header allocated by
    /// [`alloc`](Self::alloc) with no outstanding references.
    unsafe fn dealloc(ptr: NonNull<Header>) {
        let hd = ThinRef::new(ptr);
        let layout = Self::layout(hd.len()).unwrap();
        dealloc_infallible(ptr.cast::<u8>(), layout);
    }

    /// Interns a string in the global cache, returning the aligned header pointer
    /// (with the reference count already bumped). This is how `IValue::new_string`
    /// stores a string that is too long to fit inline.
    pub(crate) fn intern(s: &str) -> NonNull<u8> {
        // A string this short belongs *inline*, not here. The whole reason `IString`
        // equality and hashing can be a bare pointer comparison (see `IString::eq`/`hash`,
        // and `ObjectRepr`'s hash table, which keys on the raw word) is that a given string
        // has exactly one representation: short ones inline, long ones interned. Intern a
        // short string as well and there are now two `IString`s spelling it that compare
        // unequal — objects would miss keys and could hold one twice. `IValue::new_string`
        // is the boundary that decides inline-vs-interned; nothing should reach past it to
        // intern a string that fits inline.
        debug_assert!(
            s.len() > crate::value::inline::string::CAPACITY,
            "a string that fits inline must not be interned: it would exist in two forms",
        );
        let cache = &*STRING_CACHE;
        let shard_index = cache.determine_map(s);

        // Safety: `determine_map` should only return valid shard indices
        let shard = unsafe { cache.shards().get_unchecked(shard_index) };
        let mut guard = shard.write();
        if let Some((k, _)) = guard.get_key_value(s) {
            k.upgrade()
        } else {
            let k = WeakIString {
                ptr: Self::alloc(s, shard_index),
            };
            let res = k.upgrade();
            guard.insert(k, SharedValue::new(()));
            res
        }
    }

    /// Views the interned header behind a tagged value pointer.
    ///
    /// Safety: `ptr` must be the aligned header pointer of a live interned string, and
    /// the returned reference must not outlive it.
    unsafe fn as_header<'a>(ptr: NonNull<u8>) -> ThinRef<'a, Header> {
        ThinRef::new(ptr.cast())
    }

    /// The UTF-8 bytes of an interned string.
    ///
    /// Safety: `ptr` must be the aligned header pointer of a live interned string; the
    /// returned slice borrows that allocation and must not outlive it.
    unsafe fn bytes<'a>(ptr: NonNull<u8>) -> &'a [u8] {
        Self::as_header(ptr).bytes()
    }

    /// Clones an interned string by bumping its reference count.
    ///
    /// Safety: `ptr` must be the aligned header pointer of a live interned string.
    unsafe fn bump_rc(ptr: NonNull<u8>) {
        Self::as_header(ptr)
            .rc
            .fetch_add(1, AtomicOrdering::Relaxed);
    }

    /// Releases a reference to an interned string, freeing it when the reference
    /// count reaches zero.
    ///
    /// Safety: `ptr` must be the aligned header pointer of a live interned string; the
    /// caller gives up one reference and must not use `ptr` afterwards unless it holds
    /// another.
    unsafe fn release(ptr: NonNull<u8>) {
        let hd = Self::as_header(ptr);

        // If the reference count is greater than 1, we can safely decrement it without
        // locking the string cache. `Release`: a fast-path decrementer never takes the
        // shard lock, so this is the only thing ordering its reads of the immutable bytes
        // before it gives up the reference. It pairs with the `Acquire` on the freeing
        // `fetch_sub` below, so the thread that eventually frees sees those reads completed.
        let mut rc = hd.rc.load(AtomicOrdering::Relaxed);
        while rc > 1 {
            match hd.rc.compare_exchange_weak(
                rc,
                rc - 1,
                AtomicOrdering::Release,
                AtomicOrdering::Relaxed,
            ) {
                Ok(_) => return,
                Err(new_rc) => rc = new_rc,
            }
        }

        // Slow path: we observed a reference count of 1, so we need to lock the string cache
        let cache = &*STRING_CACHE;
        // Safety: the number of shards is fixed
        let shard = cache.shards().get_unchecked(hd.shard_index());
        let mut guard = shard.write();
        // `Acquire`, to synchronize with the fast-path `Release` decrements above so that
        // every lock-free owner's reads of the bytes happen-before the free below. It does
        // not also need `Release`: the slow path and the free both run under this shard
        // lock, so a slow-path decrementer's own reads are already published to the freeing
        // thread by the lock's release/acquire (this also covers the resurrection case,
        // where a concurrent `intern` upgraded the entry and `fetch_sub` returns > 1).
        if hd.rc.fetch_sub(1, AtomicOrdering::Acquire) == 1 {
            // Reference count reached zero, free the string
            assert!(guard.remove(hd.str()).is_some());

            // Shrink the shard if it's mostly empty.
            // The second condition is necessary because `HashMap` sometimes
            // reports a capacity of zero even when it's still backed by an
            // allocation.
            if guard.len() * 3 < guard.capacity() || guard.is_empty() {
                guard.shrink_to_fit();
            }
            drop(guard);

            Self::dealloc(ptr.cast());
        }
    }
}

impl ValueRepr for InternedRepr {
    fn value_type(&self, _v: &IValue) -> ValueType {
        ValueType::String
    }
    unsafe fn partial_cmp(&self, a: &IValue, b: &IValue) -> Option<Ordering> {
        Some(string_cmp(a, b))
    }
    unsafe fn debug(&self, v: &IValue, f: &mut Formatter<'_>) -> fmt::Result {
        string_debug(v, f)
    }
    fn destructure(&self, v: IValue) -> Destructured {
        Destructured::String(IString(v))
    }
    unsafe fn destructure_ref<'a>(&self, v: &'a IValue) -> DestructuredRef<'a> {
        DestructuredRef::String(v.as_string_unchecked())
    }
    unsafe fn destructure_mut<'a>(&self, v: &'a mut IValue) -> DestructuredMut<'a> {
        DestructuredMut::String(v.as_string_unchecked_mut())
    }
    unsafe fn clone(&self, v: &IValue) -> IValue {
        Self::bump_rc(v.ptr());
        v.raw_copy()
    }
    unsafe fn drop(&self, v: &mut IValue) {
        Self::release(v.ptr());
    }
    // hash/eq use the defaults (pointer word / `raw_eq`): interning deduplicates,
    // so equal interned strings share one allocation and compare by pointer.
    /// The interned UTF-8 bytes. Safety: `v` must be a live interned string.
    unsafe fn as_bytes<'a>(&self, v: &'a IValue) -> Option<&'a [u8]> {
        Some(Self::bytes(v.ptr()))
    }
}