j2k-native 0.13.0

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

use alloc::vec;
use alloc::vec::Vec;

use super::super::build::SubBandType;
use super::super::codestream::CodeBlockStyle;
use super::super::coefficient_view::CoefficientBlockView;
use super::distortion::segment_distortion_delta;
use super::flag_passes::try_encode_column_passes;
use super::passes::{
    clear_coded_in_current_pass, mark_coded_in_current_pass, CODED_IN_CURRENT_PASS, NEGATIVE,
    SIGNIFICANT,
};
use super::preparation::prepare_padded_coefficients;
use super::{
    classic_worker_allocation, coded_bitplanes, encode_code_block, encode_code_block_i64,
    encode_code_block_segments_with_style, encode_code_block_segments_with_style_i64,
    encode_code_block_view, pack_classic_selective_bypass_tier1_tokens, try_encode_code_block,
    try_encode_neighbor_byte_passes, ClassicTier1TokenSegment, CodedBlock, EncodedCodeBlock,
};

fn assert_encoded_blocks_equal(actual: &EncodedCodeBlock, expected: &EncodedCodeBlock) {
    assert_eq!(actual.data, expected.data);
    assert_eq!(actual.num_coding_passes, expected.num_coding_passes);
    assert_eq!(actual.num_zero_bitplanes, expected.num_zero_bitplanes);
    assert_eq!(actual.ht_cleanup_length, expected.ht_cleanup_length);
    assert_eq!(actual.ht_refinement_length, expected.ht_refinement_length);
}

#[test]
fn classic_i32_strided_block_is_byte_exact_with_contiguous_adapter() {
    const WIDTH: usize = 7;
    const HEIGHT: usize = 5;
    const STRIDE: usize = 11;
    const OFFSET: usize = 13;
    let contiguous = (0_i32..i32::try_from(WIDTH * HEIGHT).expect("test size fits i32"))
        .map(|index| match index % 5 {
            0 => 0,
            1 => index * 3,
            2 => -(index * 2),
            3 => 17 - index,
            _ => index / 2,
        })
        .collect::<Vec<_>>();
    let mut padded = vec![i32::MIN; OFFSET + STRIDE * HEIGHT + 9];
    for y in 0..HEIGHT {
        padded[OFFSET + y * STRIDE..OFFSET + y * STRIDE + WIDTH]
            .copy_from_slice(&contiguous[y * WIDTH..(y + 1) * WIDTH]);
    }
    let view = CoefficientBlockView::try_new(&padded, OFFSET, WIDTH, HEIGHT, STRIDE)
        .expect("valid strided classic block");

    let expected = encode_code_block(
        &contiguous,
        u32::try_from(WIDTH).expect("test width fits u32"),
        u32::try_from(HEIGHT).expect("test height fits u32"),
        SubBandType::HighLow,
        12,
    );
    let actual = encode_code_block_view(view, SubBandType::HighLow, 12);
    assert_encoded_blocks_equal(&actual, &expected);
}

#[test]
fn classic_i64_strided_block_is_byte_exact_at_high_precision() {
    const WIDTH: usize = 5;
    const HEIGHT: usize = 5;
    const STRIDE: usize = 9;
    const OFFSET: usize = 7;
    let contiguous = (0_usize..WIDTH * HEIGHT)
        .map(|index| {
            let magnitude =
                (1_i64 << 36) - i64::try_from(index * 1_003).expect("test offset fits i64");
            if index.is_multiple_of(3) {
                -magnitude
            } else {
                magnitude
            }
        })
        .collect::<Vec<_>>();
    let mut padded = vec![i64::MIN; OFFSET + STRIDE * HEIGHT + 3];
    for y in 0..HEIGHT {
        padded[OFFSET + y * STRIDE..OFFSET + y * STRIDE + WIDTH]
            .copy_from_slice(&contiguous[y * WIDTH..(y + 1) * WIDTH]);
    }
    let view = CoefficientBlockView::try_new(&padded, OFFSET, WIDTH, HEIGHT, STRIDE)
        .expect("valid strided exact classic block");

    let expected = encode_code_block_i64(
        &contiguous,
        u32::try_from(WIDTH).expect("test width fits u32"),
        u32::try_from(HEIGHT).expect("test height fits u32"),
        SubBandType::HighHigh,
        40,
    );
    let actual = encode_code_block_view(view, SubBandType::HighHigh, 40);
    assert_encoded_blocks_equal(&actual, &expected);
}

