jiter 0.18.0

use std::sync::atomic::{AtomicPtr, Ordering};

use ahash::random_state::RandomState;
use pyo3::exceptions::{PyTypeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::{PyBool, PyString};

use crate::string_decoder::StringOutput;

#[derive(Debug, Clone, Copy, Default)]
pub enum StringCacheMode {
    #[default]
    All,
    Keys,
    None,
}

impl<'py> FromPyObject<'_, 'py> for StringCacheMode {
    type Error = PyErr;

    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult<StringCacheMode> {
        if let Ok(bool_mode) = ob.cast::<PyBool>() {
            Ok(bool_mode.is_true().into())
        } else if let Ok(str_mode) = ob.extract::<&str>() {
            match str_mode {
                "all" => Ok(Self::All),
                "keys" => Ok(Self::Keys),
                "none" => Ok(Self::None),
                _ => Err(PyValueError::new_err(
                    "Invalid string cache mode, should be `'all'`, '`keys`', `'none`' or a `bool`",
                )),
            }
        } else {
            Err(PyTypeError::new_err(
                "Invalid string cache mode, should be `'all'`, '`keys`', `'none`' or a `bool`",
            ))
        }
    }
}

impl From<bool> for StringCacheMode {
    fn from(mode: bool) -> Self {
        if mode { Self::All } else { Self::None }
    }
}

/// The string cache a parse works on, taken from the pool on the first cacheable string and
/// returned to it when the parse ends.
#[derive(Default)]
pub(crate) struct StringCacheGuard(Option<Box<PyStringCache>>);

impl Drop for StringCacheGuard {
    fn drop(&mut self) {
        if let Some(cache) = self.0.take() {
            return_string_cache(cache);
        }
    }
}

pub trait StringMaybeCache {
    fn get_key<'py>(
        py: Python<'py>,
        guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString>;

    fn get_value<'py>(
        py: Python<'py>,
        guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString> {
        Self::get_key(py, guard, string_output)
    }
}

/// # Safety
///
/// Caller must match the ascii_only flag to the string passed in.
#[inline]
unsafe fn guarded_py_string<'py>(
    py: Python<'py>,
    guard: &mut StringCacheGuard,
    string_output: &StringOutput<'_, '_>,
) -> Bound<'py, PyString> {
    let s = string_output.as_str();
    let ascii_only = string_output.ascii_only();
    if (2..64).contains(&s.len()) {
        let cache = guard.0.get_or_insert_with(|| take_pooled_cache().unwrap_or_default());
        unsafe { cache.get_or_insert(py, s, ascii_only) }
    } else {
        unsafe { pystring_fast_new_maybe_ascii(py, s, ascii_only) }
    }
}

pub struct StringCacheAll;

impl StringMaybeCache for StringCacheAll {
    fn get_key<'py>(
        py: Python<'py>,
        guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString> {
        // SAFETY: `StringOutput` guarantees that its ASCII flag matches its contents.
        unsafe { guarded_py_string(py, guard, &string_output) }
    }
}

pub struct StringCacheKeys;

impl StringMaybeCache for StringCacheKeys {
    fn get_key<'py>(
        py: Python<'py>,
        guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString> {
        // SAFETY: `StringOutput` guarantees that its ASCII flag matches its contents.
        unsafe { guarded_py_string(py, guard, &string_output) }
    }

    fn get_value<'py>(
        py: Python<'py>,
        _guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString> {
        // SAFETY: `StringOutput` guarantees that its ASCII flag matches its contents.
        unsafe { pystring_fast_new_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
    }
}

pub struct StringNoCache;

impl StringMaybeCache for StringNoCache {
    fn get_key<'py>(
        py: Python<'py>,
        _guard: &mut StringCacheGuard,
        string_output: StringOutput<'_, '_>,
    ) -> Bound<'py, PyString> {
        // SAFETY: `StringOutput` guarantees that its ASCII flag matches its contents.
        unsafe { pystring_fast_new_maybe_ascii(py, string_output.as_str(), string_output.ascii_only()) }
    }
}

