Code shared from the Rust Playground

playground.rs

//! Proof of concept for a reference counted [str]
//! Primary innovation over a plain Rc<str> is that is one less level of
//! indirection and heap allocation

use core::ptr::NonNull;
use std::alloc::Layout;
use std::alloc::{GlobalAlloc, System};
use std::fmt::Formatter;
use std::fmt::Error;

/// Header fields for a reference counted [str]
/// The [str] payload is not included so as to avoid the fat pointer in the
/// handle and instead use the length from the header.
struct ArcStrHeader {
    ref_cnt: usize, // TODO: should use AtomicUsize and or check single thread semantics
    len: usize, // could save some bytes by making this smaller than a usize. I.e. u16
                // saves 6 bytes per string trading off space for a more modest max len.
    // not pictured: the [str] payload
}

/// A handle to a reference counted [str] buffer.
struct ArcStr {
    p: NonNull<ArcStrHeader>
}

impl ArcStr {
    pub fn from_str(s: &str) -> ArcStr {
        // first, build a composite layout with header and str combined.
        let layout = Layout::new::<ArcStrHeader>();
        let (layout, offset) = layout.extend(Layout::for_value(s)).unwrap();
        
        // allocate and initialize block
        let header = unsafe {
            let header = NonNull::new(System::alloc(&System, layout)).unwrap(); // panic on OOM
            let len = s.len();
            std::ptr::write(header.cast::<ArcStrHeader>().as_ptr(), ArcStrHeader { ref_cnt: 1, len });
            let dst_ptr = header.as_ptr().add(offset) as *mut u8;
            let dst_slice = std::slice::from_raw_parts_mut(dst_ptr, len);
            dst_slice.copy_from_slice(s.as_bytes());
            header.cast::<ArcStrHeader>()
        };
        
        ArcStr {
            p: header
        }
    }
    
    pub fn ref_cnt(&self) -> usize {
        unsafe {
            self.p.as_ref().ref_cnt
        }
    }
    
    pub fn len(&self) -> usize {
        unsafe {
            self.p.as_ref().len
        }
    }

    pub fn as_str(&self) -> &str {
        // get str offset and len
        let layout = Layout::new::<ArcStrHeader>();
        let (_layout, offset) = layout.extend(Layout::for_value("")).unwrap();
        let len = self.len();
        
        // convert [str] payload to slice
        unsafe {
            let u8_slice = std::slice::from_raw_parts(self.p.cast::<u8>().as_ptr().add(offset), len);
            std::str::from_utf8_unchecked(u8_slice)
        }
    }
    
    pub fn print_raw_payload(&self) {
        let layout = Layout::new::<ArcStrHeader>();
        let (_layout, offset) = layout.extend(Layout::for_value("")).unwrap();
        let len = self.len();
        
        // convert &[u8] raw slice
        let raw = unsafe {
            std::slice::from_raw_parts(self.p.cast::<u8>().as_ptr(), offset + len)
        };
        
        println!("raw: {raw:x?}");
    }
}

impl Clone for ArcStr {
    fn clone(&self) -> Self {
        // bump the ref_cnt and return a new handle
        {
            let header = unsafe { self.p.as_ptr().as_mut().unwrap() };
            header.ref_cnt += 1;
        }
        
        let new_handle = Self {
            p: self.p
        };
        
        println!("Cloned the ArcStr");
        
        new_handle
    }
}

impl std::fmt::Debug for ArcStr {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        f.debug_struct("ArcStr")
            .field("ref_cnt", &self.ref_cnt())
            .field("len", &self.len())
            .field("payload", &self.as_str())
            .finish()
    }
}

impl Drop for ArcStr {
    fn drop(&mut self) {
        {
            let header = unsafe { self.p.as_mut() };
            header.ref_cnt -= 1;
            if header.ref_cnt > 0 {
                return;
            }
        }
    
        // first, build a composite layout with header and str combined.
        let layout = Layout::new::<ArcStrHeader>();
        let (layout, _offset) = layout.extend(Layout::for_value(self.as_str())).unwrap();
        unsafe {
            // then dealloc
            System::dealloc(&System, self.p.cast::<u8>().as_ptr(), layout);
        }
        
        println!("Dealloc'd the ArcStr");
    }
}


fn main() {
    println!("layout: {:?}", Layout::for_value("bar"));
    println!("size: {}", std::mem::size_of::<ArcStr>());
    let x = ArcStr::from_str("abcd");
    println!("x: {x:?}");
    let cloned_x = x.clone();
    println!("cloned_x: {cloned_x:?}");
    x.print_raw_payload();
    drop(cloned_x);
    println!("x: {x:?}");
}