#[test]
fn test_encode_all_zeros() {
    for (width, height, total_bitplanes) in [(4, 4, 8), (64, 64, 16), (3, 5, 0)] {
        let coeffs = vec![0i32; width as usize * height as usize];
        let result =
            encode_code_block(&coeffs, width, height, SubBandType::LowLow, total_bitplanes);
        assert_eq!(result.num_coding_passes, 0);
        assert_eq!(result.data, [] as [u8; 0]);
        assert_eq!(result.num_zero_bitplanes, total_bitplanes);
    }
}

#[test]
fn test_encode_single_nonzero() {
    let mut coeffs = vec![0i32; 16];
    coeffs[0] = 128;
    let result = encode_code_block(&coeffs, 4, 4, SubBandType::LowLow, 8);
    assert!(result.num_coding_passes > 0);
    assert_ne!(result.data, [] as [u8; 0]);
    assert_eq!(result.num_zero_bitplanes, 0);
}

#[test]
fn pack_classic_selective_bypass_tokens_matches_scalar_single_cleanup_block() {
    let style = CodeBlockStyle {
        selective_arithmetic_coding_bypass: true,
        reset_context_probabilities: false,
        termination_on_each_pass: false,
        vertically_causal_context: false,
        segmentation_symbols: false,
        high_throughput_block_coding: false,
        mixed_block_coding: false,
    };
    let coefficients = [1i32];
    let scalar =
        encode_code_block_segments_with_style(&coefficients, 1, 1, SubBandType::LowLow, 1, style);
    let token_bytes = pack_mq_test_tokens(&[(0, 1), (9, 0)]);
    let packed = pack_classic_selective_bypass_tier1_tokens(
        &token_bytes,
        &[ClassicTier1TokenSegment {
            token_bit_offset: 0,
            token_bit_count: 12,
            start_coding_pass: 0,
            end_coding_pass: 1,
            use_arithmetic: true,
        }],
        scalar.num_coding_passes,
        scalar.num_zero_bitplanes,
    )
    .expect("tokens pack");

    assert_eq!(packed.data, scalar.data);
    assert_eq!(packed.num_coding_passes, scalar.num_coding_passes);
    assert_eq!(packed.num_zero_bitplanes, scalar.num_zero_bitplanes);
    assert_eq!(packed.segments.len(), scalar.segments.len());
    for (packed_segment, scalar_segment) in packed.segments.iter().zip(&scalar.segments) {
        assert_eq!(packed_segment.data_offset, scalar_segment.data_offset);
        assert_eq!(packed_segment.data_length, scalar_segment.data_length);
        assert_eq!(
            packed_segment.start_coding_pass,
            scalar_segment.start_coding_pass
        );
        assert_eq!(
            packed_segment.end_coding_pass,
            scalar_segment.end_coding_pass
        );
        assert_eq!(packed_segment.use_arithmetic, scalar_segment.use_arithmetic);
    }
}

#[test]
fn token_segments_must_cover_coding_passes_contiguously() {
    let segment = |start_coding_pass, end_coding_pass| ClassicTier1TokenSegment {
        token_bit_offset: 0,
        token_bit_count: 0,
        start_coding_pass,
        end_coding_pass,
        use_arithmetic: false,
    };

    for (segments, passes) in [
        (vec![segment(1, 2)], 2),
        (vec![segment(0, 1), segment(2, 3)], 3),
        (vec![segment(0, 1)], 2),
    ] {
        assert!(
            pack_classic_selective_bypass_tier1_tokens(&[], &segments, passes, 0).is_err(),
            "gapped token pass schedule must be rejected: {segments:?}"
        );
    }
}

fn pack_mq_test_tokens(tokens: &[(u8, u8)]) -> Vec<u8> {
    let mut bytes = Vec::new();
    let mut current = 0u8;
    let mut bits = 0u8;
    for &(ctx, bit) in tokens {
        let value = (ctx & 0x1F) | ((bit & 1) << 5);
        for shift in (0..6).rev() {
            current = (current << 1) | ((value >> shift) & 1);
            bits += 1;
            if bits == 8 {
                bytes.push(current);
                current = 0;
                bits = 0;
            }
        }
    }

    if bits != 0 {
        bytes.push(current << (8 - bits));
    }
    bytes
}

