j2k-native 0.13.0

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

use super::{ArithmeticEncoder, ArithmeticEncoderContext};
use crate::j2c::arithmetic_decoder::{ArithmeticDecoder, ArithmeticDecoderContext};
use alloc::{vec, vec::Vec};

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_encode_decode_round_trip() {
    let symbols: Vec<u32> = vec![0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1];
    let mut encoder = ArithmeticEncoder::new();
    let mut enc_ctx = ArithmeticEncoderContext::default();

    for &s in &symbols {
        encoder.encode(s, &mut enc_ctx);
    }
    let encoded = encoder.finish();

    // Decode and verify (new() already calls initialize())
    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_ctx = ArithmeticDecoderContext::default();

    let mut decoded = Vec::new();
    for _ in 0..symbols.len() {
        decoded.push(decoder.decode(&mut dec_ctx));
    }

    assert_eq!(symbols, decoded);
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_encode_all_mps() {
    let mut encoder = ArithmeticEncoder::new();
    let mut ctx = ArithmeticEncoderContext::default();
    for _ in 0..100 {
        encoder.encode(0, &mut ctx);
    }
    let encoded = encoder.finish();

    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_ctx = ArithmeticDecoderContext::default();
    for _ in 0..100 {
        assert_eq!(decoder.decode(&mut dec_ctx), 0);
    }
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn with_capacity_preserves_round_trip_encoding() {
    let mut encoder = ArithmeticEncoder::with_capacity(128);
    let mut enc_ctx = ArithmeticEncoderContext::default();
    let symbols = [0u32, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 1];

    for &symbol in &symbols {
        encoder.encode(symbol, &mut enc_ctx);
    }
    let encoded = encoder.finish();

    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_ctx = ArithmeticDecoderContext::default();
    for &symbol in &symbols {
        assert_eq!(decoder.decode(&mut dec_ctx), symbol);
    }
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_encode_all_lps() {
    let mut encoder = ArithmeticEncoder::new();
    let mut ctx = ArithmeticEncoderContext::default();
    for _ in 0..50 {
        encoder.encode(1, &mut ctx);
    }
    let encoded = encoder.finish();

    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_ctx = ArithmeticDecoderContext::default();
    for _ in 0..50 {
        assert_eq!(decoder.decode(&mut dec_ctx), 1);
    }
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_multiple_contexts() {
    let symbols_a = [0u32, 1, 0, 0, 1, 1, 0, 1];
    let symbols_b = [1u32, 1, 0, 1, 0, 0, 1, 0];

    let mut encoder = ArithmeticEncoder::new();
    let mut ctx_a = ArithmeticEncoderContext::default();
    let mut ctx_b = ArithmeticEncoderContext::default();

    for i in 0..8 {
        encoder.encode(symbols_a[i], &mut ctx_a);
        encoder.encode(symbols_b[i], &mut ctx_b);
    }
    let encoded = encoder.finish();

    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_ctx_a = ArithmeticDecoderContext::default();
    let mut dec_ctx_b = ArithmeticDecoderContext::default();

    for i in 0..8 {
        assert_eq!(decoder.decode(&mut dec_ctx_a), symbols_a[i]);
        assert_eq!(decoder.decode(&mut dec_ctx_b), symbols_b[i]);
    }
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_many_context_round_trip() {
    let mut state = 0x1234_5678u32;
    let mut symbols = Vec::new();
    let mut labels = Vec::new();
    let mut encoder = ArithmeticEncoder::new();
    let mut enc_contexts = [ArithmeticEncoderContext::default(); 19];
    enc_contexts[0].reset_with_index(4);
    enc_contexts[17].reset_with_index(3);
    enc_contexts[18].reset_with_index(46);

    for _ in 0..100_000 {
        state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
        let label = (state % 19) as usize;
        state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
        let bit = (state >> 31) & 1;
        encoder.encode(bit, &mut enc_contexts[label]);
        labels.push(label);
        symbols.push(bit);
    }

    let encoded = encoder.finish();
    // Internal 0xFF bytes (and their stuffed successors) must survive; only a
    // terminal 0xFF is omitted.
    assert!(encoded[..encoded.len() - 1].contains(&0xFF));
    assert_ne!(encoded.last(), Some(&0xFF));
    let mut decoder = ArithmeticDecoder::new(&encoded);
    let mut dec_contexts = [ArithmeticDecoderContext::default(); 19];
    dec_contexts[0].reset_with_index(4);
    dec_contexts[17].reset_with_index(3);
    dec_contexts[18].reset_with_index(46);

    for (index, (&label, &symbol)) in labels.iter().zip(symbols.iter()).enumerate() {
        let decoded = decoder.decode(&mut dec_contexts[label]);
        assert_eq!(decoded, symbol, "mismatch at symbol {index}");
    }
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn test_context_state_identical() {
    let mut enc_ctx = ArithmeticEncoderContext::default();
    let mut dec_ctx = ArithmeticDecoderContext::default();

    let bits = [0u32, 0, 1, 0, 1, 1, 0, 0];
    let mut encoder = ArithmeticEncoder::new();
    for &b in &bits {
        encoder.encode(b, &mut enc_ctx);
    }
    let encoded = encoder.finish();

    let mut decoder = ArithmeticDecoder::new(&encoded);
    for &b in &bits {
        let decoded = decoder.decode(&mut dec_ctx);
        assert_eq!(decoded, b);
    }

    // Both contexts should be in same state
    assert_eq!(enc_ctx.index(), dec_ctx.index());
    assert_eq!(enc_ctx.mps(), dec_ctx.mps());
}

#[test]
fn checked_payload_limit_is_sticky_and_never_underflows_renormalization() {
    let mut encoder =
        ArithmeticEncoder::try_with_byte_limit(0).expect("sentinel allocation succeeds");
    let mut context = ArithmeticEncoderContext::default();
    for index in 0..10_000 {
        encoder.encode(index & 1, &mut context);
    }
    assert_eq!(
        encoder
            .finish_checked()
            .expect_err("zero-byte payload plan is exceeded"),
        crate::EncodeError::InternalInvariant {
            what: "MQ encoder exceeded its checked payload plan",
        }
    );
}

#[test]
fn finish_omits_pending_terminal_ff() {
    // T.800 B.10.7: the bytes included for a code-block must not end in 0xFF.
    let mut encoder = ArithmeticEncoder::new();
    let mut context = ArithmeticEncoderContext::default();
    encoder.encode(0, &mut context);
    assert_eq!(encoder.finish(), [0x7f]);
}

#[test]
#[cfg_attr(
    test,
    expect(clippy::similar_names, reason = "paired encoder/decoder state")
)]
fn short_sequences_never_end_in_ff_and_round_trip() {
    for len in 1..=8_u32 {
        for pattern in 0..(1_u32 << len) {
            let symbols: Vec<u32> = (0..len).map(|bit| (pattern >> bit) & 1).collect();
            let mut encoder = ArithmeticEncoder::new();
            let mut enc_ctx = ArithmeticEncoderContext::default();
            for &symbol in &symbols {
                encoder.encode(symbol, &mut enc_ctx);
            }
            let encoded = encoder.finish();
            assert!(!encoded.is_empty(), "len {len} pattern {pattern:#b}");
            assert_ne!(
                encoded.last(),
                Some(&0xFF),
                "len {len} pattern {pattern:#b}"
            );

            let mut decoder = ArithmeticDecoder::new(&encoded);
            let mut dec_ctx = ArithmeticDecoderContext::default();
            for (index, &symbol) in symbols.iter().enumerate() {
                assert_eq!(
                    decoder.decode(&mut dec_ctx),
                    symbol,
                    "len {len} pattern {pattern:#b} symbol {index}"
                );
            }
        }
    }
}

#[test]
fn encoded_bytes_match_recorded_reference_streams() {
    // Rows: seed, ones per thousand, symbols, then the length and FNV-1a of
    // the bytes from the bit-by-bit Annex C procedure; 19 contexts as in
    // classic Tier-1.
    for (seed, ones, symbols, expected) in [
        (1, 500, 4_000, (523, 0x4140_f6e8_7ab6_78ca)),
        (2, 50, 4_000, (158, 0x7732_5507_3388_6950)),
        (3, 950, 4_000, (171, 0xc394_3501_d9e2_1f73)),
        (4, 5, 20_000, (121, 0x3797_f46b_e3be_012f)),
        (5, 1_000, 1_000, (14, 0x4eae_2264_3f55_d2e9)),
    ] {
        let mut state: u64 = seed;
        let mut encoder = ArithmeticEncoder::new();
        let mut contexts = [ArithmeticEncoderContext::default(); 19];
        for _ in 0..symbols {
            state = state
                .wrapping_mul(6_364_136_223_846_793_005)
                .wrapping_add(1_442_695_040_888_963_407);
            let label = usize::try_from(state >> 59).expect("five bits fit usize") % 19;
            let bit = u32::from((state >> 33) % 1_000 < ones);
            encoder.encode(bit, &mut contexts[label]);
        }
        let bytes = encoder.finish();
        let digest = bytes.iter().fold(0xcbf2_9ce4_8422_2325_u64, |hash, &byte| {
            (hash ^ u64::from(byte)).wrapping_mul(0x100_0000_01b3)
        });
        assert_eq!((bytes.len(), digest), expected, "seed {seed}");
    }
}

#[test]
fn codewords_stay_within_fifteen_bits_per_symbol_plus_the_flush() {
    // A symbol renormalizes by at most 15 bits and a byte carries at least 7,
    // so N symbols give at most 15 N / 7 bytes plus two flush bytes. Patterns:
    // every symbol the MPS, every symbol the LPS, alternating, random at one
    // half and one in twenty, and an adversary that drives a context to the
    // smallest Qe with MPS symbols and then codes its LPS, the costliest
    // symbol there is.
    for pattern in 0..6 {
        for symbols in [1_usize, 2, 7, 15, 100, 10_000, 200_000] {
            let mut encoder = ArithmeticEncoder::new();
            let mut contexts = [ArithmeticEncoderContext::default(); 19];
            let mut state = 0x9e37_79b9_7f4a_7c15_u64;
            for index in 0..symbols {
                state ^= state << 13;
                state ^= state >> 7;
                state ^= state << 17;
                let label = index % 19;
                let context = &mut contexts[label];
                let mps = context.mps();
                let bit = match pattern {
                    0 => mps,
                    1 => 1 - mps,
                    2 => u32::from(index % 2 == 0),
                    3 => u32::from(state >> 63 == 1),
                    4 => u32::from(state.is_multiple_of(20)),
                    _ if context.index() == 45 => 1 - mps,
                    _ => mps,
                };
                encoder.encode(bit, context);
            }
            let bytes = encoder.finish().len();
            assert!(
                bytes <= (15 * symbols).div_ceil(7) + 2,
                "pattern {pattern}, {symbols} symbols: {bytes} bytes"
            );
        }
    }
}