rpdfium_codec/jpx/

decoder.rs

// Derived from PDFium's jpx codec
// Original: Copyright 2014 The PDFium Authors
// Licensed under BSD-3-Clause / Apache-2.0
// See pdfium-upstream/LICENSE for the original license.

//! JPXDecode (JPEG 2000) filter — decodes JPEG 2000 codestreams and JP2 containers
//! using the `hayro-jpeg2000` crate.

use crate::error::DecodeError;

/// Maximum number of resolution levels that can be skipped.
///
/// Ported from upstream PDFium's `CJPX_Decoder::kMaxResolutionsToSkip`.
pub const MAX_RESOLUTIONS_TO_SKIP: u8 = 32;

/// Color space handling option for JPEG 2000 decoding.
///
/// Ported from upstream PDFium's `CJPX_Decoder::ColorSpaceOption`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum JpxColorSpaceOption {
    /// Don't use JP2 color space info — treat as raw components.
    None,
    /// Normal color space handling (default).
    #[default]
    Normal,
    /// Indexed (palette) color space — resolve palette indices to colors.
    Indexed,
}

/// Options for JPEG 2000 decoding.
///
/// Maps to upstream PDFium's `CJPX_Decoder::Create()` parameters.
#[derive(Debug, Clone, Default)]
pub struct JpxDecodeOptions {
    /// Target resolution as (width, height).
    ///
    /// When set, the decoder may decode at a lower resolution level that
    /// best matches the target, avoiding full-resolution decode for thumbnail
    /// generation. Maps to upstream's `resolution_levels_to_skip` parameter
    /// via hayro-jpeg2000's `target_resolution` setting.
    pub target_resolution: Option<(u32, u32)>,

    /// Color space handling option.
    pub color_space_option: JpxColorSpaceOption,

    /// Enable strict decoding mode.
    ///
    /// When enabled, the decoder rejects non-conforming streams that would
    /// otherwise be decoded with best-effort. Defaults to `false`.
    pub strict: bool,
}

/// Information about a JPEG 2000 image's decoded result.
#[derive(Debug, Clone)]
pub struct Jpeg2000Info {
    /// Image width in pixels.
    pub width: u32,
    /// Image height in pixels.
    pub height: u32,
    /// Whether the image has an alpha channel.
    pub has_alpha: bool,
    /// The original bit depth before normalization to 8-bit.
    pub original_bit_depth: u8,
    /// Number of color components (excluding alpha).
    /// 1 = Gray, 3 = RGB, 4 = CMYK.
    pub num_components: u8,
}

/// Build hayro-jpeg2000 `DecodeSettings` from our `JpxDecodeOptions`.
fn build_decode_settings(options: &JpxDecodeOptions) -> hayro_jpeg2000::DecodeSettings {
    let mut settings = hayro_jpeg2000::DecodeSettings::default();
    // target_resolution: not available in hayro-jpeg2000 0.1.x.
    // Will be forwarded when the dependency is upgraded to a version
    // that exposes this setting (available in 0.3+).
    let _ = options.target_resolution;
    settings.strict = options.strict;
    settings.resolve_palette_indices = match options.color_space_option {
        JpxColorSpaceOption::None => false,
        JpxColorSpaceOption::Normal => true,
        JpxColorSpaceOption::Indexed => true,
    };
    settings
}

/// Decode JPEG 2000 (JPX) compressed data.
///
/// Accepts both raw JPEG 2000 codestreams (`.j2c`) and JP2-wrapped images.
/// Returns interleaved 8-bit pixel data.
///
/// The hayro-jpeg2000 library automatically normalizes component precision to 8 bits,
/// so no manual precision adjustment is needed.
pub fn decode(input: &[u8]) -> Result<Vec<u8>, DecodeError> {
    let bitmap = hayro_jpeg2000::decode(input, &hayro_jpeg2000::DecodeSettings::default())
        .map_err(|e| DecodeError::InvalidInput(format!("JPEG2000: {e}")))?;
    Ok(bitmap.data)
}

/// Decode JPEG 2000 data with custom options.
///
/// Allows specifying target resolution for thumbnail generation and
/// color space handling options. See [`JpxDecodeOptions`] for details.
pub fn decode_with_options(
    input: &[u8],
    options: &JpxDecodeOptions,
) -> Result<Vec<u8>, DecodeError> {
    let settings = build_decode_settings(options);
    let bitmap = hayro_jpeg2000::decode(input, &settings)
        .map_err(|e| DecodeError::InvalidInput(format!("JPEG2000: {e}")))?;
    Ok(bitmap.data)
}

