nigig-org/crates/apps/pdf/pdf-document/tests/jpeg_images.rs

//! JPEG decode, downsampling and filter round-trip acceptance tests.
//!
//! The merge criteria from `REVIEWS/adr/0016-pdf-image-decode-surface.md`.
//!
//! The defect these exist for was the most convincing kind. `decode_jpeg_data`
//! took five underscore-prefixed arguments, ignored all of them, and returned
//! `Some(())`. Its caller allocated a zero-filled buffer and returned it as
//! decoded pixels, so **every JPEG in every PDF rendered as a black rectangle
//! at full alpha** with no error anywhere.
//!
//! A test asserting "decoding returned a buffer of the right length" would
//! have passed against that stub. So every test here asserts **pixel
//! values**.

use std::path::PathBuf;

use nigig_pdf_cos::filter::{decode_stream, encode_ascii_hex, encode_flate};
use nigig_pdf_cos::{PdfDict, PdfObj, PdfStream};
use nigig_pdf_document::PdfDocument;
use nigig_pdf_graphics::image::ImageInfo;
use nigig_pdf_graphics::jpeg::{self, JpegError};

fn corpus(relative: &str) -> Vec<u8> {
    let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../tests/corpus")
        .join(relative);
    std::fs::read(&path).unwrap_or_else(|e| panic!("missing {}: {e}", path.display()))
}

/// Pull the raw `/DCTDecode` bytes of the first image XObject out of a
/// fixture, the way a renderer reaches them.
fn image_of(fixture: &str) -> ImageInfo {
    let bytes = corpus(fixture);
    let mut doc = PdfDocument::parse(&bytes).expect("fixture parses");
    let page = doc.page(0).expect("page 0");
    let entry = page.xobjects.get("Im1").expect("the image xobject");
    let obj = doc.resolve_ref(entry.obj_ref).expect("resolves");
    let stream = obj.as_stream().expect("images are streams").clone();
    ImageInfo::from_dict(&stream.dict, "Im1", stream.data).expect("image info")
}

fn close(a: u8, b: u8, tolerance: i32) -> bool {
    (a as i32 - b as i32).abs() <= tolerance
}

// ------------------------------------------------------------------ JPEG

/// The headline test: real pixel values, not a buffer length.
#[test]
fn a_red_jpeg_decodes_to_red_pixels() {
    let image = image_of("images/jpeg_rgb.pdf");
    assert!(image.is_jpeg());

    let rgba = image.decode_to_rgba().expect("a baseline JPEG must decode");
    assert_eq!(rgba.len(), 16 * 16 * 4);

    // JPEG is lossy; 95-quality solid colour lands within a couple of levels.
    assert!(
        close(rgba[0], 255, 3) && close(rgba[1], 0, 3) && close(rgba[2], 0, 3),
        "expected red, got ({}, {}, {}) - the stub returned all zeroes",
        rgba[0],
        rgba[1],
        rgba[2]
    );
    assert_eq!(rgba[3], 255, "opaque");
}

/// The stub's exact signature: a correctly sized, uniformly zero buffer.
/// This is the test that would have caught it.
#[test]
fn a_decoded_jpeg_is_not_a_uniformly_zero_buffer() {
    for fixture in [
        "images/jpeg_rgb.pdf",
        "images/jpeg_gray.pdf",
        "images/jpeg_subsampled.pdf",
    ] {
        let image = image_of(fixture);
        let rgba = image.decode_to_rgba().expect("decodes");
        let colour: Vec<u8> = rgba
            .chunks(4)
            .flat_map(|px| [px[0], px[1], px[2]])
            .collect();
        assert!(
            colour.iter().any(|b| *b != 0),
            "{fixture} decoded to all-zero colour, which is what the stub \
             produced for every JPEG in every document"
        );
    }
}

#[test]
fn a_grayscale_jpeg_decodes_to_its_grey_level() {
    let image = image_of("images/jpeg_gray.pdf");
    let rgba = image.decode_to_rgba().expect("decodes");
    assert!(close(rgba[0], 128, 3), "mid-grey expected, got {}", rgba[0]);
    // Grayscale expands to equal RGB channels.
    assert_eq!(rgba[0], rgba[1]);
    assert_eq!(rgba[1], rgba[2]);
}

/// With 4:2:0 the chroma planes are half resolution, so the upsampling path
/// runs. A decoder that ignored sampling factors would produce a colour
/// shift here and pass the solid-red test.
#[test]
fn a_subsampled_jpeg_decodes_with_correct_colour() {
    let image = image_of("images/jpeg_subsampled.pdf");
    let rgba = image.decode_to_rgba().expect("decodes");
    assert!(
        close(rgba[0], 0, 6) && close(rgba[1], 128, 6) && close(rgba[2], 255, 6),
        "expected (0,128,255), got ({}, {}, {})",
        rgba[0],
        rgba[1],
        rgba[2]
    );
}

/// ADR 0016 rule 2: progressive is refused by name, never approximated.
#[test]
fn a_progressive_jpeg_is_refused_by_name() {
    let bytes = corpus("images/jpeg_progressive.pdf");
    let mut doc = PdfDocument::parse(&bytes).expect("the document still opens");
    let page = doc.page(0).expect("page 0");
    let entry = page.xobjects.get("Im1").expect("image");
    let obj = doc.resolve_ref(entry.obj_ref).expect("resolves");
    let stream = obj.as_stream().expect("stream").clone();

    match jpeg::decode(&stream.data) {
        Err(JpegError::UnsupportedProfile { name, .. }) => {
            assert_eq!(name, "progressive")
        }
        Ok(_) => panic!("progressive JPEG must not decode as if it were baseline"),
        Err(other) => panic!("expected a named refusal, got {other:?}"),
    }
}

