Fix frame table key aliasing and arena handle offset truncation

This commit is contained in:
Dennis Kobert
2026-07-31 16:46:20 +02:00
parent 1fccb112cb
commit 5541769dc5
3 changed files with 50 additions and 17 deletions
+9 -6
View File
@@ -370,13 +370,16 @@ impl ProtoNetwork {
let mut combined_deps = ContextFeatures::default(); let mut combined_deps = ContextFeatures::default();
let node_index = id.0 as usize; let node_index = id.0 as usize;
let context_features = self.nodes[node_index].1.context_features.clone(); let (extract, inject) = {
let dependencies = &self.nodes[node_index].1.context_features;
(dependencies.extract, dependencies.inject)
};
let mut inputs = match &self.nodes[node_index].1.construction_args { let mut inputs = match &self.nodes[node_index].1.construction_args {
// We pretend like we have already placed context modification nodes after ourselves because value nodes don't need to be cached // We pretend like we have already placed context modification nodes after ourselves because value nodes don't need to be cached
ConstructionArgs::Value(_) => return (context_features.extract, Some(id)), ConstructionArgs::Value(_) => return (extract, Some(id)),
ConstructionArgs::Nodes(items) => items.clone(), ConstructionArgs::Nodes(items) => items.clone(),
ConstructionArgs::Inline(_) => return (context_features.extract, Some(id)), ConstructionArgs::Inline(_) => return (extract, Some(id)),
}; };
// Compute the dependencies for each branch and combine all of them // Compute the dependencies for each branch and combine all of them
@@ -389,9 +392,9 @@ impl ProtoNetwork {
let mut new_deps = combined_deps; let mut new_deps = combined_deps;
// Remove requirements which this node provides // Remove requirements which this node provides
new_deps &= !context_features.inject; new_deps &= !inject;
// Add requirements we have // Add requirements we have
new_deps |= context_features.extract; new_deps |= extract;
// If we either introduce new dependencies, we can cache all children which don't yet need that dependency // If we either introduce new dependencies, we can cache all children which don't yet need that dependency
let we_introduce_new_deps = !combined_deps.contains(new_deps); let we_introduce_new_deps = !combined_deps.contains(new_deps);
@@ -407,7 +410,7 @@ impl ProtoNetwork {
self.nodes[node_index].1.construction_args = ConstructionArgs::Nodes(inputs); self.nodes[node_index].1.construction_args = ConstructionArgs::Nodes(inputs);
// Which dependencies do we supply (and don't need ourselves)? // Which dependencies do we supply (and don't need ourselves)?
let net_injections = context_features.inject.difference(context_features.extract); let net_injections = inject.difference(extract);
// Which dependencies still need to be met after this node? // Which dependencies still need to be met after this node?
let remaining_deps_from_children = combined_deps.difference(net_injections); let remaining_deps_from_children = combined_deps.difference(net_injections);
+20 -9
View File
@@ -4,6 +4,12 @@ use std::mem::MaybeUninit;
use std::sync::Mutex; use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
/// Handle word layout: 24 generation bits above 40 offset bits, so a 1 TiB arena
/// is addressable and generations wrap after ~3 days of 60fps resets.
const OFFSET_BITS: u32 = 40;
const OFFSET_MASK: u64 = (1 << OFFSET_BITS) - 1;
const GENERATION_MASK: u64 = (1 << (64 - OFFSET_BITS)) - 1;
pub struct Arena { pub struct Arena {
generation: AtomicU64, generation: AtomicU64,
offset: AtomicUsize, offset: AtomicUsize,
@@ -70,6 +76,9 @@ impl Arena {
pub fn alloc<T>(&self, value: T) -> Option<(&T, ArenaWeak<T>)> { pub fn alloc<T>(&self, value: T) -> Option<(&T, ArenaWeak<T>)> {
let offset = self.reserve(size_of::<T>(), align_of::<T>())?; let offset = self.reserve(size_of::<T>(), align_of::<T>())?;
// Built before the write so an unencodable offset drops `value` here
// rather than stranding it in the arena without drop glue.
let weak = ArenaWeak::new(self.generation(), offset)?;
let ptr = unsafe { self.base().add(offset) }.cast::<T>(); let ptr = unsafe { self.base().add(offset) }.cast::<T>();
// SAFETY: freshly reserved, aligned, in-bounds, unaliased. // SAFETY: freshly reserved, aligned, in-bounds, unaliased.
unsafe { ptr.write(value) }; unsafe { ptr.write(value) };
@@ -80,7 +89,7 @@ impl Arena {
self.drops.lock().unwrap().push(DropEntry { offset, drop_fn: glue::<T> }); self.drops.lock().unwrap().push(DropEntry { offset, drop_fn: glue::<T> });
} }
// SAFETY: initialized above; insert-only, so no `&mut` to it can exist. // SAFETY: initialized above; insert-only, so no `&mut` to it can exist.
Some((unsafe { &*ptr }, ArenaWeak::new(self.generation(), offset))) Some((unsafe { &*ptr }, weak))
} }
pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> Option<&[T]> { pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> Option<&[T]> {
@@ -137,20 +146,22 @@ impl<T> Copy for ArenaWeak<T> {}
impl<T> ArenaWeak<T> { impl<T> ArenaWeak<T> {
pub const NULL: Self = ArenaWeak { word: 0, _marker: PhantomData }; pub const NULL: Self = ArenaWeak { word: 0, _marker: PhantomData };
fn new(generation: u64, offset: usize) -> Self { /// `None` once the offset leaves the encodable range, so an oversized arena
debug_assert!(offset < u32::MAX as usize); /// refuses to hand out a handle rather than truncating it to a live address.
Self { fn new(generation: u64, offset: usize) -> Option<Self> {
word: ((generation & 0xFFFF_FFFF) << 32) | offset as u64, let offset = u64::try_from(offset).ok().filter(|offset| *offset <= OFFSET_MASK)?;
Some(Self {
word: ((generation & GENERATION_MASK) << OFFSET_BITS) | offset,
_marker: PhantomData, _marker: PhantomData,
} })
} }
pub fn upgrade(self, arena: &Arena) -> Option<&T> { pub fn upgrade(self, arena: &Arena) -> Option<&T> {
let generation = self.word >> 32; let generation = self.word >> OFFSET_BITS;
if generation != arena.generation() & 0xFFFF_FFFF { if generation != arena.generation() & GENERATION_MASK {
return None; return None;
} }
let offset = (self.word & 0xFFFF_FFFF) as usize; let offset = (self.word & OFFSET_MASK) as usize;
// SAFETY: same generation means the entry was fully written before its // SAFETY: same generation means the entry was fully written before its
// word was published (Release) and cannot move or be overwritten within // word was published (Release) and cannot move or be overwritten within
// a generation (insert-only); the Acquire load that produced this word // a generation (insert-only); the Acquire load that produced this word
@@ -45,8 +45,8 @@ impl<T, const CAP: usize> FrameTable<T, CAP> {
} }
pub fn lookup(&self, hash: u64) -> Lookup<'_, T> { pub fn lookup(&self, hash: u64) -> Lookup<'_, T> {
// Keys are forced odd so the empty sentinel (0) can never collide. // Only hash 0 is remapped, so distinct hashes stay distinct keys.
let key = hash | 1; let key = if hash == 0 { 1 } else { hash };
for probe in 0..CAP { for probe in 0..CAP {
let slot = &self.slots[(key as usize).wrapping_add(probe) % CAP]; let slot = &self.slots[(key as usize).wrapping_add(probe) % CAP];
let stored = slot.key.load(Ordering::Acquire); let stored = slot.key.load(Ordering::Acquire);
@@ -123,6 +123,25 @@ mod tests {
assert!(matches!(table.lookup(7), Lookup::Vacant(_))); assert!(matches!(table.lookup(7), Lookup::Vacant(_)));
} }
#[test]
fn neighboring_hashes_do_not_share_an_entry() {
let table = FrameTable::<u32, 8>::new();
let Lookup::Vacant(slot) = table.lookup(6) else { unreachable!() };
slot.publish(600, Finality::AllFinal);
assert!(matches!(table.lookup(7), Lookup::Vacant(_)), "an even hash must not answer for its odd neighbor");
}
#[test]
fn the_zero_hash_round_trips() {
let table = FrameTable::<u32, 8>::new();
let Lookup::Vacant(slot) = table.lookup(0) else { unreachable!() };
slot.publish(11, Finality::AllFinal);
let Lookup::Hit(_, value) = table.lookup(0) else {
panic!("the remapped sentinel hash must still hit");
};
assert_eq!(*value, 11);
}
#[test] #[test]
fn distinct_keys_probe_past_collisions_until_full() { fn distinct_keys_probe_past_collisions_until_full() {
let table = FrameTable::<u32, 4>::new(); let table = FrameTable::<u32, 4>::new();