/// Decode JPEG 2000 data and return both pixel data and image info.
///
/// This provides access to the original bit depth and other metadata
/// that may be needed for color space interpretation.
pub fn decode_with_info(input: &[u8]) -> Result<(Vec<u8>, Jpeg2000Info), DecodeError> {
    let bitmap = hayro_jpeg2000::decode(input, &hayro_jpeg2000::DecodeSettings::default())
        .map_err(|e| DecodeError::InvalidInput(format!("JPEG2000: {e}")))?;

    // num_components should exclude alpha (per Jpeg2000Info doc comment).
    // hayro-jpeg2000's color_space.num_channels() already excludes alpha,
    // so we use it directly.
    let num_components = bitmap.color_space.num_channels();

    let info = Jpeg2000Info {
        width: bitmap.width,
        height: bitmap.height,
        has_alpha: bitmap.has_alpha,
        original_bit_depth: bitmap.original_bit_depth,
        num_components,
    };

    Ok((bitmap.data, info))
}

/// Decode JPEG 2000 data with custom options and return both pixel data and image info.
pub fn decode_with_info_and_options(
    input: &[u8],
    options: &JpxDecodeOptions,
) -> Result<(Vec<u8>, Jpeg2000Info), DecodeError> {
    let settings = build_decode_settings(options);
    let bitmap = hayro_jpeg2000::decode(input, &settings)
        .map_err(|e| DecodeError::InvalidInput(format!("JPEG2000: {e}")))?;

    let num_components = bitmap.color_space.num_channels();
    let info = Jpeg2000Info {
        width: bitmap.width,
        height: bitmap.height,
        has_alpha: bitmap.has_alpha,
        original_bit_depth: bitmap.original_bit_depth,
        num_components,
    };

    Ok((bitmap.data, info))
}

/// Per-component precision information parsed from JP2/codestream headers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ComponentPrecision {
    /// Bits per component (1-38 per spec).
    pub bits: u8,
    /// Whether the component values are signed.
    pub is_signed: bool,
}

/// Parse component precision info from a JP2 Image Header Box (ihdr).
///
/// Scans for the JP2 Image Header Box to extract precision and signedness.
/// Returns `None` if the data is not a JP2 container or the ihdr box is not found.
pub fn parse_jp2_precision(data: &[u8]) -> Option<Vec<ComponentPrecision>> {
    // JP2 files start with the signature box: 12 bytes
    // \x00\x00\x00\x0C jP  \x20\x20 \x0D\x0A\x87\x0A
    if data.len() < 12 {
        return None;
    }
    if &data[4..8] != b"jP  " && &data[4..8] != b"jP\x1a\x1a" {
        // Not a JP2 container — might be a raw codestream
        return parse_codestream_precision(data);
    }

    // Scan JP2 boxes looking for ihdr (Image Header Box) inside jp2h (JP2 Header)
    let mut pos = 0;
    while pos + 8 <= data.len() {
        let box_len =
            u32::from_be_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]) as usize;
        let box_type = &data[pos + 4..pos + 8];

        if box_len < 8 || pos + box_len > data.len() {
            break;
        }

        if box_type == b"jp2h" {
            // Search inside jp2h for ihdr
            return parse_ihdr_in_jp2h(&data[pos + 8..pos + box_len]);
        }

        pos += box_len;
    }

    None
}

/// Parse ihdr box inside a jp2h superbox.
fn parse_ihdr_in_jp2h(data: &[u8]) -> Option<Vec<ComponentPrecision>> {
    let mut pos = 0;
    while pos + 8 <= data.len() {
        let box_len =
            u32::from_be_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]) as usize;
        let box_type = &data[pos + 4..pos + 8];

        if box_len < 8 || pos + box_len > data.len() {
            break;
        }

        if box_type == b"ihdr" {
            // ihdr: 4(height) + 4(width) + 2(num_components) + 1(bpc) + ...
            let content = &data[pos + 8..pos + box_len];
            if content.len() >= 11 {
                let num_components = u16::from_be_bytes([content[8], content[9]]) as usize;
                let bpc_byte = content[10];

                if bpc_byte == 0xFF {
                    // bpc varies per component — would need bpcc box
                    // For now return None to indicate we can't determine from ihdr alone
                    return None;
                }

                let bits = (bpc_byte & 0x7F) + 1;
                let is_signed = (bpc_byte & 0x80) != 0;
                let precision = ComponentPrecision { bits, is_signed };
                return Some(vec![precision; num_components]);
            }
        }

        pos += box_len;
    }

    None
}

