#![allow(unsafe_code)]
use std::os::raw::c_void;
use std::ptr;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use dtact::c_ffi::{dtact_default_config, dtact_fiber_launch, dtact_handle_t, dtact_init};
use dtact::dtact_await;
use crate::error::{FfiError, cold_ffi_error_runtime_uninitialized};
#[derive(Clone, Copy)]
pub struct RuntimeGate {
_private: (),
}
static RUNTIME_GATE: OnceLock<RuntimeGate> = OnceLock::new();
pub fn ensure_runtime() -> RuntimeGate {
*RUNTIME_GATE.get_or_init(|| {
let cfg = dtact_default_config();
let cfg_ptr: *const _ = &raw const cfg;
unsafe {
let _ = dtact_init(cfg_ptr);
}
if let Some(rt) = dtact::GLOBAL_RUNTIME.get() {
rt.start();
}
RuntimeGate { _private: () }
})
}
pub fn runtime_gate() -> Result<RuntimeGate, FfiError> {
RUNTIME_GATE
.get()
.copied()
.map_or_else(cold_ffi_error_runtime_uninitialized, Ok)
}
extern "C" fn task_trampoline(arg: *mut c_void) {
let invoke: unsafe fn(*mut ()) = unsafe { *arg.cast::<unsafe fn(*mut ())>() };
unsafe { invoke(arg.cast::<()>()) };
}
const POOL_INLINE_CAPACITY: usize = 64;
#[repr(C, align(16))]
struct PoolNode {
word0: usize,
_word1: usize,
data: [u8; POOL_INLINE_CAPACITY],
}
static POOL_HEAD: AtomicU64 = AtomicU64::new(0);
#[inline(always)]
fn pack(ptr: *mut PoolNode, tag: u16) -> u64 {
let bits = ptr as u64;
debug_assert_eq!(bits >> 48, 0, "pool ptr exceeds 48 bits — LA57 unsupported");
bits | (u64::from(tag) << 48)
}
#[inline(always)]
const fn unpack(val: u64) -> (*mut PoolNode, u16) {
let ptr = (val & 0x0000_FFFF_FFFF_FFFF) as *mut PoolNode;
let tag = (val >> 48) as u16;
(ptr, tag)
}
#[cold]
fn alloc_pool_node() -> Box<PoolNode> {
Box::new(PoolNode {
word0: 0,
_word1: 0,
data: [0u8; POOL_INLINE_CAPACITY],
})
}
#[inline]
fn pool_acquire() -> Box<PoolNode> {
let mut head_val = POOL_HEAD.load(Ordering::Acquire);
loop {
let (head, tag) = unpack(head_val);
if head.is_null() {
return alloc_pool_node();
}
let next = unsafe { (*head).word0 as *mut PoolNode };
let new_val = pack(next, tag.wrapping_add(1));
match POOL_HEAD.compare_exchange_weak(
head_val,
new_val,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return unsafe { Box::from_raw(head) },
Err(current) => head_val = current,
}
}
}
#[inline]
fn pool_release(node: Box<PoolNode>) {
let node_ptr = Box::into_raw(node);
let mut head_val = POOL_HEAD.load(Ordering::Relaxed);
loop {
let (head, tag) = unpack(head_val);
unsafe { (*node_ptr).word0 = head as usize };
let new_val = pack(node_ptr, tag.wrapping_add(1));
match POOL_HEAD.compare_exchange_weak(
head_val,
new_val,
Ordering::Release,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(current) => head_val = current,
}
}
}
unsafe fn invoke_and_drop_pooled<F: FnOnce() + Send + 'static>(raw: *mut ()) {
let node_ptr = raw.cast::<PoolNode>();
let f = unsafe { ptr::read((*node_ptr).data.as_ptr().cast::<F>()) };
pool_release(unsafe { Box::from_raw(node_ptr) });
f();
}
#[repr(C)]
struct TaskEnvelope<F: FnOnce() + Send + 'static> {
invoke: unsafe fn(*mut ()),
closure: core::mem::ManuallyDrop<F>,
}
impl<F: FnOnce() + Send + 'static> TaskEnvelope<F> {
unsafe fn invoke_and_drop(raw: *mut ()) {
let env_ptr = raw.cast::<Self>();
let f = unsafe { core::mem::ManuallyDrop::take(&mut (*env_ptr).closure) };
unsafe { drop(Box::from_raw(env_ptr)) };
f();
}
}
#[derive(Clone, Copy)]
pub struct TaskHandle(dtact_handle_t);
impl TaskHandle {
#[must_use]
pub const fn raw_id(self) -> u64 {
self.0.0
}
#[must_use]
pub const fn from_raw(id: u64) -> Self {
Self(dtact_handle_t(id))
}
}
pub fn spawn_task<F: FnOnce() + Send + 'static>(_gate: RuntimeGate, f: F) -> TaskHandle {
if core::mem::size_of::<F>() <= POOL_INLINE_CAPACITY && core::mem::align_of::<F>() <= 16 {
let mut node = pool_acquire();
node.word0 = invoke_and_drop_pooled::<F> as *const () as usize;
unsafe {
ptr::write(node.data.as_mut_ptr().cast::<F>(), f);
}
let arg = Box::into_raw(node).cast::<c_void>();
let handle = unsafe { dtact_fiber_launch(task_trampoline, arg) };
TaskHandle(handle)
} else {
let arg: *mut TaskEnvelope<F> = Box::into_raw(Box::new(TaskEnvelope {
invoke: TaskEnvelope::<F>::invoke_and_drop,
closure: core::mem::ManuallyDrop::new(f),
}));
let handle = unsafe { dtact_fiber_launch(task_trampoline, arg.cast::<c_void>()) };
TaskHandle(handle)
}
}
pub fn join(handle: TaskHandle) {
dtact_await(handle.0);
}
pub fn parallel_for_each<I, F, T>(gate: RuntimeGate, tasks: I) -> Vec<Option<T>>
where
I: IntoIterator<Item = F>,
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
use std::cell::UnsafeCell;
use std::sync::Arc;
let tasks: Vec<F> = tasks.into_iter().collect();
let n = tasks.len();
struct SendSync<T>(T);
unsafe impl<T: Send> Send for SendSync<T> {}
unsafe impl<T: Send> Sync for SendSync<T> {}
let slots: Arc<SendSync<UnsafeCell<Vec<Option<T>>>>> =
Arc::new(SendSync(UnsafeCell::new((0..n).map(|_| None).collect())));
let mut handles: Vec<TaskHandle> = Vec::with_capacity(n);
for (i, task) in tasks.into_iter().enumerate() {
let slots_arc = Arc::clone(&slots);
handles.push(spawn_task(gate, move || {
let value = std::panic::catch_unwind(std::panic::AssertUnwindSafe(task)).ok();
unsafe {
let vec_ptr: *mut Vec<Option<T>> = slots_arc.0.get();
(&mut *vec_ptr)[i] = value;
}
}));
}
for h in handles {
join(h);
}
let inner = Arc::try_unwrap(slots)
.unwrap_or_else(|_| unreachable!("all fibers have been joined; Arc is unique"));
inner.0.into_inner()
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
fn ensure_runtime_is_idempotent() {
let g1 = ensure_runtime();
let g2 = ensure_runtime();
let _ = (g1, g2);
assert!(runtime_gate().is_ok());
}
#[test]
fn spawn_task_runs_closure() {
let gate = ensure_runtime();
let counter = Arc::new(AtomicUsize::new(0));
let c = Arc::clone(&counter);
let h = spawn_task(gate, move || {
c.fetch_add(7, Ordering::Release);
});
join(h);
assert_eq!(counter.load(Ordering::Acquire), 7);
}
#[test]
fn parallel_for_each_preserves_order_and_runs_all() {
let gate = ensure_runtime();
let results = parallel_for_each(gate, (0u32..8).map(|i| move || i * i));
assert_eq!(
results,
alloc::vec![
Some(0),
Some(1),
Some(4),
Some(9),
Some(16),
Some(25),
Some(36),
Some(49)
]
);
}
#[test]
fn runtime_gate_reports_uninit_only_before_init() {
let _ = ensure_runtime();
assert!(runtime_gate().is_ok());
}
#[test]
fn pool_recycles_nodes_across_waves() {
let gate = ensure_runtime();
let counter = Arc::new(AtomicUsize::new(0));
let wave1: Vec<_> = (0..16)
.map(|_| {
let c = Arc::clone(&counter);
spawn_task(gate, move || {
c.fetch_add(1, Ordering::Release);
})
})
.collect();
for h in wave1 {
join(h);
}
assert_eq!(
counter.load(Ordering::Acquire),
16,
"wave 1 must run all tasks"
);
let wave2: Vec<_> = (0..16)
.map(|_| {
let c = Arc::clone(&counter);
spawn_task(gate, move || {
c.fetch_add(1, Ordering::Release);
})
})
.collect();
for h in wave2 {
join(h);
}
assert_eq!(
counter.load(Ordering::Acquire),
32,
"wave 2 must run all tasks"
);
}
#[test]
fn oversized_closure_uses_fallback_path() {
let gate = ensure_runtime();
let big: [u8; 256] = [42u8; 256];
let counter = Arc::new(AtomicUsize::new(0));
let c = Arc::clone(&counter);
let h = spawn_task(gate, move || {
c.fetch_add(big[0] as usize, Ordering::Release);
});
join(h);
assert_eq!(counter.load(Ordering::Acquire), 42);
}
}
#[cfg(test)]
extern crate alloc;