melinoe 0.10.0

use crate::cell::MelinoeCell;

use super::{par_chunks::ParChunks, ShardChunks};

/// A move-only, [`Send`] write capability over a disjoint partition of a branded
/// region.
///
/// Construct one over a whole region with [`new`](Self::new), then subdivide
/// with [`split_at`](Self::split_at) / [`chunks`](Self::chunks) to obtain
/// disjoint shards for concurrent threads. The shard grants exclusive read+write
/// to *its* cells and—because it is move-only and each sub-slice is
/// non-overlapping—two shards can never reach the same cell.
///
/// `#[repr(transparent)]` over the underlying `&mut [MelinoeCell<'brand, T>]`:
/// the capability is just the slice reference, with no extra footprint. It is
/// `Send`/`Sync` exactly when `MelinoeCell<'brand, T>` is (i.e. `T: Send` /
/// `T: Send + Sync`).
#[repr(transparent)]
pub struct WriterShard<'a, 'brand, T> {
    cells: &'a mut [MelinoeCell<'brand, T>],
}

impl<'a, 'brand, T> WriterShard<'a, 'brand, T> {
    /// Wrap an exclusive borrow of a contiguous region as a single shard.
    #[inline]
    #[must_use]
    pub fn new(cells: &'a mut [MelinoeCell<'brand, T>]) -> Self {
        Self { cells }
    }

    /// Number of cells in this partition.
    #[inline]
    #[must_use]
    pub const fn len(&self) -> usize {
        self.cells.len()
    }

    /// Whether this partition is empty.
    #[inline]
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.cells.is_empty()
    }

    /// View the partition as a plain shared slice (the lower capability; needs
    /// `&self`).
    ///
    /// Zero-cost: `MelinoeCell<'brand, T>` is `#[repr(transparent)]` over
    /// `UnsafeCell<T>` over `T`, so `[MelinoeCell<'brand, T>]` shares the layout
    /// of `[T]`. Exposing a `&[T]` lets callers use ordinary slice operations
    /// (iteration, search, SIMD) over the region.
    #[inline]
    #[must_use]
    pub fn as_slice(&self) -> &[T] {
        let ptr = MelinoeCell::slice_as_unsafe_cell(self.cells).get();
        // SAFETY: the shared `&self` borrow excludes `&mut self`—the only source
        // of a `&mut T` here—and the shard's `&mut [MelinoeCell]` ownership
        // excludes all external/token access, so no `&mut T` to these cells exists
        // while the `&[T]` lives. The pointer carries whole-region provenance via
        // `UnsafeCell::get`.
        unsafe { &*ptr.cast_const() }
    }

    /// View the partition as a plain exclusive slice (the higher capability;
    /// needs `&mut self`, which also grants read).
    ///
    /// Zero-cost for the same layout reason as [`as_slice`](Self::as_slice).
    #[inline]
    #[must_use]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        let ptr = MelinoeCell::slice_as_unsafe_cell(self.cells).get();
        // SAFETY: `&mut self` grants exclusive access to the cells, so the
        // `&mut [T]` is unaliased; `UnsafeCell::get` supplies the interior-mutable
        // provenance over the whole region.
        unsafe { &mut *ptr }
    }

    /// Consume the shard into an exclusive `&mut [T]` over its whole partition.
    ///
    /// The consuming analogue of [`as_mut_slice`](Self::as_mut_slice): because
    /// the shard owns the region's exclusive borrow, taking it by value lets the
    /// returned slice carry the region's `'a` lifetime rather than borrowing the
    /// shard value. A driver that materializes one disjoint slice per buffer at
    /// the same moment — the multi-buffer chunk operators — needs every slice to
    /// outlive its own shard, which `as_mut_slice` cannot express.
    #[inline]
    #[must_use]
    pub fn into_mut_slice(self) -> &'a mut [T] {
        let ptr = MelinoeCell::slice_as_unsafe_cell(self.cells).get();
        // SAFETY: consuming `self` moves its exclusive `&'a mut
        // [MelinoeCell<'brand, T>]` borrow into this function and drops the
        // shard afterwards, so no other reference to these cells exists; the
        // `&'a mut [T]` is therefore unaliased for `'a`. Same conversion as
        // `as_mut_slice`, with the region's lifetime instead of the shard's.
        unsafe { &mut *ptr }
    }

    /// Shared read of the `index`-th cell (needs `&self`).
    #[inline]
    #[must_use]
    pub fn get(&self, index: usize) -> Option<&T> {
        self.as_slice().get(index)
    }

    /// Iterator of shared reads over the partition (needs `&self`).
    #[inline]
    pub fn iter(&self) -> core::slice::Iter<'_, T> {
        self.as_slice().iter()
    }

    /// Exclusive access to the `index`-th cell (needs `&mut self`).
    #[inline]
    #[must_use]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        self.as_mut_slice().get_mut(index)
    }

    /// Iterator of exclusive references over the partition (needs `&mut self`).
    #[inline]
    pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, T> {
        self.as_mut_slice().iter_mut()
    }

    /// Divide this shard into two disjoint shards at `mid`, consuming `self`.
    ///
    /// The two results borrow non-overlapping halves of the same region, so they
    /// may be written concurrently on different threads. This is the primitive
    /// behind recursive divide-and-conquer partitioning.
    ///
    /// # Panics
    ///
    /// Panics if `mid > self.len()` (as [`slice::split_at_mut`]).
    #[inline]
    #[must_use]
    pub fn split_at(self, mid: usize) -> (Self, Self) {
        let (left, right) = self.cells.split_at_mut(mid);
        (Self::new(left), Self::new(right))
    }

    /// Consume the shard into an iterator of disjoint shards of `chunk_size`
    /// cells each (the final shard may be shorter).
    ///
    /// `chunk_size` is clamped to at least `1`. The iterator yields strictly
    /// non-overlapping shards, suitable for distributing across a thread pool.
    #[inline]
    pub fn chunks(self, chunk_size: usize) -> ShardChunks<'a, 'brand, T> {
        ShardChunks {
            rest: Some(self.cells),
            chunk: chunk_size.max(1),
        }
    }

    /// Consume the shard into an *indexed* view of its disjoint `chunk_size`-cell
    /// partitions (the random-access counterpart to [`chunks`](Self::chunks)).
    ///
    /// Where [`chunks`](Self::chunks) yields partitions sequentially — each
    /// `next()` reborrowing the remainder — [`ParChunks`] exposes
    /// [`len`](ParChunks::len) and
    /// [`get_unchecked_chunk`](ParChunks::get_unchecked_chunk) so a caller can
    /// request partition `c` on demand without threading `&mut` state through the
    /// sequence. This is the shape a work-stealing thread pool needs (e.g.
    /// `moirai-parallel`'s partition drivers), and it is the single authoritative
    /// home for the disjoint-sub-slice range math those consumers would otherwise
    /// hand-roll.
    ///
    /// `chunk_size` is clamped to at least `1`.
    #[inline]
    pub fn par_chunks(self, chunk_size: usize) -> ParChunks<'a, 'brand, T> {
        ParChunks::new(self.cells, chunk_size)
    }

    /// Borrow the underlying cells immutably (e.g. for token-mediated reads).
    #[inline]
    #[must_use]
    pub fn as_cells(&self) -> &[MelinoeCell<'brand, T>] {
        self.cells
    }

    /// Recover the underlying exclusive slice, consuming the shard.
    #[inline]
    #[must_use]
    pub fn into_cells(self) -> &'a mut [MelinoeCell<'brand, T>] {
        self.cells
    }
}

impl<'a, 'brand, T> core::fmt::Debug for WriterShard<'a, 'brand, T> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("WriterShard")
            .field("len", &self.cells.len())
            .finish_non_exhaustive()
    }
}

impl<'s, 'a, 'brand, T> IntoIterator for &'s WriterShard<'a, 'brand, T> {
    type Item = &'s T;
    type IntoIter = core::slice::Iter<'s, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.as_slice().iter()
    }
}

impl<'s, 'a, 'brand, T> IntoIterator for &'s mut WriterShard<'a, 'brand, T> {
    type Item = &'s mut T;
    type IntoIter = core::slice::IterMut<'s, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.as_mut_slice().iter_mut()
    }
}