rximg/src/quantization.rs

use crate::{constants::TILE_SZ2, quality_settings::{self, QualitySettings}};

const ZIGZAG: [u8; TILE_SZ2] = [
    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,
];

pub fn to_quantized(
    quality_settings: QualitySettings,
    coefs: [i16; TILE_SZ2]
) -> [i16; TILE_SZ2] {
    let mut quant: [i16; TILE_SZ2] = [0; TILE_SZ2];

    for cf_ix in 0..TILE_SZ2 {
        let div = divisor(quality_settings, cf_ix);
        let qval = (coefs[cf_ix] + div / 2) / div;
        quant[ZIGZAG[cf_ix] as usize] = qval
    }

    let mut indices = [0; TILE_SZ2];
    for i in 0..indices.len() {
        indices[i] = i;
    }
    indices.sort_by_key(|i| -quant[*i].abs());
    for i in (quality_settings.values[quality_settings::QUANTIZATION_N_QUANTS_TO_KEEP] as usize)..indices.len() {
        quant[indices[i]] = 0;
    }

    quant
}

pub fn from_quantized(
    quality_settings: QualitySettings,
    quant: [i16; TILE_SZ2]
) -> [i16; TILE_SZ2] {
    let mut coefs: [i16; TILE_SZ2] = [0; TILE_SZ2];
    for cf_ix in 0..TILE_SZ2 {
        let div = divisor(quality_settings, cf_ix);
        coefs[cf_ix] = quant[ZIGZAG[cf_ix] as usize].wrapping_mul(div);
    }
    coefs
}

pub fn divisor(
    quality_settings: QualitySettings,
    cf_ix: usize
) -> i16 {
    if cf_ix == 0 { return quality_settings.values[quality_settings::QUANTIZATION_QUANT0_DIVISOR] as i16; }
    let x = cf_ix % 8;
    let y = cf_ix / 8;
    let div = 
        (quality_settings.values[quality_settings::QUANTIZATION_CONSTANT] as i16).wrapping_add(
            (x as i16 + y as i16).wrapping_mul(
                quality_settings.values[quality_settings::QUANTIZATION_LINEAR] as i16
            )
        );
    if div==32 && cf_ix != 0 {
        dbg!(cf_ix, x, y, div);
    }
    return div;
}