j2k-native 0.13.1

// SPDX-License-Identifier: MIT OR Apache-2.0

use alloc::{vec, vec::Vec};
use std::panic::catch_unwind;

use super::super::*;

mod resident_errors;
mod whole_tile;

fn codestream_fingerprint(codestream: &[u8]) -> (usize, u64) {
    let hash = codestream
        .iter()
        .fold(0xcbf2_9ce4_8422_2325_u64, |hash, byte| {
            (hash ^ u64::from(*byte)).wrapping_mul(0x0000_0100_0000_01b3)
        });
    (codestream.len(), hash)
}

#[test]
fn public_encode_distinguishes_malformed_and_unsupported_requests() {
    let malformed = encode(&[], 0, 1, 1, 8, false, &EncodeOptions::default())
        .expect_err("zero-width encode must fail");
    assert_eq!(
        malformed,
        crate::EncodeError::InvalidInput {
            what: "invalid dimensions"
        }
    );

    let sampled_multitile = encode(
        &[0; 16],
        4,
        4,
        1,
        8,
        false,
        &EncodeOptions {
            component_sampling: Some(vec![(2, 1)]),
            tile_size: Some((2, 2)),
            ..EncodeOptions::default()
        },
    )
    .expect_err("sampled multi-tile encode is not implemented");
    assert_eq!(
        sampled_multitile,
        crate::EncodeError::Unsupported {
            what: "multi-tile encode with component sampling is not implemented"
        }
    );

    let high_bit_sampled = encode(
        &[0; 4],
        1,
        1,
        1,
        25,
        false,
        &EncodeOptions {
            reversible: true,
            component_sampling: Some(vec![(2, 1)]),
            ..EncodeOptions::default()
        },
    )
    .expect_err("exact high-bit raw encode requires full-resolution components");
    assert_eq!(
        high_bit_sampled,
        crate::EncodeError::Unsupported {
            what: "25-38 bit encode currently requires full-resolution components"
        }
    );
}

#[test]
fn public_code_block_exponents_fail_typed_without_panicking() {
    let oversized_width = catch_unwind(|| {
        encode(
            &[0; 4],
            2,
            2,
            1,
            8,
            false,
            &EncodeOptions {
                code_block_width_exp: u8::MAX,
                ..EncodeOptions::default()
            },
        )
    })
    .expect("u8::MAX code-block exponent must not panic")
    .expect_err("u8::MAX code-block exponent must fail");
    assert_eq!(
        oversized_width,
        crate::EncodeError::InvalidInput {
            what: "code-block width exponent exceeds supported range"
        }
    );

    let high_bit_data = [0_u8; 4];
    let high_bit_planes = [EncodeTypedComponentPlane {
        data: &high_bit_data,
        x_rsiz: 1,
        y_rsiz: 1,
        bit_depth: 25,
        signed: false,
    }];
    let oversized_height = catch_unwind(|| {
        encode_typed_component_planes_53(
            &high_bit_planes,
            1,
            1,
            &EncodeOptions {
                code_block_height_exp: 30,
                ..EncodeOptions::default()
            },
        )
    })
    .expect("large typed code-block exponent must not panic")
    .expect_err("large typed code-block exponent must fail");
    assert_eq!(
        oversized_height,
        crate::EncodeError::InvalidInput {
            what: "code-block height exponent exceeds supported range"
        }
    );

    let oversized_area = catch_unwind(|| {
        encode(
            &[0; 4],
            2,
            2,
            1,
            8,
            false,
            &EncodeOptions {
                code_block_width_exp: 5,
                code_block_height_exp: 4,
                ..EncodeOptions::default()
            },
        )
    })
    .expect("oversized code-block area must not panic")
    .expect_err("oversized code-block area must fail");
    assert_eq!(
        oversized_area,
        crate::EncodeError::InvalidInput {
            what: "code-block dimensions exceed JPEG 2000 Part 1 area limit"
        }
    );
}

