tpt-cv-core 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Zero-copy image buffers over borrowed slices, with row-pitch (stride)
//! support for sub-image views and a caller-provided destination buffer
//! convention so pipelines never allocate per frame.

use core::ops::{Deref, DerefMut};

use crate::pixel::{Pixel, Sample};

/// A borrowed, multi-channel image view over `&[T]`.
///
/// `C` is the channel count (1, 3, or 4 in practice). Samples are stored
/// interleaved: pixel `(x, y)` occupies the `C` elements starting at
/// `data[y * row_stride + x * C]`.
///
/// The view is zero-copy: sub-images and row iterators borrow the parent
/// buffer. Allocating owned images requires the `alloc` feature
/// ([`ImageBuf`]).
pub struct Image<'a, T: Sample, const C: usize> {
    data: &'a [T],
    width: usize,
    height: usize,
    row_stride: usize,
}

impl<'a, T: Sample, const C: usize> Clone for Image<'a, T, C> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a, T: Sample, const C: usize> Copy for Image<'a, T, C> {}

impl<'a, T: Sample, const C: usize> Image<'a, T, C> {
    /// Create a contiguous image view over `data`.
    ///
    /// Returns `None` if `data` is too short for `width * height * C`.
    #[inline]
    pub fn new(data: &'a [T], width: usize, height: usize) -> Option<Self> {
        let needed = width.checked_mul(height)?.checked_mul(C)?;
        if data.len() < needed {
            return None;
        }
        Some(Self {
            data,
            width,
            height,
            row_stride: width * C,
        })
    }

    /// Create an image view over `data` with an explicit row stride
    /// (in samples). Enables row-pitch (padding/ROI) layouts.
    ///
    /// Returns `None` if the layout would read out of bounds.
    #[inline]
    pub fn with_stride(
        data: &'a [T],
        width: usize,
        height: usize,
        row_stride: usize,
    ) -> Option<Self> {
        if row_stride < width * C {
            return None;
        }
        let rows = data.len().checked_div(row_stride)?;
        if rows < height {
            return None;
        }
        Some(Self {
            data,
            width,
            height,
            row_stride,
        })
    }

    /// Image width in pixels.
    #[inline]
    pub const fn width(&self) -> usize {
        self.width
    }

    /// Image height in pixels.
    #[inline]
    pub const fn height(&self) -> usize {
        self.height
    }

    /// Channel count (`C`).
    #[inline]
    pub const fn channels(&self) -> usize {
        C
    }

    /// Row stride in samples.
    #[inline]
    pub const fn row_stride(&self) -> usize {
        self.row_stride
    }

    /// Whether rows are packed with no padding (`row_stride == width * C`).
    #[inline]
    pub fn is_contiguous(&self) -> bool {
        self.row_stride == self.width * C
    }

    /// Total number of samples (including row padding).
    #[inline]
    pub const fn len(&self) -> usize {
        self.data.len()
    }

    /// Whether the buffer contains no samples.
    #[inline]
    pub const fn is_empty(&self) -> bool {
        self.width == 0 || self.height == 0
    }

