cowstr 0.13.0

use core::mem::{align_of, size_of};
use std::alloc;
use std::alloc::Layout;
use std::ptr::{addr_of_mut, NonNull};

#[cfg(not(feature = "multithreaded"))]
use std::cell::Cell;

#[cfg(feature = "multithreaded")]
use std::sync::atomic::{fence, AtomicUsize, Ordering};

#[cfg(feature = "multithreaded")]
use threadcell::ThreadCell;

use crate::Error;

/// Wraps raw memory to a reference counted string.
///
/// Currently rust has no support for representing/reallocating DST's that contain a
/// header. To solve this problem we store the raw allocated *mut u8 pointing to our
/// allocation and cast it opportunistically into the desired data/header parts. This is safe
/// because `RcString` constructors ensure that objects are always properly initialized.
/// This approach fixes the stacked-borrows bugs miri reported in an earlier implementation.
///
/// A `RcString` must be explicitly deallocated otherwise it will leak.
#[derive(Debug, Copy, Clone)]
#[repr(transparent)]
pub(crate) struct RcString(NonNull<u8>);

unsafe impl Send for RcString {}
unsafe impl Sync for RcString {}

// ctor/dtor/allocation
impl RcString {
    // PLANNED: unsafe fn try_from_ptr(allocated: *mut u8, layout: Layout) -> Result

    /// Creates a `RcString` from a raw pointer and layout.
    /// Panics when the pointer is null.
    #[inline]
    unsafe fn from_ptr(allocated: *mut u8, layout: Layout) -> Self {
        if allocated.is_null() {
            alloc::handle_alloc_error(layout);
        }
        Self(NonNull::new_unchecked(allocated))
    }

    /// Allocated a new empty `RcString` with at least the given capacity.
    pub(crate) fn allocate(capacity: usize) -> Self {
        unsafe {
            let layout = Layout::from_size_align_unchecked(
                capacity
                    .checked_add(RCSTRING_SIZE)
                    .expect("Capacity overflow"),
                RCSTRING_ALIGN,
            );

            let allocated = Self::from_ptr(alloc::alloc(layout), layout);

            // Caveat: allocated is uninitialized, it becomes initialized below.
            // this is not a problem since we only access it via ptr.write()

            // TODO: #75402 ptr_as_uninit when stable
            let rcstring = allocated.0.as_ptr() as *mut RcStringHeader;

            #[cfg(feature = "multithreaded")]
            addr_of_mut!((*rcstring).refcount).write(AtomicUsize::new(1usize));
            #[cfg(not(feature = "multithreaded"))]
            addr_of_mut!((*rcstring).refcount).write(Cell::new(1usize));

            #[cfg(feature = "multithreaded")]
            addr_of_mut!((*rcstring).len).write(AtomicUsize::new(0));
            #[cfg(not(feature = "multithreaded"))]
            addr_of_mut!((*rcstring).len).write(Cell::new(0));

            addr_of_mut!((*rcstring).capacity).write(layout.size() - RCSTRING_SIZE);

            #[cfg(feature = "multithreaded")]
            addr_of_mut!((*rcstring).push_lock).write(ThreadCell::new_disowned(()));

            // everything safe and sound now
            allocated
        }
    }

    /// Grows an `RcString` to have free space for at least 'reserve' bytes.
    ///
    /// Panics: when the reference count is not one.
    pub(crate) fn grow(mut self, reserve: usize) -> Self {
        debug_assert_eq!(self.strong_count(), 1);
        let layout = self.current_layout();
        let min_new_size = (self.len_relaxed() + RCSTRING_SIZE)
            .checked_add(reserve)
            .expect("Capacity overflow");

        let new_size = DEFAULT_ALLOCATION_STRATEGY
            .grow(layout.size(), min_new_size)
            .expect("Capacity overflow");

        // really grow
        if new_size > layout.size() {
            self = unsafe {
                Self::from_ptr(alloc::realloc(self.0.as_ptr(), layout, new_size), layout)
            };
            self.header_mut().capacity = new_size - RCSTRING_SIZE;
        }

        self
    }

    /// Shrinks an `RcString` to at minimum `new_capacity` or its current length, whatever is
    /// larger.
    ///
    /// Panics: when the reference count is not one.
    pub(crate) fn shrink(mut self, new_capacity: usize) -> Self {
        debug_assert_eq!(self.strong_count(), 1);

        let layout = self.current_layout();

        let new_size = DEFAULT_ALLOCATION_STRATEGY
            .align((self.len_relaxed() + RCSTRING_SIZE).max(new_capacity + RCSTRING_SIZE));

        // really shrink
        if new_size < layout.size() {
            self = unsafe {
                Self::from_ptr(alloc::realloc(self.0.as_ptr(), layout, new_size), layout)
            };
            self.header_mut().capacity = new_size - RCSTRING_SIZE;
        }

        self
    }

