semiflow 0.15.0-beta

//! Parallel helpers for `Strang3D::apply` when the `parallel` feature is enabled.
//!
//! Mirrors `strang2d_parallel` for 3D: same scatter/gather approach, same
//! bit-equality contract (ADR-0018). Three passes:
//!
//! - **X-pass** — X-pencils are contiguous; plain `chunks_mut` suffices.
//! - **Y-pass** — Y-pencils have stride `nx`; gather/scatter via `temp` buffer.
//! - **Z-pass** — Z-pencils have stride `nx*ny`; same gather/scatter approach.
//!
//! Each pencil is fully local to one thread; output written to disjoint memory.
//! Ceiling-divide chunking is deterministic. No FP summation across threads.
//!
//! ## Wave 5 (ADR-0045) — generic-over-F lift
//!
//! All functions are now generic over `F: SemiflowFloat + Send + Sync`.
//! Thread-local pools are dispatched via the sealed `ParallelPool3D` trait
//! (see [`crate::parallel_pool`]). The f64 codegen path is **byte-identical** to Wave 2.
//!
//! ## Wave 2 (ADR-0042) thread-local pools
//!
//! Each worker thread owns a `thread_local!` `RefCell<ScratchPool<F>>`. The
//! per-pencil `buf` that was previously allocated with `vec![F::zero(); n]` is
//! now borrowed from the per-precision pool. Capacity is grow-only; the pool is
//! cleared only via the explicit test hook [`drain_thread_local_pools_3d`].

use alloc::{sync::Arc, vec::Vec};

use crate::step_plan::{axis_plan, ThreadPencil};
use core::cell::Cell;
use std::sync::Mutex;

// ---------------------------------------------------------------------------
// Re-export the f64 3D pool for backward compat (tests use PARALLEL_3D_POOL)
// ---------------------------------------------------------------------------
/// Per-thread grow-only scratch pool for 3D parallel pencil buffers (f64).
///
/// Not part of the stable public API; not frozen at v1.0.0.
#[doc(hidden)]
pub use crate::parallel_pool::PARALLEL_3D_POOL_F64 as PARALLEL_3D_POOL;
use crate::{
    chernoff::ChernoffFunction,
    error::SemiflowError,
    float::SemiflowFloat,
    grid::Grid1D,
    grid_fn::GridFn1D,
    parallel_pool::{drain_thread_local_pools_3d_for, ParallelPool3D},
    strang2d_parallel::chunk_count,
};

/// Drain the calling thread's 3D parallel scratch pool (f64 + f32).
///
/// Clears all free buffers (capacity released). Used in
/// `tests/parallel_scratch_drain.rs`.
///
/// Not part of the stable public API; not frozen at v1.0.0.
#[doc(hidden)]
pub fn drain_thread_local_pools_3d() {
    drain_thread_local_pools_3d_for::<f64>();
    drain_thread_local_pools_3d_for::<f32>();
}

// ---------------------------------------------------------------------------
// Constants and test hook
// ---------------------------------------------------------------------------

/// Minimum pencils per thread; below this threshold the serial path is used.
pub(crate) const MIN_PENCILS_PER_THREAD: usize = 16;

thread_local! {
    /// Test hook: pin the 3D parallel thread count to a fixed value.
    ///
    /// `None` → `available_parallelism()`. `Some(k)` → exactly `k` threads.
    /// Separate from `strang2d_parallel::FORCE_THREADS` to avoid cross-
    /// contamination between 2D and 3D test runs in the same process.
    ///
    /// Not part of the stable public API; not frozen at v1.0.0.
    #[doc(hidden)]
    pub static FORCE_THREADS_3D: Cell<Option<usize>> = const { Cell::new(None) };
}

/// Resolve thread count for `total_pencils`.
pub(crate) fn resolve_threads_3d(total_pencils: usize) -> usize {
    let raw = FORCE_THREADS_3D.with(Cell::get).unwrap_or_else(|| {
        std::thread::available_parallelism().map_or(1, std::num::NonZeroUsize::get)
    });
    raw.min(total_pencils / MIN_PENCILS_PER_THREAD).max(1)
}

// ---------------------------------------------------------------------------
// Shared types
// ---------------------------------------------------------------------------

type ErrorSlot = Arc<Mutex<Option<SemiflowError>>>;

fn record_error(slot: &ErrorSlot, e: SemiflowError) {
    let mut g = slot.lock().unwrap();
    if g.is_none() {
        *g = Some(e);
    }
}