#[test]
fn token_segments_must_follow_selective_bypass_coding_modes() {
    let segment = |start_coding_pass, end_coding_pass, use_arithmetic| ClassicTier1TokenSegment {
        token_bit_offset: 0,
        token_bit_count: 0,
        start_coding_pass,
        end_coding_pass,
        use_arithmetic,
    };
    let malformed = [
        (vec![segment(0, 1, false)], 1),
        (vec![segment(0, 10, true), segment(10, 11, true)], 11),
        (vec![segment(0, 11, true), segment(11, 12, false)], 12),
        (vec![segment(0, 10, true), segment(10, 13, false)], 13),
        (vec![segment(0, 12, true), segment(12, 13, false)], 13),
    ];

    for (segments, passes) in malformed {
        assert!(
            pack_classic_selective_bypass_tier1_tokens(&[], &segments, passes, 0).is_err(),
            "invalid selective-bypass schedule must be rejected: {segments:?}"
        );
    }
}

#[test]
fn test_encode_various_magnitudes() {
    let coeffs: Vec<i32> = (0..64)
        .map(|x| if x % 3 == 0 { x * 10 } else { -x })
        .collect();
    let result = encode_code_block(&coeffs, 8, 8, SubBandType::HighHigh, 12);
    assert!(result.num_coding_passes > 0);
    assert_ne!(result.data, [] as [u8; 0]);
}

