flense 0.2.1

use core::marker::PhantomData;

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

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

impl<'a, T> LensesSlice<'a, T> for LensSlice<'a, T>
where
    T: TupleSet,
{
    #[inline]
    fn lens_slice(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> Clone for LensSlice<'_, T>
where
    T: TupleSet,
    T::ConsSlice: Clone,
{
    #[inline]
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            len: self.len,
            _phantom: PhantomData,
        }
    }
}
impl<T> Copy for LensSlice<'_, T>
where
    T: TupleSet,
    T::ConsSlice: Copy,
{
}

impl<'a, T> LensSlice<'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 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,
    ) -> (
        LensSlice<'a, Lhs>,
        LensSlice<
            '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();
        (
            LensSlice {
                storage: lhs_storage,
                len: self.len,
                _phantom: PhantomData,
            },
            LensSlice {
                storage: rem_storage,
                len: self.len,
                _phantom: PhantomData,
            },
        )
    }

    /// Joins another [`LensSlice`] 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: LensSlice<'a, Rhs>,
    ) -> Option<
        LensSlice<
            '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 [`LensSlice`] 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: LensSlice<'a, Rhs>,
    ) -> LensSlice<
        '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,
            >,
        >,
    {
        LensSlice {
            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 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<&'a 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) -> &'a 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. The lens holds a shared borrow of the source
        // for `'a`.
        unsafe { &*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 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 LensSlice<'a, T>
where
    T: TupleSet,
    T::ConsSlice: ConsSliceToConsPtr<Result = T::ConsPtr>,
{
    type Item = <LensSliceIter<'a, T> as Iterator>::Item;
    type IntoIter = LensSliceIter<'a, T>;

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

pub struct LensSliceIter<'a, T>
where
    T: TupleSet,
{
    lens_slice: LensSlice<'a, T>,
    offset: usize,
}

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

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