pawkit-interner 0.1.9

#![feature(ptr_metadata, cold_path, layout_for_ptr)]

use core::{fmt, str};
use std::{
    alloc::{Layout, alloc, dealloc},
    fmt::{Debug, Display},
    hint::cold_path,
    ops::Deref,
    ptr::{self, NonNull, copy_nonoverlapping, slice_from_raw_parts_mut},
    sync::{
        LazyLock,
        atomic::{AtomicUsize, Ordering},
    },
};

use dashmap::{DashMap, Entry};

const MAX_REFCOUNT: usize = (isize::MAX) as usize;

struct InternInner {
    strong: AtomicUsize,
    weak: AtomicUsize,
    data: str,
}

/// Represents an interned string.
/// Similar to an Arc<str>, but it also deduplicates strings.
/// It is designed for fast forward lookup and cloning, with reverse lookup being prioritized over insertions / deletions.
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InternString {
    inner: NonNull<InternInner>,
}

/// A weak interned string.
/// Similar to Weak<str>, but references an InternString instead
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WeakInternString {
    inner: NonNull<InternInner>,
}

/// # SAFETY
/// The keys are casted to a static lifetime, and reference the underlying InternInner data,
/// but the keys are removed before the InternInner is freed.
static DATA: LazyLock<DashMap<&'static str, WeakInternString>> = LazyLock::new(Default::default);

impl InternInner {
    fn layout_for(len: usize) -> Layout {
        return Layout::new::<AtomicUsize>()
            .extend(Layout::new::<AtomicUsize>())
            .unwrap()
            .0
            .extend(Layout::array::<u8>(len).unwrap())
            .unwrap()
            .0
            .pad_to_align();
    }

    unsafe fn alloc(s: &str) -> (NonNull<Self>, &'static str) {
        unsafe {
            let layout = Self::layout_for(s.len());

            let ptr = slice_from_raw_parts_mut(alloc(layout), s.len()) as *mut Self;

            assert_eq!(layout, Layout::for_value_raw(ptr));

            let value = &mut *ptr;

            value.strong = AtomicUsize::new(1);
            value.weak = AtomicUsize::new(1);

            copy_nonoverlapping(s.as_ptr(), value.data.as_mut_ptr(), s.len());

            return (NonNull::new_unchecked(ptr), &value.data);
        }
    }

    unsafe fn dealloc(value: NonNull<Self>) {
        unsafe {
            let ptr = value.as_ptr();
            let layout = Layout::for_value_raw(ptr);

            dealloc(ptr as *mut u8, layout);
        }
    }
}

impl InternString {
    fn inner(&self) -> &InternInner {
        unsafe {
            return &*self.inner.as_ptr();
        }
    }

    /// Creates a new intern string from the provided string
    /// If an intern string exists with the same content, it returns a clone of that.
    pub fn new(s: &str) -> Self {
        if let Some(weak) = DATA.get(s) {
            let value = Self { inner: weak.inner };

            value.inner().strong.fetch_add(1, Ordering::Relaxed);

            return value;
        }

        let (ptr, str) = unsafe { InternInner::alloc(s) };

        match DATA.entry(str) {
            Entry::Occupied(occupied) => {
                cold_path();

                unsafe {
                    InternInner::dealloc(ptr);
                }

                let value = Self {
                    inner: occupied.get().inner,
                };

                value.inner().strong.fetch_add(1, Ordering::Relaxed);

                return value;
            }
            Entry::Vacant(vacant) => {
                vacant.insert(WeakInternString { inner: ptr });

                let value = Self { inner: ptr };

                return value;
            }
        }
    }

    /// Retruns the underlying string data.
    pub fn as_str<'a>(&'a self) -> &'a str {
        return &self.inner().data;
    }

    /// Converts into a weak reference
    pub fn into_weak(&self) -> WeakInternString {
        let weak = WeakInternString { inner: self.inner };

        self.inner().weak.fetch_add(1, Ordering::Relaxed);

        return weak;
    }
}