#[test]
fn test_zero_bitplanes_count() {
    // Max value is 7 (3 bits), so with Mb=8 we have 8 - 3 = 5 zero bitplanes.
    let coeffs = vec![7i32, -3, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
    let result = encode_code_block(&coeffs, 4, 4, SubBandType::LowLow, 8);
    assert_eq!(result.num_zero_bitplanes, 5);
}

#[test]
fn padded_coefficient_preparation_stores_sign_in_state_flags() {
    let coeffs = vec![7i64, -3, 0, -9];
    let (magnitudes, states) = prepare_padded_coefficients(&coeffs, 2, 2, 4);

    assert_eq!(magnitudes[5], 7);
    assert_eq!(magnitudes[6], 3);
    assert_eq!(magnitudes[9], 0);
    assert_eq!(magnitudes[10], 9);
    assert_eq!(states[5] & NEGATIVE, 0);
    assert_ne!(states[6] & NEGATIVE, 0);
    assert_eq!(states[9] & NEGATIVE, 0);
    assert_ne!(states[10] & NEGATIVE, 0);
}

#[test]
fn clear_coded_in_current_pass_keeps_every_other_state_flag() {
    let mut states = vec![SIGNIFICANT, NEGATIVE, 0, SIGNIFICANT | NEGATIVE];
    mark_coded_in_current_pass(0, &mut states);
    mark_coded_in_current_pass(2, &mut states);

    clear_coded_in_current_pass(&mut states);

    assert_eq!(states, [SIGNIFICANT, NEGATIVE, 0, SIGNIFICANT | NEGATIVE]);
    assert!(states
        .iter()
        .all(|state| state & CODED_IN_CURRENT_PASS == 0));
}

#[test]
fn pcrd_distortion_delta_reflects_residual_error_reduction() {
    let sparse_delta = segment_distortion_delta(&[8], 0, 1, 4);
    let dense_delta = segment_distortion_delta(&[15], 0, 1, 4);

    assert!(
            dense_delta > sparse_delta,
            "coefficients with the same MSB but larger residual error should have larger PCRD distortion reduction"
        );
}

#[test]
fn str011a_classic_tier1_byte_baseline() {
    let coefficients = (0..64)
        .map(|index| {
            let magnitude = i64::from((index * 37) % 4096);
            if index % 3 == 0 {
                -magnitude
            } else {
                magnitude
            }
        })
        .collect::<Vec<_>>();
    let styles = [
        CodeBlockStyle::default(),
        CodeBlockStyle {
            termination_on_each_pass: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            selective_arithmetic_coding_bypass: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            reset_context_probabilities: true,
            vertically_causal_context: true,
            segmentation_symbols: true,
            ..CodeBlockStyle::default()
        },
    ];
    let expected = [
        (99, 0xfcce_40ac_5a7f_501d, 34, 0, 1, 0x64c9_46a4_01c3_99fb),
        (130, 0xcdf7_5e12_515b_dd00, 34, 0, 34, 0x5029_9d58_2988_06ce),
        (112, 0x2828_c2eb_1fa9_8008, 34, 0, 17, 0x6706_112a_c1da_7d36),
        (104, 0x05e7_b41b_483d_1290, 34, 0, 1, 0xdef6_5a04_ddeb_b7c8),
    ];

    for (index, style) in styles.iter().enumerate() {
        let encoded = encode_code_block_segments_with_style_i64(
            &coefficients,
            8,
            8,
            SubBandType::HighHigh,
            12,
            *style,
        );
        let digest = encoded
            .data
            .iter()
            .fold(0xcbf2_9ce4_8422_2325u64, |hash, byte| {
                hash.wrapping_mul(0x0100_0000_01b3) ^ u64::from(*byte)
            });
        let mut segment_bytes = Vec::new();
        for segment in &encoded.segments {
            segment_bytes.extend_from_slice(&segment.data_offset.to_le_bytes());
            segment_bytes.extend_from_slice(&segment.data_length.to_le_bytes());
            segment_bytes.push(segment.start_coding_pass);
            segment_bytes.push(segment.end_coding_pass);
            segment_bytes.extend_from_slice(&segment.distortion_delta.to_bits().to_le_bytes());
            segment_bytes.push(u8::from(segment.use_arithmetic));
        }
        let segment_digest = segment_bytes
            .iter()
            .fold(0xcbf2_9ce4_8422_2325u64, |hash, byte| {
                hash.wrapping_mul(0x0100_0000_01b3) ^ u64::from(*byte)
            });
        assert_eq!(
            (
                encoded.data.len(),
                digest,
                encoded.num_coding_passes,
                encoded.num_zero_bitplanes,
                encoded.segments.len(),
                segment_digest,
            ),
            expected[index],
            "classic Tier-1 bytes or segment accounting changed for style {index}",
        );
    }
}

#[test]
fn worst_case_blocks_stay_within_the_worker_payload_plan() {
    // Rows: width, height and the total bitplanes to try. Small blocks try
    // every count up to 64, the widest u64 magnitude; large ones the widest,
    // tallest and largest shapes at 1, 16 and 31 bitplanes.
    let every: Vec<u8> = (1..=64).collect();
    let few = [1_u8, 16, 31];
    let shapes: [(usize, usize, &[u8]); 8] = [
        (1, 1, &every),
        (4, 4, &every),
        (3, 5, &every),
        (64, 64, &few),
        (7, 64, &few),
        (64, 7, &few),
        (1024, 4, &few),
        (4, 1024, &few),
    ];
    let styles = [
        CodeBlockStyle::default(),
        CodeBlockStyle {
            vertically_causal_context: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            selective_arithmetic_coding_bypass: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            termination_on_each_pass: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            reset_context_probabilities: true,
            segmentation_symbols: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            selective_arithmetic_coding_bypass: true,
            reset_context_probabilities: true,
            termination_on_each_pass: true,
            vertically_causal_context: true,
            segmentation_symbols: true,
            ..CodeBlockStyle::default()
        },
    ];
    let mut state = 0x2545_f491_4f6c_dd1d_u64;
    for &(width, height, bitplane_counts) in &shapes {
        for &total_bitplanes in bitplane_counts {
            let top = (1_u64 << (total_bitplanes - 1)) - 1 + (1_u64 << (total_bitplanes - 1));
            // Draws: uniform noise over every bitplane, one bit at a random
            // plane, every coefficient at the top plane, which makes each
            // full cleanup column a run-length interruption followed by three
            // significance and sign pairs, and alternating extremes.
            for draw in 0..4 {
                let coefficients: Vec<i64> = (0..width * height)
                    .map(|index| {
                        state ^= state << 13;
                        state ^= state >> 7;
                        state ^= state << 17;
                        let magnitude = match draw {
                            0 => state & top,
                            1 => 1 << (state % u64::from(total_bitplanes)),
                            2 => 1 << (total_bitplanes - 1),
                            _ => top,
                        };
                        let negative =
                            (draw == 3 && index % 2 == 0) || (draw != 3 && state >> 63 == 1);
                        match i64::try_from(magnitude) {
                            Ok(value) if negative => -value,
                            Ok(value) => value,
                            Err(_) => i64::MIN,
                        }
                    })
                    .collect();
                let plan = classic_worker_allocation(width, height, total_bitplanes)
                    .expect("valid code-block plan");
                for (index, style) in styles.iter().enumerate() {
                    let sub_band_type = [
                        SubBandType::LowLow,
                        SubBandType::HighLow,
                        SubBandType::LowHigh,
                        SubBandType::HighHigh,
                    ][(index + draw) % 4];
                    let encoded = super::try_encode_code_block_segments_with_style_i64(
                        &coefficients,
                        u32::try_from(width).expect("width fits u32"),
                        u32::try_from(height).expect("height fits u32"),
                        sub_band_type,
                        total_bitplanes,
                        style,
                    )
                    .unwrap_or_else(|error| {
                        panic!(
                            "{width}x{height}, {total_bitplanes} bitplanes, draw {draw}: {error}"
                        )
                    });
                    assert!(
                        encoded.data.len() <= plan.payload_bytes,
                        "{width}x{height}, {total_bitplanes} bitplanes, draw {draw}, style {index}: \
                         {} bytes over the {} byte plan",
                        encoded.data.len(),
                        plan.payload_bytes
                    );
                }
            }
        }
    }
}

#[test]
fn column_word_passes_match_neighbor_byte_passes() {
    // Rows: width, height and total bitplanes; partial stripes, single rows
    // and columns, and 31 bitplanes, the most the column words hold.
    let shapes = [
        (1, 1, 8),
        (1, 4, 8),
        (3, 7, 10),
        (63, 65, 13),
        (1, 7, 9),
        (5, 3, 10),
        (4, 4, 12),
        (7, 5, 11),
        (13, 9, 14),
        (64, 1, 9),
        (1, 64, 9),
        (63, 61, 13),
        (32, 64, 12),
        (1024, 4, 10),
        (64, 64, 16),
        (64, 64, 31),
    ];
    let styles = [
        CodeBlockStyle::default(),
        CodeBlockStyle {
            reset_context_probabilities: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            segmentation_symbols: true,
            ..CodeBlockStyle::default()
        },
        CodeBlockStyle {
            reset_context_probabilities: true,
            segmentation_symbols: true,
            ..CodeBlockStyle::default()
        },
    ];
    let sub_bands = [
        SubBandType::LowLow,
        SubBandType::HighLow,
        SubBandType::LowHigh,
        SubBandType::HighHigh,
    ];
    let mut state = 0x9e37_79b9_7f4a_7c15_u64;
    for (case, &(width, height, total_bitplanes)) in shapes.iter().enumerate() {
        let full = (1_u64 << total_bitplanes) - 1;
        // Draws: zero, sparse, uniform, mostly small, and alternating bits.
        for draw in 0..5 {
            let coefficients: Vec<i32> = (0..width * height)
                .map(|_| {
                    state = state
                        .wrapping_mul(6_364_136_223_846_793_005)
                        .wrapping_add(1_442_695_040_888_963_407);
                    let random = state >> 11;
                    let magnitude = match draw {
                        0 => 0,
                        4 => full / 3,
                        1 if !random.is_multiple_of(16) => 0,
                        3 => (random & full) >> (random % u64::from(total_bitplanes)),
                        _ => random & full,
                    };
                    let magnitude = i32::try_from(magnitude).expect("31-bit magnitude fits i32");
                    if draw == 4 || random & (1 << 52) == 0 {
                        magnitude
                    } else {
                        -magnitude
                    }
                })
                .collect();
            let view = CoefficientBlockView::try_contiguous(&coefficients, width, height)
                .expect("contiguous test block");
            let Some(num_bitplanes) = coded_bitplanes(view, total_bitplanes).expect("fits") else {
                continue;
            };
            let allocation = classic_worker_allocation(width, height, total_bitplanes)
                .expect("valid code-block plan");
            for (band, &sub_band_type) in sub_bands.iter().enumerate() {
                let style = &styles[(case + draw + band) % styles.len()];
                let column = try_encode_column_passes(
                    view,
                    sub_band_type,
                    num_bitplanes,
                    style,
                    allocation.payload_bytes,
                )
                .expect("column-word passes");
                let neighbor = try_encode_neighbor_byte_passes(
                    view,
                    sub_band_type,
                    num_bitplanes,
                    style,
                    allocation,
                )
                .expect("neighbor-byte passes");
                let summary = |block: CodedBlock| (block.data, block.num_coding_passes);
                assert_eq!(
                    summary(column),
                    summary(neighbor),
                    "{width}x{height}, {total_bitplanes} bitplanes, draw {draw}, {sub_band_type:?}"
                );
            }
        }
    }
}

#[test]
fn classic_encode_rejects_magnitudes_above_the_bitplane_count() {
    for (coefficient, total_bitplanes) in [(2, 1), (i32::MIN, 31)] {
        assert_eq!(
            try_encode_code_block(&[coefficient], 1, 1, SubBandType::LowLow, total_bitplanes)
                .unwrap_err(),
            crate::EncodeError::InvalidInput {
                what: "classic code-block magnitude exceeds configured bitplane count",
            }
        );
    }
}