    /// The raw backing slice (all rows, including stride padding).
    #[inline]
    pub fn as_slice(&self) -> &'a [T] {
        self.data
    }

    /// The backing slice truncated to the exact samples used by the view,
    /// if the image is contiguous.
    #[inline]
    pub fn as_contiguous_slice(&self) -> Option<&'a [T]> {
        self.is_contiguous()
            .then(|| &self.data[..self.width * self.height * C])
    }

    /// Sample offset of pixel `(x, y)`.
    #[inline]
    pub fn offset(&self, x: usize, y: usize) -> usize {
        y * self.row_stride + x * C
    }

    /// Read the pixel at `(x, y)`.
    #[inline]
    pub fn pixel(&self, x: usize, y: usize) -> Pixel<T, C> {
        let off = self.offset(x, y);
        let mut channels = [T::ZERO; C];
        channels.copy_from_slice(&self.data[off..off + C]);
        Pixel::new(channels)
    }

    /// Row `y` as a slice of `C * width` samples.
    #[inline]
    pub fn row(&self, y: usize) -> &'a [T] {
        let start = y * self.row_stride;
        &self.data[start..start + self.width * C]
    }

    /// A sub-image (region of interest) view; zero-copy.
    ///
    /// Returns `None` if the region lies outside the image.
    #[inline]
    pub fn sub_image(&self, x: usize, y: usize, width: usize, height: usize) -> Option<Self> {
        if x.checked_add(width)? > self.width || y.checked_add(height)? > self.height {
            return None;
        }
        if width == 0 || height == 0 {
            return None;
        }
        Some(Self {
            data: &self.data[y * self.row_stride + x * C..],
            width,
            height,
            row_stride: self.row_stride,
        })
    }

    /// Iterate over rows (as `&[T]` sample slices).
    pub fn rows(&self) -> impl Iterator<Item = &'a [T]> + 'a {
        let data = self.data;
        let stride = self.row_stride;
        let row_len = self.width * C;
        (0..self.height).map(move |y| &data[y * stride..y * stride + row_len])
    }

    /// Iterate over rows as pixel-row iterators.
    pub fn iter_rows(&self) -> Rows<'a, T, C> {
        Rows {
            image: *self,
            row: 0,
        }
    }

    /// Iterate over all pixels in raster order.
    pub fn iter(&self) -> Pixels<'a, T, C> {
        Pixels {
            image: *self,
            index: 0,
        }
    }

    /// Map over pixels, writing into `out` (a matching-sized contiguous
    /// destination). Returns `false` on size mismatch (caller's job).
    pub fn map_into<S: Sample, const D: usize>(
        &self,
        out: &mut ImageMut<'_, S, D>,
        f: impl Fn(Pixel<T, C>) -> Pixel<S, D>,
    ) -> bool {
        if out.width != self.width || out.height != self.height {
            return false;
        }
        for y in 0..self.height {
            for x in 0..self.width {
                out.set_pixel(x, y, f(self.pixel(x, y)));
            }
        }
        true
    }
}

/// Iterator over rows as [`PixelRow`]s.
pub struct Rows<'a, T: Sample, const C: usize> {
    image: Image<'a, T, C>,
    row: usize,
}

impl<'a, T: Sample, const C: usize> Iterator for Rows<'a, T, C> {
    type Item = PixelRow<'a, T, C>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.row >= self.image.height {
            return None;
        }
        let row = PixelRow {
            data: self.image.row(self.row),
        };
        self.row += 1;
        Some(row)
    }
}

/// A single row of pixels borrowing the image buffer.
#[derive(Clone, Copy)]
pub struct PixelRow<'a, T: Sample, const C: usize> {
    data: &'a [T],
}

impl<'a, T: Sample, const C: usize> PixelRow<'a, T, C> {
    /// Number of pixels in the row.
    #[inline]
    pub const fn len(&self) -> usize {
        self.data.len() / C
    }

    /// Whether the row is empty.
    #[inline]
    pub const fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// The pixel at column `x`.
    #[inline]
    pub fn pixel(&self, x: usize) -> Pixel<T, C> {
        let off = x * C;
        let mut channels = [T::ZERO; C];
        channels.copy_from_slice(&self.data[off..off + C]);
        Pixel::new(channels)
    }

    /// Raw sample slice for the row.
    #[inline]
    pub fn as_slice(&self) -> &'a [T] {
        self.data
    }
}

impl<'a, T: Sample, const C: usize> Deref for PixelRow<'a, T, C> {
    type Target = [T];

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.data
    }
}

/// Iterator over all pixels in raster order.
pub struct Pixels<'a, T: Sample, const C: usize> {
    image: Image<'a, T, C>,
    index: usize,
}