impl WeakInternString {
    fn inner(&self) -> &InternInner {
        unsafe {
            return &*self.inner.as_ptr();
        }
    }

    /// Returns true if the string is still alive
    /// The string is considered dead if there are no strong references
    pub fn is_alive(&self) -> bool {
        return self.inner().strong.load(Ordering::Acquire) != 0;
    }

    /// Gets the underlying string if it's still alive
    pub fn as_str<'a>(&'a self) -> Option<&'a str> {
        if !self.is_alive() {
            return None;
        }

        return Some(&self.inner().data);
    }

    /// Converts to a strong reference if the string is still alive
    pub fn into_strong(&self) -> Option<InternString> {
        #[inline]
        fn checked_increment(n: usize) -> Option<usize> {
            if n == 0 {
                return None;
            }
            assert!(n <= MAX_REFCOUNT);
            return Some(n + 1);
        }

        self.inner()
            .strong
            .fetch_update(Ordering::Acquire, Ordering::Relaxed, checked_increment)
            .ok()?;

        return Some(InternString { inner: self.inner });
    }
}

unsafe impl Send for InternString {}
unsafe impl Sync for InternString {}

unsafe impl Send for WeakInternString {}
unsafe impl Sync for WeakInternString {}

impl Drop for InternString {
    fn drop(&mut self) {
        if self.inner().strong.fetch_sub(1, Ordering::Release) != 1 {
            return;
        }

        self.inner().strong.load(Ordering::Acquire);

        DATA.remove(self.as_str());
    }
}

impl Drop for WeakInternString {
    fn drop(&mut self) {
        if self.inner().weak.fetch_sub(1, Ordering::Release) == 1 {
            self.inner().weak.load(Ordering::Acquire);

            unsafe {
                InternInner::dealloc(self.inner);
            }
        }
    }
}

impl Clone for InternString {
    fn clone(&self) -> Self {
        let ref_count = self.inner().strong.fetch_add(1, Ordering::Relaxed);

        assert!(ref_count < MAX_REFCOUNT);

        return Self {
            inner: self.inner.clone(),
        };
    }
}

impl Clone for WeakInternString {
    fn clone(&self) -> Self {
        let ref_count = self.inner().weak.fetch_add(1, Ordering::Relaxed);

        assert!(ref_count < MAX_REFCOUNT);

        return Self {
            inner: self.inner.clone(),
        };
    }
}

impl PartialEq<WeakInternString> for InternString {
    fn eq(&self, other: &WeakInternString) -> bool {
        return ptr::addr_eq(self.inner.as_ptr(), other.inner.as_ptr());
    }
}

impl PartialEq<InternString> for WeakInternString {
    fn eq(&self, other: &InternString) -> bool {
        return ptr::addr_eq(self.inner.as_ptr(), other.inner.as_ptr());
    }
}

impl PartialEq<&str> for InternString {
    fn eq(&self, other: &&str) -> bool {
        return self.as_str() == *other;
    }
}

impl PartialEq<&str> for WeakInternString {
    fn eq(&self, other: &&str) -> bool {
        let Some(str) = self.as_str() else {
            return false;
        };

        return str == *other;
    }
}

impl Deref for InternString {
    type Target = str;

    fn deref(&self) -> &Self::Target {
        return self.as_str();
    }
}

impl Debug for InternString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        return f.debug_tuple("InternString").field(&self.as_str()).finish();
    }
}

impl Display for InternString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        return f.write_str(self.as_str());
    }
}

impl Debug for WeakInternString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut t = f.debug_tuple("WeakInternString");

        let Some(s) = self.as_str() else {
            return t.field(&"<dead>").finish();
        };

        return t.field(&s).finish();
    }
}

