flense 0.4.1

//! Shared unsafe pointer-arithmetic cores for [`LensBase`](super::LensBase) and
//! [`LensBaseMut`](super::LensBaseMut).
//!
//! These helpers operate on a borrowed storage chain (`&T::Cons<N>`) so both
//! the shared and mutable lenses can call them without duplicating the
//! SAFETY-critical math. Each helper carries the pointer-arithmetic
//! justification; callers forward their own preconditions and the
//! aliasing/lifetime reasoning unique to the shared-vs-unique split.

use core::{
    array,
    ptr::NonNull,
};

use crate::{
    Field,
    dim::Dim,
    type_lists::{
        GetMetadata,
        GetPtr,
        Mutate,
        TupleSet,
    },
};

/// Returns a pointer to the lensed element of field `Elt` at `pos`.
///
/// # Safety
/// 1. `dimensions.contains(pos)` must hold for the lens that owns `storage`;
///    the `LensBase`/`LensBaseMut` contract bounds `dimensions` to
///    `isize::MAX`, so `pos` fits in `isize` and the offset `sum(strides *
///    pos)` lands inside the original allocation.
#[inline]
#[must_use]
pub unsafe fn elem_ptr<T, Elt, Index, const N: usize>(
    storage: &T::Cons<N>,
    pos: Dim<usize, N>,
) -> *mut Elt::Type
where
    T: TupleSet,
    Elt: Field,
    T::Cons<N>: GetPtr<N, Elt, Index> + GetMetadata<N, Elt, Index>,
{
    let base: *mut Elt::Type = storage.get_ptr().as_ptr();
    let strides = *storage.get_metadata();
    // SAFETY: By contract `dimensions.contains(pos)`, which bounds `pos` to
    // `isize::MAX`.
    let pos: Dim<isize, N> = unsafe { pos.try_into().unwrap_unchecked() };
    // SAFETY: By contract the offset `sum(strides * pos)` lands inside the
    // original slice. `Adapter` guarantees the resulting pointer is valid and
    // properly aligned for `Elt::Type`.
    unsafe { base.byte_offset((strides * pos).sum()) }
}

/// Projects `storage` to a zero-dimensional chain pointing at `pos`.
///
/// # Safety
/// 1. `dimensions.contains(pos)` must hold for the lens that owns `storage`, so
///    the mutated pointers stay within the original allocation.
#[inline]
#[must_use]
pub unsafe fn project_point<T, const N: usize>(
    storage: &T::Cons<N>,
    pos: Dim<usize, N>,
) -> T::Cons<0>
where
    T: TupleSet,
    T::Cons<N>: Mutate<N, 0, Result = T::Cons<0>>,
{
    // SAFETY: By contract `dimensions.contains(pos)`, which bounds `pos` to
    // `isize::MAX`.
    let pos: Dim<isize, N> = unsafe { pos.try_into().unwrap_unchecked() };
    // SAFETY: By contract the mutated pointers stay within the original slice
    // allocation.
    unsafe {
        storage
            .clone()
            .mutate(|ptr, _size, strides| (ptr.byte_offset((*strides * pos).sum()), Dim([])))
    }
}

/// Projects `storage` to a sub-slice whose first element is at `start`,
/// preserving all `N` axes and their strides.
///
/// # Safety
/// 1. `start` must be compatible with the slice's dimensions; see
///    [`Dim::is_compatible_with`].
#[inline]
#[must_use]
pub unsafe fn project_slice<T, const N: usize>(
    storage: &T::Cons<N>,
    start: Dim<usize, N>,
) -> T::Cons<N>
where
    T: TupleSet,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
    // SAFETY: By contract `start` is compatible with the dimensions, which
    // bounds it to `isize::MAX`.
    let pos: Dim<isize, N> = unsafe { start.try_into().unwrap_unchecked() };
    // SAFETY: By contract the mutated pointers stay within the original slice.
    // However, because a dimension's length could become zero by moving the
    // start to the very end of the allocation, `wrapping_byte_offset` must be
    // used.
    unsafe {
        storage.clone().mutate(|ptr, _size, strides| {
            (
                NonNull::new_unchecked(ptr.as_ptr().wrapping_byte_offset((*strides * pos).sum())),
                *strides,
            )
        })
    }
}

