edgefirst_codec/

decoder.rs

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

//! [`ImageDecoder`] — reusable decoder state for zero-allocation hot loops.

use crate::error::CodecError;
#[cfg(not(feature = "png"))]
use crate::error::UnsupportedFeature;
use crate::options::ImageInfo;
use crate::pixel::ImagePixel;
use edgefirst_tensor::{DType, Tensor, TensorDyn};
use std::io::Read;

/// IDCT accuracy/speed selection for the software JPEG decoder.
///
/// [`Accurate`](Self::Accurate) is the default and the kernel every published
/// benchmark quotes: the `islow`-class Loeffler IDCT, bit-comparable to
/// libjpeg-turbo's default. [`Fast`](Self::Fast) opts into the AAN
/// `ifast`-class kernel — roughly an eighth of the multiplies, at a small,
/// bounded pixel accuracy cost (and the documented `ifast` degeneration at
/// very high quality factors). Fast is advisory: paths without a fast kernel
/// (non-NEON tiers, hardware V4L2/nvJPEG decoders, PNG) use their normal
/// accurate path.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum DctMethod {
    /// Accurate `islow`-class IDCT (default).
    #[default]
    Accurate,
    /// Fast AAN `ifast`-class IDCT (opt-in).
    Fast,
}

/// Chroma upsampling filter for a fused native-4:2:0 `Rgb` decode (see
/// [`ImageDecoder::set_output_format`]).
///
/// [`Box`](Self::Box) — 2×2 nearest-neighbour replication — is the only
/// filter implemented today; the type exists so that a future
/// fancy/triangle-filtered mode (libjpeg-turbo's default, and a materially
/// more expensive one — see BENCHMARKS.md) can be added as a new variant
/// without a breaking API change to
/// [`ImageDecoder::set_chroma_upsample`]. `#[non_exhaustive]` for the same
/// reason: adding a variant later must not be a semver break for callers who
/// match on this type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum ChromaUpsample {
    /// 2×2 nearest-neighbour (box) chroma replication.
    #[default]
    Box,
}

/// Reusable image decoder with internal scratch buffers.
///
/// Create one `ImageDecoder` at program initialisation and pass it to
/// [`ImageLoad::load_image`](crate::ImageLoad::load_image) in the hot loop.
/// The scratch buffers grow to the high-water mark and are reused across calls
/// — no per-frame allocations after the first few frames.
///
/// By default the decoder produces the source's native format and configures
/// the destination tensor's dimensions + pixel format accordingly (JPEG →
/// `Nv12`/`Nv16`/`Nv24`/`Grey`, PNG → `Rgb`/`Rgba`/`Grey`).
/// [`set_output_format`](Self::set_output_format) opts a JPEG decode into a
/// fused `Rgb` or `Nv12` output, where the colour conversion or chroma
/// downsample happens at the MCU write stage instead of in a second pass. The
/// decoder never rotates or resizes; use `ImageProcessor::convert()` for that.
///
/// # Example
///
/// ```rust,no_run
/// use edgefirst_codec::{ImageDecoder, ImageLoad};
/// use edgefirst_tensor::{CpuAccess, Tensor, TensorTrait, TensorMemory, PixelFormat};
///
/// let mut decoder = ImageDecoder::new();
/// // Allocate a buffer large enough for the biggest expected image. The
/// // decoder CPU-writes the pixels, so declare `CpuAccess::Write`.
/// let mut tensor = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12, Some(TensorMemory::Mem),
///                                       CpuAccess::Write)
///     .expect("alloc");
///
/// for _ in 0..100 {
///     let frame = std::fs::read("frame.jpg").unwrap();
///     let _info = tensor.load_image(&mut decoder, &frame);
/// }
/// ```
pub struct ImageDecoder {
    /// Reusable JPEG decoder state (Huffman tables, MCU scratch buffers).
    pub(crate) jpeg_state: crate::jpeg::JpegDecoderState,
    /// Scratch buffer for PNG decode output.
    pub(crate) scratch: Vec<u8>,
    /// Buffer for `Read` → `&[u8]` conversion.
    pub(crate) input_buffer: Vec<u8>,
}

