cowstr/

rcstring.rs

#![allow(clippy::inline_always)]

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

use std::sync::atomic::{fence, AtomicBool, AtomicUsize, Ordering};

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(Copy, Clone)]
#[repr(transparent)]
pub(crate) struct RcString(NonNull<u8>);

// # Safety
//
// RcString is for most parts immutable when shared. Mutations can happen only when the
// refcount is one or the unsafe try_push* API's are used.
//
//  - Reference counting is atomic.
//  - `try_push()` is Sync by using a ThreadCell and only increment the .len atomically
//    once data is in place
unsafe impl Send for RcString {}
unsafe impl Sync for RcString {}

// ctor/dtor/allocation
impl RcString {
    /// 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 {
            // Safety: using const RCSTRING_ALIGN which is defined by align_of the header.
            let layout = Layout::from_size_align_unchecked(
                capacity
                    .checked_add(RCSTRING_HEADER_SIZE)
                    .expect("Capacity overflow"),
                RCSTRING_ALIGN,
            );

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

            // Safety: `allocated` is uninitialized, it becomes initialized below.
            // this is not a problem since we only access it via addr_of_mut/ptr.write()
            let rcstring = allocated.0.cast::<RcStringHeader>().as_ptr();
            addr_of_mut!((*rcstring).refcount).write(AtomicUsize::new(1usize));
            addr_of_mut!((*rcstring).len).write(AtomicUsize::new(0));
            addr_of_mut!((*rcstring).capacity).write(layout.size() - RCSTRING_HEADER_SIZE);
            addr_of_mut!((*rcstring).guard_active).write(AtomicBool::new(false));

