From 925075337ba01f56d46240cd2a31152c68d539b8 Mon Sep 17 00:00:00 2001 From: Dennis Kobert Date: Sat, 25 Jul 2026 17:05:40 +0000 Subject: [PATCH] Add the generation-scoped arena and the FrameTable frame memo store --- node-graph/libraries/core-types/src/arena.rs | 230 ++++++++++++++++++ .../libraries/core-types/src/frame_table.rs | 155 ++++++++++++ node-graph/libraries/core-types/src/lib.rs | 2 + 3 files changed, 387 insertions(+) create mode 100644 node-graph/libraries/core-types/src/arena.rs create mode 100644 node-graph/libraries/core-types/src/frame_table.rs diff --git a/node-graph/libraries/core-types/src/arena.rs b/node-graph/libraries/core-types/src/arena.rs new file mode 100644 index 0000000000..2018691018 --- /dev/null +++ b/node-graph/libraries/core-types/src/arena.rs @@ -0,0 +1,230 @@ +use std::cell::UnsafeCell; +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; + +pub struct Arena { + generation: AtomicU64, + offset: AtomicUsize, + buf: Box<[UnsafeCell>]>, + drops: Mutex>, +} + +struct DropEntry { + offset: usize, + drop_fn: unsafe fn(*mut u8), +} + +// SAFETY: disjoint regions are handed out by an atomic bump; a region is written +// only by its allocating caller before publication, and cross-thread hand-off is +// ordered by the Release/Acquire pair on the published handle word. +unsafe impl Sync for Arena {} +unsafe impl Send for Arena {} + +impl Arena { + pub fn new(capacity: usize) -> Self { + let buf = (0..capacity).map(|_| UnsafeCell::new(MaybeUninit::uninit())).collect(); + Self { + // Starts at 1 so the null handle word (0) can never upgrade. + generation: AtomicU64::new(1), + offset: AtomicUsize::new(0), + buf, + drops: Mutex::new(Vec::new()), + } + } + + pub fn generation(&self) -> u64 { + self.generation.load(Ordering::Acquire) + } + + fn base(&self) -> *mut u8 { + self.buf.as_ptr() as *mut u8 + } + + fn reserve(&self, size: usize, align: usize) -> Option { + debug_assert!(align.is_power_of_two()); + let base = self.base() as usize; + let mut start = 0; + self.offset + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + // Alignment is computed on the absolute address; the backbone + // allocation itself has no alignment guarantee. + let addr = (base.checked_add(current)?.checked_add(align - 1)?) & !(align - 1); + start = addr - base; + let end = start.checked_add(size)?; + (end <= self.buf.len()).then_some(end) + }) + .ok()?; + Some(start) + } + + pub fn alloc(&self, value: T) -> Option<(&T, ArenaWeak)> { + let offset = self.reserve(size_of::(), align_of::())?; + let ptr = unsafe { self.base().add(offset) }.cast::(); + // SAFETY: freshly reserved, aligned, in-bounds, unaliased. + unsafe { ptr.write(value) }; + if std::mem::needs_drop::() { + unsafe fn glue(p: *mut u8) { + unsafe { p.cast::().drop_in_place() } + } + self.drops.lock().unwrap().push(DropEntry { offset, drop_fn: glue:: }); + } + // SAFETY: initialized above; insert-only, so no `&mut` to it can exist. + Some((unsafe { &*ptr }, ArenaWeak::new(self.generation(), offset))) + } + + pub fn alloc_slice_copy(&self, src: &[T]) -> Option<&[T]> { + let buf = self.alloc_scratch::(src.len())?; + for (slot, &value) in buf.iter_mut().zip(src) { + slot.write(value); + } + // SAFETY: every lane written above from `src`. + Some(unsafe { std::slice::from_raw_parts(buf.as_ptr().cast::(), src.len()) }) + } + + // Not drop-tracked: callers must either consume every written lane or + // restrict themselves to `Copy` payloads (leak, not UB, otherwise). + #[allow(clippy::mut_from_ref)] + pub fn alloc_scratch(&self, len: usize) -> Option<&mut [MaybeUninit]> { + let size = size_of::().checked_mul(len)?; + let offset = self.reserve(size, align_of::())?; + let ptr = unsafe { self.base().add(offset) }.cast::>(); + // SAFETY: exclusive region; lifetime tied to `&self`, and `reset` takes + // `&mut self`, so the slice cannot outlive the generation. + Some(unsafe { std::slice::from_raw_parts_mut(ptr, len) }) + } + + pub fn reset(&mut self) { + let base = self.base(); + for entry in self.drops.get_mut().unwrap().drain(..) { + // SAFETY: registered at alloc time; insert-only means the region was + // never overwritten within this generation. + unsafe { (entry.drop_fn)(base.add(entry.offset)) } + } + *self.offset.get_mut() = 0; + self.generation.fetch_add(1, Ordering::Release); + } +} + +impl Drop for Arena { + fn drop(&mut self) { + self.reset(); + } +} + +pub struct ArenaWeak { + word: u64, + _marker: PhantomData<*const T>, +} + +impl Clone for ArenaWeak { + fn clone(&self) -> Self { + *self + } +} +impl Copy for ArenaWeak {} + +impl ArenaWeak { + pub const NULL: Self = ArenaWeak { word: 0, _marker: PhantomData }; + + fn new(generation: u64, offset: usize) -> Self { + debug_assert!(offset < u32::MAX as usize); + Self { + word: ((generation & 0xFFFF_FFFF) << 32) | offset as u64, + _marker: PhantomData, + } + } + + pub fn upgrade(self, arena: &Arena) -> Option<&T> { + let generation = self.word >> 32; + if generation != arena.generation() & 0xFFFF_FFFF { + return None; + } + let offset = (self.word & 0xFFFF_FFFF) as usize; + // SAFETY: same generation means the entry was fully written before its + // word was published (Release) and cannot move or be overwritten within + // a generation (insert-only); the Acquire load that produced this word + // ordered the payload writes. + Some(unsafe { &*arena.base().add(offset).cast::() }) + } +} + +pub struct ArenaCell { + word: AtomicU64, + _marker: PhantomData T>, +} + +impl Default for ArenaCell { + fn default() -> Self { + Self { + word: AtomicU64::new(0), + _marker: PhantomData, + } + } +} + +impl ArenaCell { + pub fn new() -> Self { + Self::default() + } + + pub fn load<'e>(&self, arena: &'e Arena) -> Option<&'e T> { + let weak = ArenaWeak:: { + word: self.word.load(Ordering::Acquire), + _marker: PhantomData, + }; + weak.upgrade(arena) + } + + pub fn store(&self, weak: ArenaWeak) { + self.word.store(weak.word, Ordering::Release); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicU32; + + #[test] + fn alloc_upgrade_reset_miss() { + let mut arena = Arena::new(1024); + let cell = ArenaCell::new(); + let (value, weak) = arena.alloc(41u32).unwrap(); + assert_eq!(*value, 41); + cell.store(weak); + assert_eq!(cell.load(&arena), Some(&41)); + arena.reset(); + assert_eq!(cell.load(&arena), None, "stale handle must miss"); + } + + #[test] + fn capacity_survives_reset() { + let mut arena = Arena::new(64 + align_of::() - 1); + for _ in 0..10 { + for _ in 0..16 { + assert!(arena.alloc(0u32).is_some()); + } + assert!(arena.alloc(0u32).is_none(), "exhausted within generation"); + arena.reset(); + } + } + + #[test] + fn drop_glue_runs_on_reset() { + static DROPS: AtomicU32 = AtomicU32::new(0); + struct Probe(#[allow(dead_code)] String); + impl Drop for Probe { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::Relaxed); + } + } + let mut arena = Arena::new(1024); + arena.alloc(Probe("owns heap".into())).unwrap(); + arena.alloc(Probe("me too".into())).unwrap(); + assert_eq!(DROPS.load(Ordering::Relaxed), 0); + arena.reset(); + assert_eq!(DROPS.load(Ordering::Relaxed), 2); + } +} diff --git a/node-graph/libraries/core-types/src/frame_table.rs b/node-graph/libraries/core-types/src/frame_table.rs new file mode 100644 index 0000000000..d35081d057 --- /dev/null +++ b/node-graph/libraries/core-types/src/frame_table.rs @@ -0,0 +1,155 @@ +use crate::gpoll::Finality; +use std::cell::UnsafeCell; +use std::mem::MaybeUninit; +use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; + +const SLOT_EMPTY: u8 = 0; +const SLOT_FINAL: u8 = 1; +const SLOT_PARTIAL: u8 = 2; + +pub struct FrameTable { + slots: [FrameSlot; CAP], +} + +struct FrameSlot { + key: AtomicU64, + state: AtomicU8, + value: UnsafeCell>, +} + +pub enum Lookup<'t, T> { + Hit(Finality, &'t T), + Vacant(VacantSlot<'t, T>), + Full, +} + +pub struct VacantSlot<'t, T> { + slot: &'t FrameSlot, +} + +impl Default for FrameTable { + fn default() -> Self { + Self::new() + } +} + +impl FrameTable { + pub fn new() -> Self { + Self { + slots: std::array::from_fn(|_| FrameSlot { + key: AtomicU64::new(0), + state: AtomicU8::new(SLOT_EMPTY), + value: UnsafeCell::new(MaybeUninit::uninit()), + }), + } + } + + pub fn lookup(&self, hash: u64) -> Lookup<'_, T> { + // Keys are forced odd so the empty sentinel (0) can never collide. + let key = hash | 1; + for probe in 0..CAP { + let slot = &self.slots[(key as usize).wrapping_add(probe) % CAP]; + let stored = slot.key.load(Ordering::Acquire); + if stored == key { + return match slot.state.load(Ordering::Acquire) { + // SAFETY: a published state was stored with Release after the + // value write; the Acquire load above ordered that write. + SLOT_FINAL => Lookup::Hit(Finality::AllFinal, unsafe { (*slot.value.get()).assume_init_ref() }), + SLOT_PARTIAL => Lookup::Hit(Finality::Partial, unsafe { (*slot.value.get()).assume_init_ref() }), + _ => Lookup::Full, + }; + } + if stored == 0 && slot.key.compare_exchange(0, key, Ordering::AcqRel, Ordering::Acquire).is_ok() { + return Lookup::Vacant(VacantSlot { slot }); + } + } + Lookup::Full + } +} + +impl Drop for FrameTable { + fn drop(&mut self) { + for slot in &self.slots { + if slot.state.load(Ordering::Acquire) != SLOT_EMPTY { + // SAFETY: a non-empty state is only ever stored after the value + // write in `publish`. + unsafe { (*slot.value.get()).assume_init_drop() } + } + } + } +} + +impl<'t, T> VacantSlot<'t, T> { + pub fn publish(self, value: T, finality: Finality) -> &'t T { + // SAFETY: the CAS in `lookup` reserved this slot exclusively for us and + // its state is still SLOT_EMPTY, so nobody reads the value yet. + let lent = unsafe { &*(*self.slot.value.get()).write(value) }; + let state = match finality { + Finality::AllFinal => SLOT_FINAL, + Finality::Partial => SLOT_PARTIAL, + }; + self.slot.state.store(state, Ordering::Release); + lent + } + + pub fn release(self) { + self.slot.key.store(0, Ordering::Release); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicU32; + + #[test] + fn publish_then_hit_with_finality() { + let table = FrameTable::::new(); + let Lookup::Vacant(slot) = table.lookup(7) else { + panic!("fresh table must be vacant"); + }; + assert_eq!(*slot.publish(41, Finality::Partial), 41); + let Lookup::Hit(finality, value) = table.lookup(7) else { + panic!("published key must hit"); + }; + assert_eq!((finality, *value), (Finality::Partial, 41)); + } + + #[test] + fn released_slot_is_vacant_again() { + let table = FrameTable::::new(); + let Lookup::Vacant(slot) = table.lookup(7) else { unreachable!() }; + slot.release(); + assert!(matches!(table.lookup(7), Lookup::Vacant(_))); + } + + #[test] + fn distinct_keys_probe_past_collisions_until_full() { + let table = FrameTable::::new(); + for hash in [2, 4, 6, 8] { + let Lookup::Vacant(slot) = table.lookup(hash) else { + panic!("hash {hash} should find a vacant slot"); + }; + slot.publish(hash as u32, Finality::AllFinal); + } + assert!(matches!(table.lookup(100), Lookup::Full)); + } + + #[test] + fn drop_runs_glue_for_published_values_only() { + static DROPS: AtomicU32 = AtomicU32::new(0); + struct Probe; + impl Drop for Probe { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::Relaxed); + } + } + let table = FrameTable::::new(); + let Lookup::Vacant(slot) = table.lookup(1) else { unreachable!() }; + slot.publish(Probe, Finality::AllFinal); + let Lookup::Vacant(reserved_but_unpublished) = table.lookup(2) else { unreachable!() }; + reserved_but_unpublished.release(); + drop(table); + assert_eq!(DROPS.load(Ordering::Relaxed), 1); + } +} diff --git a/node-graph/libraries/core-types/src/lib.rs b/node-graph/libraries/core-types/src/lib.rs index 36a0fac6e0..f7a9b25bab 100644 --- a/node-graph/libraries/core-types/src/lib.rs +++ b/node-graph/libraries/core-types/src/lib.rs @@ -1,8 +1,10 @@ extern crate log; +pub mod arena; pub mod bounds; pub mod consts; pub mod context; +pub mod frame_table; pub mod generic; pub mod gnode; pub mod gpoll;