oxihuman-export 0.2.1

// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
// SPDX-License-Identifier: Apache-2.0
#![allow(dead_code)]

//! AV1 default scan orders for transform coefficient coding.
//!
//! Scan orders define the traversal sequence through a square transform block's
//! coefficient grid.  AV1 uses diagonal-stripe (zig-zag) scans for most block sizes.

/// The 4×4 default scan order from the AV1 specification.
/// Coefficients are traversed along anti-diagonals, top-left first.
pub const DEFAULT_SCAN_4X4: [u16; 16] =
    [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];

/// Default 8×8 diagonal scan.
/// Generated by `zigzag_scan(8)` at compile time.
pub const DEFAULT_SCAN_8X8: [u16; 64] = generate_zigzag_8x8();

/// Default 16×16 diagonal scan.
/// Generated by `zigzag_scan(16)` at compile time.
pub const DEFAULT_SCAN_16X16: [u16; 256] = generate_zigzag_16x16();

/// Return the default scan order for a given `log2(transform_size)`.
///
/// - `log2_size == 2` → 4×4  (16 coefficients)
/// - `log2_size == 3` → 8×8  (64 coefficients)
/// - `log2_size == 4` → 16×16 (256 coefficients)
///
/// Panics if `log2_size` is not in 2..=4.
pub fn default_scan(log2_size: usize) -> &'static [u16] {
    match log2_size {
        2 => &DEFAULT_SCAN_4X4,
        3 => &DEFAULT_SCAN_8X8,
        4 => &DEFAULT_SCAN_16X16,
        _ => panic!("default_scan: unsupported log2_size {log2_size} (must be 2, 3, or 4)"),
    }
}

/// Generate a diagonal zig-zag scan for a square block of side `size`.
///
/// `size` must be a power of two in [4, 8, 16, 32].
/// The returned `Vec<u16>` has length `size * size`, listing linear indices
/// `(row * size + col)` in the AV1 default zig-zag anti-diagonal order.
///
/// The traversal alternates direction on successive anti-diagonals to form the
/// classic zig-zag pattern (identical to the JPEG / AV1 spec default scan):
/// - Even diagonals: row decreases (bottom-left → top-right).
/// - Odd  diagonals: row increases (top-right → bottom-left).
pub fn zigzag_scan(size: usize) -> Vec<u16> {
    debug_assert!(
        (4..=32).contains(&size) && size.is_power_of_two(),
        "zigzag_scan: size must be 4, 8, 16, or 32; got {size}"
    );
    let n = size * size;
    let mut out = Vec::with_capacity(n);

    for d in 0..(2 * size - 1) {
        let row_start = if d < size { 0 } else { d + 1 - size };
        let row_end   = if d < size { d } else { size - 1 };

        if d.is_multiple_of(2) {
            // Even diagonal: row decreases (bottom-left → top-right).
            let mut row = row_end;
            loop {
                let col = d - row;
                out.push((row * size + col) as u16);
                if row == row_start {
                    break;
                }
                row -= 1;
            }
        } else {
            // Odd diagonal: row increases (top-right → bottom-left).
            for row in row_start..=row_end {
                let col = d - row;
                out.push((row * size + col) as u16);
            }
        }
    }

    debug_assert_eq!(out.len(), n);
    out
}

// ---- Const-friendly generation helpers ----
// We need const fns to initialise DEFAULT_SCAN_8X8 / DEFAULT_SCAN_16X16 at compile time.
// Rust stable const fns do not support loops over range iterators yet, so we use
// an iterative while-loop form.

const fn generate_zigzag_8x8() -> [u16; 64] {
    let size: usize = 8;
    let mut out = [0u16; 64];
    let mut idx: usize = 0;
    let mut d: usize = 0;
    while d < 2 * size - 1 {
        let row_start: usize = if d < size { 0 } else { d + 1 - size };
        let row_end: usize   = if d < size { d } else { size - 1 };
        if d.is_multiple_of(2) {
            // Even: row decreasing.
            let mut row = row_end;
            loop {
                let col = d - row;
                out[idx] = (row * size + col) as u16;
                idx += 1;
                if row == row_start { break; }
                row -= 1;
            }
        } else {
            // Odd: row increasing.
            let mut row = row_start;
            while row <= row_end {
                let col = d - row;
                out[idx] = (row * size + col) as u16;
                idx += 1;
                row += 1;
            }
        }
        d += 1;
    }
    out
}