fn extract_error_3d(slot: ErrorSlot) -> Result<(), SemiflowError> {
    match Arc::try_unwrap(slot) {
        Ok(mutex) => match mutex.into_inner().unwrap() {
            Some(e) => Err(e),
            None => Ok(()),
        },
        Err(_) => Ok(()),
    }
}

// ---------------------------------------------------------------------------
// X-pass (3D) — contiguous pencils, no gather/scatter needed
// ---------------------------------------------------------------------------

/// Apply one X-pass to a 3D `state` (x-fastest) using `n_threads`.
///
/// Generic over `F: SemiflowFloat + Send + Sync` (Wave 5, ADR-0045 §5.3).
/// f64 codegen path is byte-identical to Wave 2.
pub(crate) fn parallel_x_pass_3d<X, F>(
    state: &mut [F],
    gx: Grid1D<F>,
    n_threads: usize,
    op: &X,
    tau: F,
) -> Result<(), SemiflowError>
where
    F: SemiflowFloat + Send + Sync + ParallelPool3D,
    X: ChernoffFunction<F, S = GridFn1D<F>> + Clone + Send + Sync,
{
    let nx = gx.n;
    let total_rows = state.len() / nx;
    let chunk_size = chunk_count(total_rows, n_threads);
    let error_slot: ErrorSlot = Arc::new(Mutex::new(None));
    let plan = axis_plan(op, tau, gx)?;
    let plan = plan.as_ref();

    std::thread::scope(|s| {
        for row_chunk in state.chunks_mut(chunk_size * nx) {
            let rows = row_chunk.len() / nx;
            let err_arc = Arc::clone(&error_slot);
            s.spawn(move || x_pass_chunk_3d(row_chunk, (gx, rows), op, tau, &err_arc, plan));
        }
    });

    extract_error_3d(error_slot)
}

fn x_pass_chunk_3d<X, F>(
    row_chunk: &mut [F],
    (gx, rows): (Grid1D<F>, usize),
    op: &X,
    tau: F,
    err: &ErrorSlot,
    plan: Option<&(Vec<F>, usize)>,
) where
    F: SemiflowFloat + Send + Sync + ParallelPool3D,
    X: ChernoffFunction<F, S = GridFn1D<F>>,
{
    let nx = gx.n;
    // Borrow pencil buffer from thread-local pool (Wave 2, ADR-0042).
    let mut buf = F::with_pool_3d(|pool| pool.take_vec(nx));
    let mut tp = ThreadPencil::new(plan);
    for r in 0..rows {
        let start = r * nx;
        buf.copy_from_slice(&row_chunk[start..start + nx]);
        let row_fn = GridFn1D {
            values: core::mem::take(&mut buf),
            grid: gx,
        };
        match tp.apply(op, tau, &row_fn) {
            Ok(ev) => {
                row_chunk[start..start + nx].copy_from_slice(ev);
                buf = row_fn.values;
            }
            Err(e) => {
                F::with_pool_3d(|pool| pool.return_vec(row_fn.values));
                record_error(err, e);
                return;
            }
        }
    }
    F::with_pool_3d(|pool| pool.return_vec(buf));
}

// ---------------------------------------------------------------------------
// Y-pass (3D) — stride-nx pencils, gather/scatter over (i,k) pairs
// ---------------------------------------------------------------------------

/// Shared context for Y-pass threads.
#[derive(Clone, Copy)]
struct YCtx3D<'p, F: SemiflowFloat> {
    nx: usize,
    nz: usize,
    gy: Grid1D<F>,
    plan: Option<&'p (Vec<F>, usize)>,
}