#[test]
fn reference_codestreams_remain_byte_exact() {
    let rgb_pixels = (0..8 * 8 * 3)
        .map(|index| u8::try_from((index * 29 + 17) & 0xff).expect("sample fits u8"))
        .collect::<Vec<_>>();
    let classic = encode(
        &rgb_pixels,
        8,
        8,
        3,
        8,
        false,
        &EncodeOptions {
            num_decomposition_levels: 1,
            reversible: true,
            use_mct: true,
            ..EncodeOptions::default()
        },
    )
    .expect("classic reference encode");

    let gray_pixels = (0..8 * 8)
        .map(|index| u8::try_from((index * 37 + 11) & 0xff).expect("sample fits u8"))
        .collect::<Vec<_>>();
    let ht = encode_htj2k(
        &gray_pixels,
        8,
        8,
        1,
        8,
        false,
        &EncodeOptions {
            num_decomposition_levels: 1,
            reversible: false,
            guard_bits: 2,
            ..EncodeOptions::default()
        },
    )
    .expect("HT reference encode");

    let mut high_bit_pixels = Vec::with_capacity(4 * 4 * 4);
    for index in 0..4 * 4_u32 {
        high_bit_pixels.extend_from_slice(&(index * 1_048_583 + 19).to_le_bytes());
    }
    let high_bit = encode(
        &high_bit_pixels,
        4,
        4,
        1,
        25,
        false,
        &EncodeOptions {
            num_decomposition_levels: 1,
            reversible: true,
            ..EncodeOptions::default()
        },
    )
    .expect("high-bit reference encode");

    assert_eq!(
        [
            codestream_fingerprint(&classic),
            codestream_fingerprint(&ht),
            codestream_fingerprint(&high_bit),
        ],
        [
            (289, 10_919_473_942_675_803_716),
            // Irreversible: detail steps follow 9/7 synthesis norms.
            (205, 2_029_000_814_618_401_571),
            (120, 14_442_175_691_548_962_802),
        ]
    );
}

#[test]
fn ht_accelerator_hooks_keep_pipeline_order() {
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    enum Stage {
        WholeTile,
        Deinterleave,
        ForwardDwt,
        Packetization,
    }

    #[derive(Default)]
    struct RecordingAccelerator {
        stages: Vec<Stage>,
    }

    impl crate::J2kEncodeStageAccelerator for RecordingAccelerator {
        fn encode_htj2k_tile(
            &mut self,
            _job: crate::J2kHtj2kTileEncodeJob<'_>,
        ) -> crate::J2kEncodeStageResult<Option<Vec<u8>>> {
            self.stages.push(Stage::WholeTile);
            Ok(None)
        }

        fn encode_deinterleave(
            &mut self,
            _job: crate::J2kDeinterleaveToF32Job<'_>,
        ) -> crate::J2kEncodeStageResult<Option<Vec<Vec<f32>>>> {
            self.stages.push(Stage::Deinterleave);
            Ok(None)
        }

        fn encode_forward_dwt53(
            &mut self,
            _job: crate::J2kForwardDwt53Job<'_>,
        ) -> crate::J2kEncodeStageResult<Option<crate::J2kForwardDwt53Output>> {
            self.stages.push(Stage::ForwardDwt);
            Ok(None)
        }

        fn encode_packetization(
            &mut self,
            _job: crate::J2kPacketizationEncodeJob<'_>,
        ) -> crate::J2kEncodeStageResult<Option<Vec<u8>>> {
            self.stages.push(Stage::Packetization);
            Ok(None)
        }
    }

    let pixels = (0..8 * 8)
        .map(|index| u8::try_from((index * 41 + 3) & 0xff).expect("sample fits u8"))
        .collect::<Vec<_>>();
    let options = EncodeOptions {
        num_decomposition_levels: 1,
        reversible: true,
        use_ht_block_coding: true,
        ..EncodeOptions::default()
    };
    let expected = encode(&pixels, 8, 8, 1, 8, false, &options).expect("CPU reference encode");
    let mut accelerator = RecordingAccelerator::default();
    let actual = encode_with_accelerator(&pixels, 8, 8, 1, 8, false, &options, &mut accelerator)
        .expect("fallback accelerator encode");

    assert_eq!(actual, expected);
    assert_eq!(
        accelerator.stages,
        [
            Stage::WholeTile,
            Stage::Deinterleave,
            Stage::ForwardDwt,
            Stage::Packetization,
        ]
    );
}

fn textured_samples(count: usize, mask: u32) -> Vec<u32> {
    let mut state = 0x2545_f491_u32;
    (0..count)
        .map(|index| {
            state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            let index = u32::try_from(index).expect("fixture index fits u32");
            (index.wrapping_mul(37) ^ (state >> 20)) & mask
        })
        .collect()
}

