diff --git a/crates/aac/src/decoder.rs b/crates/aac/src/decoder.rs --- a/crates/aac/src/decoder.rs +++ b/crates/aac/src/decoder.rs @@ -5,8 +5,8 @@ use filmcraft_bitstream::BitReader; use crate::config::{AudioSpecificConfig, ElementType, ProgramConfig}; use crate::huffman::{self, INTENSITY_HCB, INTENSITY_HCB2, NOISE_HCB, ZERO_HCB}; use crate::ics::{IcsInfo, WindowSequence}; -use crate::mdct::{WindowShape, mdct_long, mdct_short, window}; -use crate::tables::{pow43, sf_gain}; +use crate::mdct::{MdctScratch, WindowShape, mdct_long, mdct_short, window}; +use crate::tables::{pow43, pow43_table, sf_gain}; use crate::tns::{self, TnsData}; use crate::{Error, Result}; @@ -41,6 +41,23 @@ struct ChannelState { prev_shape: WindowShape, } +/// Buffers reused by every frame so the steady state does not allocate per channel. +struct Scratch { + /// Quantised values of the channel being parsed. + q: Vec, + /// Windowed output of the current block (2048) and the raw inverse transform (2048 / 256). + buf: Vec, + y: Vec, + long: MdctScratch, + short: MdctScratch, +} + +impl Scratch { + fn new() -> Scratch { + Scratch { q: vec![0; 1024], buf: vec![0.0; 2048], y: vec![0.0; 2048], long: mdct_long().scratch(), short: mdct_short().scratch() } + } +} + /// AAC-LC decoder. pub struct Decoder { asc: AudioSpecificConfig, @@ -49,6 +66,7 @@ pub struct Decoder { offsets: Vec, channels: Vec, noise_state: u32, + scratch: Scratch, /// A fill element carried SBR data (implicitly signalled HE-AAC). sbr_seen: bool, } @@ -73,7 +91,8 @@ impl Decoder { if layout.is_empty() && asc.channel_config != 0 { return Err(Error::InvalidConfig("unsupported channel configuration")); } - let mut d = Decoder { asc, layout: Vec::new(), offsets: Vec::new(), channels: Vec::new(), noise_state: 0x2545_F491, sbr_seen: false }; + let mut d = + Decoder { asc, layout: Vec::new(), offsets: Vec::new(), channels: Vec::new(), noise_state: 0x2545_F491, scratch: Scratch::new(), sbr_seen: false }; d.set_layout(layout); Ok(d) } @@ -389,7 +408,9 @@ impl Decoder { } // spectral_data - let mut q = vec![0i32; 1024]; + let mut q = std::mem::take(&mut self.scratch.q); + q.clear(); + q.resize(1024, 0); let mut vals = [0i32; 4]; for g in 0..info.num_groups { let w0 = info.group_start(g); @@ -421,6 +442,7 @@ impl Decoder { } // dequantise + let t43 = pow43_table(); for g in 0..info.num_groups { let w0 = info.group_start(g); for sfb in 0..info.max_sfb { @@ -430,14 +452,16 @@ impl Decoder { } let gain = sf_gain(ics.sf[g][sfb]); for w in w0..w0 + info.group_len[g] as usize { - for k in swb[sfb] as usize..swb[sfb + 1] as usize { - let v = q[w * 128 + k]; - let m = pow43(v.unsigned_abs()) * gain; - ics.spec[w * 128 + k] = if v < 0 { -m } else { m }; + let r = w * 128 + swb[sfb] as usize..w * 128 + swb[sfb + 1] as usize; + for (o, &v) in ics.spec[r.clone()].iter_mut().zip(&q[r]) { + let a = v.unsigned_abs(); + let m = t43.get(a as usize).copied().unwrap_or_else(|| pow43(a)) * gain; + *o = if v < 0 { -m } else { m }; } } } } + self.scratch.q = q; Ok(ics) } @@ -487,23 +511,24 @@ impl Decoder { if let Some(t) = &ics.tns { tns::apply_decoder(&mut ics.spec, &ics.info, self.sf_index(), t); } - synth(&mut self.channels[ch], &ics.info, &ics.spec) + synth(&mut self.channels[ch], &mut self.scratch, &ics.info, &ics.spec) } } /// Inverse MDCT, windowing and overlap-add (§4.6.18), output scaled to ±1.0 full scale. -fn synth(state: &mut ChannelState, info: &IcsInfo, spec: &[f32]) -> Vec { - let mut buf = vec![0f32; 2048]; +fn synth(state: &mut ChannelState, sc: &mut Scratch, info: &IcsInfo, spec: &[f32]) -> Vec { + let Scratch { buf, y, long, short, .. } = sc; + buf.fill(0.0); let prev = state.prev_shape; let cur = info.window_shape; let norm = 1.0 / 32768.0; match info.window_sequence { WindowSequence::EightShort => { let t = mdct_short(); - let mut y = vec![0f32; 256]; + let y = &mut y[..256]; let wc = window(cur, false); for w in 0..8 { - t.inverse(&spec[w * 128..(w + 1) * 128], &mut y); + t.inverse_with(&spec[w * 128..(w + 1) * 128], y, short); let wl = window(if w == 0 { prev } else { cur }, false); let off = 448 + 128 * w; let s = norm / 128.0; @@ -514,8 +539,7 @@ fn synth(state: &mut ChannelState, info: &IcsInfo, spec: &[f32]) -> Vec { } } seq => { - let mut y = vec![0f32; 2048]; - mdct_long().inverse(&spec[..1024], &mut y); + mdct_long().inverse_with(&spec[..1024], y, long); let s = norm / 1024.0; match seq { WindowSequence::LongStop => { diff --git a/crates/aac/src/huffman.rs b/crates/aac/src/huffman.rs --- a/crates/aac/src/huffman.rs +++ b/crates/aac/src/huffman.rs @@ -53,60 +53,92 @@ pub fn book(cb: u8) -> &'static Book { &BOOKS[(cb as usize).clamp(1, 11) - 1] } -const LEAF: u32 = 1 << 31; +/// Bits resolved by the first table level. +const L1: u32 = 9; +/// Flag of a first-level entry that points to a sub-table. +const SUB: u32 = 1 << 31; -/// Binary decoding tree for one codebook. -struct Tree { - nodes: Vec<[u32; 2]>, +/// Two-level decoding table for one prefix code. A leaf entry is `len << 24 | symbol`; a first-level +/// entry with `SUB` set is `SUB | offset << 5 | width` and selects a sub-table of `1 << width` +/// leaves indexed by the next `width` bits. An entry of 0 is a bit pattern that no codeword starts +/// with (the codebooks are complete, so this does not occur). +struct Lut { + t: Vec, } -impl Tree { - fn build(codes: &[(u32, u8)]) -> Tree { - let mut nodes: Vec<[u32; 2]> = vec![[0, 0]]; +impl Lut { + fn build(codes: &[(u32, u8)]) -> Lut { + let mut t = vec![0u32; 1 << L1]; + // sub-table width per first-level prefix + let mut width = vec![0u32; 1 << L1]; + for &(code, len) in codes { + let len = len as u32; + if len > L1 { + let p = (code >> (len - L1)) as usize; + width[p] = width[p].max(len - L1); + } + } + for (p, &w) in width.iter().enumerate() { + if w > 0 { + t[p] = SUB | (t.len() as u32) << 5 | w; + t.resize(t.len() + (1 << w), 0); + } + } for (sym, &(code, len)) in codes.iter().enumerate() { - let mut node = 0usize; - for i in (0..len).rev() { - let bit = ((code >> i) & 1) as usize; - if i == 0 { - nodes[node][bit] = LEAF | sym as u32; - } else { - if nodes[node][bit] == 0 { - nodes.push([0, 0]); - nodes[node][bit] = (nodes.len() - 1) as u32; - } - node = nodes[node][bit] as usize; - } + let len = len as u32; + let leaf = len << 24 | sym as u32; + if len <= L1 { + let first = (code << (L1 - len)) as usize; + t[first..first + (1 << (L1 - len))].fill(leaf); + } else { + let e = t[(code >> (len - L1)) as usize]; + let (off, w) = (((e & !SUB) >> 5) as usize, e & 31); + let rest = code & ((1 << (len - L1)) - 1); + let first = off + (rest << (w - (len - L1))) as usize; + t[first..first + (1 << (w - (len - L1)))].fill(leaf); } } - Tree { nodes } + Lut { t } } + /// The symbol of the codeword at the reader position; `Eof` if the codeword is cut off. #[inline] fn decode(&self, br: &mut BitReader) -> Result { - let mut node = 0usize; - loop { - let bit = br.read_bits(1)? as usize; - let next = self.nodes[node][bit]; - if next & LEAF != 0 { - return Ok(next & !LEAF); - } - if next == 0 { - return Err(Error::Bitstream("invalid Huffman codeword")); - } - node = next as usize; + let bits = peek32(br); + let mut e = self.t[(bits >> (32 - L1)) as usize]; + if e & SUB != 0 { + let w = e & 31; + e = self.t[((e & !SUB) >> 5) as usize + ((bits << L1) >> (32 - w)) as usize]; } + let len = e >> 24; + if len == 0 { + return Err(Error::Bitstream("invalid Huffman codeword")); + } + br.skip(len as usize).map_err(|_| Error::Eof)?; + Ok(e & 0xFF_FFFF) + } +} + +/// The next 32 bits (zero padded past the end), from one 8-byte load when possible. +#[inline] +fn peek32(br: &BitReader) -> u32 { + let pos = br.position(); + let byte = pos >> 3; + match br.data().get(byte..byte.saturating_add(8)).and_then(|b| <[u8; 8]>::try_from(b).ok()) { + Some(b) => ((u64::from_be_bytes(b) << (pos & 7)) >> 32) as u32, + None => br.peek(32), } } -fn trees() -> &'static [Tree; 12] { - static TREES: OnceLock<[Tree; 12]> = OnceLock::new(); - TREES.get_or_init(|| std::array::from_fn(|i| if i == 0 { Tree::build(&SCALEFACTOR) } else { Tree::build(BOOKS[i - 1].codes) })) +fn luts() -> &'static [Lut; 12] { + static LUTS: OnceLock<[Lut; 12]> = OnceLock::new(); + LUTS.get_or_init(|| std::array::from_fn(|i| if i == 0 { Lut::build(&SCALEFACTOR) } else { Lut::build(BOOKS[i - 1].codes) })) } /// Decode one scalefactor difference (`-60..=60`). #[inline] pub fn decode_sf(br: &mut BitReader) -> Result { - Ok(trees()[0].decode(br)? as i32 - 60) + Ok(luts()[0].decode(br)? as i32 - 60) } /// Write one scalefactor difference (`-60..=60`). @@ -125,7 +157,7 @@ pub fn sf_bits(diff: i32) -> u32 { /// Decode one codeword of codebook `cb` (1..=11) into `out[..dim]`, including sign bits and escapes. pub fn decode_spectral(br: &mut BitReader, cb: u8, out: &mut [i32]) -> Result<()> { let b = book(cb); - let mut idx = trees()[cb as usize].decode(br)?; + let mut idx = luts()[cb as usize].decode(br)?; for i in (0..b.dim).rev() { out[i] = (idx % b.modulo) as i32; idx /= b.modulo; @@ -259,6 +291,53 @@ mod tests { } } + /// Every codeword, followed by arbitrary bits, decodes to its own symbol and length. + #[test] + fn lut_decodes_every_codeword() { + let mut seed = 7u32; + for (i, lut) in luts().iter().enumerate() { + let codes: &[(u32, u8)] = if i == 0 { &SCALEFACTOR } else { BOOKS[i - 1].codes }; + for (sym, &(code, len)) in codes.iter().enumerate() { + for _ in 0..3 { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + let mut bw = BitWriter::new(); + bw.write_bits(code, len as u32); + bw.write_bits(seed >> 8, 24); + let data = bw.finish(); + let mut br = BitReader::new(&data); + assert_eq!(lut.decode(&mut br).unwrap(), sym as u32, "book {i}"); + assert_eq!(br.position(), len as usize); + } + } + } + } + + /// A codeword cut off by the end of the data is `Eof`, also at an unaligned position. + #[test] + fn truncated_codeword_is_eof() { + for (i, lut) in luts().iter().enumerate() { + let codes: &[(u32, u8)] = if i == 0 { &SCALEFACTOR } else { BOOKS[i - 1].codes }; + let &(code, len) = codes.iter().max_by_key(|c| c.1).unwrap(); + for lead in 0..8u32 { + let mut bw = BitWriter::new(); + bw.write_bits(0, lead); + bw.write_bits(code, len as u32); + let total = lead as usize + len as usize; + let data = bw.finish(); + // drop the last byte if that cuts the codeword, else the codeword is whole + let short = &data[..data.len() - 1]; + let mut br = BitReader::new(short); + br.skip(lead as usize).unwrap(); + let r = lut.decode(&mut br); + if short.len() * 8 < total { + assert!(matches!(r, Err(Error::Eof)), "book {i} lead {lead}"); + } else { + assert!(r.is_ok()); + } + } + } + } + #[test] fn spectral_roundtrip_every_book() { for cb in 1..=11u8 { diff --git a/crates/aac/src/mdct.rs b/crates/aac/src/mdct.rs --- a/crates/aac/src/mdct.rs +++ b/crates/aac/src/mdct.rs @@ -7,7 +7,11 @@ use std::sync::OnceLock; /// Radix-2 complex FFT (forward, `e^{-2πi nk/P}`). struct Fft { n: usize, - twiddle: Vec<[f32; 2]>, + /// Twiddle of the length-2 stage, `e^{-2πi 0/n}`. + w0: [f32; 2], + /// Twiddles of the stages of length 4, 8, ..., n, one contiguous run of `len / 2` per stage. + stage_re: Vec, + stage_im: Vec, rev: Vec, } @@ -16,13 +20,23 @@ impl Fft { assert!(n.is_power_of_two()); let bits = n.trailing_zeros(); let rev = (0..n as u32).map(|i| if bits == 0 { 0 } else { i.reverse_bits() >> (32 - bits) }).collect(); - let twiddle = (0..n / 2) + let twiddle: Vec<[f32; 2]> = (0..n / 2) .map(|k| { let a = -2.0 * PI * k as f64 / n as f64; [a.cos() as f32, a.sin() as f32] }) .collect(); - Fft { n, twiddle, rev } + let (mut stage_re, mut stage_im) = (Vec::new(), Vec::new()); + let mut len = 4; + while len <= n { + let step = n / len; + for k in 0..len / 2 { + stage_re.push(twiddle[k * step][0]); + stage_im.push(twiddle[k * step][1]); + } + len <<= 1; + } + Fft { n, w0: twiddle.first().copied().unwrap_or([1.0, 0.0]), stage_re, stage_im, rev } } fn run(&self, re: &mut [f32], im: &mut [f32]) { @@ -34,24 +48,80 @@ impl Fft { im.swap(i, j); } } - let mut len = 2; - while len <= n { + if n >= 2 { + let [wr, wi] = self.w0; + let ((re2, _), (im2, _)) = (re.as_chunks_mut::<2>(), im.as_chunks_mut::<2>()); + for (r, i) in re2.iter_mut().zip(im2) { + let xr = r[1] * wr - i[1] * wi; + let xi = r[1] * wi + i[1] * wr; + r[1] = r[0] - xr; + i[1] = i[0] - xi; + r[0] += xr; + i[0] += xi; + } + } + let mut len = 4; + let mut off = 0; + // Two radix-2 stages (len, 2 * len) per pass over the data: the same butterflies in the same + // order per element, so the result is bit-identical to running the stages one by one. + while len * 2 <= n { + let h = len / 2; + let (t1r, t1i) = (&self.stage_re[off..off + h], &self.stage_im[off..off + h]); + let (t2r, t2i) = (&self.stage_re[off + h..off + 3 * h], &self.stage_im[off + h..off + 3 * h]); + off += 3 * h; + for (rc, ic) in re.chunks_exact_mut(4 * h).zip(im.chunks_exact_mut(4 * h)) { + let (ra, rest) = rc.split_at_mut(h); + let (rb, rest) = rest.split_at_mut(h); + let (rc2, rd) = rest.split_at_mut(h); + let (ia, rest) = ic.split_at_mut(h); + let (ib, rest) = rest.split_at_mut(h); + let (ic2, id) = rest.split_at_mut(h); + for k in 0..h { + let (w1r, w1i) = (t1r[k], t1i[k]); + // stage `len`: (a, b) and (c, d) + let xr = rb[k] * w1r - ib[k] * w1i; + let xi = rb[k] * w1i + ib[k] * w1r; + let (br, bi) = (ra[k] - xr, ia[k] - xi); + let (ar, ai) = (ra[k] + xr, ia[k] + xi); + let xr = rd[k] * w1r - id[k] * w1i; + let xi = rd[k] * w1i + id[k] * w1r; + let (dr, di) = (rc2[k] - xr, ic2[k] - xi); + let (cr, ci) = (rc2[k] + xr, ic2[k] + xi); + // stage `2 * len`: (a, c) with w2[k], (b, d) with w2[k + h] + let (w2r, w2i) = (t2r[k], t2i[k]); + let xr = cr * w2r - ci * w2i; + let xi = cr * w2i + ci * w2r; + rc2[k] = ar - xr; + ic2[k] = ai - xi; + ra[k] = ar + xr; + ia[k] = ai + xi; + let (w2r, w2i) = (t2r[k + h], t2i[k + h]); + let xr = dr * w2r - di * w2i; + let xi = dr * w2i + di * w2r; + rd[k] = br - xr; + id[k] = bi - xi; + rb[k] = br + xr; + ib[k] = bi + xi; + } + } + len <<= 2; + } + if len <= n { let half = len / 2; - let step = n / len; - for start in (0..n).step_by(len) { + let (twr, twi) = (&self.stage_re[off..off + half], &self.stage_im[off..off + half]); + for (rc, ic) in re.chunks_exact_mut(len).zip(im.chunks_exact_mut(len)) { + let (ra, rb) = rc.split_at_mut(half); + let (ia, ib) = ic.split_at_mut(half); for k in 0..half { - let [wr, wi] = self.twiddle[k * step]; - let a = start + k; - let b = a + half; - let xr = re[b] * wr - im[b] * wi; - let xi = re[b] * wi + im[b] * wr; - re[b] = re[a] - xr; - im[b] = im[a] - xi; - re[a] += xr; - im[a] += xi; + let (wr, wi) = (twr[k], twi[k]); + let xr = rb[k] * wr - ib[k] * wi; + let xi = rb[k] * wi + ib[k] * wr; + rb[k] = ra[k] - xr; + ib[k] = ia[k] - xi; + ra[k] += xr; + ia[k] += xi; } } - len <<= 1; } } } @@ -83,19 +153,22 @@ impl Mdct { Mdct { m, fft: Fft::new(q), pre, post } } - /// DCT-IV: `out[k] = Σ u[n] cos(π/M (n+½)(k+½))`. - fn dct4(&self, u: &[f32], out: &mut [f32]) { + /// Working buffers for this transform size. + pub fn scratch(&self) -> MdctScratch { + MdctScratch { re: vec![0.0; self.m / 2], im: vec![0.0; self.m / 2], u: vec![0.0; self.m] } + } + + /// DCT-IV: `out[k] = Σ u[n] cos(π/M (n+½)(k+½))`; `re` and `im` (length M/2) are overwritten. + fn dct4(&self, u: &[f32], out: &mut [f32], re: &mut [f32], im: &mut [f32]) { let m = self.m; let q = m / 2; - let mut re = vec![0f32; q]; - let mut im = vec![0f32; q]; for n in 0..q { let (a, b) = (u[2 * n], u[m - 1 - 2 * n]); let [c, s] = self.pre[n]; re[n] = a * c - b * s; im[n] = a * s + b * c; } - self.fft.run(&mut re, &mut im); + self.fft.run(re, im); for k in 0..q { let [c, s] = self.post[k]; let wr = re[k] * c - im[k] * s; @@ -109,21 +182,28 @@ impl Mdct { pub fn forward(&self, x: &[f32], out: &mut [f32]) { let m = self.m; let h = m / 2; - let mut u = vec![0f32; m]; + let MdctScratch { re, im, mut u } = self.scratch(); for n in 0..h { // quarters a, b, c, d of length M/2 u[n] = -x[m + h - 1 - n] - x[m + h + n]; u[h + n] = x[n] - x[m - 1 - n]; } - self.dct4(&u, out); + let (mut re, mut im) = (re, im); + self.dct4(&u, out, &mut re, &mut im); } /// Inverse MDCT (unscaled): `y[n] = Σ X[k] cos(π/M (n + ½ + M/2)(k + ½))`, `y` has length 2M. + #[cfg(test)] pub fn inverse(&self, spec: &[f32], y: &mut [f32]) { + self.inverse_with(spec, y, &mut self.scratch()); + } + + /// [`Self::inverse`] with caller-provided working buffers from [`Self::scratch`]. + pub fn inverse_with(&self, spec: &[f32], y: &mut [f32], sc: &mut MdctScratch) { let m = self.m; let h = m / 2; - let mut u = vec![0f32; m]; - self.dct4(spec, &mut u); + let MdctScratch { re, im, u } = sc; + self.dct4(spec, u, re, im); for n in 0..h { y[n] = u[h + n]; y[m - 1 - n] = -u[h + n]; @@ -133,6 +213,13 @@ impl Mdct { } } +/// Working buffers of one [`Mdct`] size, reusable across transforms. +pub struct MdctScratch { + re: Vec, + im: Vec, + u: Vec, +} + /// Cached transforms for 1024 and 128 lines. pub fn mdct_long() -> &'static Mdct { static M: OnceLock = OnceLock::new(); diff --git a/crates/aac/src/tables.rs b/crates/aac/src/tables.rs --- a/crates/aac/src/tables.rs +++ b/crates/aac/src/tables.rs @@ -99,12 +99,22 @@ pub fn tns_max_order(short: bool) -> usize { if short { 7 } else { 12 } } -/// `x^(4/3)` for the integer quantised magnitudes 0..=8191. -pub fn pow43(q: u32) -> f32 { +/// The `x^(4/3)` table for the integer quantised magnitudes 0..=8191. +pub fn pow43_table() -> &'static [f32; 8192] { use std::sync::OnceLock; - static TABLE: OnceLock> = OnceLock::new(); - let t = TABLE.get_or_init(|| (0..8192u32).map(|i| (i as f64).powf(4.0 / 3.0) as f32).collect()); - match t.get(q as usize) { + static TABLE: OnceLock> = OnceLock::new(); + TABLE.get_or_init(|| { + let mut t = Box::new([0f32; 8192]); + for (i, v) in t.iter_mut().enumerate() { + *v = (i as f64).powf(4.0 / 3.0) as f32; + } + t + }) +} + +/// `x^(4/3)` for an integer quantised magnitude (table lookup up to 8191). +pub fn pow43(q: u32) -> f32 { + match pow43_table().get(q as usize) { Some(v) => *v, None => (q as f64).powf(4.0 / 3.0) as f32, } @@ -113,6 +123,16 @@ pub fn pow43(q: u32) -> f32 { /// `2^(0.25 * (sf - 100))`, the dequantisation gain of scalefactor `sf`. #[inline] pub fn sf_gain(sf: i32) -> f32 { + use std::sync::OnceLock; + static TABLE: OnceLock<[f32; 256]> = OnceLock::new(); + let t = TABLE.get_or_init(|| std::array::from_fn(|i| sf_gain_exact(i as i32))); + match usize::try_from(sf).ok().and_then(|i| t.get(i)) { + Some(g) => *g, + None => sf_gain_exact(sf), + } +} + +fn sf_gain_exact(sf: i32) -> f32 { (0.25 * (sf as f64 - 100.0)).exp2() as f32 }