#![allow(dead_code)]
use jxl_grid::{MutableSubgrid, SharedSubgrid};
use jxl_modular::ChannelShift;
use jxl_vardct::{BlockInfo, TransformType};
use crate::vardct::{
dct_common::DctDirection,
transform_common::{AFV_BASIS, transform_varblocks_inner},
};
use super::dct_2d;
#[inline(always)]
pub(crate) fn aux_idct2_in_place_2(block: &mut MutableSubgrid<'_>) {
let c00 = block.get(0, 0);
let c01 = block.get(1, 0);
let c10 = block.get(0, 1);
let c11 = block.get(1, 1);
*block.get_mut(0, 0) = c00 + c01 + c10 + c11;
*block.get_mut(1, 0) = c00 + c01 - c10 - c11;
*block.get_mut(0, 1) = c00 - c01 + c10 - c11;
*block.get_mut(1, 1) = c00 - c01 - c10 + c11;
}
#[inline(always)]
pub(crate) fn aux_idct2_in_place<const SIZE: usize>(block: &mut MutableSubgrid<'_>) {
debug_assert!(SIZE.is_power_of_two());
let num_2x2 = SIZE / 2;
let mut scratch = [[0.0f32; SIZE]; SIZE];
for y in 0..num_2x2 {
for x in 0..num_2x2 {
let c00 = block.get(x, y);
let c01 = block.get(x + num_2x2, y);
let c10 = block.get(x, y + num_2x2);
let c11 = block.get(x + num_2x2, y + num_2x2);
scratch[2 * y][2 * x] = c00 + c01 + c10 + c11;
scratch[2 * y][2 * x + 1] = c00 + c01 - c10 - c11;
scratch[2 * y + 1][2 * x] = c00 - c01 + c10 - c11;
scratch[2 * y + 1][2 * x + 1] = c00 - c01 - c10 + c11;
}
}
for (y, scratch_row) in scratch.into_iter().enumerate() {
block.get_row_mut(y)[..SIZE].copy_from_slice(&scratch_row);
}
}
pub(crate) fn transform_dct2(coeff: &mut MutableSubgrid<'_>) {
aux_idct2_in_place::<2>(coeff);
aux_idct2_in_place::<4>(coeff);
aux_idct2_in_place::<8>(coeff);
}
pub(crate) fn transform_dct4(coeff: &mut MutableSubgrid<'_>) {
aux_idct2_in_place::<2>(coeff);
let mut scratch = [0.0f32; 64];
for y in 0..2 {
for x in 0..2 {
let mut scratch = MutableSubgrid::from_buf(&mut scratch[(y * 2 + x) * 16..], 4, 4, 4);
for iy in 0..4 {
for ix in 0..4 {
*scratch.get_mut(iy, ix) = coeff.get(x + ix * 2, y + iy * 2);
}
}
dct_2d(&mut scratch, DctDirection::Inverse);
}
}
for y in 0..2 {
for x in 0..2 {
let scratch = &scratch[(y * 2 + x) * 16..][..16];
for iy in 0..4 {
for ix in 0..4 {
*coeff.get_mut(x * 4 + ix, y * 4 + iy) = scratch[iy * 4 + ix];
}
}
}
}
}
pub(crate) fn transform_hornuss(coeff: &mut MutableSubgrid<'_>) {
aux_idct2_in_place::<2>(coeff);
let mut scratch = [0.0f32; 64];
for y in 0..2 {
for x in 0..2 {
let scratch = &mut scratch[(y * 2 + x) * 16..][..16];
for iy in 0..4 {
for ix in 0..4 {
scratch[iy * 4 + ix] = coeff.get(x + ix * 2, y + iy * 2);
}
}
let residual_sum: f32 = scratch[1..].iter().copied().sum();
let avg = scratch[0] - residual_sum / 16.0;
scratch[0] = scratch[5];
scratch[5] = 0.0;
for s in scratch.iter_mut() {
*s += avg;
}
}
}
for y in 0..2 {
for x in 0..2 {
let scratch = &mut scratch[(y * 2 + x) * 16..][..16];
for iy in 0..4 {
for ix in 0..4 {
*coeff.get_mut(x * 4 + ix, y * 4 + iy) = scratch[iy * 4 + ix];
}
}
}
}
}
pub(crate) fn transform_dct4x8<const TR: bool>(coeff: &mut MutableSubgrid<'_>) {
let coeff0 = coeff.get(0, 0);
let coeff1 = coeff.get(0, 1);
*coeff.get_mut(0, 0) = coeff0 + coeff1;
*coeff.get_mut(0, 1) = coeff0 - coeff1;
let mut scratch = [0.0f32; 64];
for idx in [0, 1] {
let mut scratch = MutableSubgrid::from_buf(&mut scratch[(idx * 32)..], 8, 4, 8);
for iy in 0..4 {
for ix in 0..8 {
*scratch.get_mut(ix, iy) = coeff.get(ix, iy * 2 + idx);
}
}
dct_2d(&mut scratch, DctDirection::Inverse);
}
if TR {
for y in 0..8 {
for x in 0..8 {
*coeff.get_mut(y, x) = scratch[y * 8 + x];
}
}
} else {
for y in 0..8 {
coeff.get_row_mut(y)[..8].copy_from_slice(&scratch[y * 8..][..8]);
}
}
}
pub(crate) fn transform_afv<const N: usize>(coeff: &mut MutableSubgrid<'_>) {
assert!(N < 4);
let flip_x = N % 2;
let flip_y = N / 2;
let mut coeff_afv = [0.0f32; 16];
coeff_afv[0] = (coeff.get(0, 0) + coeff.get(1, 0) + coeff.get(0, 1)) * 4.0;
for (idx, v) in coeff_afv.iter_mut().enumerate().skip(1) {
let iy = idx / 4;
let ix = idx % 4;
*v = coeff.get(2 * ix, 2 * iy);
}
let mut samples_afv = [0.0f32; 16];
for (coeff, basis) in coeff_afv.into_iter().zip(AFV_BASIS) {
for (sample, basis) in samples_afv.iter_mut().zip(basis) {
*sample = coeff.mul_add(basis, *sample);
}
}
let mut scratch_4x4 = [0.0f32; 16];
let mut scratch_4x8 = [0.0f32; 32];
scratch_4x4[0] = coeff.get(0, 0) - coeff.get(1, 0) + coeff.get(0, 1);
for iy in 0..4 {
for ix in 0..4 {
if ix | iy == 0 {
continue;
}
scratch_4x4[ix * 4 + iy] = coeff.get(2 * ix + 1, 2 * iy);
}
}
dct_2d(
&mut MutableSubgrid::from_buf(&mut scratch_4x4, 4, 4, 4),
DctDirection::Inverse,
);
scratch_4x8[0] = coeff.get(0, 0) - coeff.get(0, 1);
for iy in 0..4 {
for ix in 0..8 {
if ix | iy == 0 {
continue;
}
scratch_4x8[iy * 8 + ix] = coeff.get(ix, 2 * iy + 1);
}
}
dct_2d(
&mut MutableSubgrid::from_buf(&mut scratch_4x8, 8, 4, 8),
DctDirection::Inverse,
);
for iy in 0..4 {
let afv_y = if flip_y == 0 { iy } else { 3 - iy };
for ix in 0..4 {
let afv_x = if flip_x == 0 { ix } else { 3 - ix };
*coeff.get_mut(flip_x * 4 + ix, flip_y * 4 + iy) = samples_afv[afv_y * 4 + afv_x];
}
}
for iy in 0..4 {
let y = flip_y * 4 + iy;
for ix in 0..4 {
let x = (1 - flip_x) * 4 + ix;
*coeff.get_mut(x, y) = scratch_4x4[iy * 4 + ix];
}
}
for iy in 0..4 {
let y = (1 - flip_y) * 4 + iy;
coeff.get_row_mut(y)[..8].copy_from_slice(&scratch_4x8[iy * 8..][..8]);
}
}
fn transform_dct(coeff: &mut MutableSubgrid<'_>) {
dct_2d(coeff, DctDirection::Inverse);
}
fn transform(coeff: &mut MutableSubgrid<'_>, dct_select: TransformType) {
use TransformType::*;
match dct_select {
Dct2 => transform_dct2(coeff),
Dct4 => transform_dct4(coeff),
Hornuss => transform_hornuss(coeff),
Dct4x8 => transform_dct4x8::<false>(coeff),
Dct8x4 => transform_dct4x8::<true>(coeff),
Afv0 => transform_afv::<0>(coeff),
Afv1 => transform_afv::<1>(coeff),
Afv2 => transform_afv::<2>(coeff),
Afv3 => transform_afv::<3>(coeff),
_ => transform_dct(coeff),
}
}
pub fn transform_varblocks(
lf: &[SharedSubgrid<f32>; 3],
coeff_out: &mut [MutableSubgrid<'_, f32>; 3],
shifts_cbycr: [ChannelShift; 3],
block_info: &SharedSubgrid<BlockInfo>,
) {
unsafe {
transform_varblocks_inner(
lf,
coeff_out,
shifts_cbycr,
block_info,
super::dct::dct_2d,
transform,
);
}
}