melinoe 0.9.0

//! The single generic engine behind every partition driver.
//!
//! Both the mutable-shard ([`super::map`]) and shared-slice ([`super::read_map`])
//! drivers reduce to the same shape: *run `num_chunks` independent tasks, each
//! producing an `R`, and collect the results in index order, propagating the
//! first panic*. The only thing that varies between them is how task `index`
//! turns into an `R` — the mutable path builds a
//! [`WriterShard`](crate::region::WriterShard) from a disjoint sub-slice and
//! calls the user closure; the shared path builds a `&[T]`. That variation is
//! captured entirely by the `run: Fn(usize) -> R` closure each driver supplies,
//! so the delicate machinery — the `MaybeUninit` out-buffer, the
//! [`ExecutorDropGuard`] unwind handling, the panic mutex, and the
//! `Vec::from_raw_parts` teardown — lives here exactly once.

use std::vec::Vec;

use super::executor::registered_parallel_executor;

/// Guards the raw `MaybeUninit` out-buffer while the registered executor runs.
///
/// If a task panics and unwinds *through* the executor call (rather than being
/// caught by the per-task `catch_unwind`), this guard runs on unwind and (a)
/// drops exactly the slots the tasks marked initialized and (b) frees the
/// backing allocation, so no leak or double-drop occurs. On the normal path the
/// driver sets `active = false` after the executor returns and takes ownership of
/// the buffer itself.
struct ExecutorDropGuard<R> {
    out_ptr: *mut core::mem::MaybeUninit<R>,
    capacity: usize,
    successful_ptr: *const bool,
    num_chunks: usize,
    active: bool,
}

impl<R> Drop for ExecutorDropGuard<R> {
    fn drop(&mut self) {
        if self.active {
            for index in 0..self.num_chunks {
                // SAFETY: `successful_ptr` points at a live `[bool; num_chunks]`
                // (the driver's `successful` vec, borrowed for the executor call),
                // and `out_ptr` at a `MaybeUninit<R>` allocation of `capacity >=
                // num_chunks`. A `true` flag was set only after the matching slot
                // was `write`-initialized with a valid `R`, so dropping it in
                // place is sound and happens at most once.
                unsafe {
                    if *self.successful_ptr.add(index) {
                        self.out_ptr.add(index).cast::<R>().drop_in_place();
                    }
                }
            }
            // SAFETY: `out_ptr`/`capacity` are the pointer and capacity of the
            // `Vec<MaybeUninit<R>>` the driver `forget`-leaked; reconstructing it
            // with length `0` frees the backing allocation without dropping any
            // element (initialized ones were dropped above).
            unsafe {
                let _ = Vec::from_raw_parts(self.out_ptr, 0, self.capacity);
            }
        }
    }
}

/// Context handed to the raw task wrapper through the executor's `*mut ()`.
///
/// Holds only shared/read-only state plus the per-slot output pointers; each
/// task writes solely to its own `index` slot and flag, so no field is a shared
/// mutable aliasing hazard.
struct TaskContext<'a, R, Run> {
    run: &'a Run,
    out_ptr: *mut core::mem::MaybeUninit<R>,
    successful_ptr: *mut bool,
    panic_payload: &'a std::sync::Mutex<Option<std::boxed::Box<dyn std::any::Any + Send>>>,
}

/// The raw per-task entry point handed to a registered executor.
///
/// # Safety
///
/// `data` must be a live `*mut TaskContext<'_, R, Run>` valid for the whole
/// executor call, and the executor must invoke each `index` in
/// `0..num_chunks` at most once (the [`ParallelExecutor`](super::ParallelExecutor)
/// contract), so no two invocations touch the same `out_ptr`/`successful_ptr`
/// slot.
unsafe fn task_wrapper<R, Run>(index: usize, data: *mut ())
where
    R: Send,
    Run: Fn(usize) -> R + Sync,
{
    // SAFETY: by the function contract `data` is a live `TaskContext<'_, R, Run>`
    // borrowed for the duration of the executor call; the fields read here
    // (`run`, `panic_payload`) are shared-immutable, and the per-slot writes
    // below target this task's unique `index`.
    let ctx = unsafe { &*(data as *const TaskContext<'_, R, Run>) };

    let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (ctx.run)(index)));

    match result {
        Ok(val) => {
            // SAFETY: `index` is unique to this task (executor contract), so this
            // is the sole writer of slot `index`; the slot is uninitialized
            // `MaybeUninit<R>` and is written exactly once, after which its flag
            // is set so teardown knows to drop it.
            unsafe {
                ctx.out_ptr
                    .add(index)
                    .write(core::mem::MaybeUninit::new(val));
                *ctx.successful_ptr.add(index) = true;
            }
        }
        Err(payload) => {
            if let Ok(mut g) = ctx.panic_payload.lock() {
                if g.is_none() {
                    *g = Some(payload);
                }
            }
        }
    }
}