    /// Creates a new `RcString` by copying from a &[u8] with some spare capacity.
    ///
    /// # Safety
    ///
    /// The source must be valid utf-8
    unsafe fn from_bytes_unchecked(source: &[u8], reserve: usize) -> Self {
        let mut allocated = Self::allocate(
            source
                .len()
                .checked_add(reserve)
                .expect("Capacity overflow"),
        );

        std::ptr::copy_nonoverlapping(source.as_ptr(), allocated.data_mut(), source.len());
        allocated.len_set_release(source.len());

        allocated
    }

    /// Creates a new `RcString` by copying from a &str with some spare capacity.
    #[inline]
    pub(crate) fn from_str(source: &str, reserve: usize) -> Self {
        // Safety: source is a valid utf-8 string
        unsafe { Self::from_bytes_unchecked(source.as_bytes(), reserve) }
    }

    /// Unconditionally deallocates a `RcString`.
    ///
    /// # Safety
    ///
    /// Must only be called with a valid `RcString` pointer whose refcount is zero.
    // The side effect here is that the memory becomes freed, which we can't easily test
    #[mutants::skip]
    pub(crate) unsafe fn dealloc(self) {
        debug_assert_eq!(self.strong_count(), 0);
        let layout = self.current_layout();
        alloc::dealloc(self.0.as_ptr(), layout);
    }

    /// Returns the current layout.
    #[inline]
    fn current_layout(&self) -> Layout {
        unsafe {
            // Safety: &self is a valid RcString
            Layout::from_size_align_unchecked(
                self.header().capacity + RCSTRING_SIZE,
                RCSTRING_ALIGN,
            )
        }
    }
}

// push etc
impl RcString {
    /// Appends the given `char` to the end of this `RcString`.
    ///
    /// # Panics
    ///
    /// There is not enough space available. The capacity has to be extended before calling
    /// this.
    #[inline]
    pub(crate) fn push(&mut self, ch: char) {
        let mut buf = [0u8; 4];
        let bytes = ch.encode_utf8(&mut buf).as_bytes();

        assert!(self.spare_capacity() >= bytes.len());

        unsafe {
            std::ptr::copy_nonoverlapping(
                bytes.as_ptr(),
                self.data_mut().add(self.header().len_relaxed()),
                bytes.len(),
            );
        }
        self.len_add_release(bytes.len());
    }

    /// Tries to append the given `char` to the end of this `RcString`.
    /// This locks the string for appending and will fail when there is not enough spare capacity.
    pub(crate) unsafe fn try_push(&self, ch: char) -> Result<(), Error> {
        let mut buf = [0u8; 4];
        let bytes = ch.encode_utf8(&mut buf).as_bytes();

        #[cfg(feature = "multithreaded")]
        if !self.header().push_lock.is_owned() {
            return Err(Error::ThreadOwnership);
        }
        if self.spare_capacity() >= bytes.len() {
            std::ptr::copy_nonoverlapping(
                bytes.as_ptr(),
                // can be safely cast to *mut here because we write into the uninitialized
                // locked part
                (self.data() as *mut u8).add(self.header().len_relaxed()),
                bytes.len(),
            );

            self.len_add_release(bytes.len());
            Ok(())
        } else {
            Err(Error::Capacity)
        }
    }

    /// Appends the given `str` to the end of this `RcString`.
    ///
    /// # Panics
    ///
    /// There is not enough space available. The capacity has to be extended before calling
    /// this.
    #[inline]
    pub(crate) fn push_str(&mut self, s: &str) {
        assert!(self.spare_capacity() >= s.len());

        unsafe {
            std::ptr::copy_nonoverlapping(
                s.as_ptr(),
                self.data_mut().add(self.header().len_relaxed()),
                s.len(),
            );
        }
        self.len_add_release(s.len());
    }

    /// Tries to appends the given `str` to the end of this `RcString`.
    /// This locks the string for appending and will fail when there is not enough spare capacity.
    pub(crate) unsafe fn try_push_str(&self, s: &str) -> Result<(), Error> {
        #[cfg(feature = "multithreaded")]
        if !self.header().push_lock.is_owned() {
            return Err(Error::ThreadOwnership);
        }
        if self.spare_capacity() >= s.len() {
            std::ptr::copy_nonoverlapping(
                s.as_ptr(),
                // can be safely cast to *mut here because we write into the uninitialized
                // locked part
                (self.data() as *mut u8).add(self.header().len_relaxed()),
                s.len(),
            );
            self.len_add_release(s.len());
            Ok(())
        } else {
            Err(Error::Capacity)
        }
    }

