flense 0.2.0

use core::marker::PhantomData;

use crate::{
    Field,
    lens::{
        Lens,
        LensMut,
    },
    lenses::LensesSliceMut,
    type_lists::{
        Concat,
        ConsSet,
        ConsSliceToConsPtr,
        GetPtr,
        GetStride,
        SculptSlice,
        TupleSet,
    },
};

/// Mutable reified lens into contiguous memory.
///
/// Refers to a mutable slice of a set of [`Field`].
///
/// This is to lensing what a `&mut [_]` slice is to slices.
#[derive(Debug)]
pub struct LensSliceMut<'a, T>
where
    T: TupleSet,
{
    storage: T::ConsSlice,
    len: usize,
    _phantom: PhantomData<&'a mut ()>,
}

// The lack of Copy + Clone for LensSliceMut is for correctness; ConsSlice is
// always Copy + Clone since it is just pointers and strides, but allowing a
// copy of the reference-like LensSliceMut would be a soundness issue.

impl<'a, T> LensesSliceMut<'a, T> for LensSliceMut<'a, T>
where
    T: TupleSet,
{
    #[inline]
    fn lens_slice_mut(self) -> Self {
        self
    }
}

impl<'a, T> LensSliceMut<'a, T>
where
    T: TupleSet,
{
    #[inline]
    #[must_use]
    pub(crate) const fn new(storage: T::ConsSlice, len: usize) -> Self {
        Self {
            storage,
            len,
            _phantom: PhantomData,
        }
    }

    /// Split `self` into two mutable slice lenses.
    ///
    /// `Lhs` is the tuple of fields that should appear in the left-hand
    /// returned lens; the right-hand lens contains the remaining fields.
    ///
    /// `Indices` is a single tuple-nested list of shrinking-list positions (one
    /// per element of `Lhs`) and should be inferred; pass `_` at the call site.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn split<Lhs, Indices>(
        self,
    ) -> (
        LensSliceMut<'a, Lhs>,
        LensSliceMut<
            'a,
            <<T::ConsSlice as SculptSlice<Lhs::ConsSlice, Indices>>::Remainder as ConsSet>::Tuple,
        >,
    )
    where
        Lhs: TupleSet,
        T::ConsSlice: SculptSlice<Lhs::ConsSlice, Indices>,
        <T::ConsSlice as SculptSlice<Lhs::ConsSlice, Indices>>::Remainder: ConsSet<
            Tuple: TupleSet<
                ConsSlice = <T::ConsSlice as SculptSlice<Lhs::ConsSlice, Indices>>::Remainder,
            >,
        >,
    {
        let (lhs_storage, rem_storage) = self.storage.sculpt_slice();
        (
            LensSliceMut {
                storage: lhs_storage,
                len: self.len,
                _phantom: PhantomData,
            },
            LensSliceMut {
                storage: rem_storage,
                len: self.len,
                _phantom: PhantomData,
            },
        )
    }

    /// Joins another [`LensSliceMut`] into `self`, forming one.
    ///
    /// The joined slice is rectangular, truncating `other` if needed. All
    /// elements from `self` are retained.
    ///
    /// # Returns
    /// Returns [`None`] if `other` has a length lesser than `self.len()`.
    /// Returns the joined slices otherwise.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn join<Rhs>(
        self,
        other: LensSliceMut<'a, Rhs>,
    ) -> Option<
        LensSliceMut<
            'a,
            <<T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result as ConsSet>::Tuple,
        >,
    >
    where
        Rhs: TupleSet,
        T::ConsSlice: Concat<T::ConsSlice, Rhs::ConsSlice>,
        <T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result: ConsSet<
            Tuple: TupleSet<
                ConsSlice = <T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result,
            >,
        >,
    {
        if other.len >= self.len {
            // SAFETY: `other.len >= self.len` checked above.
            Some(unsafe { self.join_unchecked(other) })
        } else {
            None
        }
    }

    /// Joins another [`LensSliceMut`] into `self`, forming one without checking
    /// slice lengths.
    ///
    /// # Safety
    /// 1. `other` must have a length greater than or equal to `self.len()`.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub unsafe fn join_unchecked<Rhs>(
        self,
        other: LensSliceMut<'a, Rhs>,
    ) -> LensSliceMut<
        'a,
        <<T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result as ConsSet>::Tuple,
    >
    where
        Rhs: TupleSet,
        T::ConsSlice: Concat<T::ConsSlice, Rhs::ConsSlice>,
        <T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result: ConsSet<
            Tuple: TupleSet<
                ConsSlice = <T::ConsSlice as Concat<T::ConsSlice, Rhs::ConsSlice>>::Result,
            >,
        >,
    {
        LensSliceMut {
            storage: self.storage.concat(other.storage),
            len: self.len,
            _phantom: PhantomData,
        }
    }

    /// Returns a pointer to the first lensed element of a field.
    #[inline]
    #[must_use]
    pub fn as_ptr<Elt, Index>(&self) -> *const Elt::Type
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index>,
    {
        self.storage.get_ptr().as_ptr()
    }

    /// Returns a mutable pointer to the first lensed element of a field.
    #[inline]
    #[must_use]
    pub fn as_mut_ptr<Elt, Index>(&mut self) -> *mut Elt::Type
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index>,
    {
        self.storage.get_ptr().as_ptr()
    }

    /// Returns a shared reference to the `i`-th lensed element of a field, or
    /// `None` if `i >= self.len()`.
    #[inline]
    #[must_use]
    pub fn get<Elt, Index>(&self, i: usize) -> Option<&Elt::Type>
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index> + GetStride<Elt, Index>,
    {
        if i < self.len {
            // SAFETY: `i < self.len` checked above.
            Some(unsafe { self.get_unchecked(i) })
        } else {
            None
        }
    }

    /// Returns a shared reference to the `i`-th lensed element of a field,
    /// without bounds checking.
    ///
    /// # Safety
    /// `i` must be `< self.len()`.
    #[inline]
    #[must_use]
    pub unsafe fn get_unchecked<Elt, Index>(&self, i: usize) -> &Elt::Type
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index> + GetStride<Elt, Index>,
    {
        let base: *const Elt::Type = self.storage.get_ptr().as_ptr();
        let stride = self.storage.get_stride();
        // SAFETY: By the caller's contract `i < self.len()`, so the offset
        // `stride * i` lands inside the original slice. `Adapter` guarantees
        // the resulting pointer is valid for reads of `Elt::Type` and
        // properly aligned. Tying the returned reference to `&self`
        // reborrows the lens shared for the call's lifetime, so no
        // `&mut Elt::Type` to the same field can coexist.
        unsafe { &*base.byte_add(stride * i) }
    }

    /// Returns a unique reference to the `i`-th lensed element of a field, or
    /// `None` if `i >= self.len()`.
    #[inline]
    #[must_use]
    pub fn get_mut<Elt, Index>(&mut self, i: usize) -> Option<&mut Elt::Type>
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index> + GetStride<Elt, Index>,
    {
        if i < self.len {
            // SAFETY: `i < self.len` checked above.
            Some(unsafe { self.get_mut_unchecked(i) })
        } else {
            None
        }
    }

    /// Returns a unique reference to the `i`-th lensed element of a field,
    /// without bounds checking.
    ///
    /// # Safety
    /// `i` must be `< self.len()`.
    #[inline]
    #[must_use]
    pub unsafe fn get_mut_unchecked<Elt, Index>(&mut self, i: usize) -> &mut Elt::Type
    where
        Elt: Field,
        T::ConsSlice: GetPtr<Elt, Index> + GetStride<Elt, Index>,
    {
        let base: *mut Elt::Type = self.storage.get_ptr().as_ptr();
        let stride = self.storage.get_stride();
        // SAFETY: By the caller's contract `i < self.len()`, so the offset
        // `stride * i` lands inside the original slice. `Adapter` guarantees
        // the resulting pointer is valid for reads and writes of `Elt::Type`
        // and properly aligned. `&mut self` reborrows the lens uniquely for
        // the call's lifetime, so the returned `&mut` cannot alias any
        // other reference reachable through the lens; the disjointness
        // check at construction rules out aliasing through other fields.
        unsafe { &mut *base.byte_add(stride * i) }
    }

    /// Returns a [`Lens`] to the `i`-th lensed fields, or `None` if `i >=
    /// self.len()`.
    #[inline]
    #[must_use]
    pub fn get_all(&self, i: usize) -> Option<Lens<'a, T>>
    where
        T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
    {
        if i < self.len {
            // SAFETY: `i < self.len` checked above.
            Some(unsafe { self.get_all_unchecked(i) })
        } else {
            None
        }
    }

    /// Returns a [`Lens`] to the `i`-th lensed fields, without bounds checking.
    ///
    /// # Safety
    /// `i` must be `< self.len()`.
    #[inline]
    #[must_use]
    pub unsafe fn get_all_unchecked(&self, i: usize) -> Lens<'a, T>
    where
        T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
    {
        // SAFETY: Caller guarantees `i < self.len()`, so `to_cons_ptr` stays within the
        // original slice allocation.
        Lens::new(unsafe { self.storage.to_cons_ptr(i) })
    }

    /// Returns a [`LensMut`] to the `i`-th lensed fields, or `None` if `i >=
    /// self.len()`.
    #[inline]
    #[must_use]
    pub fn get_all_mut(&mut self, i: usize) -> Option<LensMut<'a, T>>
    where
        T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
    {
        if i < self.len {
            // SAFETY: `i < self.len` checked above.
            Some(unsafe { self.get_all_mut_unchecked(i) })
        } else {
            None
        }
    }

    /// Returns a [`LensMut`] to the `i`-th lensed fields, without bounds
    /// checking.
    ///
    /// # Safety
    /// `i` must be `< self.len()`.
    #[inline]
    #[must_use]
    pub unsafe fn get_all_mut_unchecked(&mut self, i: usize) -> LensMut<'a, T>
    where
        T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
    {
        // SAFETY: Caller guarantees `i < self.len()`, so `to_cons_ptr` stays within the
        // original slice allocation.
        LensMut::new(unsafe { self.storage.to_cons_ptr(i) })
    }

    /// Returns the number of lensed elements.
    #[inline]
    #[must_use]
    pub const fn len(&self) -> usize {
        self.len
    }

    /// Returns the number of bytes between lensed elements of a field.
    #[inline]
    #[must_use]
    pub fn stride<Elt, Index>(&self) -> usize
    where
        Elt: Field,
        T::ConsSlice: GetStride<Elt, Index>,
    {
        self.storage.get_stride()
    }

    /// Returns `true` if the lens has a length of 0.
    #[inline]
    #[must_use]
    pub const fn is_empty(&self) -> bool {
        self.len == 0
    }
}

impl<'a, T> IntoIterator for LensSliceMut<'a, T>
where
    T: TupleSet,
    T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
{
    type Item = <LensSliceMutIterMut<'a, T> as Iterator>::Item;
    type IntoIter = LensSliceMutIterMut<'a, T>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        LensSliceMutIterMut {
            lens_slice_mut: self,
            offset: 0,
        }
    }
}

pub struct LensSliceMutIterMut<'a, T>
where
    T: TupleSet,
    T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
{
    lens_slice_mut: LensSliceMut<'a, T>,
    offset: usize,
}

impl<'a, T> Iterator for LensSliceMutIterMut<'a, T>
where
    T: TupleSet,
    T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
{
    type Item = LensMut<'a, T>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let offset = self.offset;
        self.offset += 1;
        self.lens_slice_mut.get_all_mut(offset)
    }
}