edgefirst-codec 0.26.0

// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
// SPDX-License-Identifier: Apache-2.0

//! Custom baseline JPEG decoder with a zero-allocation hot loop and strided
//! native output.
//!
//! The decoder emits the source's native format only — `Grey` for greyscale
//! (1-component) JPEGs, and the matching semi-planar format for colour
//! (3-component) JPEGs by subsampling: `Nv12` for 4:2:0, `Nv16` for 4:2:2,
//! `Nv24` for 4:4:4. It never converts to RGB and never rotates: colour and
//! geometry are applied downstream by `ImageProcessor::convert()`. EXIF
//! orientation is reported in [`ImageInfo`].

pub mod bitstream;
pub mod huffman;
pub mod idct;
pub mod markers;
pub mod mcu;
pub mod types;

/// Optional V4L2 hardware JPEG-decoder backend (Linux only, `v4l2` feature).
#[cfg(all(target_os = "linux", feature = "v4l2"))]
mod v4l2;

/// Optional nvJPEG GPU JPEG-decoder backend (Linux + CUDA, `nvjpeg` feature).
#[cfg(all(target_os = "linux", feature = "nvjpeg"))]
mod nvjpeg;

/// Runtime probe: `true` when the nvJPEG GPU decoder is available.
#[cfg(all(target_os = "linux", feature = "nvjpeg"))]
pub(crate) use nvjpeg::is_nvjpeg_available as nvjpeg_available;

/// Parse a boolean backend-gate environment variable: `true` only for
/// `1`/`true`/`yes` (case-insensitive, trimmed); absent or anything else is
/// `false`. Shared by the V4L2 opt-out (`EDGEFIRST_DISABLE_V4L2`) and the nvJPEG
/// opt-in (`EDGEFIRST_ENABLE_NVJPEG`) gates.
#[cfg(all(target_os = "linux", any(feature = "v4l2", feature = "nvjpeg")))]
pub(crate) fn env_flag(name: &str) -> bool {
    std::env::var(name)
        .map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
        .unwrap_or(false)
}

use crate::error::{CodecError, UnsupportedFeature};
use crate::exif::read_exif_orientation;
use crate::options::ImageInfo;
use crate::pixel::ImagePixel;
use edgefirst_tensor::{PixelFormat, Tensor, TensorTrait};

/// Reusable JPEG decoder state.
///
/// Holds scratch buffers that grow to the high-water mark and are reused across
/// decode calls. After the first frame at a given resolution, subsequent
/// decodes perform zero heap allocations.
pub struct JpegDecoderState {
    /// MCU scratch buffers (per-component IDCT output for one MCU row band).
    mcu_scratch: Option<mcu::McuScratch>,
    /// V4L2 hardware decoder, lazily probed on first decode. Probed at most
    /// once; a ready context is reused (and amortises per-image setup) across
    /// decodes.
    #[cfg(all(target_os = "linux", feature = "v4l2"))]
    v4l2: v4l2::V4l2Probe,
    /// nvJPEG GPU decoder, lazily probed on first decode. Preferred over V4L2
    /// when a CUDA device + libnvjpeg are present and the destination is a
    /// CUDA-backed tensor; reuses one nvJPEG handle/state/stream across decodes.
    #[cfg(all(target_os = "linux", feature = "nvjpeg"))]
    nvjpeg: nvjpeg::NvJpegProbe,
}

impl JpegDecoderState {
    pub fn new() -> Self {
        Self {
            mcu_scratch: None,
            #[cfg(all(target_os = "linux", feature = "v4l2"))]
            v4l2: v4l2::V4l2Probe::default(),
            #[cfg(all(target_os = "linux", feature = "nvjpeg"))]
            nvjpeg: nvjpeg::NvJpegProbe::default(),
        }
    }
}

impl Default for JpegDecoderState {
    fn default() -> Self {
        Self::new()
    }
}

