diff --git a/src/sync.rs b/src/sync.rs index 8db61871..e25743e1 100644 --- a/src/sync.rs +++ b/src/sync.rs @@ -66,6 +66,12 @@ pub mod shim { /// A polyfill for `std::sync::OnceLock`. pub struct OnceLock(Mutex, UnsafeCell>); + impl Default for OnceLock { + fn default() -> Self { + OnceLock::new() + } + } + impl OnceLock { pub const fn new() -> OnceLock { OnceLock(Mutex::new(false), UnsafeCell::new(MaybeUninit::uninit())) diff --git a/src/table/memo.rs b/src/table/memo.rs index 416b6e78..0daaca6c 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -5,6 +5,7 @@ use std::{ ptr::{self, NonNull}, }; +use portable_atomic::hint::spin_loop; use thin_vec::ThinVec; use crate::sync::atomic::{AtomicPtr, Ordering}; @@ -47,6 +48,7 @@ struct MemoEntry { atomic_memo: AtomicPtr, } +#[derive(Default)] pub struct MemoEntryType { data: OnceLock, } @@ -120,22 +122,34 @@ impl MemoTableTypes { memo_type: &MemoEntryType, ) { let memo_ingredient_index = memo_ingredient_index.as_usize(); - while memo_ingredient_index >= self.types.count() { - self.types.push(MemoEntryType { - data: OnceLock::new(), - }); + + // Try to create our entry if it has not already been created. + if memo_ingredient_index >= self.types.count() { + while self.types.push(MemoEntryType::default()) < memo_ingredient_index {} + } + + loop { + let Some(memo_entry_type) = self.types.get(memo_ingredient_index) else { + // It's possible that someone else began pushing to our index but has not + // completed the entry's initialization yet, as `boxcar` is lock-free. This + // is extremely unlikely given initialization is just a handful of instructions. + // Additionally, this function is generally only called on startup, so we can + // just spin here. + spin_loop(); + continue; + }; + + memo_entry_type + .data + .set( + *memo_type.data.get().expect( + "cannot provide an empty `MemoEntryType` for `MemoEntryType::set()`", + ), + ) + .ok() + .expect("memo type should only be set once"); + break; } - let memo_entry_type = self.types.get(memo_ingredient_index).unwrap(); - memo_entry_type - .data - .set( - *memo_type - .data - .get() - .expect("cannot provide an empty `MemoEntryType` for `MemoEntryType::set()`"), - ) - .ok() - .expect("memo type should only be set once"); } /// # Safety