/// Apply one Y-pass to `state` (3D x-fastest) using `n_threads`.
///
/// Total Y-pencils = `nx * nz`. Two-phase gather/scatter via `temp`.
/// Generic over `F: SemiflowFloat + Send + Sync` (Wave 5, ADR-0045 §5.3).
/// Phase 2 of Y-pass: scatter `temp` back into `state` (parallel, row-major).
#[allow(clippy::too_many_arguments)]
fn scatter_y_temp<F: SemiflowFloat + Send + Sync>(
    state: &mut [F],
    temp: &[F],
    nx: usize,
    ny: usize,
    nz: usize,
    row_chunk: usize,
    col_chunk: usize,
) {
    std::thread::scope(|s| {
        for (ridx, sr) in state.chunks_mut(row_chunk * nx).enumerate() {
            let rs = ridx * row_chunk;
            let rows = sr.len() / nx;
            s.spawn(move || y_scatter(sr, temp, nx, ny, nz, col_chunk, rs, rows));
        }
    });
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn parallel_y_pass_3d<Y, F>(
    state: &mut [F],
    nx: usize,
    gy: Grid1D<F>,
    nz: usize,
    n_threads: usize,
    op: &Y,
    tau: F,
    temp: &mut Vec<F>,
) -> Result<(), SemiflowError>
where
    F: SemiflowFloat + Send + Sync + ParallelPool3D,
    Y: ChernoffFunction<F, S = GridFn1D<F>> + Clone + Send + Sync,
{
    let ny = gy.n;
    let total = nx * nz;
    let col_chunk = chunk_count(total, n_threads);
    let row_chunk = chunk_count(total, n_threads);
    let plan = axis_plan(op, tau, gy)?;
    let ctx = YCtx3D {
        nx,
        nz,
        gy,
        plan: plan.as_ref(),
    };
    temp.resize(total * ny, F::zero());
    let err: ErrorSlot = Arc::new(Mutex::new(None));

    // Phase 1: gather + apply each (i,k) Y-pencil.
    let state_ref: &[F] = state;
    std::thread::scope(|s| {
        for (cidx, tc) in temp.chunks_mut(col_chunk * ny).enumerate() {
            let ps = cidx * col_chunk;
            let p_count = tc.len() / ny;
            let e2 = Arc::clone(&err);
            s.spawn(move || y_apply_chunk(state_ref, ctx, op, tau, ps, p_count, tc, &e2));
        }
    });
    first_error(&err)?;

    // Phase 2: scatter temp → state.
    scatter_y_temp(state, temp, nx, ny, nz, row_chunk, col_chunk);
    Ok(())
}

/// Gather + apply Y-pencils `ps..ps+p_count`.
///
/// Pencil index `p = i + k*nx` → `(i = p%nx, k = p/nx)`.
#[allow(clippy::too_many_arguments)]
fn y_apply_chunk<Y, F>(
    state: &[F],
    ctx: YCtx3D<'_, F>,
    op: &Y,
    tau: F,
    ps: usize,
    p_count: usize,
    tc: &mut [F],
    err: &ErrorSlot,
) where
    F: SemiflowFloat + Send + Sync + ParallelPool3D,
    Y: ChernoffFunction<F, S = GridFn1D<F>>,
{
    let ny = ctx.gy.n;
    let nx = ctx.nx;
    let total = nx * ctx.nz;
    let mut buf = F::with_pool_3d(|pool| pool.take_vec(ny));
    let mut tp = ThreadPencil::new(ctx.plan);
    for p in 0..p_count {
        let pi = ps + p;
        if pi >= total {
            break;
        }
        let i = pi % nx;
        let k = pi / nx;
        for j in 0..ny {
            buf[j] = state[k * nx * ny + j * nx + i];
        }
        let col_fn = GridFn1D {
            values: core::mem::take(&mut buf),
            grid: ctx.gy,
        };
        match tp.apply(op, tau, &col_fn) {
            Ok(ev) => {
                let ts = p * ny;
                tc[ts..ts + ny].copy_from_slice(ev);
                buf = col_fn.values;
            }
            Err(e) => {
                F::with_pool_3d(|pool| pool.return_vec(col_fn.values));
                record_error(err, e);
                return;
            }
        }
    }
    F::with_pool_3d(|pool| pool.return_vec(buf));
}

/// Scatter Y-pass results: for row `(j,k)` in `state`, read from temp at
/// `(pencil_idx / col_chunk)*col_chunk*ny + (pencil_idx % col_chunk)*ny + j`
/// where `pencil_idx = i + k*nx`.
#[allow(clippy::too_many_arguments)]
fn y_scatter<F: SemiflowFloat + Send>(
    sr: &mut [F],
    temp: &[F],
    nx: usize,
    ny: usize,
    nz: usize,
    col_chunk: usize,
    rs: usize,
    rows: usize,
) {
    for r in 0..rows {
        let gr = rs + r;
        if gr >= ny * nz {
            break;
        }
        let j = gr % ny;
        let k = gr / ny;
        for i in 0..nx {
            let pi = i + k * nx;
            let ti = (pi / col_chunk) * col_chunk * ny + (pi % col_chunk) * ny + j;
            sr[r * nx + i] = temp[ti];
        }
    }
}

#[path = "strang3d_parallel_z.rs"]
mod z_pass;
pub(crate) use z_pass::parallel_z_pass_3d;

/// The first error recorded by a pass's worker threads, if any.
fn first_error(slot: &ErrorSlot) -> Result<(), SemiflowError> {
    match &*slot.lock().unwrap() {
        Some(e) => Err(e.clone()),
        None => Ok(()),
    }
}