/// Parse component precision from a raw JPEG 2000 codestream SIZ marker.
fn parse_codestream_precision(data: &[u8]) -> Option<Vec<ComponentPrecision>> {
    // SOC marker: 0xFF4F, then SIZ marker: 0xFF51
    if data.len() < 4 || data[0] != 0xFF || data[1] != 0x4F {
        return None;
    }
    if data[2] != 0xFF || data[3] != 0x51 {
        return None;
    }

    // SIZ marker: length(2) + Rsiz(2) + Xsiz(4) + Ysiz(4) + XOsiz(4) + YOsiz(4)
    //             + XTsiz(4) + YTsiz(4) + XTOsiz(4) + YTOsiz(4) + Csiz(2)
    //             + per-component: Ssiz(1) + XRsiz(1) + YRsiz(1)
    if data.len() < 6 {
        return None;
    }
    let siz_len = u16::from_be_bytes([data[4], data[5]]) as usize;
    if data.len() < 4 + siz_len {
        return None;
    }

    // Offset within SIZ data (after length field)
    let siz_data = &data[6..4 + siz_len];
    // Rsiz(2) + Xsiz(4) + Ysiz(4) + XOsiz(4) + YOsiz(4)
    // + XTsiz(4) + YTsiz(4) + XTOsiz(4) + YTOsiz(4) + Csiz(2) = 36 bytes
    if siz_data.len() < 36 {
        return None;
    }

    let csiz = u16::from_be_bytes([siz_data[34], siz_data[35]]) as usize;
    let comp_start = 36;

    let mut result = Vec::with_capacity(csiz);
    for i in 0..csiz {
        let offset = comp_start + i * 3;
        if offset >= siz_data.len() {
            break;
        }
        let ssiz = siz_data[offset];
        let bits = (ssiz & 0x7F) + 1;
        let is_signed = (ssiz & 0x80) != 0;
        result.push(ComponentPrecision { bits, is_signed });
    }

    if result.is_empty() {
        None
    } else {
        Some(result)
    }
}

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

    #[test]
    fn test_decode_invalid_data_returns_error() {
        let result = decode(&[1, 2, 3]);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.to_string().contains("JPEG2000"),
            "error should mention JPEG2000: {err}"
        );
    }

    #[test]
    fn test_decode_empty_data_returns_error() {
        let result = decode(&[]);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(
            err.to_string().contains("JPEG2000"),
            "error should mention JPEG2000: {err}"
        );
    }

    #[test]
    fn test_decode_truncated_jp2_header_returns_error() {
        // JP2 signature prefix, but truncated
        let data = b"\x00\x00\x00\x0C\x6A\x50\x20\x20\x0D\x0A";
        let result = decode(data);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_truncated_codestream_returns_error() {
        // Codestream SOC + SIZ markers, but truncated
        let data = b"\xFF\x4F\xFF\x51\x00\x00";
        let result = decode(data);
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_codestream_precision_basic() {
        // Build a minimal SIZ marker: SOC + SIZ
        let mut data = Vec::new();
        // SOC
        data.extend_from_slice(&[0xFF, 0x4F]);
        // SIZ marker
        data.extend_from_slice(&[0xFF, 0x51]);
        // Length = 2 (self) + 36 (header) + 3*2 (components) = 44
        data.extend_from_slice(&[0x00, 44u8]);
        // Rsiz (2)
        data.extend_from_slice(&[0x00, 0x00]);
        // Xsiz (4) = 100
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x64]);
        // Ysiz (4) = 200
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0xC8]);
        // XOsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
        // YOsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
        // XTsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x64]);
        // YTsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0xC8]);
        // XTOsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
        // YTOsiz (4)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
        // Csiz (2) = 2 components
        data.extend_from_slice(&[0x00, 0x02]);
        // Component 0: Ssiz=0x07 (8-bit unsigned), XRsiz=1, YRsiz=1
        data.extend_from_slice(&[0x07, 0x01, 0x01]);
        // Component 1: Ssiz=0x8F (16-bit signed), XRsiz=1, YRsiz=1
        data.extend_from_slice(&[0x8F, 0x01, 0x01]);

        let result = parse_jp2_precision(&data).unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result[0].bits, 8);
        assert!(!result[0].is_signed);
        assert_eq!(result[1].bits, 16);
        assert!(result[1].is_signed);
    }

    #[test]
    fn test_parse_jp2_precision_with_ihdr() {
        // Build a minimal JP2 file with signature + jp2h/ihdr boxes
        let mut data = Vec::new();

        // Signature box (12 bytes)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x0C]);
        data.extend_from_slice(b"jP  ");
        data.extend_from_slice(&[0x0D, 0x0A, 0x87, 0x0A]);

        // File type box (20 bytes)
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x14]);
        data.extend_from_slice(b"ftyp");
        data.extend_from_slice(b"jp2 ");
        data.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
        data.extend_from_slice(b"jp2 ");

        // jp2h box containing ihdr
        // ihdr box: 8 (box header) + 14 (content) = 22 bytes
        // jp2h box: 8 (box header) + 22 (ihdr) = 30 bytes
        let ihdr_len: u32 = 22;
        let jp2h_len: u32 = 8 + ihdr_len;

        data.extend_from_slice(&jp2h_len.to_be_bytes());
        data.extend_from_slice(b"jp2h");

        data.extend_from_slice(&ihdr_len.to_be_bytes());
        data.extend_from_slice(b"ihdr");
        // Height (4) = 480
        data.extend_from_slice(&480u32.to_be_bytes());
        // Width (4) = 640
        data.extend_from_slice(&640u32.to_be_bytes());
        // NumComponents (2) = 3
        data.extend_from_slice(&3u16.to_be_bytes());
        // BPC byte: 0x07 = 8-bit unsigned (bits-1=7, sign=0)
        data.push(0x07);
        // Compression type (1)
        data.push(0x07);
        // UnkC (1) = 0
        data.push(0x00);
        // IPR (1) = 0
        data.push(0x00);

        let result = parse_jp2_precision(&data).unwrap();
        assert_eq!(result.len(), 3);
        for comp in &result {
            assert_eq!(comp.bits, 8);
            assert!(!comp.is_signed);
        }
    }

    #[test]
    fn test_parse_jp2_precision_not_jp2() {
        // Random data — not JP2 or codestream
        let data = [0x00, 0x01, 0x02, 0x03, 0x04, 0x05];
        assert!(parse_jp2_precision(&data).is_none());
    }

    #[test]
    fn test_parse_jp2_precision_empty() {
        assert!(parse_jp2_precision(&[]).is_none());
    }

    #[test]
    fn test_component_precision_eq() {
        let a = ComponentPrecision {
            bits: 8,
            is_signed: false,
        };
        let b = ComponentPrecision {
            bits: 8,
            is_signed: false,
        };
        assert_eq!(a, b);
    }

    // -----------------------------------------------------------------------
    // JpxDecodeOptions tests (R3 + R4)
    // -----------------------------------------------------------------------

    #[test]
    fn test_jpx_decode_options_default() {
        let opts = JpxDecodeOptions::default();
        assert_eq!(opts.target_resolution, None);
        assert_eq!(opts.color_space_option, JpxColorSpaceOption::Normal);
        assert!(!opts.strict);
    }

    #[test]
    fn test_jpx_color_space_option_default() {
        assert_eq!(JpxColorSpaceOption::default(), JpxColorSpaceOption::Normal);
    }

    #[test]
    fn test_build_settings_default() {
        let opts = JpxDecodeOptions::default();
        let settings = build_decode_settings(&opts);
        assert!(!settings.strict);
        assert!(settings.resolve_palette_indices);
    }

    #[test]
    fn test_build_settings_with_target_resolution() {
        // target_resolution is accepted in our API but not yet forwarded to
        // hayro-jpeg2000 v0.1 (which lacks the field). Verify it doesn't panic.
        let opts = JpxDecodeOptions {
            target_resolution: Some((320, 240)),
            ..Default::default()
        };
        let settings = build_decode_settings(&opts);
        // Settings should still be valid
        assert!(!settings.strict);
    }

    #[test]
    fn test_build_settings_color_space_none() {
        let opts = JpxDecodeOptions {
            color_space_option: JpxColorSpaceOption::None,
            ..Default::default()
        };
        let settings = build_decode_settings(&opts);
        assert!(!settings.resolve_palette_indices);
    }

    #[test]
    fn test_build_settings_color_space_indexed() {
        let opts = JpxDecodeOptions {
            color_space_option: JpxColorSpaceOption::Indexed,
            ..Default::default()
        };
        let settings = build_decode_settings(&opts);
        assert!(settings.resolve_palette_indices);
    }

    #[test]
    fn test_build_settings_strict() {
        let opts = JpxDecodeOptions {
            strict: true,
            ..Default::default()
        };
        let settings = build_decode_settings(&opts);
        assert!(settings.strict);
    }

    #[test]
    fn test_decode_with_options_invalid_data() {
        let opts = JpxDecodeOptions::default();
        let result = decode_with_options(&[1, 2, 3], &opts);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_with_info_and_options_invalid_data() {
        let opts = JpxDecodeOptions {
            target_resolution: Some((100, 100)),
            ..Default::default()
        };
        let result = decode_with_info_and_options(&[1, 2, 3], &opts);
        assert!(result.is_err());
    }

    #[test]
    fn test_max_resolutions_to_skip_constant() {
        assert_eq!(MAX_RESOLUTIONS_TO_SKIP, 32);
    }

    // --- Upstream ports: jpx_unittest.cpp ---
    //
    // The upstream C++ tests exercise OpenJPEG stream adapter callbacks
    // (DecodeData read/skip/seek) and YUV420ToRGB conversion. rpdfium
    // delegates to hayro-jpeg2000, which handles stream I/O internally.
    // These tests verify equivalent error-handling behavior at the Rust API
    // boundary.

    /// Upstream: TEST(fxcodec, DecodeDataNullDecodeData)
    ///
    /// Null/missing data should produce an error, not a crash.
    #[test]
    fn test_jpx_decode_null_data() {
        // Empty slice (null-equivalent) returns error
        let result = decode(&[]);
        assert!(result.is_err());
        let result2 = decode_with_options(&[], &JpxDecodeOptions::default());
        assert!(result2.is_err());
    }

    /// Upstream: TEST(fxcodec, DecodeDataNullStream)
    ///
    /// Zero-length reads on null stream should error.
    #[test]
    fn test_jpx_decode_null_stream_equivalent() {
        // Attempt decode with completely invalid data
        let result = decode_with_info(&[]);
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(err.to_string().contains("JPEG2000"));
    }

    /// Upstream: TEST(fxcodec, DecodeDataZeroSize)
    ///
    /// Zero-size data should error.
    #[test]
    fn test_jpx_decode_zero_size_data() {
        let result = decode_with_info_and_options(&[], &JpxDecodeOptions::default());
        assert!(result.is_err());
    }

    /// Upstream: TEST(fxcodec, DecodeDataReadInBounds)
    #[test]
    #[ignore = "requires OpenJPEG stream adapter internals not exposed in hayro-jpeg2000"]
    fn test_jpx_decode_data_read_in_bounds() {
        // Upstream tests opj_read_from_memory with exact-sized reads.
        // hayro-jpeg2000 manages stream I/O internally.
    }

    /// Upstream: TEST(fxcodec, DecodeDataReadBeyondBounds)
    #[test]
    #[ignore = "requires OpenJPEG stream adapter internals not exposed in hayro-jpeg2000"]
    fn test_jpx_decode_data_read_beyond_bounds() {
        // Upstream tests opj_read_from_memory beyond data length.
    }

    /// Upstream: TEST(fxcodec, DecodeDataSkip)
    #[test]
    #[ignore = "requires OpenJPEG stream adapter internals not exposed in hayro-jpeg2000"]
    fn test_jpx_decode_data_skip() {
        // Upstream tests opj_skip_from_memory with various offsets.
    }

    /// Upstream: TEST(fxcodec, DecodeDataSeek)
    #[test]
    #[ignore = "requires OpenJPEG stream adapter internals not exposed in hayro-jpeg2000"]
    fn test_jpx_decode_data_seek() {
        // Upstream tests opj_seek_from_memory with various positions.
    }

    /// Upstream: TEST(fxcodec, YUV420ToRGB)
    #[test]
    #[ignore = "requires CJPX_Decoder::Sycc420ToRgbForTesting internals not exposed in hayro-jpeg2000"]
    fn test_jpx_yuv420_to_rgb() {
        // Upstream tests YUV420 -> RGB color conversion with various image widths.
        // hayro-jpeg2000 handles color conversion internally.
    }
}