/// Final byte of every non-empty code-block contribution, per packet.
fn contribution_final_bytes(codestream: &[u8]) -> Vec<u8> {
    use crate::j2c::{build, decode::DecompositionStorage, segment, tile};

    let parsed = crate::j2c::parse_raw(codestream, &crate::DecodeSettings::default())
        .expect("parse encoded codestream");
    let tiles = tile::parse(
        &mut crate::reader::BitReader::new(parsed.data),
        &parsed.header,
        0,
    )
    .expect("parse encoded tiles");
    let mut storage = DecompositionStorage::default();
    let mut final_bytes = Vec::new();
    for tile in tiles.iter() {
        storage.reset_for_next_tile();
        build::build(
            tile,
            &mut storage,
            tiles.structural_workspace_bytes(),
            false,
            build::BuildWorkspace::MetadataOnly,
            None,
        )
        .expect("plan encoded code-blocks");
        segment::parse(tile, &parsed.header, &mut storage).expect("parse encoded packets");
        for code_block in &storage.code_blocks {
            for layer in &storage.layers[code_block.layers.clone()] {
                let Some(segments) = layer.segments.clone() else {
                    continue;
                };
                if let Some(&last) = storage.segments[segments]
                    .iter()
                    .rev()
                    .find_map(|segment| segment.data.last())
                {
                    final_bytes.push(last);
                }
            }
        }
    }
    final_bytes
}

struct TerminationCase<'a> {
    name: &'static str,
    pixels: &'a [u8],
    components: u16,
    bit_depth: u8,
    options: &'a EncodeOptions,
}

#[test]
fn classic_code_block_contributions_never_end_in_ff() {
    // T.800 B.10.7: the bytes included for a code-block must not end in 0xFF.
    let rgb8 = textured_samples(64 * 64 * 3, 0xff)
        .into_iter()
        .map(|sample| u8::try_from(sample).expect("masked sample fits u8"))
        .collect::<Vec<_>>();
    let rgb16 = textured_samples(64 * 64 * 3, 0xffff)
        .into_iter()
        .flat_map(|sample| {
            u16::try_from(sample)
                .expect("masked sample fits u16")
                .to_le_bytes()
        })
        .collect::<Vec<_>>();
    let gray8 = textured_samples(64 * 64, 0xff)
        .into_iter()
        .map(|sample| u8::try_from(sample).expect("masked sample fits u8"))
        .collect::<Vec<_>>();
    let lossless_rgb = EncodeOptions {
        num_decomposition_levels: 3,
        reversible: true,
        use_mct: true,
        ..EncodeOptions::default()
    };
    let layered_gray = EncodeOptions {
        num_decomposition_levels: 3,
        reversible: true,
        num_layers: 3,
        quality_layer_byte_targets: vec![256, 1_024, 8_192],
        ..EncodeOptions::default()
    };
    let cases = [
        TerminationCase {
            name: "rgb8 lossless",
            pixels: &rgb8,
            components: 3,
            bit_depth: 8,
            options: &lossless_rgb,
        },
        TerminationCase {
            name: "rgb16 lossless",
            pixels: &rgb16,
            components: 3,
            bit_depth: 16,
            options: &lossless_rgb,
        },
        TerminationCase {
            name: "gray8 three layers",
            pixels: &gray8,
            components: 1,
            bit_depth: 8,
            options: &layered_gray,
        },
    ];

    for case in cases {
        let name = case.name;
        let codestream = encode(
            case.pixels,
            64,
            64,
            case.components,
            case.bit_depth,
            false,
            case.options,
        )
        .unwrap_or_else(|error| panic!("{name} encode: {error:?}"));
        let final_bytes = contribution_final_bytes(&codestream);
        let ending_in_ff: usize = final_bytes
            .iter()
            .map(|&byte| usize::from(byte == 0xFF))
            .sum();
        assert!(!final_bytes.is_empty(), "{name} has no contributions");
        assert_eq!(
            ending_in_ff,
            0,
            "{name}: {ending_in_ff}/{} contributions end in 0xFF",
            final_bytes.len()
        );

        let image = crate::Image::new(&codestream, &crate::DecodeSettings::default())
            .unwrap_or_else(|error| panic!("{name} parse: {error:?}"));
        let decoded = image
            .decode_native()
            .unwrap_or_else(|error| panic!("{name} decode: {error:?}"));
        assert_eq!(decoded.data, case.pixels, "{name} lossless round trip");
    }
}