/// Each cache is a quarter of a megabyte, so a pool that grew to a large thread count would hold
/// on to that memory for good; when every slot is taken, a returned cache is dropped instead.
const POOL_SLOTS: usize = 8;

/// The string caches no parse is using, parked one per slot as `Box::into_raw`. A parse swaps one
/// out, or builds its own if every slot is empty, and puts it back when it is done, so a cache is
/// only ever reachable from one parse at a time and nothing is held while Python code can run.
///
/// Taking a cache is a single exchange, which is what makes taking one per string affordable for
/// callers of `cached_py_string` that cannot hold onto one. Under the GIL parses never overlap and
/// only the first slot is ever used; on free-threaded builds overlapping parses take different
/// slots and never wait on each other. A slot is only ever written when it is empty, so putting a
/// cache back cannot displace another parse's.
static STRING_CACHE: [AtomicPtr<PyStringCache>; POOL_SLOTS] =
    [const { AtomicPtr::new(std::ptr::null_mut()) }; POOL_SLOTS];

/// Take a cache out of the pool, or `None` if every slot is empty.
fn take_pooled_cache() -> Option<Box<PyStringCache>> {
    for slot in &STRING_CACHE {
        let cache = slot.swap(std::ptr::null_mut(), Ordering::Acquire);
        if !cache.is_null() {
            // SAFETY: a non-null slot holds `Box::into_raw` of a cache no parse is using, and
            // swapping in null is the only way to reach it, so the box is now ours alone.
            return Some(unsafe { Box::from_raw(cache) });
        }
    }
    None
}

/// Put a cache back in the first empty slot, or drop it if the pool is full.
fn return_string_cache(cache: Box<PyStringCache>) {
    let cache = Box::into_raw(cache);
    for slot in &STRING_CACHE {
        if slot
            .compare_exchange(std::ptr::null_mut(), cache, Ordering::Release, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
    }
    // SAFETY: no slot took the pointer, so it is still ours, and it came from `Box::into_raw`.
    drop(unsafe { Box::from_raw(cache) });
}

/// Hand each pooled cache to `f`, one slot at a time. A cache a parse is using is not in the pool,
/// so it is left alone.
///
/// Each cache goes back into the slot it came from, and only that one slot is empty while `f` runs,
/// so a parse starting meanwhile still finds the other seven rather than building its own.
fn for_each_pooled_cache(mut f: impl FnMut(&mut PyStringCache)) {
    for slot in &STRING_CACHE {
        let cache = slot.swap(std::ptr::null_mut(), Ordering::Acquire);
        if cache.is_null() {
            continue;
        }
        // SAFETY: as in `take_pooled_cache`.
        let mut cache = unsafe { Box::from_raw(cache) };
        f(&mut cache);
        let cache = Box::into_raw(cache);
        if slot
            .compare_exchange(std::ptr::null_mut(), cache, Ordering::Release, Ordering::Relaxed)
            .is_err()
        {
            // a parse put its own cache in this slot while we held ours
            // SAFETY: no slot took the pointer, so it is still ours, and it came from `Box::into_raw`.
            drop(unsafe { Box::from_raw(cache) });
        }
    }
}

/// The number of entries in the string caches no parse is using.
pub fn cache_usage() -> usize {
    let mut usage = 0;
    for_each_pooled_cache(|cache| usage += cache.usage());
    usage
}

/// Clear the string caches no parse is using; a cache in use by a parse is left as it is.
pub fn cache_clear() {
    for_each_pooled_cache(PyStringCache::clear);
}

/// Create a cached Python `str` from a string slice
#[inline]
pub fn cached_py_string<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
    // SAFETY: the ASCII fast path is disabled.
    unsafe { cached_py_string_maybe_ascii(py, s, false) }
}

/// Create a cached Python `str` from a string slice.
///
/// # Safety
///
/// Caller must pass ascii-only string.
#[inline]
pub unsafe fn cached_py_string_ascii<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
    // SAFETY: the caller guarantees that `s` is ASCII only.
    unsafe { cached_py_string_maybe_ascii(py, s, true) }
}

