flense 0.4.1

use core::{
    array,
    iter::FusedIterator,
    marker::PhantomData,
    ops::RangeBounds,
};

use crate::{
    Field,
    dim::Dim,
    lens::{
        LensBase,
        raw,
    },
    lenses::LensesBaseMut,
    type_lists::{
        Concat,
        ConsSet,
        GetMetadata,
        GetPtr,
        Mutate,
        Sculpt,
        TupleSet,
    },
};

/// Mutable reified lens.
///
/// Refers to a mutable set of [`Field`].
pub struct LensBaseMut<'a, T, const N: usize>
where
    T: TupleSet,
    T::Cons<N>: 'a,
{
    storage: T::Cons<N>,
    dimensions: Dim<usize, N>,
    _phantom: PhantomData<&'a mut T>,
}

// The lack of Copy + Clone for LensMut is for correctness; Cons<()> is always
// Copy + Clone since it is just pointers, but allowing a copy of the
// reference-like LensMut would be a soundness issue.

impl<'a, T, const N: usize> LensesBaseMut<'a, T, N> for LensBaseMut<'a, T, N>
where
    T: TupleSet,
{
    #[inline]
    fn lens_base_mut(self) -> Self {
        self
    }
}

impl<'a, T> LensBaseMut<'a, T, 0>
where
    T: TupleSet,
{
    /// Returns a shared reference to the lensed field.
    #[expect(
        clippy::should_implement_trait,
        reason = "signature is generic over the field type, not interchangeable with `AsRef`"
    )]
    #[inline]
    #[must_use]
    pub fn as_ref<'b, Elt, Index>(&'b self) -> &'b Elt::Type
    where
        'a: 'b,
        Elt: Field,
        T::Cons<0>: GetPtr<0, Elt, Index>,
    {
        // SAFETY: `Adapter` guarantees the stored 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` can
        // coexist.
        unsafe { &*self.storage.get_ptr().as_ptr() }
    }

    /// Returns a unique reference to the lensed field.
    #[expect(
        clippy::should_implement_trait,
        reason = "signature is generic over the field type, not interchangeable with `AsMut`"
    )]
    #[inline]
    #[must_use]
    pub fn as_mut<'b, Elt, Index>(&'b mut self) -> &'b mut Elt::Type
    where
        'a: 'b,
        Elt: Field,
        T::Cons<0>: GetPtr<0, Elt, Index>,
    {
        // SAFETY: `Adapter` guarantees the stored 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.
        unsafe { &mut *self.storage.get_ptr().as_ptr() }
    }

    /// Joins another [`LensBaseMut`] into `self`, forming one.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn join<Rhs>(
        self,
        other: impl LensesBaseMut<'a, Rhs, 0>,
    ) -> LensBaseMut<
        'a,
        <<T::Cons<0> as Concat<0, T::Cons<0>, Rhs::Cons<0>>>::Result as ConsSet<0>>::Tuple,
        0,
    >
    where
        Rhs: 'a + TupleSet,
        Rhs::Cons<0>: 'a,
        T::Cons<0>: Concat<
                0,
                T::Cons<0>,
                Rhs::Cons<0>,
                Result: ConsSet<
                    0,
                    Tuple: TupleSet<
                        Cons<0> = <T::Cons<0> as Concat<0, T::Cons<0>, Rhs::Cons<0>>>::Result,
                    >,
                >,
            >,
    {
        LensBaseMut {
            storage: self.storage.concat(other.lens_base_mut().storage),
            dimensions: self.dimensions,
            _phantom: PhantomData,
        }
    }
}