/// The codec's native output format for a JPEG: `Grey` for 1-component
/// (greyscale) images, and the matching semi-planar format for 3-component
/// (YCbCr) images by subsampling — `Nv24` (4:4:4), `Nv16` (4:2:2), or `Nv12`
/// (4:2:0); non-standard subsamplings downsample to `Nv12`.
fn native_format(headers: &markers::JpegHeaders) -> crate::Result<PixelFormat> {
    let comps = &headers.header.components;
    match comps.len() {
        1 => Ok(PixelFormat::Grey),
        3 => {
            // Emit the JPEG's NATIVE chroma format so no resampling is needed:
            //   4:4:4 (full-res chroma)        → NV24
            //   4:2:2 (half-width chroma)      → NV16
            //   4:2:0 (quarter-res chroma)     → NV12
            // The downsample path (avg_block in `write_nv12_rows`) then only
            // runs for non-standard subsamplings (4:1:1, 4:1:0, mismatched
            // Cb/Cr), which fall back to NV12.
            let max_h = headers.header.max_h_samp as usize;
            let max_v = headers.header.max_v_samp as usize;
            let cb = comps[1].sampling;
            let cr = comps[2].sampling;
            // Both chroma components must share sampling for a clean native
            // mapping; otherwise fall back to NV12.
            if cb.h == cr.h && cb.v == cr.v && cb.h as usize > 0 && cb.v as usize > 0 {
                let h_ratio = max_h / cb.h as usize;
                let v_ratio = max_v / cb.v as usize;
                return Ok(match (h_ratio, v_ratio) {
                    (1, 1) => PixelFormat::Nv24, // 4:4:4
                    (2, 1) => PixelFormat::Nv16, // 4:2:2
                    (2, 2) => PixelFormat::Nv12, // 4:2:0
                    _ => PixelFormat::Nv12,      // exotic → downsample to 4:2:0
                });
            }
            Ok(PixelFormat::Nv12)
        }
        n => Err(CodecError::Unsupported(
            UnsupportedFeature::JpegComponentCount {
                components: n as u8,
            },
        )),
    }
}

/// Native luma/primary-plane row stride in bytes for a freshly decoded image.
/// GREY and all semi-planar luma/UV planes are 1 byte/pixel.  Semi-planar grids
/// require at least `even(width)` bytes per row (chroma alignment); the stride
/// is further rounded up to 64 bytes so `ImageInfo.row_stride` honours the
/// same 64-byte-alignment invariant as `Tensor::image()`.
///
/// This MUST stay equal to the image crate's `align_width_for_gpu_pitch(width, 1)`
/// and `PixelFormat::semi_planar_surface_dims(..)` pitch (both produce the same
/// 64-aligned even width). The codec crate can't depend on `edgefirst-image`, so
/// the value is recomputed here — keep the two in lockstep if either changes.
fn native_row_stride(width: usize) -> usize {
    width.next_multiple_of(2).next_multiple_of(64)
}

/// Parse JPEG headers and return native dimensions, format, and EXIF
/// orientation without decoding pixels. The codec does not rotate, so `width`
/// and `height` are the source's true dimensions and the orientation is
/// reported for the caller to apply via `convert()`.
pub fn peek_jpeg_info(data: &[u8]) -> crate::Result<ImageInfo> {
    let headers = markers::parse_markers(data)?;
    let img_w = headers.header.width as usize;
    let img_h = headers.header.height as usize;
    let format = native_format(&headers)?;
    let (rotation_degrees, flip_horizontal) = headers
        .exif_data
        .as_deref()
        .map(read_exif_orientation)
        .unwrap_or((0, false));
    Ok(ImageInfo {
        width: img_w,
        height: img_h,
        format,
        row_stride: native_row_stride(img_w),
        rotation_degrees,
        flip_horizontal,
    })
}