/// # Safety
///
/// Caller must match the ascii_only flag to the string passed in.
unsafe fn cached_py_string_maybe_ascii<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
    // SAFETY: this function's caller guarantees that `ascii_only` matches `s`.
    unsafe {
        // from tests, 0 and 1 character strings are faster not cached
        if (2..64).contains(&s.len()) {
            let mut cache = take_pooled_cache().unwrap_or_default();
            let py_string = cache.get_or_insert(py, s, ascii_only);
            return_string_cache(cache);
            py_string
        } else {
            pystring_fast_new_maybe_ascii(py, s, ascii_only)
        }
    }
}

// Capacity should be a power of 2 so the compiler can convert `%` to a right shift below
// Using a smaller number here (e.g. 1024) seems to be faster in many cases than a larger number (like 65536)
const CAPACITY: usize = 16_384;
type Entry = Option<(u64, Py<PyString>)>;

/// This is a Fully associative cache with LRU replacement policy.
/// See https://en.wikipedia.org/wiki/Cache_placement_policies#Fully_associative_cache
#[derive(Debug)]
struct PyStringCache {
    entries: Box<[Entry; CAPACITY]>,
    hash_builder: RandomState,
}

const ARRAY_REPEAT_VALUE: Entry = None;

impl Default for PyStringCache {
    fn default() -> Self {
        Self {
            // Make sure we don't allocate a large array on the stack (e.g. using `Box::new([ARRAY_REPEAT_VALUE; CAPACITY])`)
            // to avoid potential stack overflows.
            // Note: we might want to use `Box::<[Entry; CAPACITY]>::new_zeroed()` if we ever bump the MSRV above 1.92,
            // given that `Entry` can be used as a niche optimization.
            entries: std::iter::repeat_with(|| ARRAY_REPEAT_VALUE)
                .take(CAPACITY)
                .collect::<Vec<_>>()
                .into_boxed_slice()
                .try_into()
                .unwrap(),
            hash_builder: RandomState::default(),
        }
    }
}

impl PyStringCache {
    /// Lookup the cache for an entry with the given string. If it exists, return it.
    /// If it is not set or has a different string, insert it and return it.
    ///
    /// # Safety
    ///
    /// `ascii_only` must only be set to `true` if the string is guaranteed to be ASCII only.
    unsafe fn get_or_insert<'py>(&mut self, py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
        let hash = self.hash_builder.hash_one(s);

        let hash_index = hash as usize % CAPACITY;

        let set_entry = |entry: &mut Entry| {
            // SAFETY: caller upholds invariant
            let py_str = unsafe { pystring_fast_new_maybe_ascii(py, s, ascii_only) };
            if let Some((_, old_py_str)) = entry.replace((hash, py_str.clone().unbind())) {
                // micro-optimization: bind the old entry before dropping it so that PyO3 can
                // fast-path the drop (Bound::drop is faster than Py::drop)
                drop(old_py_str.into_bound(py));
            }
            py_str
        };

        // we try up to 5 contiguous slots to find a match or an empty slot
        for index in hash_index..hash_index.wrapping_add(5) {
            if let Some(entry) = self.entries.get_mut(index) {
                if let Some((entry_hash, py_str_ob)) = entry {
                    // to avoid a string comparison, we first compare the hashes
                    if *entry_hash == hash {
                        // if the hashes match, we compare the strings to be absolutely sure - as a hashmap would do
                        if py_str_ob.bind(py) == s {
                            // the strings matched, return the cached string object
                            return py_str_ob.bind(py).to_owned();
                        }
                    }
                } else {
                    // we got to an empty entry, use it
                    return set_entry(entry);
                }
            } else {
                // we reached the end of entries, break
                break;
            }
        }
        // we tried all 5 slots (or got to the end of entries) without finding a match
        // or an empty slot, make this LRU by replacing the first entry
        let entry = self.entries.get_mut(hash_index).unwrap();
        set_entry(entry)
    }

    /// get the number of entries in the cache that are set
    fn usage(&self) -> usize {
        self.entries.iter().filter(|e| e.is_some()).count()
    }

    /// clear the cache by resetting all entries to `None`
    fn clear(&mut self) {
        self.entries.fill_with(|| None);
    }
}