impl<'a, T: Sample, const C: usize> Iterator for Pixels<'a, T, C> {
    type Item = Pixel<T, C>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let total = self.image.width * self.image.height;
        if self.index >= total {
            return None;
        }
        let x = self.index % self.image.width;
        let y = self.index / self.image.width;
        self.index += 1;
        Some(self.image.pixel(x, y))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let rem = self.image.width * self.image.height - self.index;
        (rem, Some(rem))
    }
}

impl<'a, T: Sample, const C: usize> ExactSizeIterator for Pixels<'a, T, C> {}

/// A mutable, multi-channel image view over `&mut [T]`.
pub struct ImageMut<'a, T: Sample, const C: usize> {
    data: &'a mut [T],
    width: usize,
    height: usize,
    row_stride: usize,
}

impl<'a, T: Sample, const C: usize> ImageMut<'a, T, C> {
    /// Create a contiguous mutable image view over `data`.
    #[inline]
    pub fn new(data: &'a mut [T], width: usize, height: usize) -> Option<Self> {
        let needed = width.checked_mul(height)?.checked_mul(C)?;
        if data.len() < needed {
            return None;
        }
        Some(Self {
            data,
            width,
            height,
            row_stride: width * C,
        })
    }

    /// Create a mutable image view with an explicit row stride.
    #[inline]
    pub fn with_stride(
        data: &'a mut [T],
        width: usize,
        height: usize,
        row_stride: usize,
    ) -> Option<Self> {
        if row_stride < width * C {
            return None;
        }
        let rows = data.len().checked_div(row_stride)?;
        if rows < height {
            return None;
        }
        Some(Self {
            data,
            width,
            height,
            row_stride,
        })
    }

    /// Image width in pixels.
    #[inline]
    pub const fn width(&self) -> usize {
        self.width
    }

    /// Image height in pixels.
    #[inline]
    pub const fn height(&self) -> usize {
        self.height
    }

    /// Channel count (`C`).
    #[inline]
    pub const fn channels(&self) -> usize {
        C
    }

    /// Row stride in samples.
    #[inline]
    pub const fn row_stride(&self) -> usize {
        self.row_stride
    }

    /// Whether rows are packed with no padding.
    #[inline]
    pub fn is_contiguous(&self) -> bool {
        self.row_stride == self.width * C
    }

