Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use crate::{
BLOCK_DIM, BLOCK_SIZE,
bit_reader::BitReader,
entropy_coding::decode::SymbolReader,
error::Result,
frame::Histograms,
headers::permutation::Permutation,
util::{CeilLog2, tracing_wrappers::*},
};
use jxl_transforms::transform_map::*;
use std::borrow::Cow;
use std::mem;
use std::sync::OnceLock;
pub const NUM_ORDERS: usize = 13;
pub const TRANSFORM_TYPE_LUT: [HfTransformType; NUM_ORDERS] = [
HfTransformType::DCT,
HfTransformType::IDENTITY, // a.k.a. "Hornuss"
HfTransformType::DCT16X16,
HfTransformType::DCT32X32,
HfTransformType::DCT8X16,
HfTransformType::DCT8X32,
HfTransformType::DCT16X32,
HfTransformType::DCT64X64,
HfTransformType::DCT32X64,
HfTransformType::DCT128X128,
HfTransformType::DCT64X128,
HfTransformType::DCT256X256,
HfTransformType::DCT128X256,
];
pub const NUM_PERMUTATION_CONTEXTS: usize = 8;
/// Cached natural coefficient orders per transform type.
/// Each entry is computed lazily on first access, avoiding computation
/// of orders for large transforms that are never used.
static NATURAL_COEFF_ORDERS: [OnceLock<Vec<u32>>; NUM_ORDERS] = [
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
OnceLock::new(),
];
/// Get cached natural coefficient order for a transform type index.
/// Computes the order lazily on first access.
fn get_natural_coeff_order(idx: usize) -> &'static Vec<u32> {
NATURAL_COEFF_ORDERS[idx].get_or_init(|| natural_coeff_order(TRANSFORM_TYPE_LUT[idx]))
}
pub fn natural_coeff_order(transform: HfTransformType) -> Vec<u32> {
let cx = covered_blocks_x(transform) as usize;
let cy = covered_blocks_y(transform) as usize;
let xsize: usize = cx * BLOCK_DIM;
assert!(cx >= cy);
// We compute the zigzag order for a cx x cx block, then discard all the
// lines that are not multiple of the ratio between cx and cy.
let xs = cx / cy;
let xsm = xs - 1;
let xss = xs.ceil_log2();
let mut out: Vec<u32> = vec![0; cx * cy * BLOCK_SIZE];
// First half of the block
let mut cur = cx * cy;
for i in 0..xsize {
for j in 0..(i + 1) {
let mut x = j;
let mut y = i - j;
if i % 2 != 0 {
mem::swap(&mut x, &mut y);
}
if (y & xsm) != 0 {
continue;
}
y >>= xss;
let val;
if x < cx && y < cy {
val = y * cx + x;
} else {
val = cur;
cur += 1;
}
out[val] = (y * xsize + x) as u32;
}
}
// Second half
for ir in 1..xsize {
let ip = xsize - ir;
let i = ip - 1;
for j in 0..(i + 1) {
let mut x = xsize - 1 - (i - j);
let mut y = xsize - 1 - j;
if !i.is_multiple_of(2) {
mem::swap(&mut x, &mut y);
}
if (y & xsm) != 0 {
continue;
}
y >>= xss;
let val = cur;
cur += 1;
out[val] = (y * xsize + x) as u32;
}
}
out
}
pub fn decode_coeff_orders(used_orders: u32, br: &mut BitReader) -> Result<Vec<Permutation>> {
// Use cached natural coefficient orders instead of recomputing
let all_component_orders = 3 * NUM_ORDERS;
let mut permutations: Vec<Permutation> = (0..all_component_orders)
.map(|o| Permutation(Cow::Borrowed(get_natural_coeff_order(o / 3))))
.collect();
if used_orders == 0 {
return Ok(permutations);
}
let histograms = Histograms::decode(NUM_PERMUTATION_CONTEXTS, br, true)?;
let mut reader = SymbolReader::new(&histograms, br, None)?;
for (ord, transform_type) in TRANSFORM_TYPE_LUT.iter().enumerate() {
if used_orders & (1 << ord) == 0 {
continue;
}
debug!(?transform_type);
let num_blocks = covered_blocks_x(*transform_type) * covered_blocks_y(*transform_type);
for c in 0..3 {
let size = num_blocks * BLOCK_SIZE as u32;
let permutation = Permutation::decode(size, num_blocks, &histograms, br, &mut reader)?;
let index = 3 * ord + c;
permutations[index].compose(&permutation);
}
}
reader.check_final_state(&histograms, br)?;
Ok(permutations)
}
#[cfg(test)]
mod tests {
use super::*;
// Golden data generated from libjxl's `ComputeNaturalCoeffOrder` for DCT (8x8).
const COEFF_ORDER_1X1: [u32; 64] = [
0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27,
20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
];
// Golden data generated from libjxl's `ComputeNaturalCoeffOrder` for DCT8X16 (16x8).
const COEFF_ORDER_2X1: [u32; 128] = [
0, 1, 16, 2, 3, 17, 32, 18, 4, 5, 19, 33, 48, 34, 20, 6, 7, 21, 35, 49, 64, 50, 36, 22, 8,
9, 23, 37, 51, 65, 80, 66, 52, 38, 24, 10, 11, 25, 39, 53, 67, 81, 96, 82, 68, 54, 40, 26,
12, 13, 27, 41, 55, 69, 83, 97, 112, 98, 84, 70, 56, 42, 28, 14, 15, 29, 43, 57, 71, 85,
99, 113, 114, 100, 86, 72, 58, 44, 30, 31, 45, 59, 73, 87, 101, 115, 116, 102, 88, 74, 60,
46, 47, 61, 75, 89, 103, 117, 118, 104, 90, 76, 62, 63, 77, 91, 105, 119, 120, 106, 92, 78,
79, 93, 107, 121, 122, 108, 94, 95, 109, 123, 124, 110, 111, 125, 126, 127,
];
#[test]
fn test_natural_coeff_order() {
let order_1x1 = natural_coeff_order(HfTransformType::DCT);
assert_eq!(order_1x1, COEFF_ORDER_1X1);
let order_2x1 = natural_coeff_order(HfTransformType::DCT8X16);
assert_eq!(order_2x1, COEFF_ORDER_2X1);
}
}