impl<'a, T, const N: usize> LensBaseMut<'a, T, N>
where
    T: TupleSet,
{
    #[inline]
    #[must_use]
    pub(crate) const fn new(storage: T::Cons<N>, dimensions: Dim<usize, N>) -> Self {
        Self {
            storage,
            dimensions,
            _phantom: PhantomData,
        }
    }

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

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

    /// Returns the number of lensed elements in all dimensions.
    #[inline]
    #[must_use]
    pub fn sizes(&self) -> Dim<usize, N> {
        self.dimensions
    }

    /// Returns the number of bytes between lensed elements of a field in all
    /// dimensions.
    #[inline]
    #[must_use]
    pub fn strides<Elt, Index>(&self) -> Dim<isize, N>
    where
        Elt: Field,
        T::Cons<N>: GetMetadata<N, Elt, Index>,
    {
        *self.storage.get_metadata()
    }

    /// 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::Cons<N>: GetPtr<N, 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::Cons<N>: GetPtr<N, Elt, Index>,
    {
        self.storage.get_ptr().as_ptr()
    }

    /// Converts this mutable lens into a shared [`LensBase`] over the same
    /// fields, preserving the full source lifetime `'a`.
    ///
    /// Consuming the unique borrow trades away mutation for the shared lens's
    /// abilities: the result is freely [`Copy`]able and its reads live for `'a`
    /// rather than being reborrow-scoped to a borrow of `self`.
    #[inline]
    #[must_use]
    pub fn into_const(self) -> LensBase<'a, T, N> {
        LensBase::new(self.storage, self.dimensions)
    }

    /// Returns a shared reference to the lensed element at `pos`, or `None` if
    /// `!self.dimensions.contains(pos)`.
    #[inline]
    #[must_use]
    pub fn get<'b, Elt, Index>(&'b self, pos: impl Into<Dim<usize, N>>) -> Option<&'b Elt::Type>
    where
        'a: 'b,
        Elt: Field,
        T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
    {
        let pos = pos.into();
        if self.dimensions.contains(&pos) {
            // SAFETY: `self.dimensions.contains(pos)` checked above.
            Some(unsafe { self.get_unchecked(pos) })
        } else {
            None
        }
    }

    /// Returns a shared reference to the lensed element at `pos`, without
    /// bounds checking.
    ///
    /// # Safety
    /// 1. `self.dimensions().contains(pos)` must be true.
    #[inline]
    #[must_use]
    pub unsafe fn get_unchecked<'b, Elt, Index>(
        &'b self,
        pos: impl Into<Dim<usize, N>>,
    ) -> &'b Elt::Type
    where
        'a: 'b,
        Elt: Field,
        T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
    {
        // SAFETY: Forwarded from this method's contract `dimensions.contains(pos)`.
        // Tying the returned reference to `&self` reborrows the lens shared for the
        // call's lifetime, so no `&mut Elt::Type` can coexist.
        unsafe { &*raw::elem_ptr::<T, Elt, Index, N>(&self.storage, pos.into()) }
    }

    /// Returns a unique reference to the lensed element at `pos`, or `None` if
    /// `!self.dimensions.contains(pos)`.
    #[inline]
    #[must_use]
    pub fn get_mut<'b, Elt, Index>(
        &'b mut self,
        pos: impl Into<Dim<usize, N>>,
    ) -> Option<&'b mut Elt::Type>
    where
        'a: 'b,
        Elt: Field,
        T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
    {
        let pos = pos.into();
        if self.dimensions.contains(&pos) {
            // SAFETY: `self.dimensions.contains(pos)` checked above.
            Some(unsafe { self.get_mut_unchecked(pos) })
        } else {
            None
        }
    }

    /// Returns a unique reference to the lensed element at `pos`, without
    /// bounds checking.
    ///
    /// # Safety
    /// 1. `self.dimensions().contains(pos)` must be true.
    #[inline]
    #[must_use]
    pub unsafe fn get_mut_unchecked<'b, Elt, Index>(
        &'b mut self,
        pos: impl Into<Dim<usize, N>>,
    ) -> &'b mut Elt::Type
    where
        'a: 'b,
        Elt: Field,
        T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
    {
        // SAFETY: Forwarded from this method's contract `dimensions.contains(pos)`.
        // `&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 *raw::elem_ptr::<T, Elt, Index, N>(&self.storage, pos.into()) }
    }

    /// Returns a [`Lens`](crate::lens::Lens) to the lensed fields at `pos`, or
    /// `None` if `!self.dimensions.contains(pos)`.
    #[inline]
    #[must_use]
    pub fn get_all<'b>(&'b self, pos: impl Into<Dim<usize, N>>) -> Option<LensBase<'b, T, 0>>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, 0, Result = T::Cons<0>>,
    {
        let pos = pos.into();
        if self.dimensions.contains(&pos) {
            // SAFETY: `self.dimensions.contains(pos)` checked above.
            Some(unsafe { self.get_all_unchecked(pos) })
        } else {
            None
        }
    }

    /// Returns a [`Lens`](crate::lens::Lens) to the lensed fields at `pos`,
    /// without bounds checking.
    ///
    /// # Safety
    /// 1. `self.dimensions().contains(pos)` must be true.
    #[inline]
    #[must_use]
    pub unsafe fn get_all_unchecked<'b>(
        &'b self,
        pos: impl Into<Dim<usize, N>>,
    ) -> LensBase<'b, T, 0>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, 0, Result = T::Cons<0>>,
    {
        LensBase::new(
            // SAFETY: Forwarded from this method's contract `dimensions.contains(pos)`,
            // so the mutated pointers stay within the original slice allocation. The
            // returned lens reborrows `self` shared for `'b`, so no `&mut` can coexist
            // with it.
            unsafe { raw::project_point::<T, N>(&self.storage, pos.into()) },
            Dim([]),
        )
    }

    /// Returns a [`LensMut`](crate::lens::LensMut) to the lensed fields at
    /// `pos`, or `None` if `!self.dimensions.contains(pos)`.
    #[inline]
    #[must_use]
    pub fn get_all_mut<'b>(
        &'b mut self,
        pos: impl Into<Dim<usize, N>>,
    ) -> Option<LensBaseMut<'b, T, 0>>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, 0, Result = T::Cons<0>>,
    {
        let pos = pos.into();
        if self.dimensions.contains(&pos) {
            // SAFETY: `self.dimensions.contains(pos)` checked above.
            Some(unsafe { self.get_all_mut_unchecked(pos) })
        } else {
            None
        }
    }

    /// Returns a [`LensMut`](crate::lens::LensMut) to the lensed fields at
    /// `pos`, without bounds checking.
    ///
    /// # Safety
    /// 1. `self.dimensions().contains(pos)` must be true.
    #[inline]
    #[must_use]
    pub unsafe fn get_all_mut_unchecked<'b>(
        &'b mut self,
        pos: impl Into<Dim<usize, N>>,
    ) -> LensBaseMut<'b, T, 0>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, 0, Result = T::Cons<0>>,
    {
        LensBaseMut::new(
            // SAFETY: Forwarded from this method's contract `dimensions.contains(pos)`,
            // so the mutated pointers stay within the original slice allocation. `&mut
            // self` reborrows the lens uniquely for `'b`, so the returned mutable lens
            // cannot alias any other reference reachable through the lens.
            unsafe { raw::project_point::<T, N>(&self.storage, pos.into()) },
            Dim([]),
        )
    }

    /// Slice this lens by the given range.
    #[inline]
    #[must_use]
    pub fn slice<'b, R>(&'b self, range: impl Into<Dim<R, N>>) -> Option<LensBase<'b, T, N>>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
        R: RangeBounds<usize>,
    {
        let (start, end) = range.into().to_indices(self.dimensions)?;
        if start.is_compatible_with(&end) && end.is_compatible_with(&self.dimensions) {
            // SAFETY: Verified the range compatibility above.
            Some(unsafe { self.slice_from_indices(start, end) })
        } else {
            None
        }
    }

    /// Slice this lens by the given range.
    ///
    /// # Safety
    /// 1. The range's start bounds must be compatible with its end bounds; see
    ///    [`Dim::is_compatible_with`].
    /// 2. The range's end bounds must be compatible with `self.dimensions()`.
    #[inline]
    #[must_use]
    pub unsafe fn slice_unchecked<'b, R>(
        &'b self,
        range: impl Into<Dim<R, N>>,
    ) -> LensBase<'b, T, N>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
        R: RangeBounds<usize>,
    {
        // SAFETY: By contract 2 the end bounds fit `self.dimensions()` (bounded
        // to `isize::MAX`); an overflowing bound could not, so `to_indices`
        // returns `Some`.
        let (start, end) = unsafe { range.into().to_indices(self.dimensions).unwrap_unchecked() };
        // SAFETY: Forwarded from this method's contract.
        unsafe { self.slice_from_indices(start, end) }
    }

    /// Build a shared sub-lens from already-resolved `[start, end)` indices.
    ///
    /// Takes indices, not a `RangeBounds`: a `RangeBounds` could use interior
    /// mutability to report different bounds on re-evaluation, so it must be
    /// resolved exactly once before reaching here.
    ///
    /// # Safety
    /// 1. `start` must be compatible with `end`; see
    ///    [`Dim::is_compatible_with`].
    /// 2. `end` must be compatible with `self.dimensions()`.
    #[inline]
    #[must_use]
    unsafe fn slice_from_indices<'b>(
        &'b self,
        start: Dim<usize, N>,
        end: Dim<usize, N>,
    ) -> LensBase<'b, T, N>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
    {
        LensBase::new(
            // SAFETY: Forwarded from this method's contract that `start` is compatible
            // with the dimensions, so the mutated pointers stay within the original
            // slice. The returned lens reborrows `self` shared for `'b`, so no `&mut`
            // can coexist with it.
            unsafe { raw::project_slice::<T, N>(&self.storage, start) },
            end - start,
        )
    }

    /// Slice this lens by the given range.
    #[inline]
    #[must_use]
    pub fn slice_mut<'b, R>(
        &'b mut self,
        range: impl Into<Dim<R, N>>,
    ) -> Option<LensBaseMut<'b, T, N>>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
        R: RangeBounds<usize>,
    {
        let (start, end) = range.into().to_indices(self.dimensions)?;
        if start.is_compatible_with(&end) && end.is_compatible_with(&self.dimensions) {
            // SAFETY: Verified the range compatibility above.
            Some(unsafe { self.slice_mut_from_indices(start, end) })
        } else {
            None
        }
    }

    /// Slice this lens by the given range.
    ///
    /// # Safety
    /// 1. The range's start bounds must be compatible with its end bounds; see
    ///    [`Dim::is_compatible_with`].
    /// 2. The range's end bounds must be compatible with `self.dimensions()`.
    #[inline]
    #[must_use]
    pub unsafe fn slice_mut_unchecked<'b, R>(
        &'b mut self,
        range: impl Into<Dim<R, N>>,
    ) -> LensBaseMut<'b, T, N>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
        R: RangeBounds<usize>,
    {
        // SAFETY: By contract 2 the end bounds fit `self.dimensions()` (bounded
        // to `isize::MAX`); an overflowing bound could not, so `to_indices`
        // returns `Some`.
        let (start, end) = unsafe { range.into().to_indices(self.dimensions).unwrap_unchecked() };
        // SAFETY: Forwarded from this method's contract.
        unsafe { self.slice_mut_from_indices(start, end) }
    }

    /// Build a mutable sub-lens from already-resolved `[start, end)` indices.
    ///
    /// Takes indices, not a `RangeBounds`: a `RangeBounds` could use interior
    /// mutability to report different bounds on re-evaluation, so it must be
    /// resolved exactly once before reaching here.
    ///
    /// # Safety
    /// 1. `start` must be compatible with `end`; see
    ///    [`Dim::is_compatible_with`].
    /// 2. `end` must be compatible with `self.dimensions()`.
    #[inline]
    #[must_use]
    unsafe fn slice_mut_from_indices<'b>(
        &'b mut self,
        start: Dim<usize, N>,
        end: Dim<usize, N>,
    ) -> LensBaseMut<'b, T, N>
    where
        'a: 'b,
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
    {
        LensBaseMut::new(
            // SAFETY: Forwarded from this method's contract that `start` is compatible
            // with the dimensions, so the mutated pointers stay within the original
            // slice. `&mut self` reborrows the lens uniquely for `'b`, so the returned
            // mutable lens cannot alias any other reference reachable through the lens.
            unsafe { raw::project_slice::<T, N>(&self.storage, start) },
            end - start,
        )
    }

    /// Materialize a mutable slice of a user-defined type for field access.
    ///
    /// `f` receives a pointer to each element, its size, and its strides.
    #[inline]
    pub fn materialize<U, URet>(self, f: U) -> (Dim<usize, N>, impl AsMut<[URet]> + 'a)
    where
        U: 'a + Fn(*mut (), usize, Dim<isize, N>) -> URet,
        URet: 'a,
    {
        (
            self.dimensions,
            self.storage
                .materialize(move |ptr, size, strides| f(ptr.as_ptr(), size, strides)),
        )
    }

    /// Split `self` into two mutable 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,
    ) -> (
        LensBaseMut<'a, Lhs, N>,
        LensBaseMut<
            'a,
            <<T::Cons<N> as Sculpt<Lhs::Cons<N>, Indices>>::Remainder as ConsSet<N>>::Tuple,
            N,
        >,
    )
    where
        Lhs: TupleSet,
        T::Cons<N>: Sculpt<
                Lhs::Cons<N>,
                Indices,
                Remainder: ConsSet<
                    N,
                    Tuple: TupleSet<
                        Cons<N> = <T::Cons<N> as Sculpt<Lhs::Cons<N>, Indices>>::Remainder,
                    >,
                >,
            >,
    {
        let (lhs_storage, rem_storage) = self.storage.sculpt();
        (
            LensBaseMut {
                storage: lhs_storage,
                dimensions: self.dimensions,
                _phantom: PhantomData,
            },
            LensBaseMut {
                storage: rem_storage,
                dimensions: self.dimensions,
                _phantom: PhantomData,
            },
        )
    }

    /// Joins another [`LensBaseMut`] into `self`, forming one.
    ///
    /// Truncates `other` to `self` if `other` is larger; if it is smaller,
    /// returns `None`.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn try_join<Rhs>(
        self,
        other: impl LensesBaseMut<'a, Rhs, N>,
    ) -> Option<
        LensBaseMut<
            'a,
            <<T::Cons<N> as Concat<N, T::Cons<N>, Rhs::Cons<N>>>::Result as ConsSet<N>>::Tuple,
            N,
        >,
    >
    where
        Rhs: 'a + TupleSet,
        Rhs::Cons<N>: 'a,
        T::Cons<N>: Concat<
                N,
                T::Cons<N>,
                Rhs::Cons<N>,
                Result: ConsSet<
                    N,
                    Tuple: TupleSet<
                        Cons<N> = <T::Cons<N> as Concat<N, T::Cons<N>, Rhs::Cons<N>>>::Result,
                    >,
                >,
            >,
    {
        let other = other.lens_base_mut();
        if self.dimensions.is_compatible_with(&other.dimensions) {
            Some(LensBaseMut {
                storage: self.storage.concat(other.storage),
                dimensions: self.dimensions,
                _phantom: PhantomData,
            })
        } else {
            None
        }
    }

    /// Joins another [`LensBaseMut`] into `self`, forming one.
    ///
    /// Truncates to the smaller of `self` and `other`.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn zip<Rhs>(
        self,
        other: impl LensesBaseMut<'a, Rhs, N>,
    ) -> LensBaseMut<
        'a,
        <<T::Cons<N> as Concat<N, T::Cons<N>, Rhs::Cons<N>>>::Result as ConsSet<N>>::Tuple,
        N,
    >
    where
        Rhs: 'a + TupleSet,
        Rhs::Cons<N>: 'a,
        T::Cons<N>: Concat<
                N,
                T::Cons<N>,
                Rhs::Cons<N>,
                Result: ConsSet<
                    N,
                    Tuple: TupleSet<
                        Cons<N> = <T::Cons<N> as Concat<N, T::Cons<N>, Rhs::Cons<N>>>::Result,
                    >,
                >,
            >,
    {
        let other = other.lens_base_mut();
        LensBaseMut {
            storage: self.storage.concat(other.storage),
            dimensions: self.dimensions.min(&other.dimensions),
            _phantom: PhantomData,
        }
    }
}