impl ImageDecoder {
    /// Create a new decoder with empty scratch buffers.
    pub fn new() -> Self {
        Self {
            jpeg_state: crate::jpeg::JpegDecoderState::new(),
            scratch: Vec::new(),
            input_buffer: Vec::new(),
        }
    }

    /// Request a fused decode output format instead of the source's native
    /// format (JPEG only; PNG is unaffected).
    ///
    /// This is a **pure CPU, single-pass** path inside the software JPEG
    /// decoder — colour conversion / chroma downsample happens at the MCU
    /// write stage. It is **not** a GPU hybrid or nvJPEG path; V4L2/nvJPEG
    /// are bypassed whenever the resolved output differs from native.
    ///
    /// - `Some(Rgb)`: 4:4:4 colour JPEGs decode straight to interleaved RGB
    ///   (no chroma resampling); native 4:2:0 colour JPEGs also decode
    ///   straight to RGB, with a 2×2 nearest-neighbour (box) chroma upsample
    ///   fused into the write — not libjpeg's fancy/triangle upsampling, a
    ///   deliberate speed tradeoff (see BENCHMARKS.md). Other sources (4:2:2,
    ///   greyscale, non-standard subsamplings) cannot honour an RGB request:
    ///   the decode returns [`CodecError::UnsupportedFormat`] rather than
    ///   silently substituting the native semi-planar format. Check the
    ///   source's subsampling first (e.g. via [`crate::peek_info`]) if it
    ///   isn't known to be 4:4:4 or 4:2:0.
    /// - `Some(Nv12)`: colour JPEGs decode to NV12, downsampling chroma at
    ///   the write stage (2×2 average for 4:4:4, vertical for 4:2:2).
    /// - `None` (default): native format (`Nv12`/`Nv16`/`Nv24`/`Grey`).
    ///
    /// Callers still run `ImageProcessor::convert()` on the result for
    /// model-input preprocessing (letterbox, resize, EXIF orientation). With
    /// fused RGB that convert step is typically a pure resize.
    pub fn set_output_format(&mut self, format: Option<edgefirst_tensor::PixelFormat>) {
        self.jpeg_state.preferred_format = format;
    }

    /// Select the software JPEG IDCT kernel class — see [`DctMethod`].
    /// Off (Accurate) by default; `EDGEFIRST_CODEC_DCT=fast` in the
    /// environment flips a **new** decoder's default for A/B runs.
    pub fn set_dct_method(&mut self, method: DctMethod) {
        self.jpeg_state.fast_dct = method == DctMethod::Fast;
    }

    /// Select the chroma upsampling filter for a fused native-4:2:0 `Rgb`
    /// decode — see [`ChromaUpsample`]. [`ChromaUpsample::Box`] (the
    /// default) is the only filter implemented today; this setter exists so
    /// a future filter can be selected without changing
    /// [`set_output_format`](Self::set_output_format)'s signature. Has no
    /// effect on 4:4:4 sources (no chroma resampling occurs there) or when
    /// the output format isn't `Rgb`.
    pub fn set_chroma_upsample(&mut self, mode: ChromaUpsample) {
        self.jpeg_state.chroma_upsample = mode;
    }