    #[cfg(feature = "multithreaded")]
    #[inline]
    pub(crate) fn acquire(&self) -> bool {
        self.header().push_lock.try_acquire_once()
    }

    #[cfg(feature = "multithreaded")]
    #[inline]
    pub(crate) fn release(&self) {
        self.header().push_lock.release();
    }
}

// conversions
impl RcString {
    /// Returns self as &str
    #[inline]
    pub(crate) fn as_str(&self) -> &str {
        unsafe {
            std::str::from_utf8_unchecked(std::slice::from_raw_parts(
                self.data(),
                self.len_acquire(),
            ))
        }
    }

    /// Returns self as &mut str
    /// Safety: must be called with refcount equal one.
    #[inline]
    pub(crate) unsafe fn as_mut_str(&mut self) -> &mut str {
        debug_assert_eq!(self.strong_count(), 1);
        std::str::from_utf8_unchecked_mut(std::slice::from_raw_parts_mut(
            self.data_mut(),
            self.len_acquire(),
        ))
    }
}

// accessors
impl RcString {
    #[inline(always)]
    fn header(&self) -> &RcStringHeader {
        unsafe { std::mem::transmute::<_, &RcStringHeader>(self.0) }
    }

    #[inline(always)]
    fn header_mut(&mut self) -> &mut RcStringHeader {
        unsafe { std::mem::transmute::<_, &mut RcStringHeader>(self.0) }
    }

    #[inline(always)]
    fn data(&self) -> *const u8 {
        unsafe { self.0.as_ptr().add(RCSTRING_SIZE) }
    }

    #[inline(always)]
    fn data_mut(&mut self) -> *mut u8 {
        unsafe { self.0.as_ptr().add(RCSTRING_SIZE) }
    }

    #[inline(always)]
    pub(crate) fn spare_capacity(&self) -> usize {
        self.header().spare_capacity()
    }

    #[inline(always)]
    pub(crate) fn capacity(&self) -> usize {
        self.header().capacity()
    }

    #[inline(always)]
    pub(crate) fn strong_count(&self) -> usize {
        self.header().strong_count()
    }

    #[inline(always)]
    pub(crate) fn increment_strong_count(&self) {
        self.header().increment_strong_count();
    }

    #[inline(always)]
    #[must_use = "Would leak when 'true' is ignored"]
    pub(crate) fn decrement_strong_count(&self) -> bool {
        self.header().decrement_strong_count()
    }

    /// Get len from atomic or Cell with relaxed semantic.
    #[inline(always)]
    #[mutants::skip]
    fn len_relaxed(&self) -> usize {
        self.header().len_relaxed()
    }

    /// Get len from atomic or Cell with Acquire semantic.
    #[inline(always)]
    fn len_acquire(&self) -> usize {
        self.header().len_acquire()
    }

    /// Add 'n' to the length with Release semantic.
    #[inline(always)]
    fn len_add_release(&self, n: usize) {
        self.header().len_add_release(n);
    }

    /// Sets a new length with Release semantic.
    #[inline(always)]
    fn len_set_release(&self, n: usize) {
        self.header().len_set_release(n);
    }
}

#[cfg(feature = "multithreaded")]
#[derive(Debug)]
#[repr(C, align(32))]
pub(crate) struct RcStringHeader {
    refcount: AtomicUsize,
    len: AtomicUsize,
    capacity: usize,
    push_lock: ThreadCell<()>, // protect data and len on try_push
}

#[cfg(not(feature = "multithreaded"))]
#[derive(Debug)]
#[repr(C, align(32))]
pub(crate) struct RcStringHeader {
    refcount: Cell<usize>,
    len: Cell<usize>,
    capacity: usize,
}

const RCSTRING_SIZE: usize = size_of::<RcStringHeader>();
const RCSTRING_ALIGN: usize = align_of::<RcStringHeader>();

impl RcStringHeader {
    /// Get len from atomic or Cell with relaxed semantic.
    #[cfg(feature = "multithreaded")]
    #[inline(always)]
    fn len_relaxed(&self) -> usize {
        self.len.load(Ordering::Relaxed)
    }

    #[doc(hidden)]
    #[cfg(not(feature = "multithreaded"))]
    #[mutants::skip]
    #[inline(always)]
    fn len_relaxed(&self) -> usize {
        self.len.get()
    }