/// Joins two mutable lenses, forming one.
///
/// See [`LensBaseMut::join`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn join_lens_mut<'a, Lhs, Rhs>(
    lhs: impl LensesBaseMut<'a, Lhs, 0>,
    rhs: impl LensesBaseMut<'a, Rhs, 0>,
) -> LensBaseMut<
    'a,
    <<Lhs::Cons<0> as Concat<0, Lhs::Cons<0>, Rhs::Cons<0>>>::Result as ConsSet<0>>::Tuple,
    0,
>
where
    Lhs: 'a + TupleSet,
    Lhs::Cons<0>: 'a
        + Concat<
            0,
            Lhs::Cons<0>,
            Rhs::Cons<0>,
            Result: ConsSet<
                0,
                Tuple: TupleSet<
                    Cons<0> = <Lhs::Cons<0> as Concat<0, Lhs::Cons<0>, Rhs::Cons<0>>>::Result,
                >,
            >,
        >,
    Rhs: 'a + TupleSet,
    Rhs::Cons<0>: 'a,
{
    lhs.lens_base_mut().join(rhs.lens_base_mut())
}

/// Attempt to join two lenses, forming one.
///
/// See [`LensBaseMut::try_join`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn try_join_lens_mut<'a, Lhs, Rhs, const N: usize>(
    lhs: impl LensesBaseMut<'a, Lhs, N>,
    rhs: impl LensesBaseMut<'a, Rhs, N>,
) -> Option<
    LensBaseMut<
        'a,
        <<Lhs::Cons<N> as Concat<N, Lhs::Cons<N>, Rhs::Cons<N>>>::Result as ConsSet<N>>::Tuple,
        N,
    >,
>
where
    Lhs: 'a + TupleSet,
    Lhs::Cons<N>: 'a
        + Concat<
            N,
            Lhs::Cons<N>,
            Rhs::Cons<N>,
            Result: ConsSet<
                N,
                Tuple: TupleSet<
                    Cons<N> = <Lhs::Cons<N> as Concat<N, Lhs::Cons<N>, Rhs::Cons<N>>>::Result,
                >,
            >,
        >,
    Rhs: 'a + TupleSet,
    Rhs::Cons<N>: 'a,
{
    lhs.lens_base_mut().try_join(rhs.lens_base_mut())
}

/// Join two lenses, forming one.
///
/// See [`LensBaseMut::zip`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn zip_lens_mut<'a, Lhs, Rhs, const N: usize>(
    lhs: impl LensesBaseMut<'a, Lhs, N>,
    rhs: impl LensesBaseMut<'a, Rhs, N>,
) -> LensBaseMut<
    'a,
    <<Lhs::Cons<N> as Concat<N, Lhs::Cons<N>, Rhs::Cons<N>>>::Result as ConsSet<N>>::Tuple,
    N,
>
where
    Lhs: 'a + TupleSet,
    Lhs::Cons<N>: 'a
        + Concat<
            N,
            Lhs::Cons<N>,
            Rhs::Cons<N>,
            Result: ConsSet<
                N,
                Tuple: TupleSet<
                    Cons<N> = <Lhs::Cons<N> as Concat<N, Lhs::Cons<N>, Rhs::Cons<N>>>::Result,
                >,
            >,
        >,
    Rhs: 'a + TupleSet,
    Rhs::Cons<N>: 'a,
{
    lhs.lens_base_mut().zip(rhs.lens_base_mut())
}

// `IntoIterator` is implemented per concrete dimension count because stable
// Rust cannot express `M = N - 1` as a const-generic bound; these impls are
// the only constructors of `LensBaseMutIter`, so they alone pin that
// relationship.
macro_rules! impl_into_iterator {
    ($($n:literal => $m:literal),* $(,)?) => {$(
        impl<'a, T> IntoIterator for LensBaseMut<'a, T, $n>
        where
            T: TupleSet,
            T::Cons<$m>: 'a,
            T::Cons<$n>: Mutate<$n, $m, Result = T::Cons<$m>>,
            T::Cons<$n>: Mutate<$n, $n, Result = T::Cons<$n>>,
        {
            type Item = <LensBaseMutIter<'a, T, $n, $m> as Iterator>::Item;
            type IntoIter = LensBaseMutIter<'a, T, $n, $m>;

            #[inline]
            fn into_iter(self) -> Self::IntoIter {
                LensBaseMutIter { lens: self }
            }
        }
    )*};
}

impl_into_iterator!(
    1 => 0,
    2 => 1,
    3 => 2,
    4 => 3,
    5 => 4,
    6 => 5,
    7 => 6,
    8 => 7,
);

/// Iterator over the last axis of a [`LensBaseMut`], yielding mutable lenses
/// of one fewer dimension.
///
/// `M` is always `N - 1`; the relationship is enforced by the
/// [`IntoIterator`] impls, which are this type's only constructors.
///
/// `lens` doubles as both cursors. `lens.dimensions[N - 1]` is the number of
/// hyperplanes still to yield; both ends shrink it in place. `lens.storage` is
/// the front cursor: `next` reads the hyperplane there and then slides it
/// forward one step, while `next_back` indexes from that same sliding base.
/// Distinct cursor positions select disjoint hyperplanes, so the yielded
/// mutable lenses never alias. Advancing by a pointer step rather than `base +
/// stride * i` keeps LLVM from versioning the loop and inhibiting the unroller.
/// See [`raw::drop_last_axis`].
pub struct LensBaseMutIter<'a, T, const N: usize, const M: usize>
where
    T: TupleSet,
{
    lens: LensBaseMut<'a, T, N>,
}

impl<'a, T, const N: usize, const M: usize> LensBaseMutIter<'a, T, N, M>
where
    T: TupleSet,
    T::Cons<M>: 'a,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
{
    /// Returns the back cursor: the current length of the iterated-over
    /// dimension.
    #[expect(
        clippy::indexing_slicing,
        reason = "N is at least 1, enforced by the `IntoIterator` impls"
    )]
    #[inline]
    #[must_use]
    const fn end(&self) -> usize {
        self.lens.dimensions.0[N - 1]
    }

    /// Returns a mutable lens to the hyperplane at `index` of the last axis.
    ///
    /// # Safety
    /// 1. `index < self.lens.dimensions.0[N - 1]` must be true.
    /// 2. `index` must never be repeated across calls on the same iterator, so
    ///    the returned mutable lenses never alias.
    #[expect(
        clippy::indexing_slicing,
        reason = "axis indices are bounded by N, and M is always N - 1"
    )]
    #[inline]
    #[must_use]
    unsafe fn get_unchecked(&self, index: usize) -> LensBaseMut<'a, T, M> {
        #[expect(
            clippy::cast_possible_wrap,
            reason = "index is less than a dimension length, which must fit in `isize`"
        )]
        let index = index as isize;
        LensBaseMut::new(
            // SAFETY: Forwarded from this method's contract that `index` is in bounds of the last
            // axis. Hyperplanes at distinct indices are disjoint, and the caller never repeats an
            // index, so the yielded mutable lenses never alias each other.
            unsafe { raw::project_hyperplane::<T, N, M>(&self.lens.storage, index) },
            Dim(array::from_fn(|axis| self.lens.dimensions.0[axis])),
        )
    }
}

impl<'a, T, const N: usize, const M: usize> Iterator for LensBaseMutIter<'a, T, N, M>
where
    T: TupleSet,
    T::Cons<M>: 'a,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
    type Item = LensBaseMut<'a, T, M>;

    #[expect(
        clippy::indexing_slicing,
        reason = "N is at least 1, enforced by the `IntoIterator` impls"
    )]
    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.end() > 0 {
            self.lens.dimensions.0[N - 1] -= 1;
            // SAFETY: `end() > 0`, so the front cursor sits on a valid hyperplane; the
            // decrement above only touches the last axis. The cursor slides past this
            // hyperplane below, so successive yields select disjoint hyperplanes and the
            // mutable lenses never alias.
            let item = LensBaseMut::new(
                unsafe { raw::drop_last_axis::<T, N, M>(&self.lens.storage) },
                Dim(array::from_fn(|axis| self.lens.dimensions.0[axis])),
            );
            // SAFETY: the front cursor moves at most one past the last element of the axis,
            // which `advance_last_axis` tolerates.
            self.lens.storage = unsafe { raw::advance_last_axis::<T, N>(&self.lens.storage) };
            Some(item)
        } else {
            None
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.end();
        (len, Some(len))
    }
}

impl<'a, T, const N: usize, const M: usize> ExactSizeIterator for LensBaseMutIter<'a, T, N, M>
where
    T: TupleSet,
    T::Cons<M>: 'a,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
    #[inline]
    fn len(&self) -> usize {
        self.end()
    }
}

impl<'a, T, const N: usize, const M: usize> DoubleEndedIterator for LensBaseMutIter<'a, T, N, M>
where
    T: TupleSet,
    T::Cons<M>: 'a,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
    #[expect(
        clippy::indexing_slicing,
        reason = "N is at least 1, enforced by the `IntoIterator` impls"
    )]
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.end() > 0 {
            self.lens.dimensions.0[N - 1] -= 1;
            // SAFETY: The shrunken last dimension is a valid index relative to
            // the front cursor (`storage` may have slid forward via `next`);
            // front and back meet when `end()` reaches 0, so the yielded mutable
            // lens never aliases one from `next`.
            Some(unsafe { self.get_unchecked(self.end()) })
        } else {
            None
        }
    }
}

impl<'a, T, const N: usize, const M: usize> FusedIterator for LensBaseMutIter<'a, T, N, M>
where
    T: TupleSet,
    T::Cons<M>: 'a,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
}