    /// Number of samples in the backing buffer.
    #[inline]
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Whether the buffer contains no samples.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.width == 0 || self.height == 0
    }

    /// The raw backing slice.
    #[inline]
    pub fn as_slice(&self) -> &[T] {
        self.data
    }

    /// The mutable raw backing slice.
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        self.data
    }

    /// The backing slice truncated to the exact samples of the view, if
    /// contiguous.
    #[inline]
    pub fn as_contiguous_slice(&self) -> Option<&[T]> {
        self.is_contiguous()
            .then(|| &self.data[..self.width * self.height * C])
    }

    /// The mutable backing slice truncated to the exact samples of the view,
    /// if contiguous.
    #[inline]
    pub fn as_contiguous_slice_mut(&mut self) -> Option<&mut [T]> {
        self.is_contiguous()
            .then(|| &mut self.data[..self.width * self.height * C])
    }

    /// Sample offset of pixel `(x, y)`.
    #[inline]
    pub fn offset(&self, x: usize, y: usize) -> usize {
        y * self.row_stride + x * C
    }

    /// Read the pixel at `(x, y)`.
    #[inline]
    pub fn pixel(&self, x: usize, y: usize) -> Pixel<T, C> {
        let off = self.offset(x, y);
        let mut channels = [T::ZERO; C];
        channels.copy_from_slice(&self.data[off..off + C]);
        Pixel::new(channels)
    }

    /// Write the pixel at `(x, y)`.
    #[inline]
    pub fn pixel_mut(&mut self, x: usize, y: usize) -> PixelMut<'_, T, C> {
        let off = self.offset(x, y);
        PixelMut {
            data: &mut self.data[off..off + C],
        }
    }

    /// Set the pixel at `(x, y)`.
    #[inline]
    pub fn set_pixel(&mut self, x: usize, y: usize, value: Pixel<T, C>) {
        let off = self.offset(x, y);
        self.data[off..off + C].copy_from_slice(&value.channels);
    }

    /// Row `y` as an immutable sample slice.
    #[inline]
    pub fn row(&self, y: usize) -> &[T] {
        let start = y * self.row_stride;
        &self.data[start..start + self.width * C]
    }

    /// Row `y` as a mutable sample slice.
    #[inline]
    pub fn row_mut(&mut self, y: usize) -> &mut [T] {
        let start = y * self.row_stride;
        &mut self.data[start..start + self.width * C]
    }

    /// A sub-image (ROI) view; zero-copy.
    #[inline]
    pub fn sub_image(
        &self,
        x: usize,
        y: usize,
        width: usize,
        height: usize,
    ) -> Option<Image<'_, T, C>> {
        if x.checked_add(width)? > self.width || y.checked_add(height)? > self.height {
            return None;
        }
        if width == 0 || height == 0 {
            return None;
        }
        Some(Image {
            data: &self.data[y * self.row_stride..],
            width,
            height,
            row_stride: self.row_stride,
        })
    }

    /// A mutable sub-image (ROI) view; zero-copy. The view borrows `self`
    /// for the returned lifetime.
    #[inline]
    pub fn sub_image_mut(
        &mut self,
        x: usize,
        y: usize,
        width: usize,
        height: usize,
    ) -> Option<ImageMut<'_, T, C>> {
        if x.checked_add(width)? > self.width || y.checked_add(height)? > self.height {
            return None;
        }
        if width == 0 || height == 0 {
            return None;
        }
        let (row_stride, data) = (self.row_stride, &mut *self.data);
        Some(ImageMut {
            data: &mut data[y * row_stride + x * C..],
            width,
            height,
            row_stride,
        })
    }

    /// Immutable view of this mutable image.
    #[inline]
    pub fn as_image(&self) -> Image<'_, T, C> {
        Image {
            data: self.data,
            width: self.width,
            height: self.height,
            row_stride: self.row_stride,
        }
    }

    /// Fill every pixel with `value`.
    pub fn fill(&mut self, value: Pixel<T, C>) {
        for y in 0..self.height {
            let row = self.row_mut(y);
            for px in row.chunks_exact_mut(C) {
                px.copy_from_slice(&value.channels);
            }
        }
    }

    /// Copy all pixels from `src` into this image. Returns `false` on a
    /// size mismatch.
    pub fn copy_from(&mut self, src: &Image<'_, T, C>) -> bool {
        if src.width != self.width || src.height != self.height {
            return false;
        }
        let width = self.width;
        let height = self.height;
        for y in 0..height {
            let dst_row = self.row_mut(y);
            let src_row = src.row(y);
            for x in 0..width {
                let off = x * C;
                dst_row[off..off + C].copy_from_slice(&src_row[off..off + C]);
            }
        }
        true
    }
}

impl<'a, T: Sample, const C: usize> Deref for ImageMut<'a, T, C> {
    type Target = Image<'a, T, C>;

    #[inline]
    fn deref(&self) -> &Self::Target {
        // SAFETY: Image is a read-only projection of the same buffer/layout.
        unsafe { &*(self as *const Self as *const Image<'a, T, C>) }
    }
}

impl<'a, T: Sample, const C: usize> DerefMut for ImageMut<'a, T, C> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        // SAFETY: The reborrowed view is used only immutably; the struct is
        // never constructed through the pointer (its fields are identical).
        unsafe { &mut *(self as *mut Self as *mut Image<'a, T, C>) }
    }
}

/// A mutable borrow of a single pixel's channels.
pub struct PixelMut<'a, T: Sample, const C: usize> {
    data: &'a mut [T],
}

impl<'a, T: Sample, const C: usize> PixelMut<'a, T, C> {
    /// Read the current value.
    #[inline]
    pub fn get(&self) -> Pixel<T, C> {
        let mut channels = [T::ZERO; C];
        channels.copy_from_slice(self.data);
        Pixel::new(channels)
    }