/// Run `num_chunks` tasks that each turn an index into an `R`, collecting the
/// results in index order and re-raising the first task panic.
///
/// `num_chunks` must be `>= 1` (callers return early for the empty region).
/// `run(index)` is invoked once per `index` in `0..num_chunks`; the closure owns
/// whatever per-index shard construction the driver needs. When a parallel
/// executor is registered it drives the tasks; otherwise the work runs on scoped
/// OS threads (one task inline on the current thread).
///
/// # Panics
///
/// Re-raises (via [`resume_unwind`](std::panic::resume_unwind)) the first panic
/// observed from any `run` invocation, after dropping every successfully
/// produced result so nothing leaks.
pub(super) fn drive<R, Run>(num_chunks: usize, run: Run) -> Vec<R>
where
    R: Send,
    Run: Fn(usize) -> R + Sync,
{
    debug_assert!(num_chunks >= 1, "drive requires at least one chunk");

    if let Some(executor) = registered_parallel_executor() {
        let mut out: Vec<core::mem::MaybeUninit<R>> = Vec::with_capacity(num_chunks);
        // SAFETY: `capacity == num_chunks`; the elements are `MaybeUninit`, so
        // extending the length to expose the (uninitialized) slots is sound —
        // each slot is initialized exactly once by its task before being read.
        unsafe {
            out.set_len(num_chunks);
        }
        let out_ptr = out.as_mut_ptr();
        let capacity = out.capacity();
        core::mem::forget(out);

        let mut successful = std::vec![false; num_chunks];
        let panic_payload = std::sync::Mutex::new(None);

        let mut guard = ExecutorDropGuard {
            out_ptr,
            capacity,
            successful_ptr: successful.as_ptr(),
            num_chunks,
            active: true,
        };

        let mut ctx = TaskContext {
            run: &run,
            out_ptr,
            successful_ptr: successful.as_mut_ptr(),
            panic_payload: &panic_payload,
        };

        // SAFETY: `task_wrapper::<R, Run>` matches the executor's task signature;
        // `&mut ctx` is a live `TaskContext<'_, R, Run>` valid for the whole call
        // (the executor blocks until every task completes). The executor contract
        // guarantees each `index` runs at most once, upholding `task_wrapper`'s
        // own safety requirement.
        unsafe {
            executor.execute(
                num_chunks,
                task_wrapper::<R, Run>,
                &mut ctx as *mut TaskContext<'_, R, Run> as *mut (),
            );
        }

        // The executor has returned: every task completed or unwound into the
        // panic mutex. Take manual ownership of the buffer away from the guard.
        guard.active = false;

        if let Some(payload) = panic_payload.into_inner().unwrap() {
            for (index, &success) in successful.iter().enumerate() {
                if success {
                    // SAFETY: `success` is set only after slot `index` was
                    // initialized with a valid `R`; drop it once to avoid leaking
                    // the results of the tasks that finished before the panic.
                    unsafe {
                        out_ptr.add(index).cast::<R>().drop_in_place();
                    }
                }
            }
            // SAFETY: reconstruct the leaked buffer with length `0` to free its
            // backing allocation; initialized slots were just dropped above.
            unsafe {
                let _ = Vec::from_raw_parts(out_ptr, 0, capacity);
            }
            std::panic::resume_unwind(payload);
        }

        // SAFETY: on the panic-free path every one of the `num_chunks` slots was
        // written exactly once (executor ran each index), so the buffer is a
        // fully-initialized `[R; num_chunks]`; `MaybeUninit<R>` shares `R`'s
        // layout, and `capacity` is the original allocation capacity, so this
        // reconstitutes the owning `Vec<R>` without copy.
        return unsafe { Vec::from_raw_parts(out_ptr.cast::<R>(), num_chunks, capacity) };
    }

    if num_chunks == 1 {
        return std::vec![run(0)];
    }

    std::thread::scope(|scope| {
        let run = &run;
        let mut handles = Vec::with_capacity(num_chunks - 1);
        for index in 0..(num_chunks - 1) {
            handles.push(scope.spawn(move || run(index)));
        }

        let last = run(num_chunks - 1);

        let mut results = Vec::with_capacity(num_chunks);
        for h in handles {
            match h.join() {
                Ok(value) => results.push(value),
                Err(payload) => std::panic::resume_unwind(payload),
            }
        }
        results.push(last);
        results
    })
}