use crate::foundation::consts::{
BASE_QUANT_MATRIX_XYB, BASE_QUANT_MATRIX_YCBCR, GLOBAL_SCALE_XYB, GLOBAL_SCALE_YCBCR,
};
use crate::quant::{
FREQUENCY_EXPONENT, ZERO_BIAS_MUL_XYB, ZERO_BIAS_MUL_YCBCR_HQ, ZERO_BIAS_MUL_YCBCR_LQ,
ZERO_BIAS_OFFSET_XYB, ZERO_BIAS_OFFSET_YCBCR_AC, ZERO_BIAS_OFFSET_YCBCR_DC,
};
#[derive(Clone, Debug, PartialEq)]
pub struct PerComponent<T> {
pub c0: T,
pub c1: T,
pub c2: T,
}
pub type YCbCrComponents<T> = PerComponent<T>;
pub type XybComponents<T> = PerComponent<T>;
impl<T> PerComponent<T> {
pub fn new(c0: T, c1: T, c2: T) -> Self {
Self { c0, c1, c2 }
}
pub fn map<U, F: Fn(&T) -> U>(&self, f: F) -> PerComponent<U> {
PerComponent {
c0: f(&self.c0),
c1: f(&self.c1),
c2: f(&self.c2),
}
}
pub fn get(&self, index: usize) -> &T {
match index {
0 => &self.c0,
1 => &self.c1,
_ => &self.c2,
}
}
pub fn get_mut(&mut self, index: usize) -> &mut T {
match index {
0 => &mut self.c0,
1 => &mut self.c1,
_ => &mut self.c2,
}
}
}
impl<T> YCbCrComponents<T> {
pub fn y(&self) -> &T {
&self.c0
}
pub fn cb(&self) -> &T {
&self.c1
}
pub fn cr(&self) -> &T {
&self.c2
}
pub fn y_mut(&mut self) -> &mut T {
&mut self.c0
}
pub fn cb_mut(&mut self) -> &mut T {
&mut self.c1
}
pub fn cr_mut(&mut self) -> &mut T {
&mut self.c2
}
}
impl<T> XybComponents<T> {
pub fn x(&self) -> &T {
&self.c0
}
pub fn luma(&self) -> &T {
&self.c1
}
pub fn b(&self) -> &T {
&self.c2
}
pub fn x_mut(&mut self) -> &mut T {
&mut self.c0
}
pub fn luma_mut(&mut self) -> &mut T {
&mut self.c1
}
pub fn b_mut(&mut self) -> &mut T {
&mut self.c2
}
}
impl<T: Copy> PerComponent<[T; 64]> {
pub fn coeff(&self, component: usize, k: usize) -> T {
self.get(component)[k]
}
pub fn set_coeff(&mut self, component: usize, k: usize, value: T) {
self.get_mut(component)[k] = value;
}
}
impl PerComponent<[f32; 64]> {
pub fn scale_coeff(&mut self, component: usize, k: usize, factor: f32) {
self.get_mut(component)[k] *= factor;
}
pub fn scale_component(&mut self, component: usize, factor: f32) {
for v in self.get_mut(component).iter_mut() {
*v *= factor;
}
}
pub fn scale_all(&mut self, factor: f32) {
for c in 0..3 {
self.scale_component(c, factor);
}
}
pub fn blend(&self, other: &Self, t: f32) -> Self {
let blend_arr = |a: &[f32; 64], b: &[f32; 64]| {
let mut result = [0.0f32; 64];
for i in 0..64 {
result[i] = a[i] * (1.0 - t) + b[i] * t;
}
result
};
Self {
c0: blend_arr(&self.c0, &other.c0),
c1: blend_arr(&self.c1, &other.c1),
c2: blend_arr(&self.c2, &other.c2),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum ScalingParams {
Exact,
Scaled {
global_scale: f32,
frequency_exponents: Box<[f32; 64]>,
},
}
impl ScalingParams {
#[must_use]
pub fn default_ycbcr() -> Self {
Self::Scaled {
global_scale: GLOBAL_SCALE_YCBCR,
frequency_exponents: Box::new(FREQUENCY_EXPONENT),
}
}
#[must_use]
pub fn default_xyb() -> Self {
Self::Scaled {
global_scale: GLOBAL_SCALE_XYB,
frequency_exponents: Box::new(FREQUENCY_EXPONENT),
}
}
#[must_use]
pub fn exact() -> Self {
Self::Exact
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct EncodingTables {
pub quant: PerComponent<[f32; 64]>,
pub zero_bias_mul: PerComponent<[f32; 64]>,
pub zero_bias_offset_dc: [f32; 3],
pub zero_bias_offset_ac: [f32; 3],
pub scaling: ScalingParams,
}
impl EncodingTables {
#[must_use]
pub fn default_ycbcr() -> Self {
let quant = PerComponent {
c0: std::array::from_fn(|i| BASE_QUANT_MATRIX_YCBCR[i]),
c1: std::array::from_fn(|i| BASE_QUANT_MATRIX_YCBCR[64 + i]),
c2: std::array::from_fn(|i| BASE_QUANT_MATRIX_YCBCR[128 + i]),
};
let zero_bias_mul = PerComponent {
c0: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[i]),
c1: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[64 + i]),
c2: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[128 + i]),
};
Self {
quant,
zero_bias_mul,
zero_bias_offset_dc: ZERO_BIAS_OFFSET_YCBCR_DC,
zero_bias_offset_ac: ZERO_BIAS_OFFSET_YCBCR_AC,
scaling: ScalingParams::default_ycbcr(),
}
}
#[must_use]
pub fn default_xyb() -> Self {
let quant = PerComponent {
c0: std::array::from_fn(|i| BASE_QUANT_MATRIX_XYB[i]),
c1: std::array::from_fn(|i| BASE_QUANT_MATRIX_XYB[64 + i]),
c2: std::array::from_fn(|i| BASE_QUANT_MATRIX_XYB[128 + i]),
};
let make_xyb_mul = || {
let mut arr = [ZERO_BIAS_MUL_XYB; 64];
arr[0] = 0.0; arr
};
let zero_bias_mul = PerComponent {
c0: make_xyb_mul(),
c1: make_xyb_mul(),
c2: make_xyb_mul(),
};
Self {
quant,
zero_bias_mul,
zero_bias_offset_dc: [0.0, 0.0, 0.0],
zero_bias_offset_ac: [
ZERO_BIAS_OFFSET_XYB,
ZERO_BIAS_OFFSET_XYB,
ZERO_BIAS_OFFSET_XYB,
],
scaling: ScalingParams::default_xyb(),
}
}
#[must_use]
pub fn ycbcr_hq_zero_bias_mul() -> PerComponent<[f32; 64]> {
PerComponent {
c0: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_HQ[i]),
c1: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_HQ[64 + i]),
c2: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_HQ[128 + i]),
}
}
#[must_use]
pub fn ycbcr_lq_zero_bias_mul() -> PerComponent<[f32; 64]> {
PerComponent {
c0: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[i]),
c1: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[64 + i]),
c2: std::array::from_fn(|i| ZERO_BIAS_MUL_YCBCR_LQ[128 + i]),
}
}
pub fn blend_zero_bias_mul(&mut self, hq: &PerComponent<[f32; 64]>, t: f32) {
let lq = Self::ycbcr_lq_zero_bias_mul();
self.zero_bias_mul = lq.blend(hq, t);
}
pub fn scale_quant(&mut self, component: usize, k: usize, factor: f32) {
self.quant.scale_coeff(component, k, factor);
}
pub fn scale_mul(&mut self, component: usize, k: usize, factor: f32) {
self.zero_bias_mul.scale_coeff(component, k, factor);
}
pub fn perturb_quant(&mut self, component: usize, k: usize, delta: f32) {
self.quant.get_mut(component)[k] += delta;
}
pub fn perturb_mul(&mut self, component: usize, k: usize, delta: f32) {
self.zero_bias_mul.get_mut(component)[k] += delta;
}
#[must_use]
pub fn blend(&self, other: &Self, t: f32) -> Self {
let blend_3 = |a: &[f32; 3], b: &[f32; 3]| {
[
a[0] * (1.0 - t) + b[0] * t,
a[1] * (1.0 - t) + b[1] * t,
a[2] * (1.0 - t) + b[2] * t,
]
};
Self {
quant: self.quant.blend(&other.quant, t),
zero_bias_mul: self.zero_bias_mul.blend(&other.zero_bias_mul, t),
zero_bias_offset_dc: blend_3(&self.zero_bias_offset_dc, &other.zero_bias_offset_dc),
zero_bias_offset_ac: blend_3(&self.zero_bias_offset_ac, &other.zero_bias_offset_ac),
scaling: self.scaling.clone(),
}
}
#[must_use]
pub fn generate_quant_table(
&self,
component: usize,
distance: f32,
is_420: bool,
) -> crate::quant::QuantTable {
use crate::foundation::consts::{GLOBAL_SCALE_420, K420_RESCALE};
use crate::quant::{create_quant_table, DIST_THRESHOLD};
let c = component.min(2);
let base = self.quant.get(c);
match &self.scaling {
ScalingParams::Exact => {
let mut values = [0u16; 64];
for i in 0..64 {
values[i] = base[i].round().clamp(1.0, 65535.0) as u16;
}
create_quant_table(values, true)
}
ScalingParams::Scaled {
global_scale,
frequency_exponents,
} => {
let mut values = [0u16; 64];
let mut scale_factor = *global_scale;
if is_420 {
scale_factor *= GLOBAL_SCALE_420;
}
let is_chroma_420 = is_420 && component > 0;
for i in 0..64 {
let freq_scale = if distance < DIST_THRESHOLD {
distance
} else {
let exp = frequency_exponents[i];
let mul = DIST_THRESHOLD.powf(1.0 - exp);
(0.5 * distance).max(mul * distance.powf(exp))
};
let mut scale = freq_scale * scale_factor;
if is_chroma_420 {
scale *= K420_RESCALE[i];
}
let q = (base[i] * scale).round();
values[i] = q as u16;
}
create_quant_table(values, true)
}
}
}
#[must_use]
pub fn generate_quant_tables(
&self,
distance: f32,
is_420: bool,
) -> (
crate::quant::QuantTable,
crate::quant::QuantTable,
crate::quant::QuantTable,
) {
(
self.generate_quant_table(0, distance, is_420),
self.generate_quant_table(1, distance, is_420),
self.generate_quant_table(2, distance, is_420),
)
}
#[must_use]
pub fn generate_zero_bias_params(&self, component: usize) -> crate::quant::ZeroBiasParams {
let c = component.min(2);
let mul = *self.zero_bias_mul.get(c);
let mut offset = [0.0f32; 64];
offset[0] = self.zero_bias_offset_dc[c];
for i in 1..64 {
offset[i] = self.zero_bias_offset_ac[c];
}
crate::quant::ZeroBiasParams { mul, offset }
}
#[must_use]
pub fn generate_zero_bias_all(
&self,
) -> (
crate::quant::ZeroBiasParams,
crate::quant::ZeroBiasParams,
crate::quant::ZeroBiasParams,
) {
(
self.generate_zero_bias_params(0),
self.generate_zero_bias_params(1),
self.generate_zero_bias_params(2),
)
}
#[must_use]
pub fn is_exact(&self) -> bool {
matches!(self.scaling, ScalingParams::Exact)
}
}
pub mod dct {
#[inline]
#[must_use]
pub const fn freq_distance(k: usize) -> usize {
let row = k / 8;
let col = k % 8;
row + col
}
#[inline]
#[must_use]
pub const fn row_col(k: usize) -> (usize, usize) {
(k / 8, k % 8)
}
#[must_use]
pub const fn to_zigzag(k: usize) -> usize {
ZIGZAG_ORDER[k]
}
#[must_use]
pub const fn from_zigzag(z: usize) -> usize {
INVERSE_ZIGZAG[z]
}
pub const IMPORTANCE_ORDER: [usize; 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, ];
const ZIGZAG_ORDER: [usize; 64] = [
0, 1, 5, 6, 14, 15, 27, 28, 2, 4, 7, 13, 16, 26, 29, 42, 3, 8, 12, 17, 25, 30, 41, 43, 9,
11, 18, 24, 31, 40, 44, 53, 10, 19, 23, 32, 39, 45, 52, 54, 20, 22, 33, 38, 46, 51, 55, 60,
21, 34, 37, 47, 50, 56, 59, 61, 35, 36, 48, 49, 57, 58, 62, 63,
];
const INVERSE_ZIGZAG: [usize; 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,
];
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_freq_distance() {
assert_eq!(freq_distance(0), 0); assert_eq!(freq_distance(1), 1); assert_eq!(freq_distance(8), 1); assert_eq!(freq_distance(63), 14); }
#[test]
fn test_zigzag_roundtrip() {
for k in 0..64 {
let z = to_zigzag(k);
let back = from_zigzag(z);
assert_eq!(back, k, "Roundtrip failed for k={}", k);
}
}
#[test]
fn test_importance_order_complete() {
let mut seen = [false; 64];
for &k in &IMPORTANCE_ORDER {
assert!(!seen[k], "Duplicate index {} in IMPORTANCE_ORDER", k);
seen[k] = true;
}
assert!(
seen.iter().all(|&s| s),
"Missing indices in IMPORTANCE_ORDER"
);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_ycbcr_tables() {
let tables = EncodingTables::default_ycbcr();
assert!(tables.quant.c0[0] > 0.0, "DC quant should be positive");
assert!(tables.quant.c1[0] > 0.0);
assert!(tables.quant.c2[0] > 0.0);
assert_eq!(tables.zero_bias_offset_dc, [0.0, 0.0, 0.0]);
assert!(tables.zero_bias_offset_ac[0] > 0.0);
}
#[test]
fn test_default_xyb_tables() {
let tables = EncodingTables::default_xyb();
assert_eq!(tables.zero_bias_offset_ac, [0.5, 0.5, 0.5]);
assert_eq!(tables.zero_bias_offset_dc, [0.0, 0.0, 0.0]);
assert_eq!(tables.zero_bias_mul.c0[0], 0.0);
assert_eq!(tables.zero_bias_mul.c1[0], 0.0);
assert_eq!(tables.zero_bias_mul.c2[0], 0.0);
assert_eq!(tables.zero_bias_mul.c0[1], 0.5);
}
#[test]
fn test_blend() {
let a = EncodingTables::default_ycbcr();
let b = EncodingTables::default_xyb();
let mid = a.blend(&b, 0.5);
let a_dc = a.quant.c0[0];
let b_dc = b.quant.c0[0];
let mid_dc = mid.quant.c0[0];
assert!(
(mid_dc - (a_dc + b_dc) / 2.0).abs() < 0.01,
"Blend should be midpoint"
);
}
#[test]
fn test_scale_quant() {
let mut tables = EncodingTables::default_ycbcr();
let original = tables.quant.c0[5];
tables.scale_quant(0, 5, 2.0);
assert!((tables.quant.c0[5] - original * 2.0).abs() < 0.001);
}
#[test]
fn test_ycbcr_accessors() {
let tables = EncodingTables::default_ycbcr();
let y = tables.quant.y();
let cb = tables.quant.cb();
let cr = tables.quant.cr();
assert_eq!(y, &tables.quant.c0);
assert_eq!(cb, &tables.quant.c1);
assert_eq!(cr, &tables.quant.c2);
}
#[test]
fn test_xyb_accessors() {
let tables = EncodingTables::default_xyb();
let x = tables.quant.x();
let luma = tables.quant.luma();
let b = tables.quant.b();
assert_eq!(x, &tables.quant.c0);
assert_eq!(luma, &tables.quant.c1);
assert_eq!(b, &tables.quant.c2);
}
}