    /// Overwrite the pixel.
    #[inline]
    pub fn set(&mut self, value: Pixel<T, C>) {
        self.data.copy_from_slice(&value.channels);
    }

    /// Mutable access to channel `index`.
    #[inline]
    pub fn channel_mut(&mut self, index: usize) -> &mut T {
        &mut self.data[index]
    }
}

impl<'a, T: Sample, const C: usize> Deref for PixelMut<'a, T, C> {
    type Target = [T];

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.data
    }
}

impl<'a, T: Sample, const C: usize> DerefMut for PixelMut<'a, T, C> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.data
    }
}

#[cfg(feature = "alloc")]
mod owned {
    use super::*;
    use alloc::vec::Vec;

    /// An owned, heap-allocated image. Available with the `alloc` feature.
    ///
    /// Stores the backing buffer plus the same layout parameters as
    /// [`Image`], so the owned form can be created once and reused across
    /// frames without re-allocation.
    pub struct ImageBuf<T: Sample, const C: usize> {
        data: Vec<T>,
        width: usize,
        height: usize,
        row_stride: usize,
    }

    impl<T: Sample, const C: usize> ImageBuf<T, C> {
        /// Allocate a zero-filled `width`×`height` contiguous image.
        pub fn new(width: usize, height: usize) -> Self {
            let len = width.checked_mul(height).unwrap_or(0) * C;
            Self {
                data: alloc::vec![T::ZERO; len],
                width,
                height,
                row_stride: width * C,
            }
        }

        /// Allocate a `width`×`height` image filled with `value`.
        pub fn with_value(width: usize, height: usize, value: T) -> Self {
            let len = width.checked_mul(height).unwrap_or(0) * C;
            Self {
                data: alloc::vec![value; len],
                width,
                height,
                row_stride: width * C,
            }
        }

        /// Wrap an existing vector as an image, returning `None` if its
        /// length does not match `width * height * C`.
        pub fn from_vec(data: Vec<T>, width: usize, height: usize) -> Option<Self> {
            if data.len() != width.checked_mul(height)? * C {
                return None;
            }
            Some(Self {
                row_stride: width * C,
                data,
                width,
                height,
            })
        }

        /// Image width.
        #[inline]
        pub const fn width(&self) -> usize {
            self.width
        }

        /// Image height.
        #[inline]
        pub const fn height(&self) -> usize {
            self.height
        }

        /// Borrow as an immutable view.
        #[inline]
        pub fn as_image(&self) -> Image<'_, T, C> {
            Image {
                data: &self.data,
                width: self.width,
                height: self.height,
                row_stride: self.row_stride,
            }
        }

        /// Borrow as a mutable view.
        #[inline]
        pub fn as_image_mut(&mut self) -> ImageMut<'_, T, C> {
            ImageMut {
                data: &mut self.data,
                width: self.width,
                height: self.height,
                row_stride: self.row_stride,
            }
        }

        /// Consume into the backing vector.
        #[inline]
        pub fn into_vec(self) -> Vec<T> {
            self.data
        }

        /// Reinterpret as raw bytes.
        ///
        /// Only available for `u8`-backed images: viewing a `u16`/`f32` buffer
        /// as raw bytes would be unsound, so this is gated on the sealed
        /// [`crate::pixel::ByteRepr`] marker implemented solely for `u8`
        /// (see [`ImageBuf::as_bytes`]).
        #[inline]
        pub fn as_bytes(&self) -> &[u8]
        where
            T: crate::pixel::ByteRepr,
        {
            let (prefix, bytes, suffix) = unsafe { self.data.align_to::<u8>() };
            debug_assert!(prefix.is_empty() && suffix.is_empty());
            bytes
        }
    }
}

#[cfg(feature = "alloc")]
pub use owned::ImageBuf;

#[cfg(test)]
mod tests {
    use super::*;

