edgefirst-codec 0.38.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 by default — `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. Callers may opt into a **pure CPU** fused output
//! via [`crate::ImageDecoder::set_output_format`]: `Rgb` (4:4:4 or native
//! 4:2:0 sources — the latter via a fused 2×2 box chroma upsample) or
//! `Nv12` (any colour source), converted at the MCU write stage in a single
//! software-decode pass (not a GPU hybrid). The decoder never rotates:
//! geometry and any remaining preprocessing are applied downstream by
//! `ImageProcessor::convert()`. EXIF orientation is reported in
//! [`ImageInfo`].

pub mod bitstream;
pub mod color;
pub mod cpu;
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;

/// Runtime probe: `true` when a V4L2 JPEG decoder device is present (and not
/// opted out via `EDGEFIRST_DISABLE_V4L2`). Opens and drops the device once.
#[cfg(all(target_os = "linux", feature = "v4l2"))]
pub(crate) fn v4l2_available() -> bool {
    v4l2::device_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};

/// 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>,
    /// Huffman LUTs reused across frames when the DHT payload is unchanged.
    /// Tables are moved into [`markers::JpegHeaders`] for the decode and
    /// restored afterwards so a cache hit allocates nothing.
    huffman_cache: huffman::HuffmanCache,
    /// Optional fused-output preference (see [`resolve_output_format`]).
    pub(crate) preferred_format: Option<PixelFormat>,
    /// Chroma upsampling filter for a fused native-4:2:0 `Rgb` decode — see
    /// [`crate::ChromaUpsample`]. Only `Box` is implemented; the field
    /// exists so the type is threaded through end to end before a second
    /// filter is added, rather than retrofitted later.
    pub(crate) chroma_upsample: crate::decoder::ChromaUpsample,
    /// Opt-in fast (AAN `ifast`-class) IDCT — see [`crate::DctMethod`].
    /// Default false; `EDGEFIRST_CODEC_DCT=fast` flips the constructor
    /// default so benchmarks can A/B without code changes.
    pub(crate) fast_dct: bool,
    /// 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,
            huffman_cache: huffman::HuffmanCache::default(),
            preferred_format: None,
            chroma_upsample: crate::decoder::ChromaUpsample::default(),
            fast_dct: std::env::var("EDGEFIRST_CODEC_DCT")
                .map(|v| v.trim().eq_ignore_ascii_case("fast"))
                .unwrap_or(false),
            #[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,
            },
        )),
    }
}

/// Resolve the decode output format from the native format and an optional
/// caller preference (fused decode outputs):
///
/// - `Rgb`: honoured for 4:4:4 colour JPEGs (native `Nv24`, straight per-row
///   YCbCr→RGB, no chroma resampling) and for native 4:2:0 colour JPEGs
///   (`headers` sampling factors match [`mcu::is_native_420`] exactly — a 2×2
///   nearest-neighbour chroma upsample fused into the write). `native ==
///   Nv12` alone is not sufficient: non-standard subsamplings (4:1:1, mismatched
///   Cb/Cr) also downsample to `Nv12` in [`native_format`] but aren't the
///   shape the fused RGB write targets. An explicit `Rgb` request on any
///   other source (4:2:2, greyscale, or a non-standard subsampling) returns
///   [`CodecError::UnsupportedFormat`] rather than silently substituting the
///   native format — a caller who asks for RGB and doesn't check
///   [`ImageInfo::format`](crate::ImageInfo) must not be handed a
///   differently-shaped buffer with no indication anything changed. See
///   `docs/BENCHMARKS_JPEG_REVIEW_v3.13.md` item #3.
/// - `Nv12`: honoured for any colour JPEG — the MCU writer downsamples
///   chroma (2×2 for 4:4:4, vertical for 4:2:2; native 4:2:0 passes through).
/// - Anything else (or `None`): native format.
fn resolve_output_format(
    headers: &markers::JpegHeaders,
    native: PixelFormat,
    preferred: Option<PixelFormat>,
) -> crate::Result<PixelFormat> {
    Ok(match preferred {
        Some(PixelFormat::Rgb) if native == PixelFormat::Nv24 => PixelFormat::Rgb,
        Some(PixelFormat::Rgb)
            if native == PixelFormat::Nv12 && mcu::is_native_420(&headers.header) =>
        {
            PixelFormat::Rgb
        }
        Some(PixelFormat::Rgb) => return Err(CodecError::UnsupportedFormat(PixelFormat::Rgb)),
        Some(PixelFormat::Nv12)
            if matches!(
                native,
                PixelFormat::Nv12 | PixelFormat::Nv16 | PixelFormat::Nv24
            ) =>
        {
            PixelFormat::Nv12
        }
        _ => native,
    })
}

/// 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 mut headers = {
        let _s = tracing::trace_span!("codec.decode_jpeg.parse_markers").entered();
        markers::parse_markers_with_cache(data, Some(&mut state.huffman_cache))?
    };
    let result = decode_jpeg_into_parsed(data, dst, state, &headers);
    state
        .huffman_cache
        .restore(&mut headers.dc_tables, &mut headers.ac_tables);
    result
}

/// Decode using already-parsed headers. Huffman tables live in `headers` for
/// the duration of this call; the caller restores them to the cache after.
fn decode_jpeg_into_parsed<T: ImagePixel>(
    data: &[u8],
    dst: &mut Tensor<T>,
    state: &mut JpegDecoderState,
    headers: &markers::JpegHeaders,
) -> crate::Result<ImageInfo> {
    let img_w = headers.header.width as usize;
    let img_h = headers.header.height as usize;
    let native_fmt = native_format(headers)?;
    let output_fmt = resolve_output_format(headers, native_fmt, state.preferred_format)?;
    // Fused (non-native) outputs are a CPU-decoder feature: the hardware
    // decoders produce fixed formats, so bypass them when the resolved output
    // differs from native.
    #[cfg(all(target_os = "linux", any(feature = "v4l2", feature = "nvjpeg")))]
    let allow_hw = output_fmt == native_fmt;

    // NV12 supports odd dimensions. Odd *height* gives a `H + ceil(H/2)`
    // combined-plane height (`PixelFormat::allocation_shape`). Odd *width* rounds the
    // tensor's buffer width up to even (also via `allocation_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"))]
    if allow_hw {
        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"))]
    if allow_hw {
        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.
    dst.cpu_access()
        .require(edgefirst_tensor::CpuAccess::Write)?;
    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 = crate::pixel::map_decode_dst(dst)?;
        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,
            state.fast_dct,
        )?;
    }

    // RGB output carries no chroma colorimetry (matches the nvJPEG path);
    // clear any stale value left on a recycled pool tensor.
    if output_fmt == PixelFormat::Rgb {
        dst.set_colorimetry(None);
    } else {
        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,
    })
}