const fn generate_zigzag_16x16() -> [u16; 256] {
    let size: usize = 16;
    let mut out = [0u16; 256];
    let mut idx: usize = 0;
    let mut d: usize = 0;
    while d < 2 * size - 1 {
        let row_start: usize = if d < size { 0 } else { d + 1 - size };
        let row_end: usize   = if d < size { d } else { size - 1 };
        if d.is_multiple_of(2) {
            let mut row = row_end;
            loop {
                let col = d - row;
                out[idx] = (row * size + col) as u16;
                idx += 1;
                if row == row_start { break; }
                row -= 1;
            }
        } else {
            let mut row = row_start;
            while row <= row_end {
                let col = d - row;
                out[idx] = (row * size + col) as u16;
                idx += 1;
                row += 1;
            }
        }
        d += 1;
    }
    out
}

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

    fn validate_scan(scan: &[u16], size: usize) {
        let n = size * size;
        assert_eq!(scan.len(), n, "scan length should be {n} for size={size}");

        // Every index 0..n must appear exactly once.
        let mut seen = vec![false; n];
        for &idx in scan {
            let i = idx as usize;
            assert!(i < n, "index {i} out of range for size={size}");
            assert!(!seen[i], "duplicate index {i} in scan for size={size}");
            seen[i] = true;
        }
        for (i, &s) in seen.iter().enumerate() {
            assert!(s, "index {i} missing from scan for size={size}");
        }
    }

    #[test]
    fn test_default_scan_4x4_spec_values() {
        assert_eq!(
            DEFAULT_SCAN_4X4,
            [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15]
        );
        validate_scan(&DEFAULT_SCAN_4X4, 4);
    }

    #[test]
    fn test_default_scan_8x8_valid() {
        validate_scan(&DEFAULT_SCAN_8X8, 8);
    }

    #[test]
    fn test_default_scan_16x16_valid() {
        validate_scan(&DEFAULT_SCAN_16X16, 16);
    }

    #[test]
    fn test_zigzag_scan_4x4_matches_spec() {
        let scan = zigzag_scan(4);
        assert_eq!(scan.as_slice(), &DEFAULT_SCAN_4X4[..]);
    }

    #[test]
    fn test_zigzag_scan_8x8_valid() {
        let scan = zigzag_scan(8);
        validate_scan(&scan, 8);
        // Must match the const table.
        assert_eq!(scan.as_slice(), &DEFAULT_SCAN_8X8[..]);
    }

    #[test]
    fn test_zigzag_scan_16x16_valid() {
        let scan = zigzag_scan(16);
        validate_scan(&scan, 16);
        assert_eq!(scan.as_slice(), &DEFAULT_SCAN_16X16[..]);
    }

    #[test]
    fn test_zigzag_scan_32x32_valid() {
        let scan = zigzag_scan(32);
        validate_scan(&scan, 32);
    }

    #[test]
    fn test_default_scan_dispatch() {
        assert_eq!(default_scan(2).len(), 16);
        assert_eq!(default_scan(3).len(), 64);
        assert_eq!(default_scan(4).len(), 256);
    }

    #[test]
    #[should_panic]
    fn test_default_scan_invalid_log2() {
        let _ = default_scan(5);
    }

    #[test]
    fn test_zigzag_first_and_last() {
        // DC coefficient (0,0) is always first.
        for size in [4usize, 8, 16, 32] {
            let scan = zigzag_scan(size);
            assert_eq!(scan[0], 0, "DC not first for size={size}");
            // Last entry is the bottom-right corner: (size-1)*size + (size-1) = size^2 - 1.
            assert_eq!(
                *scan.last().unwrap(),
                (size * size - 1) as u16,
                "last coefficient wrong for size={size}"
            );
        }
    }
}