/// Create a new Python `str` from a string slice, with a fast path for ASCII strings
///
/// # Safety
///
/// `ascii_only` must only be set to `true` if the string is guaranteed to be ASCII only.
unsafe fn pystring_fast_new_maybe_ascii<'py>(py: Python<'py>, s: &str, ascii_only: bool) -> Bound<'py, PyString> {
    if ascii_only {
        // SAFETY: caller upholds invariant
        unsafe { pystring_ascii_new(py, s) }
    } else {
        PyString::new(py, s)
    }
}

/// Faster creation of PyString from an ASCII string, inspired by
/// <https://github.com/ijl/orjson/blob/3.10.0/src/str/create.rs#L41>
///
/// # Safety
///
/// `s` must be ASCII only
pub unsafe fn pystring_ascii_new<'py>(py: Python<'py>, s: &str) -> Bound<'py, PyString> {
    unsafe {
        #[cfg(not(any(PyPy, GraalPy, Py_LIMITED_API)))]
        {
            if s.len() <= 1 {
                return PyString::new(py, s);
            }
            // SAFETY: `PyUnicode_New` returns a new owned reference or null. Converting it to a
            // `Bound` immediately ensures an allocation failure is handled before dereferencing it.
            let py_string = Bound::from_owned_ptr(py, pyo3::ffi::PyUnicode_New(s.len() as isize, 127));
            let ptr = py_string.as_ptr();
            // see https://github.com/pydantic/jiter/pull/72#discussion_r1545485907
            debug_assert_eq!(pyo3::ffi::PyUnicode_KIND(ptr), pyo3::ffi::PyUnicode_1BYTE_KIND);
            let data_ptr = pyo3::ffi::PyUnicode_DATA(ptr).cast();
            // SAFETY: the caller guarantees that `s` is ASCII, so `PyUnicode_New` allocated a
            // writable one-byte buffer containing `s.len()` bytes followed by a null terminator.
            core::ptr::copy_nonoverlapping(s.as_ptr(), data_ptr, s.len());
            core::ptr::write(data_ptr.add(s.len()), 0);
            py_string.cast_into_unchecked()
        }

        #[cfg(any(PyPy, GraalPy, Py_LIMITED_API))]
        {
            PyString::new(py, s)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The pool must hand back every cache it is given, keep only as many as it has slots, and be
    /// empty again once they are all taken out. Nothing here needs the interpreter: a cache that
    /// has held no strings owns no Python objects.
    #[test]
    fn pool_keeps_at_most_its_slots() {
        // more caches than there are slots, all live at once
        let taken: Vec<Box<PyStringCache>> = (0..POOL_SLOTS + 2)
            .map(|_| take_pooled_cache().unwrap_or_default())
            .collect();
        assert!(
            take_pooled_cache().is_none(),
            "the pool should be empty while they are out"
        );

        for cache in taken {
            return_string_cache(cache);
        }
        let held = STRING_CACHE
            .iter()
            .filter(|slot| !slot.load(Ordering::Relaxed).is_null())
            .count();
        assert_eq!(held, POOL_SLOTS, "the pool keeps a slotful and drops the rest");

        // and what it kept still works: each one interns a string and hands the same object back
        let mut again: Vec<Box<PyStringCache>> = (0..POOL_SLOTS).map(|_| take_pooled_cache().unwrap()).collect();
        assert!(take_pooled_cache().is_none(), "taking them all leaves it empty");
        Python::initialize();
        Python::attach(|py| {
            for (i, cache) in again.iter_mut().enumerate() {
                let text = format!("slot number {i}");
                // SAFETY: `text` is ASCII only.
                let first = unsafe { cache.get_or_insert(py, &text, true) };
                assert_eq!(first.to_str().unwrap(), text);
                // SAFETY: as above.
                let second = unsafe { cache.get_or_insert(py, &text, true) };
                assert!(first.is(&second), "a second lookup should hit the entry just made");
                assert_eq!(cache.usage(), 1);
            }
        });
        for cache in again {
            return_string_cache(cache);
        }
        assert_eq!(cache_usage(), POOL_SLOTS, "every cache came back holding its string");
        cache_clear();
        assert_eq!(cache_usage(), 0);
    }
}