    /// Get len from atomic or Cell with Acquire semantic.
    #[cfg(feature = "multithreaded")]
    #[inline(always)]
    fn len_acquire(&self) -> usize {
        self.len.load(Ordering::Acquire)
    }

    #[doc(hidden)]
    #[inline(always)]
    #[cfg(not(feature = "multithreaded"))]
    #[mutants::skip]
    fn len_acquire(&self) -> usize {
        self.len.get()
    }

    /// Add 'n' to the length with Release semantic.
    #[inline(always)]
    fn len_add_release(&self, n: usize) {
        #[cfg(feature = "multithreaded")]
        self.len.fetch_add(n, Ordering::Release);
        #[cfg(not(feature = "multithreaded"))]
        self.len.set(self.len.get() + n);
    }

    /// Sets a new length with Release semantic.
    #[inline(always)]
    fn len_set_release(&self, n: usize) {
        #[cfg(feature = "multithreaded")]
        self.len.store(n, Ordering::Release);
        #[cfg(not(feature = "multithreaded"))]
        self.len.set(n);
    }

    /// Returns the amount of bytes that can be pushed without reallocation.
    /// This is an acquire fenced load allowing any further access to be relaxed.
    #[inline]
    pub(crate) fn spare_capacity(&self) -> usize {
        self.capacity - self.len_acquire()
    }

    /// Returns the capacity of the allocation.
    #[inline]
    pub(crate) fn capacity(&self) -> usize {
        self.capacity
    }

    #[inline]
    #[cfg(feature = "multithreaded")]
    pub(crate) fn strong_count(&self) -> usize {
        self.refcount.load(Ordering::Relaxed)
    }

    #[inline]
    #[cfg(not(feature = "multithreaded"))]
    #[mutants::skip]
    pub(crate) fn strong_count(&self) -> usize {
        self.refcount.get()
    }

    #[cfg(feature = "multithreaded")]
    pub(crate) fn increment_strong_count(&self) {
        self.refcount.fetch_add(1, Ordering::Relaxed);
    }

    #[cfg(not(feature = "multithreaded"))]
    #[mutants::skip]
    pub(crate) fn increment_strong_count(&self) {
        self.refcount.set(self.refcount.get() + 1);
    }

    #[cfg(not(feature = "multithreaded"))]
    #[mutants::skip]
    #[must_use = "Would leak when 'true' is ignored"]
    pub(crate) fn decrement_strong_count(&self) -> bool {
        let count = self.refcount.get();
        self.refcount.set(count - 1);
        count == 1
    }

    #[cfg(feature = "multithreaded")]
    #[must_use = "Would leak when 'true' is ignored"]
    pub(crate) fn decrement_strong_count(&self) -> bool {
        if self.refcount.fetch_sub(1, Ordering::Release) == 1 {
            fence(Ordering::Acquire);
            true
        } else {
            false
        }
    }
}

const KB: usize = 1024;
const MB: usize = KB * 1024;
const GB: usize = MB * 1024;

/// The allocation strategy when resizing an allocation
struct AllocationStrategy {
    /// Minimum allocation size returned
    pub min_allocation: usize,
    /// first cut, below this the allocations will be doubled, above until grow_const they will be increased by 50%
    pub grow_half: usize,
    /// second cut, above this allocation will grow constantly by grow_const/2
    pub grow_const: usize,
}

const DEFAULT_ALLOCATION_STRATEGY: AllocationStrategy = AllocationStrategy {
    min_allocation: 32,
    grow_half: 8 * MB,
    grow_const: /* 1 */ GB,
};

impl AllocationStrategy {
    /// returns the new size that shall be allocated, either per allocation strategy or when
    /// that is not sufficient, the `minimum_needed` size rounded up to a multiple of `min_allocation`.
    #[inline]
    pub fn grow(&self, old_size: usize, minimum_needed: usize) -> Option<usize> {
        Some(
            self.align(
                if old_size < self.min_allocation {
                    self.min_allocation
                } else if old_size < self.grow_half {
                    old_size.checked_mul(2)?
                } else if old_size < self.grow_const {
                    old_size.checked_add(old_size / 2)?
                } else {
                    old_size.checked_add(self.grow_const / 2)?
                }
                .max(minimum_needed),
            ),
        )
    }

    /// rounds/aligns the size by the `min_allocation` bytes
    #[inline]
    pub fn align(&self, size: usize) -> usize {
        if size > 0 {
            // TODO: when stable .next_multiple_of(self.min_allocation)
            size | (self.min_allocation - 1)
        } else {
            0
        }
    }
}