tpt-cv-core 0.1.0

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

//! Pixel-wise arithmetic: add, sub, mul, blend, and clamp, with saturating
//! and wrapping variants. Operations write into a caller-provided destination
//! and return `false` on size mismatch.

use crate::image::{Image, ImageMut};
use crate::pixel::{Pixel, Sample};

/// `dst = a + b` with saturating integer arithmetic.
pub fn add_sat<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.saturating_add(y))
}

/// `dst = a + b` with wrapping integer arithmetic.
pub fn add_wrap<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.wrapping_add(y))
}

/// `dst = a - b` with saturating integer arithmetic.
pub fn sub_sat<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.saturating_sub(y))
}

/// `dst = a - b` with wrapping integer arithmetic.
pub fn sub_wrap<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.wrapping_sub(y))
}

/// `dst = a * b` with saturating integer arithmetic.
pub fn mul_sat<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.saturating_mul(y))
}

/// `dst = a * b` with wrapping integer arithmetic.
pub fn mul_wrap<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| x.wrapping_mul(y))
}

/// `dst = a` scaled by scalar `s`, saturating.
pub fn scale_sat<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    s: T,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    unary_op(a, dst, |x| x.saturating_mul(s))
}

/// `dst = a` scaled by scalar `s`, wrapping.
pub fn scale_wrap<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    s: T,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    unary_op(a, dst, |x| x.wrapping_mul(s))
}

/// `dst = a * (1 - t) + b * t`, `t` in `[0, 1]`, computed in unit space.
///
/// The blend factor `t` is a normalized value in `[0, 1]`.
pub fn blend<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    t: f32,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    let t = t.clamp(0.0, 1.0);
    let u = 1.0 - t;
    bin_op(a, b, dst, |x, y| {
        T::from_unit(x.to_unit() * u + y.to_unit() * t)
    })
}

/// Clamp every channel of `img` into `[low, high]` in place.
pub fn clamp<T: Sample, const C: usize>(
    img: &mut ImageMut<'_, T, C>,
    low: Pixel<T, C>,
    high: Pixel<T, C>,
) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * C;
            for c in 0..C {
                let v = row[off + c];
                row[off + c] = v.clamp(low.channels[c], high.channels[c]);
            }
        }
    }
}

/// Pointwise subtract of `b` from `a` followed by absolute value
/// (`dst = |a - b|`), useful for difference detection.
pub fn absdiff<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
) -> bool {
    bin_op(a, b, dst, |x, y| {
        if x > y {
            x.saturating_sub(y)
        } else {
            y.saturating_sub(x)
        }
    })
}

#[inline]
fn bin_op<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    b: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
    f: impl Fn(T, T) -> T,
) -> bool {
    if a.width() != dst.width()
        || a.height() != dst.height()
        || b.width() != dst.width()
        || b.height() != dst.height()
    {
        return false;
    }
    let h = a.height();
    for y in 0..h {
        let ar = a.row(y);
        let br = b.row(y);
        let dr = dst.row_mut(y);
        for (d, (x, y)) in dr.iter_mut().zip(ar.iter().zip(br.iter())) {
            *d = f(*x, *y);
        }
    }
    true
}