    /// Decode image data into a typed tensor, configuring its dimensions and
    /// pixel format to the decoded output format — the source's native format
    /// unless [`set_output_format`](Self::set_output_format) selected a fused
    /// one.
    ///
    /// Detects the image format (JPEG or PNG) from magic bytes.
    ///
    /// # Errors
    ///
    /// - [`CodecError::InsufficientCapacity`] if the image is larger than the
    ///   tensor's allocation
    /// - [`CodecError::UnsupportedDtype`] if `T` is not valid for the output
    ///   format (JPEG NV12/NV16/NV24/RGB/GREY require `u8`)
    /// - [`CodecError::InvalidData`] if the data is not a valid JPEG or PNG,
    ///   including a JPEG whose entropy-coded scan is truncated
    pub fn decode_into<T: ImagePixel>(
        &mut self,
        data: &[u8],
        dst: &mut Tensor<T>,
    ) -> crate::Result<ImageInfo> {
        if is_jpeg(data) {
            crate::jpeg::decode_jpeg_into(data, dst, &mut self.jpeg_state)
        } else if is_png(data) {
            {
                #[cfg(feature = "png")]
                {
                    crate::png::decode_png_into(data, dst, &mut self.scratch)
                }
                #[cfg(not(feature = "png"))]
                {
                    Err(CodecError::Unsupported(UnsupportedFeature::PngNotEnabled))
                }
            }
        } else {
            Err(CodecError::InvalidData(
                "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
            ))
        }
    }

    /// Decode image data into a type-erased tensor.
    pub fn decode_into_dyn(
        &mut self,
        data: &[u8],
        dst: &mut TensorDyn,
    ) -> crate::Result<ImageInfo> {
        match dst.dtype() {
            DType::U8 => self.decode_into(data, dst.as_typed_mut::<u8>().expect("dtype checked")),
            DType::I8 => self.decode_into(data, dst.as_typed_mut::<i8>().expect("dtype checked")),
            DType::U16 => self.decode_into(data, dst.as_typed_mut::<u16>().expect("dtype checked")),
            DType::I16 => self.decode_into(data, dst.as_typed_mut::<i16>().expect("dtype checked")),
            DType::F32 => self.decode_into(data, dst.as_typed_mut::<f32>().expect("dtype checked")),
            other => Err(CodecError::UnsupportedDtype(other)),
        }
    }

    /// Read all bytes from a `Read` source into the internal input buffer, then
    /// decode. The input buffer is reused across calls.
    pub fn decode_from_reader<T: ImagePixel, R: Read>(
        &mut self,
        mut reader: R,
        dst: &mut Tensor<T>,
    ) -> crate::Result<ImageInfo> {
        self.input_buffer.clear();
        reader.read_to_end(&mut self.input_buffer)?;
        decode_into_inner(
            &mut self.jpeg_state,
            &mut self.scratch,
            &self.input_buffer,
            dst,
        )
    }

    /// Read all bytes from a `Read` source, then decode into a type-erased
    /// tensor.
    pub fn decode_from_reader_dyn<R: Read>(
        &mut self,
        mut reader: R,
        dst: &mut TensorDyn,
    ) -> crate::Result<ImageInfo> {
        self.input_buffer.clear();
        reader.read_to_end(&mut self.input_buffer)?;
        decode_into_inner_dyn(
            &mut self.jpeg_state,
            &mut self.scratch,
            &self.input_buffer,
            dst,
        )
    }
}

/// Internal decode entry point parameterised over disjoint `&mut` borrows so
/// callers can read from a `&[u8]` borrowed from one field of [`ImageDecoder`]
/// while mutably borrowing the others.
pub(crate) fn decode_into_inner<T: ImagePixel>(
    jpeg_state: &mut crate::jpeg::JpegDecoderState,
    scratch: &mut Vec<u8>,
    data: &[u8],
    dst: &mut Tensor<T>,
) -> crate::Result<ImageInfo> {
    if is_jpeg(data) {
        crate::jpeg::decode_jpeg_into(data, dst, jpeg_state)
    } else if is_png(data) {
        {
            #[cfg(feature = "png")]
            {
                crate::png::decode_png_into(data, dst, scratch)
            }
            #[cfg(not(feature = "png"))]
            {
                let _ = scratch;
                Err(CodecError::Unsupported(UnsupportedFeature::PngNotEnabled))
            }
        }
    } else {
        Err(CodecError::InvalidData(
            "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
        ))
    }
}

