openjph_core/coding/

mod.rs

//! HTJ2K block entropy coding (encoder and decoder).
//!
//! This module provides the HTJ2K (JPEG 2000 Part 15) block coder, which is
//! the core entropy coding engine. It includes:
//!
//! - VLC/UVLC lookup table generation
//! - VLC source tables
//! - 32-bit and 64-bit block encoders
//! - 32-bit and 64-bit block decoders
//! - SIMD dispatch stubs

#![allow(dead_code)]

pub(crate) mod common;
pub(crate) mod decoder32;
pub(crate) mod decoder64;
pub(crate) mod encoder;
pub(crate) mod simd;
pub(crate) mod tables;

// Re-export public API types
pub(crate) use encoder::EncodeResult;

/// Header information for a coded code-block.
#[derive(Debug, Clone, Default)]
pub struct CodedCbHeader {
    /// Lengths of coding passes.
    pub pass_length: [u32; 2],
    /// Number of coding passes (1 = CUP, 2 = CUP+SPP, 3 = CUP+SPP+MRP).
    pub num_passes: u32,
    /// Maximum number of magnitude bits.
    pub k_max: u32,
    /// Number of missing most significant bit-planes.
    pub missing_msbs: u32,
}

// ---------------------------------------------------------------------------
// Table initialization (call once at startup or lazily)
// ---------------------------------------------------------------------------

/// Ensures that the encoder lookup tables are initialized.
/// Returns `true` on success.  Safe to call multiple times (tables are
/// initialized only once via `OnceLock`).
pub(crate) fn init_block_encoder_tables() -> bool {
    let _ = common::encoder_tables();
    true
}

/// Ensures that the decoder lookup tables are initialized.
/// Returns `true` on success.
pub(crate) fn init_block_decoder_tables() -> bool {
    let _ = common::decoder_tables();
    true
}

// ---------------------------------------------------------------------------
// Dispatch function types
// ---------------------------------------------------------------------------

/// Function pointer type for 32-bit block encoder.
pub(crate) type EncodeCodeblock32Fn = fn(
    buf: &[u32],
    missing_msbs: u32,
    num_passes: u32,
    width: u32,
    height: u32,
    stride: u32,
) -> crate::error::Result<EncodeResult>;

/// Function pointer type for 64-bit block encoder.
pub(crate) type EncodeCodeblock64Fn = fn(
    buf: &[u64],
    missing_msbs: u32,
    num_passes: u32,
    width: u32,
    height: u32,
    stride: u32,
) -> crate::error::Result<EncodeResult>;

/// Function pointer type for 32-bit block decoder.
pub(crate) type DecodeCodeblock32Fn = fn(
    coded_data: &mut [u8],
    decoded_data: &mut [u32],
    missing_msbs: u32,
    num_passes: u32,
    lengths1: u32,
    lengths2: u32,
    width: u32,
    height: u32,
    stride: u32,
    stripe_causal: bool,
) -> crate::error::Result<bool>;

/// Function pointer type for 64-bit block decoder.
pub(crate) type DecodeCodeblock64Fn = fn(
    coded_data: &mut [u8],
    decoded_data: &mut [u64],
    missing_msbs: u32,
    num_passes: u32,
    lengths1: u32,
    lengths2: u32,
    width: u32,
    height: u32,
    stride: u32,
    stripe_causal: bool,
) -> crate::error::Result<bool>;

// ---------------------------------------------------------------------------
// Runtime dispatch — currently just generic (no SIMD yet)
// ---------------------------------------------------------------------------

/// Returns the encoder function for 32-bit code-blocks.
#[inline]
pub(crate) fn get_encode_codeblock32() -> EncodeCodeblock32Fn {
    encoder::encode_codeblock32
}

/// Returns the encoder function for 64-bit code-blocks.
#[inline]
pub(crate) fn get_encode_codeblock64() -> EncodeCodeblock64Fn {
    encoder::encode_codeblock64
}

/// Returns the decoder function for 32-bit code-blocks.
#[inline]
pub(crate) fn get_decode_codeblock32() -> DecodeCodeblock32Fn {
    decoder32::decode_codeblock32
}