#[inline]
fn unary_op<T: Sample, const C: usize>(
    a: &Image<'_, T, C>,
    dst: &mut ImageMut<'_, T, C>,
    f: impl Fn(T) -> T,
) -> bool {
    if a.width() != dst.width() || a.height() != dst.height() {
        return false;
    }
    let h = a.height();
    for y in 0..h {
        let ar = a.row(y);
        let dr = dst.row_mut(y);
        for (d, x) in dr.iter_mut().zip(ar.iter()) {
            *d = f(*x);
        }
    }
    true
}

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

    fn img_of<T: Sample, const C: usize>(data: &[T]) -> ImageBuf<T, C> {
        let mut buf = ImageBuf::new(data.len() / C, 1);
        buf.as_image_mut()
            .as_contiguous_slice_mut()
            .unwrap()
            .copy_from_slice(data);
        buf
    }

    fn img_1(data: &[u8]) -> ImageBuf<u8, 1> {
        img_of(data)
    }

    #[test]
    fn add_saturating() {
        let a = img_1(&[250, 0, 10]);
        let b = img_1(&[10, 255, 5]);
        let mut dst = ImageBuf::<u8, 1>::new(3, 1);
        assert!(add_sat(
            &a.as_image(),
            &b.as_image(),
            &mut dst.as_image_mut()
        ));
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 255);
        assert_eq!(dst.as_image().pixel(1, 0).scalar(), 255);
        assert_eq!(dst.as_image().pixel(2, 0).scalar(), 15);
    }

    #[test]
    fn add_wrapping() {
        let a = img_1(&[250, 0]);
        let b = img_1(&[10, 1]);
        let mut dst = ImageBuf::<u8, 1>::new(2, 1);
        add_wrap(&a.as_image(), &b.as_image(), &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 4);
        assert_eq!(dst.as_image().pixel(1, 0).scalar(), 1);
    }

    #[test]
    fn sub_and_absdiff() {
        let a = img_1(&[10, 200, 5]);
        let b = img_1(&[30, 50, 3]);
        let mut dst = ImageBuf::<u8, 1>::new(3, 1);
        sub_sat(&a.as_image(), &b.as_image(), &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 0);
        assert_eq!(dst.as_image().pixel(1, 0).scalar(), 150);
        absdiff(&a.as_image(), &b.as_image(), &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 20);
        assert_eq!(dst.as_image().pixel(1, 0).scalar(), 150);
        assert_eq!(dst.as_image().pixel(2, 0).scalar(), 2);
    }

    #[test]
    fn scale_and_clamp() {
        let a = img_1(&[100, 200, 5]);
        let mut dst = ImageBuf::<u8, 1>::new(3, 1);
        scale_sat(&a.as_image(), 3, &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 255);
        assert_eq!(dst.as_image().pixel(1, 0).scalar(), 255);
        assert_eq!(dst.as_image().pixel(2, 0).scalar(), 15);
        clamp(&mut dst.as_image_mut(), Pixel::new([0]), Pixel::new([10]));
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 10);
        assert_eq!(dst.as_image().pixel(2, 0).scalar(), 10);
    }

    #[test]
    fn blend_known() {
        let a = img_1(&[0, 100, 255]);
        let b = img_1(&[255, 100, 0]);
        let mut dst = ImageBuf::<u8, 1>::new(3, 1);
        blend(&a.as_image(), &b.as_image(), 0.5, &mut dst.as_image_mut());
        let out = dst.as_image();
        assert_eq!(out.pixel(0, 0).scalar(), 128);
        assert_eq!(out.pixel(1, 0).scalar(), 100);
        assert_eq!(out.pixel(2, 0).scalar(), 128);
    }

    #[test]
    fn size_mismatch_rejected() {
        let a = img_1(&[1, 2, 3]);
        let b = img_1(&[1, 2]);
        let mut dst = ImageBuf::<u8, 1>::new(3, 1);
        assert!(!add_sat(
            &a.as_image(),
            &b.as_image(),
            &mut dst.as_image_mut()
        ));
        assert!(!add_sat(
            &a.as_image(),
            &a.as_image(),
            &mut ImageBuf::<u8, 1>::new(2, 1).as_image_mut()
        ));
    }

    #[test]
    fn multi_channel_ops() {
        let a = ImageBuf::<u8, 3>::with_value(1, 1, 200);
        let b = ImageBuf::<u8, 3>::with_value(1, 1, 100);
        let mut dst = ImageBuf::<u8, 3>::new(1, 1);
        add_sat(&a.as_image(), &b.as_image(), &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).channels, [255, 255, 255]);
    }

    #[test]
    fn f32_ops() {
        let a = ImageBuf::<f32, 1>::with_value(2, 1, 0.25);
        let mut dst = ImageBuf::<f32, 1>::new(2, 1);
        scale_wrap(&a.as_image(), 2.0, &mut dst.as_image_mut());
        assert_eq!(dst.as_image().pixel(0, 0).scalar(), 0.5);
    }
}