Files
scryer-prolog/src/offset_table.rs
2025-07-07 22:44:02 -07:00

423 lines
12 KiB
Rust

use std::cell::UnsafeCell;
use std::sync::Arc;
use std::{fmt, mem, ptr};
use arcu::atomic::Arcu;
use arcu::epoch_counters::GlobalEpochCounterPool;
use arcu::rcu_ref::RcuRef;
use arcu::Rcu;
use parking_lot::RwLock;
use crate::machine::machine_indices::IndexPtr;
use crate::raw_block::RawBlock;
use crate::raw_block::RawBlockTraits;
use ordered_float::OrderedFloat;
const F64_TABLE_INIT_SIZE: usize = 1 << 16;
const F64_TABLE_ALIGN: usize = 8;
const CODE_INDEX_TABLE_INIT_SIZE: usize = 1 << 16;
const CODE_INDEX_TABLE_ALIGN: usize = 8;
impl RawBlockTraits for OrderedFloat<f64> {
#[inline]
fn init_size() -> usize {
F64_TABLE_INIT_SIZE
}
#[inline]
fn align() -> usize {
F64_TABLE_ALIGN
}
}
impl RawBlockTraits for IndexPtr {
#[inline]
fn init_size() -> usize {
CODE_INDEX_TABLE_INIT_SIZE
}
#[inline]
fn align() -> usize {
CODE_INDEX_TABLE_ALIGN
}
}
#[derive(Debug)]
pub struct OffsetTableImpl<T: RawBlockTraits>(InnerOffsetTableImpl<T>);
impl<T: RawBlockTraits> From<Arc<ConcurrentOffsetTable<T>>> for OffsetTableImpl<T> {
#[inline]
fn from(value: Arc<ConcurrentOffsetTable<T>>) -> Self {
OffsetTableImpl(InnerOffsetTableImpl::Concurrent(value))
}
}
impl<T: fmt::Debug + RawBlockTraits> OffsetTableImpl<T> {
#[inline(always)]
pub fn new() -> Self {
Self(InnerOffsetTableImpl::Serial(SerialOffsetTable::new()))
}
#[must_use = "the returned concurrent table must be absorbed into the owned OffsetTable"]
pub fn single_to_concurrent(&mut self) -> Arc<ConcurrentOffsetTable<T>> {
match &mut self.0 {
InnerOffsetTableImpl::Serial(serial_tbl) => {
let empty_serial_tbl = SerialOffsetTable {
block: RawBlock::empty_block(),
};
let serial_tbl = mem::replace(serial_tbl, empty_serial_tbl);
let num_tbl_entries = serial_tbl.block.size() / size_of::<T>();
let block = Arcu::new(serial_tbl.block, GlobalEpochCounterPool);
let offset_locks: Vec<RwLock<()>> =
(0..num_tbl_entries).map(|_| RwLock::new(())).collect();
let concurrent_tbl = Arc::new(ConcurrentOffsetTable {
block,
growth_lock: RwLock::new(()),
offset_locks: RwLock::new(offset_locks),
});
self.0 = InnerOffsetTableImpl::Concurrent(concurrent_tbl.clone());
concurrent_tbl
}
InnerOffsetTableImpl::Concurrent(concurrent_tbl) => concurrent_tbl.clone(),
}
}
#[must_use = "the transition to a single-threaded offset table may fail if the concurrent table is held from multiple places"]
pub fn concurrent_to_single(&mut self) -> Result<(), ()> {
match &mut self.0 {
InnerOffsetTableImpl::Serial(_serial_tbl) => Ok(()),
InnerOffsetTableImpl::Concurrent(concurrent_tbl) => {
let table_arc = std::mem::replace(
concurrent_tbl,
Arc::new(ConcurrentOffsetTable {
block: Arcu::new(RawBlock::empty_block(), GlobalEpochCounterPool),
growth_lock: RwLock::new(()),
offset_locks: RwLock::new(vec![]),
}),
);
match Arc::try_unwrap(table_arc) {
Ok(table) => {
// this was the only instance of the concurrent table, as such
// at this point no build_with/with_entry{_mut} call can be in-progress/made
// this shouldn't be able to fail
let raw_block =
Arc::try_unwrap(table.block.replace(RawBlock::empty_block())).unwrap();
self.0 =
InnerOffsetTableImpl::Serial(SerialOffsetTable { block: raw_block });
Ok(())
}
Err(table_arc) => {
// restore the concurrent_tbl
*concurrent_tbl = table_arc;
Err(())
}
}
}
}
}
#[inline]
pub fn get_entry(&self, offset: <Self as OffsetTable<T>>::Offset) -> T
where
Self: OffsetTable<T>,
T: Copy,
{
self.with_entry(offset, |value| *value)
}
}
impl<T: fmt::Debug + RawBlockTraits> Default for OffsetTableImpl<T> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
struct SerialOffsetTable<T: RawBlockTraits> {
block: RawBlock<T>,
}
#[derive(Debug)]
pub struct ConcurrentOffsetTable<T: RawBlockTraits> {
block: Arcu<RawBlock<T>, GlobalEpochCounterPool>,
growth_lock: RwLock<()>,
offset_locks: RwLock<Vec<RwLock<()>>>,
}
#[derive(Debug)]
enum InnerOffsetTableImpl<T: RawBlockTraits> {
Serial(SerialOffsetTable<T>),
#[allow(dead_code)]
Concurrent(Arc<ConcurrentOffsetTable<T>>),
}
impl<T: RawBlockTraits> InnerOffsetTableImpl<T> {
#[inline(always)]
fn build_with(&mut self, value: T) -> usize {
match self {
Self::Concurrent(concurrent_tbl) => concurrent_tbl.build_with(value),
Self::Serial(serial_tbl) => unsafe { serial_tbl.build_with(value) },
}
}
#[inline(always)]
fn with_entry<R, F: FnOnce(&T) -> R>(&self, offset: usize, f: F) -> R {
match self {
Self::Concurrent(concurrent_tbl) => concurrent_tbl.with_entry(offset, f),
Self::Serial(serial_tbl) => f(unsafe { serial_tbl.lookup(offset) }),
}
}
#[inline(always)]
fn with_entry_mut<R, F: FnOnce(&mut T) -> R>(&mut self, offset: usize, f: F) -> R {
match self {
Self::Concurrent(concurrent_tbl) => concurrent_tbl.with_entry_mut(offset, f),
Self::Serial(serial_tbl) => f(unsafe { serial_tbl.lookup_mut(offset) }),
}
}
}
pub trait OffsetTable<T: RawBlockTraits> {
type Offset: Copy + Into<usize>;
fn build_with(&mut self, value: T) -> Self::Offset;
fn with_entry<R, F: FnOnce(&T) -> R>(&self, offset: Self::Offset, f: F) -> R;
fn with_entry_mut<R, F: FnOnce(&mut T) -> R>(&mut self, offset: Self::Offset, f: F) -> R;
}
impl OffsetTable<OrderedFloat<f64>> for OffsetTableImpl<OrderedFloat<f64>> {
type Offset = F64Offset;
fn build_with(&mut self, value: OrderedFloat<f64>) -> F64Offset {
F64Offset(self.0.build_with(value))
}
#[inline]
fn with_entry<R, F: FnOnce(&OrderedFloat<f64>) -> R>(&self, offset: F64Offset, f: F) -> R {
self.0.with_entry(offset.into(), f)
}
#[inline]
fn with_entry_mut<R, F: FnOnce(&mut OrderedFloat<f64>) -> R>(
&mut self,
offset: F64Offset,
f: F,
) -> R {
self.0.with_entry_mut(offset.into(), f)
}
}
impl OffsetTable<IndexPtr> for OffsetTableImpl<IndexPtr> {
type Offset = CodeIndexOffset;
fn build_with(&mut self, value: IndexPtr) -> CodeIndexOffset {
CodeIndexOffset(self.0.build_with(value))
}
#[inline]
fn with_entry<R, F: FnOnce(&IndexPtr) -> R>(&self, offset: CodeIndexOffset, f: F) -> R {
self.0.with_entry(offset.into(), f)
}
#[inline]
fn with_entry_mut<R, F: FnOnce(&mut IndexPtr) -> R>(
&mut self,
offset: CodeIndexOffset,
f: F,
) -> R {
self.0.with_entry_mut(offset.into(), f)
}
}
impl<T: RawBlockTraits> SerialOffsetTable<T> {
#[inline]
fn new() -> Self {
Self {
block: RawBlock::new(),
}
}
unsafe fn build_with(&mut self, value: T) -> usize {
let mut ptr;
loop {
ptr = self.block.alloc(size_of::<T>());
if ptr.is_null() {
let new_block = self.block.grow_new().unwrap();
self.block = new_block;
} else {
break;
}
}
ptr::write(ptr as *mut T, value);
ptr.addr() - self.block.base.addr()
}
#[inline]
unsafe fn lookup(&self, offset: usize) -> &T {
&*self.block.base.add(offset).cast::<T>()
}
#[inline]
unsafe fn lookup_mut(&mut self, offset: usize) -> &mut T {
&mut *self.block.base.add(offset).cast::<T>().cast_mut()
}
}
impl<T: RawBlockTraits> ConcurrentOffsetTable<T> {
#[allow(clippy::missing_safety_doc)]
fn build_with(&self, value: T) -> usize {
let growth_lock = self.growth_lock.write();
// we don't have an index table for lookups as AtomTable does so
// just get the epoch after we take the upgrade lock
let mut block_epoch = self.block.read();
let mut ptr;
loop {
ptr = unsafe { block_epoch.alloc(mem::size_of::<T>()) };
if ptr.is_null() {
let new_block = unsafe { block_epoch.grow_new().unwrap() };
self.block.replace(new_block);
block_epoch = self.block.read();
} else {
break;
}
}
let new_tbl_sz = block_epoch.size() / size_of::<T>();
let mut offset_locks = self.offset_locks.write();
offset_locks.resize_with(new_tbl_sz, || RwLock::new(()));
unsafe {
ptr::write(ptr as *mut T, value);
}
let value = ptr.addr() - block_epoch.base.addr();
// AtomTable would have to update the index table at this point
// explicit drop to ensure we don't accidentally drop it early
drop(offset_locks);
drop(growth_lock);
value
}
fn with_entry<R, F: FnOnce(&T) -> R>(&self, offset: usize, f: F) -> R {
let outer_offset_lock = self.offset_locks.read();
let inner_offset_lock = outer_offset_lock[offset / size_of::<T>()].read();
let rcu_ref = RcuRef::try_map(self.block.read(), |raw_block| unsafe {
raw_block.base.add(offset).cast::<T>().as_ref()
})
.expect("offset valid");
let result = f(&*rcu_ref);
drop(inner_offset_lock);
drop(outer_offset_lock);
result
}
fn with_entry_mut<R, F: FnOnce(&mut T) -> R>(&self, offset: usize, f: F) -> R {
let growth_lock = self.growth_lock.read();
let outer_offset_lock = self.offset_locks.read();
let inner_offset_lock = outer_offset_lock[offset / size_of::<T>()].write();
let rcu_ref = RcuRef::try_map(self.block.read(), |raw_block| unsafe {
raw_block
.base
.add(offset)
.cast_mut()
.cast::<UnsafeCell<T>>()
.as_ref()
})
.expect("offset valid");
let result = f(unsafe { &mut *rcu_ref.get().as_mut().unwrap() });
drop(inner_offset_lock);
drop(outer_offset_lock);
drop(growth_lock);
result
}
}
pub type F64Table = OffsetTableImpl<OrderedFloat<f64>>;
pub type CodeIndexTable = OffsetTableImpl<IndexPtr>;
#[derive(Clone, Copy, Debug)]
pub struct F64Offset(usize);
impl From<usize> for F64Offset {
#[inline(always)]
fn from(offset: usize) -> Self {
Self(offset)
}
}
impl From<F64Offset> for usize {
fn from(val: F64Offset) -> Self {
val.0
}
}
#[derive(Debug, Clone, Copy)]
pub struct CodeIndexOffset(usize);
impl From<usize> for CodeIndexOffset {
#[inline(always)]
fn from(offset: usize) -> Self {
Self(offset)
}
}
impl From<CodeIndexOffset> for usize {
#[inline(always)]
fn from(val: CodeIndexOffset) -> Self {
val.0
}
}
impl CodeIndexOffset {
#[inline(always)]
pub fn to_u64(self) -> u64 {
self.0 as u64
}
}
impl fmt::Display for CodeIndexOffset {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "CodeIndexOffset({})", self.0)
}
}
impl F64Offset {
#[inline(always)]
pub fn to_u64(self) -> u64 {
self.0 as u64
}
}
impl fmt::Display for F64Offset {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "F64Offset({})", self.0)
}
}