/// Decode a JPEG image from `data` into `dst`, configuring `dst`'s dimensions
/// and pixel format to the decoded native format.
///
/// The destination must be a `u8` tensor with enough underlying capacity for
/// the decoded image; smaller allocations yield
/// [`CodecError::InsufficientCapacity`]. Decoded rows are written at the
/// tensor's `effective_row_stride()` offsets. The returned [`ImageInfo`]
/// reports the EXIF orientation for the caller to apply downstream.
pub fn decode_jpeg_into<T: ImagePixel>(
    data: &[u8],
    dst: &mut Tensor<T>,
    state: &mut JpegDecoderState,
) -> crate::Result<ImageInfo> {
    let _span = tracing::trace_span!(
        "codec.decode_jpeg",
        dtype = std::any::type_name::<T>(),
        n_bytes = data.len(),
    )
    .entered();

    let headers = {
        let _s = tracing::trace_span!("codec.decode_jpeg.parse_markers").entered();
        markers::parse_markers(data)?
    };
    let img_w = headers.header.width as usize;
    let img_h = headers.header.height as usize;
    let output_fmt = native_format(&headers)?;

    // NV12 supports odd dimensions. Odd *height* gives a `H + ceil(H/2)`
    // combined-plane height (`PixelFormat::image_shape`). Odd *width* rounds the
    // tensor's buffer width up to even (also via `image_shape`), so the MCU
    // writer's `ceil(width/2)` chroma columns are byte-aligned; the reported
    // `ImageInfo.width` below stays the true odd value for a downstream crop.

    // Native NV12/GREY are u8 formats; reject non-u8 destinations.
    if T::dtype() != edgefirst_tensor::DType::U8 {
        return Err(CodecError::UnsupportedDtype(T::dtype()));
    }

    // EXIF orientation is reported, never applied.
    let (rotation_degrees, flip_horizontal) = headers
        .exif_data
        .as_deref()
        .map(read_exif_orientation)
        .unwrap_or((0, false));

    // Configure the destination tensor to the decoded image (or error if its
    // allocation is too small).
    let cap_w = dst.width().unwrap_or(0);
    let cap_h = dst.height().unwrap_or(0);
    dst.configure_image(img_w, img_h, output_fmt)
        .map_err(|e| match e {
            edgefirst_tensor::Error::InsufficientCapacity { .. } => {
                CodecError::InsufficientCapacity {
                    image: (img_w, img_h),
                    tensor: (cap_w, cap_h),
                }
            }
            other => CodecError::Tensor(other),
        })?;
    let dst_stride = dst
        .effective_row_stride()
        .unwrap_or(native_row_stride(img_w));

    // Try the nvJPEG GPU decoder first when available (CUDA + libnvjpeg). It
    // decodes interleaved RGB straight into a CUDA-backed (PBO) destination for
    // zero-copy `convert()`. A non-CUDA destination or any failure cascades to
    // the V4L2/CPU decoders.
    #[cfg(all(target_os = "linux", feature = "nvjpeg"))]
    {
        match state
            .nvjpeg
            .try_decode::<T>(data, &headers, dst, output_fmt, img_w, img_h, dst_stride)
        {
            Ok(Some(mut info)) => {
                info.rotation_degrees = rotation_degrees;
                info.flip_horizontal = flip_horizontal;
                // RGB output carries no chroma colorimetry; clear any stale
                // value left on a recycled pool tensor.
                dst.set_colorimetry(None);
                return Ok(info);
            }
            Ok(None) => {}
            Err(nvjpeg::NvJpegDecode::Fallback(why)) => {
                log::debug!("nvjpeg jpeg decode fell back: {why}");
            }
            Err(nvjpeg::NvJpegDecode::Fatal(e)) => return Err(e),
        }
    }

    // Try the V4L2 hardware decoder next when available; a probed-but-failing
    // device resets itself and falls through to the CPU decoder transparently.
    #[cfg(all(target_os = "linux", feature = "v4l2"))]
    {
        match state
            .v4l2
            .try_decode::<T>(data, &headers, dst, output_fmt, img_w, img_h, dst_stride)
        {
            Ok(Some(mut info)) => {
                info.rotation_degrees = rotation_degrees;
                info.flip_horizontal = flip_horizontal;
                dst.set_colorimetry(Some(edgefirst_tensor::Colorimetry::jfif()));
                return Ok(info);
            }
            Ok(None) => {}
            Err(v4l2::V4l2Decode::Fallback(why)) => {
                log::debug!("v4l2 jpeg decode fell back to cpu: {why}");
            }
            Err(v4l2::V4l2Decode::Fatal(e)) => return Err(e),
        }
    }

    // CPU decode: MCU loop writes NV12/GREY u8 directly into the tensor.
    match &mut state.mcu_scratch {
        Some(scratch) => scratch.ensure_capacity(&headers),
        None => state.mcu_scratch = Some(mcu::McuScratch::new(&headers)),
    }
    let mcu_scratch = state.mcu_scratch.as_mut().unwrap();

    {
        let mut map = dst.map_write()?;
        let dst_bytes: &mut [T] = &mut map;
        // SAFETY: T is u8 (checked above) — layout-identical reinterpret.
        let dst_u8: &mut [u8] = unsafe {
            std::slice::from_raw_parts_mut(dst_bytes.as_mut_ptr() as *mut u8, dst_bytes.len())
        };
        let _s = tracing::trace_span!("codec.decode_jpeg.mcu_loop").entered();
        mcu::decode_image(data, &headers, mcu_scratch, dst_u8, dst_stride, output_fmt)?;
    }

    dst.set_colorimetry(Some(edgefirst_tensor::Colorimetry::jfif()));

    Ok(ImageInfo {
        width: img_w,
        height: img_h,
        format: output_fmt,
        row_stride: dst_stride,
        rotation_degrees,
        flip_horizontal,
    })
}