/// Drops the last axis of `storage` without moving any pointer, yielding the
/// `M`-dimensional hyperplane at the chain's current position (`M` is always
/// `N - 1`).
///
/// This is [`project_hyperplane`] at `index == 0` with no stride/index
/// multiply: callers walk the axis with a cursor (see [`advance_last_axis`])
/// instead of forming `base + stride * index`. That keeps LLVM from versioning
/// the loop on a `stride == 1` predicate, which would inhibit the unroller.
///
/// # Safety
/// 1. The chain's current position must be a valid element of the last axis for
///    the lens that owns `storage`.
#[expect(
    clippy::indexing_slicing,
    reason = "axis indices are bounded by N, and M is always N - 1"
)]
#[inline]
#[must_use]
pub unsafe fn drop_last_axis<T, const N: usize, const M: usize>(storage: &T::Cons<N>) -> T::Cons<M>
where
    T: TupleSet,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
{
    const { assert!(N - 1 == M, "drop_last_axis requires M == N - 1") };
    // SAFETY: The pointers are unchanged; only the last axis is dropped from the
    // metadata, which can only restrict the reachable elements.
    unsafe {
        storage
            .clone()
            .mutate(|ptr, _size, strides| (ptr, Dim(array::from_fn(|axis| strides.0[axis]))))
    }
}

/// Advances every field pointer in `storage` by one step along the last axis,
/// keeping all `N` axes and their strides intact.
///
/// Each field moves by its own last-axis byte stride, so this is the cursor
/// counterpart to [`drop_last_axis`]: walk with `advance_last_axis`, read with
/// `drop_last_axis`. The step is a single stride add with no index multiply;
/// see [`drop_last_axis`] for why that matters to codegen.
///
/// # Safety
/// 1. The resulting position must stay within the original allocation, i.e. the
///    cursor is advanced at most one past the last valid element of the axis.
#[expect(clippy::indexing_slicing, reason = "axis indices are bounded by N")]
#[inline]
#[must_use]
pub unsafe fn advance_last_axis<T, const N: usize>(storage: &T::Cons<N>) -> T::Cons<N>
where
    T: TupleSet,
    T::Cons<N>: Mutate<N, N, Result = T::Cons<N>>,
{
    // SAFETY: By contract the advanced pointers stay within (or one past) the
    // original allocation; `wrapping_byte_offset` tolerates the one-past case.
    unsafe {
        storage.clone().mutate(|ptr, _size, strides| {
            (
                NonNull::new_unchecked(ptr.as_ptr().wrapping_byte_offset(strides.0[N - 1])),
                *strides,
            )
        })
    }
}

/// Projects `storage` to the hyperplane at `index` of the last axis, dropping
/// that axis to yield an `M`-dimensional chain (`M` is always `N - 1`).
///
/// # Safety
/// 1. `index < dimensions[N - 1]` must hold for the lens that owns `storage`,
///    so offsetting by `strides[N - 1] * index` stays within the original
///    allocation.
#[expect(
    clippy::indexing_slicing,
    reason = "axis indices are bounded by N, and M is always N - 1"
)]
#[inline]
#[must_use]
pub unsafe fn project_hyperplane<T, const N: usize, const M: usize>(
    storage: &T::Cons<N>,
    index: isize,
) -> T::Cons<M>
where
    T: TupleSet,
    T::Cons<N>: Mutate<N, M, Result = T::Cons<M>>,
{
    // `M == N - 1` is relied upon below: `strides.0[N - 1]` selects the dropped
    // axis and the result chain keeps the first `M` axes. The `IntoIterator`
    // impls pin this relationship, but as a standalone helper this is now
    // reachable independently, so assert it rather than trust the caller.
    const { assert!(N - 1 == M, "project_hyperplane requires M == N - 1") };
    // SAFETY: By contract `index` is in bounds of the last axis, so offsetting
    // every pointer by `strides[N - 1] * index` stays within the original
    // allocation; dropping the last axis only restricts the reachable elements
    // further.
    unsafe {
        storage.clone().mutate(|ptr, _size, strides| {
            (
                NonNull::new_unchecked(ptr.as_ptr().wrapping_byte_offset(strides.0[N - 1] * index)),
                Dim(array::from_fn(|axis| strides.0[axis])),
            )
        })
    }
}