Files
effectcraft/crates/opusenc/src/rate.rs
T

562 lines
18 KiB
Rust

//! Pulse cache and bit allocation (RFC 6716 §4.3.3, §4.3.4.1).
//!
//! The allocation must be bit-identical to the decoder's, so this mirrors FilmCraft's
//! spec-derived `filmcraft-opus` decoder (first-party, MIT OR Apache-2.0) step by step; the only
//! difference is that the encoder *writes* the band-skip, intensity and dual-stereo symbols.
use crate::range::{BITRES, RangeEncoder, ilog};
use crate::tables::*;
const MAX_PSEUDO: i32 = 40;
const LOG_MAX_PSEUDO: i32 = 6;
pub const MAX_FINE_BITS: i32 = 8;
const FINE_OFFSET: i32 = 21;
pub const QTHETA_OFFSET: i32 = 4;
pub const QTHETA_OFFSET_TWOPHASE: i32 = 16;
const ALLOC_STEPS: i32 = 6;
/// Number of pulses represented by pseudo-pulse index `i`.
pub fn get_pulses(i: i32) -> i32 {
if i < 8 { i } else { (8 + (i & 7)) << ((i >> 3) - 1) }
}
/// Conservative (rounded up) log2 of `val` with `frac` fractional bits.
pub fn log2_frac(mut val: u32, mut frac: u32) -> i32 {
let mut l = ilog(val);
if val & val.wrapping_sub(1) != 0 {
if l > 16 {
val = ((val - 1) >> (l - 16)) + 1;
} else {
val <<= 16 - l;
}
let mut l = (l - 1) << frac;
loop {
let b = (val >> 16) as i32;
l += b << frac;
val = (val + b as u32) >> b;
val = (val * val + 0x7FFF) >> 15;
if frac == 0 {
break;
}
frac -= 1;
}
l + (val > 0x8000) as i32
} else {
l -= 1;
l << frac
}
}
/// PVQ codebook size V(N, K) (saturating far above 2^32).
pub fn pvq_v(n: usize, k: usize) -> u128 {
if k == 0 {
return 1;
}
if n == 0 {
return 0;
}
let mut row = vec![0u128; k + 1];
row[0] = 1;
for _ in 0..n {
let mut prev_diag = row[0];
for kk in 1..=k {
let up = row[kk];
let v = up.saturating_add(row[kk - 1]).saturating_add(prev_diag);
prev_diag = up;
row[kk] = v.min(u128::MAX >> 8);
}
row[0] = 1;
}
row[k]
}
fn fits_in32(n: usize, k: i32) -> bool {
pvq_v(n, k as usize) <= u32::MAX as u128
}
/// Precomputed per-mode data.
pub struct Mode {
/// log2 of band width in 1/8 bits.
pub log_n: [i32; NB_EBANDS],
cache_index: Vec<i16>,
cache_bits: Vec<u8>,
caps: Vec<u8>,
}
impl Mode {
pub fn new() -> Mode {
let mut log_n = [0i32; NB_EBANDS];
for (i, l) in log_n.iter_mut().enumerate() {
*l = log2_frac((EBANDS[i + 1] - EBANDS[i]) as u32, BITRES);
}
let (cache_index, cache_bits) = compute_pulse_cache(MAX_LM as i32);
let mut m = Mode { log_n, cache_index, cache_bits, caps: Vec::new() };
m.caps = compute_caps(&m);
m
}
#[inline]
pub fn cache(&self, lm: i32, band: usize) -> &[u8] {
let idx = self.cache_index[((lm + 1) as usize) * NB_EBANDS + band];
&self.cache_bits[idx.max(0) as usize..]
}
pub fn bits2pulses(&self, band: usize, lm: i32, bits: i32) -> i32 {
let cache = self.cache(lm, band);
let mut lo = 0i32;
let mut hi = cache[0] as i32;
let bits = bits - 1;
for _ in 0..LOG_MAX_PSEUDO {
let mid = (lo + hi + 1) >> 1;
if cache[mid as usize] as i32 >= bits {
hi = mid;
} else {
lo = mid;
}
}
let lo_bits = if lo == 0 { -1 } else { cache[lo as usize] as i32 };
if bits - lo_bits <= cache[hi as usize] as i32 - bits { lo } else { hi }
}
pub fn pulses2bits(&self, band: usize, lm: i32, pulses: i32) -> i32 {
if pulses == 0 { 0 } else { self.cache(lm, band)[pulses as usize] as i32 + 1 }
}
/// Per-band allocation caps for this frame (`init_caps`).
pub fn init_caps(&self, lm: usize, c: usize) -> [i32; NB_EBANDS] {
let mut cap = [0i32; NB_EBANDS];
for (i, cp) in cap.iter_mut().enumerate() {
let n = (EBANDS[i + 1] - EBANDS[i]) << lm;
*cp = (self.caps[NB_EBANDS * (2 * lm + c - 1) + i] as i32 + 64) * c as i32 * n >> 2;
}
cap
}
}
fn compute_pulse_cache(lm_max: i32) -> (Vec<i16>, Vec<u8>) {
let nb = NB_EBANDS;
let rows = (lm_max + 2) as usize;
let mut index = vec![-1i16; nb * rows];
let mut entries: Vec<(usize, i32, usize)> = Vec::new();
let mut curr = 0usize;
for i in 0..rows {
for j in 0..nb {
let n = (((EBANDS[j + 1] - EBANDS[j]) << i) >> 1) as usize;
index[i * nb + j] = -1;
'search: for k in 0..=i {
for nn in 0..nb {
if k == i && nn >= j {
break;
}
if n == (((EBANDS[nn + 1] - EBANDS[nn]) << k) >> 1) as usize {
index[i * nb + j] = index[k * nb + nn];
break 'search;
}
}
}
if index[i * nb + j] == -1 && n != 0 {
let mut k = 0;
while k < MAX_PSEUDO && fits_in32(n, get_pulses(k + 1)) {
k += 1;
}
index[i * nb + j] = curr as i16;
entries.push((n, k, curr));
curr += (k + 1) as usize;
}
}
}
let mut bits = vec![0u8; curr];
for &(n, k, off) in &entries {
let maxk = get_pulses(k) as usize;
let tmp = required_bits(n, maxk, BITRES);
for j in 1..=k {
bits[off + j as usize] = (tmp[get_pulses(j) as usize] - 1) as u8;
}
bits[off] = k as u8;
}
(index, bits)
}
fn required_bits(n: usize, maxk: usize, frac: u32) -> Vec<i32> {
let mut out = vec![0i32; maxk + 1];
if n == 1 {
for o in out.iter_mut().skip(1) {
*o = 1 << frac;
}
} else {
for (k, o) in out.iter_mut().enumerate().skip(1) {
*o = log2_frac(pvq_v(n, k) as u32, frac);
}
}
out
}
fn compute_caps(m: &Mode) -> Vec<u8> {
let nb = NB_EBANDS;
let mut caps = vec![0u8; (MAX_LM + 1) * 2 * nb];
for i in 0..=MAX_LM as i32 {
for c in 1..=2i32 {
for j in 0..nb {
let mut n0 = EBANDS[j + 1] - EBANDS[j];
let max_bits = if n0 << i == 1 {
c * (1 + MAX_FINE_BITS) << BITRES
} else {
let mut lm0 = 0;
if n0 > 2 {
n0 >>= 1;
lm0 -= 1;
} else if n0 <= 1 {
lm0 = i.min(1);
n0 <<= lm0;
}
let pcache = m.cache(lm0, j);
let mut mb = pcache[pcache[0] as usize] as i32 + 1;
let mut n = n0;
for k in 0..(i - lm0) {
mb <<= 1;
let offset = ((m.log_n[j] + ((lm0 + k) << BITRES)) >> 1) - QTHETA_OFFSET;
let num = 459 * ((2 * n - 1) * offset + mb);
let den = ((2 * n - 1) << 9) - 459;
let qb = ((num + (den >> 1)) / den).min(57);
mb += qb;
n <<= 1;
}
if c == 2 {
mb <<= 1;
let offset = ((m.log_n[j] + (i << BITRES)) >> 1) - if n == 2 { QTHETA_OFFSET_TWOPHASE } else { QTHETA_OFFSET };
let ndof = 2 * n - 1 - (n == 2) as i32;
let p = if n == 2 { 512 } else { 487 };
let num = p * (mb + ndof * offset);
let den = (ndof << 9) - p;
let qb = ((num + (den >> 1)) / den).min(if n == 2 { 64 } else { 61 });
mb += qb;
}
let ndof = c * n + if c == 2 && n > 2 { 1 } else { 0 };
let mut offset = ((m.log_n[j] + (i << BITRES)) >> 1) - FINE_OFFSET;
if n == 2 {
offset += (1 << BITRES) >> 2;
}
let num = mb + ndof * offset;
let den = (ndof - 1) << BITRES;
let qb = ((num + (den >> 1)) / den).min(MAX_FINE_BITS);
mb += c * qb << BITRES;
mb
};
let w = c * ((EBANDS[j + 1] - EBANDS[j]) << i);
let v = (4 * max_bits / w) - 64;
caps[((i * 2 + c - 1) as usize) * nb + j] = v.clamp(0, 255) as u8;
}
}
}
caps
}
/// Result of the allocation process.
pub struct Allocation {
pub coded_bands: usize,
pub intensity: usize,
pub dual_stereo: bool,
pub balance: i32,
pub pulses: [i32; NB_EBANDS],
pub fine_quant: [i32; NB_EBANDS],
pub fine_priority: [i32; NB_EBANDS],
}
/// `clt_compute_allocation`, encoder side. Stereo frames are always coded as dual (L/R) stereo
/// with no intensity bands, and no band is skipped voluntarily.
#[allow(clippy::too_many_arguments)]
pub fn compute_allocation(
m: &Mode,
start: usize,
end: usize,
offsets: &[i32; NB_EBANDS],
cap: &[i32; NB_EBANDS],
alloc_trim: i32,
total: i32,
c: usize,
lm: usize,
enc: &mut RangeEncoder,
) -> Allocation {
let ci = c as i32;
let lmi = lm as i32;
let mut total = total.max(0);
let mut skip_start = start;
let skip_rsv = if total >= 1 << BITRES { 1 << BITRES } else { 0 };
total -= skip_rsv;
let mut intensity_rsv = 0;
let mut dual_stereo_rsv = 0;
if c == 2 {
intensity_rsv = LOG2_FRAC_TABLE[end - start];
if intensity_rsv > total {
intensity_rsv = 0;
} else {
total -= intensity_rsv;
dual_stereo_rsv = if total >= 1 << BITRES { 1 << BITRES } else { 0 };
total -= dual_stereo_rsv;
}
}
let mut bits1 = [0i32; NB_EBANDS];
let mut bits2 = [0i32; NB_EBANDS];
let mut thresh = [0i32; NB_EBANDS];
let mut trim_offset = [0i32; NB_EBANDS];
for j in start..end {
let w = EBANDS[j + 1] - EBANDS[j];
thresh[j] = (ci << BITRES).max((3 * w << lmi << BITRES) >> 4);
trim_offset[j] = ci * w * (alloc_trim - 5 - lmi) * (end as i32 - j as i32 - 1) * (1 << (lmi + BITRES as i32)) >> 6;
if w << lmi == 1 {
trim_offset[j] -= ci << BITRES;
}
}
let nb_alloc = ALLOC.len() as i32;
let mut lo = 1;
let mut hi = nb_alloc - 1;
loop {
let mut done = false;
let mut psum = 0;
let mid = (lo + hi) >> 1;
for j in (start..end).rev() {
let n = EBANDS[j + 1] - EBANDS[j];
let mut bitsj = ci * n * (ALLOC[mid as usize][j] as i32) << lmi >> 2;
if bitsj > 0 {
bitsj = (bitsj + trim_offset[j]).max(0);
}
bitsj += offsets[j];
if bitsj >= thresh[j] || done {
done = true;
psum += bitsj.min(cap[j]);
} else if bitsj >= ci << BITRES {
psum += ci << BITRES;
}
}
if psum > total {
hi = mid - 1;
} else {
lo = mid + 1;
}
if lo > hi {
break;
}
}
hi = lo;
lo -= 1;
for j in start..end {
let n = EBANDS[j + 1] - EBANDS[j];
let mut bits1j = ci * n * (ALLOC[lo as usize][j] as i32) << lmi >> 2;
let mut bits2j = if hi >= nb_alloc { cap[j] } else { ci * n * (ALLOC[hi as usize][j] as i32) << lmi >> 2 };
if bits1j > 0 {
bits1j = (bits1j + trim_offset[j]).max(0);
}
if bits2j > 0 {
bits2j = (bits2j + trim_offset[j]).max(0);
}
if lo > 0 {
bits1j += offsets[j];
}
bits2j += offsets[j];
if offsets[j] > 0 {
skip_start = j;
}
bits2j = (bits2j - bits1j).max(0);
bits1[j] = bits1j;
bits2[j] = bits2j;
}
interp_bits2pulses(m, start, end, skip_start, &bits1, &bits2, &thresh, cap, total, skip_rsv, intensity_rsv, dual_stereo_rsv, c, lm, enc)
}
#[allow(clippy::too_many_arguments)]
fn interp_bits2pulses(
m: &Mode,
start: usize,
end: usize,
skip_start: usize,
bits1: &[i32; NB_EBANDS],
bits2: &[i32; NB_EBANDS],
thresh: &[i32; NB_EBANDS],
cap: &[i32; NB_EBANDS],
mut total: i32,
skip_rsv: i32,
mut intensity_rsv: i32,
mut dual_stereo_rsv: i32,
c: usize,
lm: usize,
enc: &mut RangeEncoder,
) -> Allocation {
let ci = c as i32;
let alloc_floor = ci << BITRES;
let stereo = (c > 1) as i32;
let log_m = (lm as i32) << BITRES;
let mut lo = 0;
let mut hi = 1 << ALLOC_STEPS;
for _ in 0..ALLOC_STEPS {
let mid = (lo + hi) >> 1;
let mut psum = 0;
let mut done = false;
for j in (start..end).rev() {
let tmp = bits1[j] + (mid * bits2[j] >> ALLOC_STEPS);
if tmp >= thresh[j] || done {
done = true;
psum += tmp.min(cap[j]);
} else if tmp >= alloc_floor {
psum += alloc_floor;
}
}
if psum > total {
hi = mid;
} else {
lo = mid;
}
}
let mut bits = [0i32; NB_EBANDS];
let mut psum = 0;
let mut done = false;
for j in (start..end).rev() {
let mut tmp = bits1[j] + (lo * bits2[j] >> ALLOC_STEPS);
if tmp < thresh[j] && !done {
tmp = if tmp >= alloc_floor { alloc_floor } else { 0 };
} else {
done = true;
}
tmp = tmp.min(cap[j]);
bits[j] = tmp;
psum += tmp;
}
let s = start;
let mut coded_bands = end;
loop {
let j = coded_bands - 1;
if j <= skip_start {
total += skip_rsv;
break;
}
let mut left = total - psum;
let span = EBANDS[coded_bands] - EBANDS[s];
let percoeff = left / span;
left -= span * percoeff;
let rem = (left - (EBANDS[j] - EBANDS[s])).max(0);
let band_width = EBANDS[coded_bands] - EBANDS[j];
let band_bits = bits[j] + percoeff * band_width + rem;
if band_bits >= thresh[j].max(alloc_floor + (1 << BITRES)) {
// Keep this band (and every band below it).
enc.bit_logp(true, 1);
break;
}
psum -= bits[j] + intensity_rsv;
if intensity_rsv > 0 {
intensity_rsv = LOG2_FRAC_TABLE[j - s];
}
psum += intensity_rsv;
if band_bits >= alloc_floor {
psum += alloc_floor;
bits[j] = alloc_floor;
} else {
bits[j] = 0;
}
coded_bands -= 1;
}
let intensity = if intensity_rsv > 0 {
// No intensity stereo: every coded band is coded in full.
enc.uint((coded_bands - s) as u32, (coded_bands + 1 - s) as u32);
coded_bands
} else {
0
};
if intensity <= s {
total += dual_stereo_rsv;
dual_stereo_rsv = 0;
}
let dual_stereo = if dual_stereo_rsv > 0 {
enc.bit_logp(true, 1);
true
} else {
false
};
let mut left = total - psum;
let span = EBANDS[coded_bands] - EBANDS[s];
let percoeff = left / span;
left -= span * percoeff;
for j in s..coded_bands {
bits[j] += percoeff * (EBANDS[j + 1] - EBANDS[j]);
}
for j in s..coded_bands {
let tmp = left.min(EBANDS[j + 1] - EBANDS[j]);
bits[j] += tmp;
left -= tmp;
}
let mut ebits = [0i32; NB_EBANDS];
let mut fine_priority = [0i32; NB_EBANDS];
let mut balance = 0;
for j in s..coded_bands {
let n0 = EBANDS[j + 1] - EBANDS[j];
let n = n0 << lm;
let bit = bits[j] + balance;
let mut excess;
if n > 1 {
excess = (bit - cap[j]).max(0);
bits[j] = bit - excess;
let den = ci * n + if c == 2 && n > 2 && !dual_stereo && j < intensity { 1 } else { 0 };
let nclogn = den * (m.log_n[j] + log_m);
let mut offset = (nclogn >> 1) - den * FINE_OFFSET;
if n == 2 {
offset += (den << BITRES) >> 2;
}
if bits[j] + offset < den * 2 << BITRES {
offset += nclogn >> 2;
} else if bits[j] + offset < den * 3 << BITRES {
offset += nclogn >> 3;
}
let mut eb = (bits[j] + offset + (den << (BITRES - 1))).max(0);
eb = (eb / den) >> BITRES;
if ci * eb > (bits[j] >> BITRES) {
eb = bits[j] >> stereo >> BITRES;
}
eb = eb.min(MAX_FINE_BITS);
ebits[j] = eb;
fine_priority[j] = (eb * (den << BITRES) >= bits[j] + offset) as i32;
bits[j] -= ci * eb << BITRES;
} else {
excess = (bit - (ci << BITRES)).max(0);
bits[j] = bit - excess;
ebits[j] = 0;
fine_priority[j] = 1;
}
if excess > 0 {
let extra_fine = (excess >> (stereo + BITRES as i32)).min(MAX_FINE_BITS - ebits[j]);
ebits[j] += extra_fine;
let extra_bits = extra_fine * ci << BITRES;
fine_priority[j] = (extra_bits >= excess - balance) as i32;
excess -= extra_bits;
}
balance = excess;
}
for j in coded_bands..end {
ebits[j] = bits[j] >> stereo >> BITRES;
bits[j] = 0;
fine_priority[j] = (ebits[j] < 1) as i32;
}
Allocation { coded_bands, intensity, dual_stereo, balance, pulses: bits, fine_quant: ebits, fine_priority }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn caps_match_decoder_reference() {
let m = Mode::new();
let first_rows: [u8; 42] = [
224, 224, 224, 224, 224, 224, 224, 224, 160, 160, 160, 160, 185, 185, 185, 178, 178, 168, 134, 61, 37, 224, 224, 224, 224, 224, 224, 224, 224, 240,
240, 240, 240, 207, 207, 207, 198, 198, 183, 144, 66, 40,
];
assert_eq!(&m.caps[..42], &first_rows[..]);
}
#[test]
fn pvq_counts() {
assert_eq!(pvq_v(5, 1), 10);
assert_eq!(pvq_v(1, 7), 2);
assert_eq!(pvq_v(2, 9), 36);
assert_eq!(pvq_v(3, 2), 18);
}
}