use super::coefficient_source::{DwtBandView, DwtComponentSource, DwtLevelView, OwnedDwtComponent};
use core::cell::Cell;

/// Checks at its component's first read that every earlier one is dropped.
struct CheckedSource<'a> {
    component: usize,
    source: OwnedDwtComponent,
    dropped: &'a Cell<usize>,
}

impl DwtComponentSource for CheckedSource<'_> {
    fn ll(&self) -> NativeEncodePipelineResult<DwtBandView<'_>> {
        assert_eq!(
            self.dropped.get(),
            self.component,
            "earlier sources still live"
        );
        self.source.ll()
    }

    fn level_count(&self) -> usize {
        self.source.level_count()
    }

    fn level(&self, index: usize) -> NativeEncodePipelineResult<Option<DwtLevelView<'_>>> {
        self.source.level(index)
    }

    fn owned_bytes(&self) -> NativeEncodePipelineResult<usize> {
        self.source.owned_bytes()
    }
}

impl Drop for CheckedSource<'_> {
    fn drop(&mut self) {
        self.dropped.set(self.dropped.get() + 1);
    }
}

#[test]
fn staged_encodes_drop_each_component_source_once_its_code_blocks_are_prepared() {
    use super::accelerator::{prepare_accelerated_components, AcceleratedComponentRequest};
    use super::ownership::prepared_transforms_retained_bytes;
    use super::plan::{build_single_tile_plan, validate_encode_request, ValidatedEncodeRoute};
    use super::tile_encode::encode_tile_packets;

    const SIDE: u32 = 32;
    let pixels = textured_samples(SIDE as usize * SIDE as usize * 3, 0xff)
        .into_iter()
        .map(|sample| u8::try_from(sample).expect("masked sample fits u8"))
        .collect::<Vec<_>>();
    for mode in [BlockCodingMode::Classic, BlockCodingMode::HighThroughput] {
        let options = EncodeOptions {
            num_decomposition_levels: 2,
            reversible: true,
            use_mct: true,
            use_ht_block_coding: mode == BlockCodingMode::HighThroughput,
            ..EncodeOptions::default()
        };
        let session = NativeEncodeSession::try_new(NativeEncodeRetainedInput::none())
            .expect("encode session");
        let ValidatedEncodeRoute::SingleTile(validated) = validate_encode_request(
            pixels.len(),
            SIDE,
            SIDE,
            3,
            8,
            &options,
            mode,
            &[],
            &session,
        )
        .expect("encode request") else {
            panic!("fixture must be single tile");
        };
        let plan = build_single_tile_plan(
            validated,
            SIDE,
            SIDE,
            3,
            8,
            false,
            &options,
            mode,
            &[],
            &[],
            &session,
        )
        .expect("single-tile plan");
        let prepared = prepare_accelerated_components(
            &AcceleratedComponentRequest {
                pixels: &pixels,
                width: SIDE,
                height: SIDE,
                num_components: 3,
                bit_depth: 8,
                signed: false,
                options: &options,
                plan: &plan,
                profile_enabled: false,
                session: &session,
            },
            &mut crate::CpuOnlyJ2kEncodeStageAccelerator,
        )
        .expect("component transforms");
        let source_bytes = prepared_transforms_retained_bytes(&prepared).expect("source bytes");
        let dropped = Cell::new(0);
        let sources = prepared
            .decompositions
            .into_iter()
            .enumerate()
            .map(|(component, source)| CheckedSource {
                component,
                source,
                dropped: &dropped,
            })
            .collect::<Vec<_>>();

        encode_tile_packets(
            SIDE,
            SIDE,
            3,
            8,
            &options,
            &[],
            &plan,
            sources.into_iter(),
            source_bytes,
            0,
            false,
            &session,
            &mut crate::CpuOnlyJ2kEncodeStageAccelerator,
        )
        .unwrap_or_else(|error| panic!("{mode:?} encode: {error:?}"));

        assert_eq!(dropped.get(), 3, "{mode:?}");
    }
}