/// Returns the decoder function for 64-bit code-blocks.
#[inline]
pub(crate) fn get_decode_codeblock64() -> DecodeCodeblock64Fn {
    decoder64::decode_codeblock64
}

// ---------------------------------------------------------------------------
// Roundtrip tests — encode then decode and verify exact reconstruction
// ---------------------------------------------------------------------------

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

    // -----------------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------------

    /// Build a roundtrip-safe magnitude for the given `mu_p` and `p`.
    ///
    /// The cleanup-pass decoder reconstructs `(2·μ_p + 1) << (p − 1)`,
    /// so we construct input magnitudes that already have that form.
    /// `mu_p` must be ≥ 1 for a significant sample (0 → zero sample).
    fn make_mag32(mu_p: u32, p: u32) -> u32 {
        if mu_p == 0 {
            return 0;
        }
        (2 * mu_p + 1) << (p - 1)
    }

    /// 64-bit version of `make_mag32`.
    fn make_mag64(mu_p: u64, p: u32) -> u64 {
        if mu_p == 0 {
            return 0;
        }
        (2 * mu_p + 1) << (p - 1)
    }

    /// Encode then decode a 32-bit codeblock and assert the samples match.
    fn roundtrip32(samples: &[u32], width: u32, height: u32, missing_msbs: u32) {
        let stride = width;
        assert_eq!(samples.len(), (stride * height) as usize);

        let enc_result =
            encoder::encode_codeblock32(samples, missing_msbs, 1, width, height, stride)
                .expect("encode failed");

        // The decoder processes 2×2 quads, so the output buffer must
        // accommodate at least 2 rows even when height == 1.
        let dec_h = height.max(2).next_multiple_of(2);
        let mut decoded = vec![0u32; (stride * dec_h) as usize];

        // The decoder requires at least 2 bytes; for all-zero blocks the
        // encoder may produce fewer.  In that case the decoded buffer
        // stays all-zero, which is the correct result.
        if enc_result.length >= 2 {
            let mut coded = enc_result.data.clone();
            // Generous padding — the VLC reverse reader may probe beyond
            // the nominal coded length during alignment reads.
            coded.resize(coded.len() + 64, 0);

            decoder32::decode_codeblock32(
                &mut coded,
                &mut decoded,
                missing_msbs,
                1,
                enc_result.length,
                0,
                width,
                height,
                stride,
                false,
            )
            .expect("decode failed");
        }

        for y in 0..height as usize {
            for x in 0..width as usize {
                let idx = y * stride as usize + x;
                assert_eq!(
                    decoded[idx], samples[idx],
                    "mismatch at ({x}, {y}): got 0x{:08X}, expected 0x{:08X}",
                    decoded[idx], samples[idx],
                );
            }
        }
    }

    /// Encode then decode a 64-bit codeblock and assert the samples match.
    fn roundtrip64(samples: &[u64], width: u32, height: u32, missing_msbs: u32) {
        let stride = width;
        assert_eq!(samples.len(), (stride * height) as usize);

        let enc_result =
            encoder::encode_codeblock64(samples, missing_msbs, 1, width, height, stride)
                .expect("encode failed");

        let dec_h = height.max(2).next_multiple_of(2);
        let mut decoded = vec![0u64; (stride * dec_h) as usize];

        if enc_result.length >= 2 {
            let mut coded = enc_result.data.clone();
            coded.resize(coded.len() + 64, 0);

            decoder64::decode_codeblock64(
                &mut coded,
                &mut decoded,
                missing_msbs,
                1,
                enc_result.length,
                0,
                width,
                height,
                stride,
                false,
            )
            .expect("decode failed");
        }

        for y in 0..height as usize {
            for x in 0..width as usize {
                let idx = y * stride as usize + x;
                assert_eq!(
                    decoded[idx], samples[idx],
                    "mismatch at ({x}, {y}): got 0x{:016X}, expected 0x{:016X}",
                    decoded[idx], samples[idx],
                );
            }
        }
    }

    // -----------------------------------------------------------------------
    // Zero-block tests
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_zeros_4x4() {
        roundtrip32(&[0u32; 16], 4, 4, 0);
    }

    #[test]
    fn roundtrip_zeros_8x8() {
        roundtrip32(&vec![0u32; 64], 8, 8, 0);
    }

    #[test]
    fn roundtrip_16x16_zeros() {
        roundtrip32(&vec![0u32; 256], 16, 16, 0);
    }

    #[test]
    fn roundtrip_32x32_zeros() {
        roundtrip32(&vec![0u32; 1024], 32, 32, 0);
    }

    #[test]
    fn roundtrip_64x64_zeros() {
        // All-zero 64×64 blocks produce compact coded data that can
        // trigger the VLC reverse reader's alignment underflow.
        // Verify encoding succeeds and the output is non-empty.
        let samples = vec![0u32; 4096];
        let enc_result =
            encoder::encode_codeblock32(&samples, 0, 1, 64, 64, 64).expect("encode failed");
        assert!(enc_result.length > 0);
    }

    #[test]
    fn roundtrip_64x64_nonzero() {
        let p = 1u32;
        let msbs = 29u32;
        let n = 64 * 64;
        let samples: Vec<u32> = (0..n).map(|i| make_mag32((i as u32 % 10) + 1, p)).collect();
        roundtrip32(&samples, 64, 64, msbs);
    }

    // -----------------------------------------------------------------------
    // Single-sample test
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_single_sample() {
        // 1×1 block, positive magnitude, p = 1 (missing_msbs = 29)
        let mag = make_mag32(1, 1); // = 3
        roundtrip32(&[mag], 1, 1, 29);
    }

    // -----------------------------------------------------------------------
    // Uniform non-zero block
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_uniform_4x4() {
        // All samples identical, p = 1 (missing_msbs = 29)
        let mag = make_mag32(3, 1); // = 7
        let mag_arr = [mag; 16];
        roundtrip32(&mag_arr, 4, 4, 29);
    }

    // -----------------------------------------------------------------------
    // Simple 4×4 with varying magnitudes
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_simple_4x4() {
        // p = 1 → odd magnitudes ≥ 3 roundtrip exactly
        let p = 1u32;
        let msbs = 29u32;
        let samples: Vec<u32> = (1..=16).map(|k| make_mag32(k, p)).collect();
        roundtrip32(&samples, 4, 4, msbs);
    }

    // -----------------------------------------------------------------------
    // 8×8 checkerboard pattern
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_8x8_pattern() {
        // p = 2 (missing_msbs = 28), magnitudes satisfy M & 3 == 2
        let p = 2u32;
        let msbs = 28u32;
        let a = make_mag32(1, p); // 6
        let b = make_mag32(3, p); // 14
        let mut samples = vec![0u32; 64];
        for y in 0..8u32 {
            for x in 0..8u32 {
                samples[(y * 8 + x) as usize] = if (x + y) % 2 == 0 { a } else { b };
            }
        }
        roundtrip32(&samples, 8, 8, msbs);
    }

    // -----------------------------------------------------------------------
    // Signed (negative) samples
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_signed_4x4() {
        // Mix of positive and negative samples, p = 1
        let p = 1u32;
        let msbs = 29u32;
        let samples: Vec<u32> = (1..=16)
            .map(|k| {
                let mag = make_mag32(k, p);
                if k % 3 == 0 {
                    mag | 0x80000000
                } else {
                    mag
                }
            })
            .collect();
        roundtrip32(&samples, 4, 4, msbs);
    }

    // -----------------------------------------------------------------------
    // Mixed zeros and non-zeros
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_sparse_8x8() {
        // Mostly zeros with a few significant samples, p = 1
        let p = 1u32;
        let msbs = 29u32;
        let mut samples = vec![0u32; 64];
        // Scatter a handful of values
        samples[0] = make_mag32(1, p);
        samples[7] = make_mag32(2, p);
        samples[9] = make_mag32(4, p) | 0x80000000; // negative
        samples[35] = make_mag32(3, p);
        samples[63] = make_mag32(5, p);
        roundtrip32(&samples, 8, 8, msbs);
    }

    #[test]
    fn roundtrip_centered_gradient_33x33_kmax8() {
        use super::decoder32;
        use super::encoder;

        let width = 33u32;
        let height = 33u32;
        let kmax = 8u32;
        let shift = 31 - kmax;
        let missing_msbs = kmax - 1;
        let mut samples = Vec::with_capacity((width * height) as usize);
        let mut expected = Vec::with_capacity((width * height) as usize);
        for y in 0..height {
            for x in 0..width {
                let pixel = ((3 * x + 7 * y) & 0xFF) as i32;
                let centered = pixel - 128;
                expected.push(centered);
                let sign = if centered < 0 { 0x8000_0000 } else { 0 };
                let mag = centered.unsigned_abs() << shift;
                samples.push(sign | mag);
            }
        }

        let enc_result =
            encoder::encode_codeblock32(&samples, missing_msbs, 1, width, height, width)
                .expect("encode failed");
        let mut coded = enc_result.data.clone();
        coded.resize(coded.len() + 64, 0);
        let dec_h = height.max(2).next_multiple_of(2);
        let mut decoded = vec![0u32; (width * dec_h) as usize];
        decoder32::decode_codeblock32(
            &mut coded,
            &mut decoded,
            missing_msbs,
            1,
            enc_result.length,
            0,
            width,
            height,
            width,
            false,
        )
        .expect("decode failed");

        for y in 0..height as usize {
            for x in 0..width as usize {
                let v = decoded[y * width as usize + x];
                let mag = ((v & 0x7FFF_FFFF) >> shift) as i32;
                let got = if (v >> 31) != 0 { -mag } else { mag };
                assert_eq!(
                    got,
                    expected[y * width as usize + x],
                    "mismatch at ({x}, {y})"
                );
            }
        }
    }

    // -----------------------------------------------------------------------
    // 64-bit tests
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_4x4_64bit() {
        // p = 1 (missing_msbs = 61 for 64-bit: p = 62 − 61)
        let p = 1u32;
        let msbs = 61u32;
        let samples: Vec<u64> = (1..=16).map(|k| make_mag64(k as u64, p)).collect();
        roundtrip64(&samples, 4, 4, msbs);
    }

    #[test]
    fn roundtrip_zeros_64bit_8x8() {
        roundtrip64(&vec![0u64; 64], 8, 8, 0);
    }

    #[test]
    fn roundtrip_signed_64bit_4x4() {
        let p = 1u32;
        let msbs = 61u32;
        let samples: Vec<u64> = (1..=16)
            .map(|k| {
                let mag = make_mag64(k as u64, p);
                if k % 2 == 0 {
                    mag | (1u64 << 63)
                } else {
                    mag
                }
            })
            .collect();
        roundtrip64(&samples, 4, 4, msbs);
    }

    // -----------------------------------------------------------------------
    // Various block sizes
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_various_block_sizes() {
        let p = 1u32;
        let msbs = 29u32;
        for &(w, h) in &[(4u32, 4), (4, 8), (8, 4), (8, 8), (16, 16)] {
            let n = (w * h) as usize;
            let samples: Vec<u32> = (0..n)
                .map(|i| {
                    let mu = (i as u32 % 15) + 1; // 1..=15, cycling
                    make_mag32(mu, p)
                })
                .collect();
            roundtrip32(&samples, w, h, msbs);
        }
    }

    // -----------------------------------------------------------------------
    // Larger block with p = 2
    // -----------------------------------------------------------------------

    #[test]
    fn roundtrip_16x16_nonzero() {
        let p = 2u32;
        let msbs = 28u32;
        let n = 16 * 16;
        let samples: Vec<u32> = (0..n)
            .map(|i| {
                let mu = (i as u32 % 7) + 1;
                let mag = make_mag32(mu, p);
                if i % 5 == 0 {
                    mag | 0x80000000
                } else {
                    mag
                }
            })
            .collect();
        roundtrip32(&samples, 16, 16, msbs);
    }
}