            // everything safe and sound now
            allocated
        }
    }

    /// Grows an `RcString` to have free space for at least 'reserve' bytes.
    ///
    /// # Safety
    ///
    /// Must be called with reference count of 1.
    pub(crate) unsafe fn grow(mut self, reserve: usize, tryreserve: bool) -> Result<Self, Error> {
        // Safety: Internal interface, the caller has to ensure that the refcount is one.
        // This is the same with all debug_asserts later down.
        debug_assert_eq!(self.strong_count(), 1);
        let layout = self.current_layout();
        let min_new_size = (self.len_relaxed() + RCSTRING_HEADER_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() {
            let ptr = alloc::realloc(self.0.as_ptr(), layout, new_size);
            if tryreserve && ptr.is_null() {
                return Err(Error::TryReserveError);
            }
            self = Self::from_ptr(ptr, layout);
            self.header_mut().capacity = new_size - RCSTRING_HEADER_SIZE;
        }

        Ok(self)
    }

    /// Shrinks an `RcString` to at minimum `new_capacity` or its current length, whatever is
    /// larger.
    ///
    /// # Safety
    ///
    /// Must be called with reference count of 1.
    pub(crate) unsafe 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().max(new_capacity) + RCSTRING_HEADER_SIZE);

        // really shrink, or keep the old one when the shrinking fails
        if new_size < layout.size() {
            let ptr = alloc::realloc(self.0.as_ptr(), layout, new_size);
            if !ptr.is_null() {
                let mut new_self = Self::from_ptr(ptr, layout);
                new_self.header_mut().capacity = new_size - RCSTRING_HEADER_SIZE;
                self = new_self;
            };
        }

        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_HEADER_SIZE,
                RCSTRING_ALIGN,
            )
        }
    }
}

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

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

        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.
    ///
    /// # Safety
    ///
    /// * Must only be called from a `CowStrPushGuard` (or single owner) calling it
    ///   concurrently appends in random order (but safely)
    /// * Violates the immutably promise, all `CowStr` cloned from this see the mutation.
    pub(crate) unsafe fn try_push(&self, ch: char) -> Result<&str, Error> {
        let mut buf = [0u8; 4];
        let bytes = ch.encode_utf8(&mut buf).as_bytes();

        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
                // threadcell protected part
                self.data().cast_mut().add(self.header().len_relaxed()),
                bytes.len(),
            );

            let start = self.len_acquire();
            self.len_add_release(bytes.len());
            Ok(&self.as_str()[start..start + bytes.len()])
        } else {
            Err(Error::Capacity)
        }
    }

    /// Appends the given `str` to the end of this `RcString`.
    ///
    /// # Safety
    ///
    /// Must have enough spare capacity available. The capacity has to be extended before
    /// calling this.
    #[inline]
    pub(crate) unsafe fn push_str(&mut self, s: &str) {
        debug_assert!(self.spare_capacity() >= s.len());

        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`.  Self should be
    /// acquired by the current thread, will fail when there is not enough spare capacity.
    ///
    /// # Safety
    ///
    /// * Must only be called from a `CowStrPushGuard` (or single owner) calling it
    ///   concurrently appends in random order (but safely)
    /// * Violates the immutably promise, all `CowStr` cloned from this see the mutation.
    pub(crate) unsafe fn try_push_str(&self, s: &str) -> Result<&str, Error> {
        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
                // threadcell protected part
                self.data().cast_mut().add(self.header().len_relaxed()),
                s.len(),
            );
            let start = self.len_acquire();
            self.len_add_release(s.len());
            Ok(&self.as_str()[start..start + s.len()])
        } else {
            Err(Error::Capacity)
        }
    }

    // helper only used by the Guard
    #[inline]
    pub(crate) fn acquire_guard(self) -> bool {
        self.header()
            .guard_active
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
    }

    // helper only used by the Guard
    #[inline]
    pub(crate) fn release_guard(self) {
        self.header().guard_active.store(false, Ordering::Release);
    }

    // helper only used by Debug
    #[inline]
    #[mutants::skip]
    pub(crate) fn has_guard(self) -> bool {
        self.header().guard_active.load(Ordering::Relaxed)
    }

    /// Removes the last character and returns it.
    #[inline]
    pub(crate) unsafe fn pop(&mut self) -> Option<char> {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);
        let ch = self.as_str().chars().next_back()?;
        self.len_sub_release(ch.len_utf8());
        Some(ch)
    }

    /// Truncates the string to zero length. keeps capacity.
    #[inline]
    pub(crate) unsafe fn clear(&mut self) {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);
        self.len_set_release(0);
    }

    /// Truncates the string to the supplied length. keeps capacity.
    #[inline]
    pub(crate) unsafe fn truncate(&mut self, new_len: usize) {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);
        // Safety: caller checks already.
        debug_assert!(self.as_str().is_char_boundary(new_len));
        self.len_set_release(new_len);
    }

    /// remove a character from a string
    pub(crate) unsafe fn remove(&mut self, idx: usize) -> char {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);

        let Some(ch) = self.as_str()[idx..].chars().next() else {
            panic!("cannot remove a char from the end of a string")
        };

        let next = idx + ch.len_utf8();
        let len = self.len_relaxed();

        let data = self.data_mut();
        std::ptr::copy(data.add(next), data.add(idx), len - next);
        self.len_sub_release(next - idx);

        ch
    }

    /// replaces a range of a rcstring with another str
    pub(crate) unsafe fn replace_range<R>(&mut self, range: R, replace_with: &str)
    where
        R: std::slice::SliceIndex<str, Output = str>,
    {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);

        // PLANNED: this impl is a hack, make it more rusty
        let removed_slice = &self.as_str()[range];
        let old_len = self.as_str().len();
        let new_len = old_len + replace_with.len() - removed_slice.len();
        debug_assert!(self.capacity() >= new_len);

        let removed_start: usize = removed_slice
            .as_ptr()
            .offset_from(self.as_str().as_ptr())
            .try_into()
            .unwrap_unchecked();

        let tail_start = removed_start + removed_slice.len();
        let insert_end = removed_start + replace_with.len();

        let tail_len = old_len - tail_start;

        self.len_set_release(new_len);
        let data = self.data_mut();

        // shift the tail
        std::ptr::copy(data.add(tail_start), data.add(insert_end), tail_len);
        // copy replace_with
        std::ptr::copy_nonoverlapping(
            replace_with.as_ptr(),
            data.add(removed_start),
            replace_with.len(),
        );
    }
}

// conversions
impl RcString {
    /// Returns self as &str
    #[inline]
    pub(crate) fn as_str(&self) -> &str {
        // Safety: RcString always holds a valid utf-8 string
        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 {
        // Safety: caller must only call this on a mutable object.
        debug_assert_eq!(self.strong_count(), 1);
        // Safety: RcString always holds a valid utf-8 string
        std::str::from_utf8_unchecked_mut(std::slice::from_raw_parts_mut(
            self.data_mut(),
            self.len_acquire(),
        ))
    }
}

// header accessors
//
// # Safety
//
// A RcString always wraps a properly aligned and initialized block of memory. We need to view
// it as a header followed by capacity*bytes, where the bytes above header.len are not yet
// initialized. miri's stacked borrows wont let us make up objects within this block easily. However
// casting this opportunistically whenever needed is sound and safe.
//
// NOTE: I would be glad if someone can point out a less hacky way to implement such kinds of DST's.
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)]
    const fn data(&self) -> *const u8 {
        unsafe { self.0.as_ptr().add(RCSTRING_HEADER_SIZE) }
    }

    #[inline(always)]
    fn data_mut(&mut self) -> *mut u8 {
        unsafe { self.0.as_ptr().add(RCSTRING_HEADER_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]
    pub(crate) 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.
    ///
    /// # Safety
    ///
    /// The new len must comply with valid utf-8 encoding.
    #[inline(always)]
    unsafe fn len_add_release(self, n: usize) {
        self.header().len_add_release(n);
    }

    /// Add 'n' to the length with Release semantic.
    ///
    /// # Safety
    ///
    /// The new len must comply with valid utf-8 encoding.
    #[inline(always)]
    unsafe fn len_sub_release(self, n: usize) {
        self.header().len_sub_release(n);
    }

    /// Sets a new length with Release semantic.
    ///
    /// # Safety
    ///
    /// The new len must comply with valid utf-8 encoding.
    #[inline(always)]
    unsafe fn len_set_release(self, n: usize) {
        self.header().len_set_release(n);
    }
}

impl std::fmt::Debug for RcString {
    #[mutants::skip]
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        f.debug_struct("RcString")
            .field("header", self.header())
            .field("str", &self.as_str())
            .finish()
    }
}

#[derive(Debug)]
#[repr(C, align(32))]
pub(crate) struct RcStringHeader {
    refcount: AtomicUsize,
    len: AtomicUsize,
    capacity: usize,
    guard_active: AtomicBool,
}

/// The size of the header of a `CowStr` allocated memory block. This is mostly a implementation
/// detail but it is exported to give the user more control about memory
/// allocations. `CowStr::with_capacity(65536 - cowstr::RCSTRING_HEADER_SIZE)` will allocate a `CowStr`
/// that uses exactly 64kb memory.
pub const RCSTRING_HEADER_SIZE: usize = size_of::<RcStringHeader>();

/// The alignment/size granularity a `CowStr` allocated memory block. This is mostly a
/// implementation detail. To avoid tiny allocations and memory fragmentation allocated
/// lengths are a multiple of this alignment. Note that an allocation includes a header.
pub const RCSTRING_ALIGN: usize = align_of::<RcStringHeader>();

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

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

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

    /// Sub 'n' from the length with Release semantic.
    #[inline(always)]
    unsafe fn len_sub_release(&self, n: usize) {
        self.len.fetch_sub(n, Ordering::Release);
    }

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

    /// Returns the amount of bytes that can be pushed without reallocation.  This is an
    /// acquire fenced load allowing any further access to len 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) const fn capacity(&self) -> usize {
        self.capacity
    }

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

    // NOTE: incrementing the strong count can always be relaxed because for any existing
    //       object it is:
    //
    //        * '1' then we own the object (mutably), no one else can decrement the strong count
    //        * '>1' the object is immutable/shared it may be concurrently decremented, but
    //          since we own it here it can't be dropped and we are going to increment it at
    //          least to 2.
    pub(crate) fn increment_strong_count(&self) {
        self.refcount.fetch_add(1, Ordering::Relaxed);
    }

    #[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]
    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
    #[cfg(not(feature = "nightly_int_roundings"))]
    #[inline]
    #[mutants::skip]
    const fn align(&self, size: usize) -> usize {
        size | (self.min_allocation - 1)
    }

    /// rounds/aligns the size by the `min_allocation` bytes
    #[cfg(feature = "nightly_int_roundings")]
    #[inline]
    #[mutants::skip]
    fn align(&self, size: usize) -> usize {
        size.next_multiple_of(self.min_allocation)
    }
}