`cfd533ff0e2c`

source

crates/jrsonnet-interner/src/inner.rs5.5 KiBsource history
1use std::{2	alloc::{self, Layout},3	borrow::Borrow,4	cmp,5	hash::{Hash, Hasher},6	mem,7	ptr::{self, NonNull},8	slice, str,9};1011const UTF8_MASK: u32 = 1 << 31;12const REFCNT_MASK: u32 = !UTF8_MASK;1314#[repr(C)]15struct InnerHeader {16	size: u32,17	// MSB is checked utf8 flag, rest - refcnt18	utf8_refcnt: u32,19}20impl InnerHeader {21	const fn new(size: u32, is_utf8: bool) -> Self {22		Self {23			size,24			utf8_refcnt: 1 | (if is_utf8 { UTF8_MASK } else { 0 }),25		}26	}2728	const fn refcnt(&self) -> u32 {29		self.utf8_refcnt & REFCNT_MASK30	}31	const fn is_utf8(&self) -> bool {32		self.utf8_refcnt & UTF8_MASK != 033	}3435	fn set_refcnt(&mut self, cnt: u32) {36		assert_eq!(cnt & UTF8_MASK, 0);37		// Reset all bits expect last38		self.utf8_refcnt &= UTF8_MASK;39		// Store refcnt40		self.utf8_refcnt |= cnt;41	}42	fn set_is_utf8(&mut self) {43		self.utf8_refcnt |= UTF8_MASK;44	}45}4647/// Similar to Rc<[u8]>, but stores all data (refcnt, size) inline, instead of being DST48pub struct Inner(NonNull<u8>);49impl Inner {50	/// # Safety51	/// `is_utf8` should only be set if data is really checked to be utf852	/// # Panics53	/// If data is larger than 4GB54	// we allocate with correct alignment55	#[allow(clippy::cast_ptr_alignment)]56	unsafe fn new_raw(bytes: &[u8], is_utf8: bool) -> Self {57		// SAFETY:58		// - layout has non-zero size, and correct align59		// - data is written right after allocation60		// - new allocation can't overlap with passed slice61		unsafe {62			let data = alloc::alloc(Layout::from_size_align_unchecked(63				mem::size_of::<InnerHeader>() + bytes.len(),64				mem::align_of::<InnerHeader>(),65			));66			assert!(!data.is_null());67			*data.cast::<InnerHeader>() =68				InnerHeader::new(bytes.len().try_into().expect("bytes > 4GB"), is_utf8);69			ptr::copy_nonoverlapping(70				bytes.as_ptr(),71				data.add(mem::size_of::<InnerHeader>()),72				bytes.len(),73			);74			Self(NonNull::new_unchecked(data))75		}76	}77	pub fn new_bytes(bytes: &[u8]) -> Self {78		// SAFETY: is_utf8 is not set79		unsafe { Self::new_raw(bytes, false) }80	}81	#[allow(dead_code)]82	pub fn new_str(str: &str) -> Self {83		// SAFETY: strings always utf884		unsafe { Self::new_raw(str.as_bytes(), true) }85	}8687	pub const fn as_slice(&self) -> &[u8] {88		let header = Self::header(self);89		// SAFETY: data is not null, and it is correctly initialized90		let size = unsafe { (*header).size };91		// SAFETY: bytes after data is allocated to be exactly data.size in length92		unsafe {93			slice::from_raw_parts(94				self.0.as_ptr().add(mem::size_of::<InnerHeader>()),95				size as usize,96			)97		}98	}99100	/// # Safety101	/// Data should be checked to be utf8 via [`check_utf8`] first102	pub const unsafe fn as_str_unchecked(&self) -> &str {103		// SAFETY: data is checked104		unsafe { str::from_utf8_unchecked(self.as_slice()) }105	}106107	/// Check data to be utf-8108	///109	/// Positive results are cached110	pub fn check_utf8(this: &Self) -> bool {111		let header = Self::header_mut(this);112		// SAFETY: header is initialized113		if unsafe { (*header).is_utf8() } {114			return true;115		}116117		if str::from_utf8(this.as_slice()).is_ok() {118			// SAFETY: header is initialized119			unsafe { (*header).set_is_utf8() };120			true121		} else {122			false123		}124	}125126	/// Marks data as utf-8127	///128	/// # Safety129	/// data should be really utf-8130	pub unsafe fn assume_utf8(this: &Self) {131		let header = Self::header_mut(this);132		// SAFETY: header is correct133		unsafe { (*header).set_is_utf8() }134	}135136	const fn header(this: &Self) -> *const InnerHeader {137		// in `new`, we allocate with correct alignment138		#![allow(clippy::cast_ptr_alignment)]139		this.0.as_ptr() as *const InnerHeader140	}141	const fn header_mut(this: &Self) -> *mut InnerHeader {142		// in `new`, we allocate with correct alignment143		#![allow(clippy::cast_ptr_alignment)]144		this.0.as_ptr().cast::<InnerHeader>()145	}146147	fn clone(this: &Self) -> Self {148		let header = Self::header_mut(this);149		// SAFETY: header is initialized150		unsafe {151			let refcnt = (*header).refcnt() + 1;152			(*header).set_refcnt(refcnt);153		}154		Self(this.0)155	}156157	pub fn ptr_eq(a: &Self, b: &Self) -> bool {158		a.0 == b.0159	}160	pub const fn as_ptr(this: &Self) -> *const u8 {161		// SAFETY: data is initialized162		unsafe { this.0.as_ptr().add(mem::size_of::<InnerHeader>()) }163	}164165	pub const fn strong_count(this: &Self) -> u32 {166		let header = Self::header(this);167		// SAFETY: header is initialized168		unsafe { (*header).refcnt() }169	}170}171172impl Clone for Inner {173	fn clone(&self) -> Self {174		Self::clone(self)175	}176}177178impl Drop for Inner {179	fn drop(&mut self) {180		#[cold]181		#[inline(never)]182		fn dealloc(val: &Inner) {183			let header = Inner::header_mut(val);184			// SAFETY: size is correct, layout is valid185			unsafe {186				alloc::dealloc(187					val.0.as_ptr(),188					Layout::from_size_align_unchecked(189						mem::size_of::<InnerHeader>() + (*header).size as usize,190						mem::align_of::<InnerHeader>(),191					),192				);193			}194		}195		let header = Self::header_mut(self);196		// SAFETY: header is initialized197		let refcnt = unsafe {198			let refcnt = (*header).refcnt() - 1;199			(*header).set_refcnt(refcnt);200			refcnt201		};202		if refcnt == 0 {203			dealloc(self);204		}205	}206}207208impl PartialEq for Inner {209	fn eq(&self, other: &Self) -> bool {210		self.0 == other.0 || self.as_slice().eq(other.as_slice())211	}212}213impl Hash for Inner {214	fn hash<H: Hasher>(&self, state: &mut H) {215		self.as_slice().hash(state);216	}217}218impl Eq for Inner {}219impl PartialOrd for Inner {220	fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {221		self.as_slice().partial_cmp(other.as_slice())222	}223}224impl Ord for Inner {225	fn cmp(&self, other: &Self) -> cmp::Ordering {226		self.as_slice().cmp(other.as_slice())227	}228}229230impl Borrow<[u8]> for Inner {231	fn borrow(&self) -> &[u8] {232		self.as_slice()233	}234}