pub(crate) fn decode_into_inner_dyn(
    jpeg_state: &mut crate::jpeg::JpegDecoderState,
    scratch: &mut Vec<u8>,
    data: &[u8],
    dst: &mut TensorDyn,
) -> crate::Result<ImageInfo> {
    match dst.dtype() {
        DType::U8 => decode_into_inner(
            jpeg_state,
            scratch,
            data,
            dst.as_typed_mut::<u8>().expect("dtype checked"),
        ),
        DType::I8 => decode_into_inner(
            jpeg_state,
            scratch,
            data,
            dst.as_typed_mut::<i8>().expect("dtype checked"),
        ),
        DType::U16 => decode_into_inner(
            jpeg_state,
            scratch,
            data,
            dst.as_typed_mut::<u16>().expect("dtype checked"),
        ),
        DType::I16 => decode_into_inner(
            jpeg_state,
            scratch,
            data,
            dst.as_typed_mut::<i16>().expect("dtype checked"),
        ),
        DType::F32 => decode_into_inner(
            jpeg_state,
            scratch,
            data,
            dst.as_typed_mut::<f32>().expect("dtype checked"),
        ),
        other => Err(CodecError::UnsupportedDtype(other)),
    }
}

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

/// Parse image headers and return the native dimensions, format, and EXIF
/// orientation without decoding pixels.
///
/// Recommended for one-shot flows that allocate a tensor sized to the image:
///
/// ```rust,no_run
/// use edgefirst_codec::{peek_info, ImageDecoder, ImageLoad};
/// use edgefirst_tensor::{CpuAccess, Tensor, TensorMemory};
///
/// let data = std::fs::read("image.jpg").unwrap();
/// let info = peek_info(&data).unwrap();
/// let mut tensor = Tensor::<u8>::image(info.width, info.height, info.format,
///                                       Some(TensorMemory::Mem),
///                                       CpuAccess::Write).unwrap();
/// let mut decoder = ImageDecoder::new();
/// tensor.load_image(&mut decoder, &data).unwrap();
/// ```
pub fn peek_info(data: &[u8]) -> crate::Result<ImageInfo> {
    if is_jpeg(data) {
        crate::jpeg::peek_jpeg_info(data)
    } else if is_png(data) {
        {
            #[cfg(feature = "png")]
            {
                crate::png::peek_png_info(data)
            }
            #[cfg(not(feature = "png"))]
            {
                Err(CodecError::Unsupported(UnsupportedFeature::PngNotEnabled))
            }
        }
    } else {
        Err(CodecError::InvalidData(
            "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
        ))
    }
}

/// Check for JPEG magic bytes (FFD8FF).
fn is_jpeg(data: &[u8]) -> bool {
    data.len() >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF
}

/// Check for PNG magic bytes (89504E47).
fn is_png(data: &[u8]) -> bool {
    data.len() >= 4 && data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47
}

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

    #[test]
    fn magic_bytes_jpeg() {
        assert!(is_jpeg(&[0xFF, 0xD8, 0xFF, 0xE0]));
        assert!(!is_jpeg(&[0x89, 0x50, 0x4E, 0x47]));
        assert!(!is_jpeg(&[]));
        assert!(!is_jpeg(&[0xFF]));
    }

    #[test]
    fn magic_bytes_png() {
        assert!(is_png(&[0x89, 0x50, 0x4E, 0x47]));
        assert!(!is_png(&[0xFF, 0xD8, 0xFF, 0xE0]));
        assert!(!is_png(&[]));
    }

    #[test]
    fn invalid_data() {
        let mut decoder = ImageDecoder::new();
        let mut tensor = Tensor::<u8>::image(
            100,
            100,
            edgefirst_tensor::PixelFormat::Nv12,
            Some(edgefirst_tensor::TensorMemory::Mem),
            edgefirst_tensor::CpuAccess::ReadWrite,
        )
        .unwrap();
        let result = decoder.decode_into(b"not an image", &mut tensor);
        assert!(matches!(result, Err(CodecError::InvalidData(_))));
    }
}