impl Display for WeakInternString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let Some(s) = self.as_str() else {
            return f.write_str(&"<dead>");
        };

        return f.write_str(&s);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::{Arc, Barrier, Mutex};
    use std::thread;

    #[test]
    fn basic_interning() {
        let s1 = InternString::new("hello");
        let s2 = InternString::new("hello");
        let s3 = InternString::new("world");

        assert_eq!(s1, s2);
        assert_ne!(s1, s3);

        assert_eq!(s1.as_str(), "hello");
        assert_eq!(s3.as_str(), "world");
    }

    #[test]
    fn weak_references() {
        let s1 = InternString::new("weak_test");
        let w1 = s1.into_weak();

        assert!(w1.is_alive());
        assert_eq!(w1.as_str(), Some("weak_test"));

        drop(s1);

        assert!(!w1.is_alive());
        assert_eq!(w1.as_str(), None);
    }

    #[test]
    fn cloning_strong() {
        let s1 = InternString::new("clone_test");
        let s2 = s1.clone();
        let s3 = s2.clone();

        assert_eq!(s1.inner.as_ptr(), s2.inner.as_ptr());
        assert_eq!(s2.inner.as_ptr(), s3.inner.as_ptr());

        assert_eq!(s1.as_str(), "clone_test");
    }

    #[test]
    fn cloning_weak() {
        let s1 = InternString::new("weak_clone");
        let w1 = s1.into_weak();
        let w2 = w1.clone();

        assert_eq!(w1.inner.as_ptr(), w2.inner.as_ptr());

        drop(s1);

        assert!(!w1.is_alive());
        assert!(!w2.is_alive());
    }

    #[test]
    fn eq_partial() {
        let s = InternString::new("eq_test");
        let w = s.into_weak();

        assert_eq!(s, w);
        assert_eq!(w, s);

        assert_eq!(s, "eq_test");
        assert_eq!(w, "eq_test");
    }

    #[test]
    fn data_deduplication() {
        let s1 = InternString::new("dedup");
        let s2 = InternString::new("dedup");

        assert!(ptr::addr_eq(s1.inner.as_ptr(), s2.inner.as_ptr()));
    }

    #[test]
    fn drop_cleans_data() {
        let s = InternString::new("cleanup_test");

        drop(s);

        assert!(DATA.get("cleanup_test").is_none());
    }

    #[test]
    fn multithreaded_usage() {
        let s = Arc::new(InternString::new("thread_test"));

        let mut handles = vec![];

        for _ in 0..10 {
            let s_clone = Arc::clone(&s);
            handles.push(thread::spawn(move || {
                let local = s_clone.clone();
                assert_eq!(local.as_str(), "thread_test");
            }));
        }

        for handle in handles {
            handle.join().unwrap();
        }

        assert_eq!(s.as_str(), "thread_test");
    }

    #[test]
    fn weak_after_drop_multithreaded() {
        let s = InternString::new("weak_thread");
        let w = s.into_weak();

        let handle = thread::spawn(move || {
            drop(s);
        });

        handle.join().unwrap();

        assert!(!w.is_alive());
        assert_eq!(w.as_str(), None);
    }

    #[test]
    fn simultaneous_intern() {
        let barrier = Arc::new(Barrier::new(2));
        let s1 = Arc::new(Mutex::new(None));
        let s2 = Arc::new(Mutex::new(None));

        let b1 = barrier.clone();
        let s1c = s1.clone();
        let t1 = thread::spawn(move || {
            b1.wait();
            *s1c.lock().unwrap() = Some(InternString::new("race_test"));
        });

        let b2 = barrier.clone();
        let s2c = s2.clone();
        let t2 = thread::spawn(move || {
            b2.wait();
            *s2c.lock().unwrap() = Some(InternString::new("race_test"));
        });

        t1.join().unwrap();
        t2.join().unwrap();

        let s1 = s1.lock().unwrap().take().unwrap();
        let s2 = s2.lock().unwrap().take().unwrap();

        assert!(ptr::addr_eq(s1.inner.as_ptr(), s2.inner.as_ptr()));
    }

    #[test]
    fn weak_upgrade() {
        let s = InternString::new("upgrade_test");
        let w = s.into_weak();

        let s2 = w.into_strong().expect("Should upgrade");
        assert_eq!(s2.as_str(), "upgrade_test");

        drop(s);
        drop(s2);

        assert!(w.into_strong().is_none());
    }
}