    #[cfg(feature = "alloc")]
    fn buf<T: Sample, const C: usize>(w: usize, h: usize, value: T) -> ImageBuf<T, C> {
        ImageBuf::with_value(w, h, value)
    }

    #[test]
    fn contiguous_view() {
        let data = [1u8, 2, 3, 4, 5, 6];
        let img = Image::<_, 3>::new(&data, 2, 1).unwrap();
        assert_eq!(img.width(), 2);
        assert_eq!(img.height(), 1);
        assert!(img.is_contiguous());
        assert_eq!(img.pixel(0, 0).channels, [1, 2, 3]);
        assert_eq!(img.pixel(1, 0).channels, [4, 5, 6]);
        assert_eq!(img.as_contiguous_slice(), Some(&data[..]));
    }

    #[test]
    fn strided_view() {
        // Two rows of 2 pixels, padded to 8 samples per row.
        let data = [0u8, 1, 2, 3, 0, 0, 0, 0, 4, 5, 6, 7, 0, 0, 0, 0];
        let img = Image::<_, 2>::with_stride(&data, 2, 2, 8).unwrap();
        assert!(!img.is_contiguous());
        assert_eq!(img.pixel(0, 0).channels, [0, 1]);
        assert_eq!(img.pixel(1, 1).channels, [6, 7]);
        assert_eq!(img.row_stride(), 8);
    }

    #[test]
    fn sub_image_zero_copy() {
        let data = [1u8, 2, 3, 4, 5, 6, 7, 8, 9];
        let img = Image::<_, 3>::new(&data, 3, 1).unwrap();
        let sub = img.sub_image(1, 0, 2, 1).unwrap();
        assert_eq!(sub.pixel(0, 0).channels, [4, 5, 6]);
        assert_eq!(sub.pixel(1, 0).channels, [7, 8, 9]);
        assert_eq!(
            sub.as_slice().as_ptr(),
            img.as_slice().as_ptr().wrapping_add(3)
        );
    }

    #[test]
    fn sub_image_bounds() {
        let data = [0u8; 12];
        let img = Image::<_, 3>::new(&data, 2, 2).unwrap();
        assert!(img.sub_image(1, 1, 2, 1).is_none());
        assert!(img.sub_image(0, 0, 0, 1).is_none());
        assert!(img.sub_image(0, 0, 1, 1).is_some());
    }

    #[test]
    fn new_rejects_short_buffer() {
        assert!(Image::<u8, 3>::new(&[0u8; 5], 2, 1).is_none());
        assert!(Image::<u8, 3>::new(&[0u8; 6], 2, 1).is_some());
    }

    #[test]
    fn iterators_match_direct_access() {
        let data = [0u8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
        let img = Image::<_, 3>::new(&data, 2, 2).unwrap();
        let collected: Vec<[u8; 3]> = img.iter().map(|p| p.channels).collect();
        assert_eq!(
            collected,
            vec![[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]]
        );
        assert_eq!(img.iter().len(), 4);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn image_buf_roundtrip() {
        let mut b = buf::<u8, 3>(2, 2, 7);
        assert_eq!(b.as_image().pixel(1, 1).channels, [7, 7, 7]);
        b.as_image_mut().set_pixel(0, 0, Pixel::new([1, 2, 3]));
        assert_eq!(b.as_image().pixel(0, 0).channels, [1, 2, 3]);
        let v = b.into_vec();
        assert_eq!(v.len(), 12);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn map_into_writes_destination() {
        let data = [10u8, 20, 30, 40, 50, 60];
        let img = Image::<_, 3>::new(&data, 2, 1).unwrap();
        let mut out = ImageBuf::<u8, 1>::new(2, 1);
        let ok = img.map_into(&mut out.as_image_mut(), |p| {
            let sum = p.channels[0] / 3 + p.channels[1] / 3 + p.channels[2] / 3;
            Pixel::new([sum])
        });
        assert!(ok);
        assert_eq!(out.as_image().pixel(0, 0).channels, [19]);
    }
}