Skip to content

Instantly share code, notes, and snippets.

@rust-play
Created February 7, 2023 05:51
Show Gist options
  • Save rust-play/253b09af296de0f77592e9c31dca865b to your computer and use it in GitHub Desktop.
Save rust-play/253b09af296de0f77592e9c31dca865b to your computer and use it in GitHub Desktop.
Code shared from the Rust Playground
//! 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:?}");
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment