flense 0.4.0

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

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

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

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

// Manual Clone/Copy: the derive would add `T: Clone + Copy` bounds even
// though `T` is purely a phantom type-set tag.
impl<T, const N: usize> Clone for LensBase<'_, T, N>
where
    T: TupleSet,
    T::Cons<N>: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            dimensions: self.dimensions,
            _phantom: PhantomData,
        }
    }
}
impl<T, const N: usize> Copy for LensBase<'_, T, N>
where
    T: TupleSet,
    T::Cons<N>: Copy,
{
}

// Implementations that only make sense on un-dimensioned lenses, which always
// contain exactly one element.
impl<'a, T> LensBase<'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<Elt, Index>(&self) -> &'a Elt::Type
    where
        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. The lens borrows the source
        // for `'a` shared, so handing out `&'a Elt::Type` aliases nothing it
        // doesn't already alias.
        unsafe { &*self.storage.get_ptr().as_ptr() }
    }

    /// Joins another [`LensBase`] into `self`, forming one.
    #[expect(clippy::type_complexity, reason = "necessary type inference")]
    #[inline]
    #[must_use]
    pub fn join<Rhs>(
        self,
        other: impl LensesBase<'a, Rhs, 0>,
    ) -> LensBase<
        '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,
                    >,
                >,
            >,
    {
        LensBase {
            storage: self.storage.concat(other.lens_base().storage),
            dimensions: self.dimensions,
            _phantom: PhantomData,
        }
    }
}

impl<'a, T, const N: usize> LensBase<'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 shared reference to the lensed element at `pos`, or `None` if
    /// `!self.dimensions.contains(pos)`.
    #[inline]
    #[must_use]
    pub fn get<Elt, Index>(&self, pos: impl Into<Dim<usize, N>>) -> Option<&'a Elt::Type>
    where
        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<Elt, Index>(&self, pos: impl Into<Dim<usize, N>>) -> &'a Elt::Type
    where
        Elt: Field,
        T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
    {
        // SAFETY: Forwarded from this method's contract `dimensions.contains(pos)`.
        // The lens holds a shared borrow of the source for `'a`, so handing out
        // `&'a Elt::Type` aliases nothing it doesn't already alias.
        unsafe { &*raw::elem_ptr::<T, Elt, Index, N>(&self.storage, pos.into()) }
    }

    /// Returns a [`LensBase`] to the `i`-th lensed fields, or `None` if `i >=
    /// self.len()`.
    #[inline]
    #[must_use]
    pub fn get_all(&self, pos: impl Into<Dim<usize, N>>) -> Option<LensBase<'a, T, 0>>
    where
        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 [`LensBase`] to the `i`-th lensed fields, without bounds
    /// checking.
    ///
    /// # Safety
    /// 1. `self.dimensions().contains(pos)` must be true.
    #[inline]
    #[must_use]
    pub unsafe fn get_all_unchecked(&self, pos: impl Into<Dim<usize, N>>) -> LensBase<'a, T, 0>
    where
        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.
            unsafe { raw::project_point::<T, N>(&self.storage, pos.into()) },
            Dim([]),
        )
    }

    /// Slice this lens by the given range.
    #[inline]
    #[must_use]
    pub fn slice<R>(&self, range: impl Into<Dim<R, N>>) -> Option<Self>
    where
        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<R>(&self, range: impl Into<Dim<R, N>>) -> Self
    where
        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 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(&self, start: Dim<usize, N>, end: Dim<usize, N>) -> Self
    where
        T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
    {
        Self::new(
            // SAFETY: Forwarded from this method's contract that `start` is compatible
            // with the dimensions, so the mutated pointers stay within the original
            // slice.
            unsafe { raw::project_slice::<T, N>(&self.storage, start) },
            end - start,
        )
    }

    /// Materialize a 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 AsRef<[URet]> + 'a)
    where
        U: 'a + Fn(*const (), usize, Dim<isize, N>) -> URet,
        URet: 'a,
    {
        (
            self.dimensions,
            self.storage
                .materialize(move |ptr, size, strides| f(ptr.as_ptr().cast_const(), size, strides)),
        )
    }

    /// Split `self` into two 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,
    ) -> (
        LensBase<'a, Lhs, N>,
        LensBase<
            '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();
        (
            LensBase {
                storage: lhs_storage,
                dimensions: self.dimensions,
                _phantom: PhantomData,
            },
            LensBase {
                storage: rem_storage,
                dimensions: self.dimensions,
                _phantom: PhantomData,
            },
        )
    }

    /// Joins another lens 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 LensesBase<'a, Rhs, N>,
    ) -> Option<
        LensBase<
            '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();
        if self.dimensions.is_compatible_with(&other.dimensions) {
            Some(LensBase {
                storage: self.storage.concat(other.storage),
                dimensions: self.dimensions,
                _phantom: PhantomData,
            })
        } else {
            None
        }
    }

    /// Zips another lens 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 LensesBase<'a, Rhs, N>,
    ) -> LensBase<
        '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();
        LensBase {
            storage: self.storage.concat(other.storage),
            dimensions: self.dimensions.min(&other.dimensions),
            _phantom: PhantomData,
        }
    }
}

/// Join two lenses, forming one.
///
/// See [`LensBase::join`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn join_lens<'a, Lhs, Rhs>(
    lhs: impl LensesBase<'a, Lhs, 0>,
    rhs: impl LensesBase<'a, Rhs, 0>,
) -> LensBase<
    '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().join(rhs.lens_base())
}

/// Attempt to join two lenses, forming one.
///
/// See [`LensBase::try_join`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn try_join_lens<'a, Lhs, Rhs, const N: usize>(
    lhs: impl LensesBase<'a, Lhs, N>,
    rhs: impl LensesBase<'a, Rhs, N>,
) -> Option<
    LensBase<
        '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().try_join(rhs.lens_base())
}

/// Join two lenses, forming one.
///
/// See [`LensBase::zip`].
#[expect(clippy::type_complexity, reason = "necessary type inference")]
#[inline]
#[must_use]
pub fn zip_lens<'a, Lhs, Rhs, const N: usize>(
    lhs: impl LensesBase<'a, Lhs, N>,
    rhs: impl LensesBase<'a, Rhs, N>,
) -> LensBase<
    '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().zip(rhs.lens_base())
}

// `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 `LensBaseIter`, so they alone pin that
// relationship.
macro_rules! impl_into_iterator {
    ($($n:literal => $m:literal),* $(,)?) => {$(
        impl<'a, T> IntoIterator for LensBase<'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 = <LensBaseIter<'a, T, $n, $m> as Iterator>::Item;
            type IntoIter = LensBaseIter<'a, T, $n, $m>;

            #[inline]
            fn into_iter(self) -> Self::IntoIter {
                LensBaseIter { 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 [`LensBase`], yielding 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.
/// 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 LensBaseIter<'a, T, const N: usize, const M: usize>
where
    T: TupleSet,
{
    lens: LensBase<'a, T, N>,
}

impl<'a, T, const N: usize, const M: usize> LensBaseIter<'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 lens to the hyperplane at `index` of the last axis.
    ///
    /// # Safety
    /// 1. `index < self.lens.dimensions.0[N - 1]` must be true.
    #[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) -> LensBase<'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;
        LensBase::new(
            // SAFETY: Forwarded from this method's contract that `index` is in
            // bounds of the last axis.
            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 LensBaseIter<'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 = LensBase<'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.
            let item = LensBase::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 LensBaseIter<'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 LensBaseIter<'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.
            Some(unsafe { self.get_unchecked(self.end()) })
        } else {
            None
        }
    }
}

impl<'a, T, const N: usize, const M: usize> FusedIterator for LensBaseIter<'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>>,
{
}