/// A page whose image cannot be decoded must still be readable.
#[test]
fn a_page_with_an_undecodable_image_still_parses() {
    let bytes = corpus("images/jpeg_progressive.pdf");
    let mut doc = PdfDocument::parse(&bytes).expect("parses");
    let page = doc.page(0).expect("page 0");
    assert!(
        !page.content_data.is_empty(),
        "one undecodable image must not cost the page its content"
    );
}

#[test]
fn a_truncated_jpeg_is_an_error_not_partial_garbage() {
    let image = image_of("images/jpeg_rgb.pdf");
    let mut truncated = image.data.clone();
    truncated.truncate(truncated.len() / 3);
    assert!(
        jpeg::decode(&truncated).is_err(),
        "a truncated JPEG must be refused rather than half-decoded"
    );
}

#[test]
fn the_dimensions_come_from_the_jpeg_itself() {
    let image = image_of("images/jpeg_subsampled.pdf");
    let decoded = jpeg::decode(&image.data).expect("decodes");
    assert_eq!((decoded.width, decoded.height), (32, 32));
    assert_eq!(decoded.components, 3);
}

// ----------------------------------------------------------- downsampling

#[test]
fn downsampling_averages_pixels_and_halves_dimensions() {
    // A 4x2 grayscale ramp, downsampled by 2: each output pixel is the mean
    // of a 2x2 block.
    let mut dict = PdfDict::new();
    dict.set("Width", PdfObj::Int(4));
    dict.set("Height", PdfObj::Int(2));
    dict.set("BitsPerComponent", PdfObj::Int(8));
    dict.set("ColorSpace", PdfObj::Name("DeviceGray".into()));
    let data = vec![0, 10, 100, 110, 0, 10, 100, 110];
    let image = ImageInfo::from_dict(&dict, "Im", data).expect("image");

    let small = image.downsample(2).expect("downsamples");
    assert_eq!((small.width, small.height), (2, 1));
    // Left block is (0+10+0+10)/4 = 5; right is (100+110+100+110)/4 = 105.
    assert_eq!(small.data, vec![5, 105]);
}

#[test]
fn downsampling_by_one_is_the_identity() {
    let image = image_of("images/jpeg_rgb.pdf");
    let same = image.downsample(1).expect("identity");
    assert_eq!(same.width, image.width);
    assert_eq!(same.data, image.data);
}

#[test]
fn downsampling_by_zero_is_refused() {
    let image = image_of("images/jpeg_rgb.pdf");
    assert!(
        image.downsample(0).is_none(),
        "factor 0 would divide by zero; it must be refused"
    );
}

/// Downsampling must not average bytes that are not pixels — a still
/// compressed stream is not raw samples.
#[test]
fn downsampling_refuses_data_that_is_not_raw_samples() {
    let image = image_of("images/jpeg_rgb.pdf");
    // `image.data` is still JPEG-compressed, far shorter than w*h*channels.
    assert!(
        image.downsample(2).is_none(),
        "compressed data must not be box-filtered as though it were pixels"
    );
}

// ------------------------------------------------------------ round trips

#[test]
fn flate_round_trips_adversarial_inputs() {
    for case in [
        Vec::new(),
        vec![0u8],
        vec![0u8; 1000],
        vec![0xFFu8; 1000],
        (0..=255u8).cycle().take(4096).collect(),
        b"the quick brown fox".to_vec(),
    ] {
        let encoded = encode_flate(&case);
        let mut dict = PdfDict::new();
        dict.set("Filter", PdfObj::Name("FlateDecode".into()));
        let decoded = decode_stream(&PdfStream {
            dict,
            data: encoded,
        })
        .expect("flate round-trips");
        assert_eq!(
            decoded,
            case,
            "flate lost data for a {}-byte input",
            case.len()
        );
    }
}

#[test]
fn ascii_hex_round_trips_adversarial_inputs() {
    for case in [
        Vec::new(),
        vec![0u8],
        vec![0xFFu8, 0x00, 0x7F],
        (0..=255u8).collect(),
    ] {
        let encoded = encode_ascii_hex(&case);
        let mut dict = PdfDict::new();
        dict.set("Filter", PdfObj::Name("ASCIIHexDecode".into()));
        let decoded = decode_stream(&PdfStream {
            dict,
            data: encoded,
        })
        .expect("hex round-trips");
        assert_eq!(
            decoded,
            case,
            "hex lost data for a {}-byte input",
            case.len()
        );
    }
}

#[test]
fn every_jpeg_fixture_parses_without_panicking() {
    for name in [
        "jpeg_rgb.pdf",
        "jpeg_gray.pdf",
        "jpeg_subsampled.pdf",
        "jpeg_progressive.pdf",
    ] {
        let bytes = corpus(&format!("images/{name}"));
        let mut doc = PdfDocument::parse(&bytes)
            .unwrap_or_else(|e| panic!("images/{name} should parse: {e}"));
        let Ok(page) = doc.page(0) else { continue };
        for entry in page.xobjects.values() {
            if let Ok(obj) = doc.resolve_ref(entry.obj_ref) {
                if let Some(stream) = obj.as_stream() {
                    if let Some(img) = ImageInfo::from_dict(&stream.dict, "x", stream.data.clone())
                    {
                        let _ = img.decode_to_rgba();
                        let _ = img.downsample(2);
                    }
                }
            }
        }
    }
}