wmv_decoder/wma/
mdct.rs

1//! MDCT/IMDCT implementation.
2//!
3
4use std::f64::consts::PI;
5
6/// FFT-based inverse MDCT. The historical name is retained for API compatibility.
7pub struct MdctNaive {
8    /// Frame size (N). Full IMDCT output is 2N.
9    pub len: usize,
10    /// Scale factor (double precision in upstream).
11    pub scale: f64,
12    plan: Option<FftPlan>,
13}
14
15impl MdctNaive {
16    pub fn new(len: usize, scale: f64) -> Self {
17        Self {
18            len,
19            scale,
20            plan: (len >= 2 && len.is_power_of_two()).then(|| FftPlan::new(len)),
21        }
22    }
23
24    /// Half-length inverse MDCT.
25    ///
26    /// Input: N coefficients.
27    /// Output: N samples (half IMDCT), matching upstream's MDCT semantics.
28    pub fn imdct_half(&self, dst: &mut [f32], src: &[f32]) {
29        if let Some(plan) = self
30            .plan
31            .as_ref()
32            .filter(|plan| plan.phase.len() * 2 == self.len)
33        {
34            plan.imdct_half(dst, src, self.scale);
35            return;
36        }
37        // Retain support for non-power-of-two lengths and callers changing the
38        // public length field. WMA uses power-of-two lengths exclusively.
39        self.imdct_half_naive(dst, src);
40    }
41
42    fn imdct_half_naive(&self, dst: &mut [f32], src: &[f32]) {
43        // Translated from `ff_tx_mdct_naive_inv`.
44        // In upstream: len = s->len >> 1; len2 = len*2 (== s->len)
45        let len = self.len >> 1;
46        let len2 = len * 2;
47        let phase = PI / (4.0 * (len2 as f64));
48
49        for i in 0..len {
50            let mut sum_d: f64 = 0.0;
51            let mut sum_u: f64 = 0.0;
52
53            let i_d = phase * ((4 * len - 2 * i - 1) as f64);
54            let i_u = phase * ((3 * len2 + 2 * i + 1) as f64);
55
56            for j in 0..len2 {
57                let a = (2 * j + 1) as f64;
58                let a_d = (a * i_d).cos();
59                let a_u = (a * i_u).cos();
60                let val = src[j] as f64;
61                sum_d += a_d * val;
62                sum_u += a_u * val;
63            }
64
65            dst[i] = (sum_d * self.scale) as f32;
66            dst[i + len] = (-(sum_u * self.scale)) as f32;
67        }
68    }
69
70    /// Full IMDCT.
71    ///
72    /// Input: N coefficients.
73    /// Output: 2N samples.
74    pub fn imdct_full(&self, dst: &mut [f32], src: &[f32]) {
75        // Translated from `ff_tx_mdct_inv_full`.
76        let len = self.len * 2;
77        let len2 = len / 2;
78        let len4 = len / 4;
79
80        // The half IMDCT is written into the middle of the output.
81        self.imdct_half(&mut dst[len4..len4 + len2], src);
82
83        for i in 0..len4 {
84            dst[i] = -dst[len2 - i - 1];
85            dst[len - i - 1] = dst[len2 + i];
86        }
87    }
88}
89
90#[derive(Clone, Copy, Default)]
91struct Complex {
92    re: f64,
93    im: f64,
94}
95
96impl Complex {
97    fn rotation(angle: f64) -> Self {
98        let (im, re) = angle.sin_cos();
99        Self { re, im }
100    }
101
102    fn mul(self, other: Self) -> Self {
103        Self {
104            re: self.re * other.re - self.im * other.im,
105            im: self.re * other.im + self.im * other.re,
106        }
107    }
108}
109
110struct FftPlan {
111    phase: Vec<Complex>,
112    twiddles: Vec<Complex>,
113    bit_reversed: Vec<usize>,
114}
115
116impl FftPlan {
117    fn new(len: usize) -> Self {
118        let fft_len = len / 2;
119        let phase_step = -PI / len as f64;
120        Self {
121            // Split the constant quarter-sample phase equally between the
122            // input and output rotations so they can share one table.
123            phase: (0..fft_len)
124                .map(|i| Complex::rotation(phase_step * (i as f64 + 0.125)))
125                .collect(),
126            twiddles: (0..fft_len / 2)
127                .map(|i| Complex::rotation(-2.0 * PI * i as f64 / fft_len as f64))
128                .collect(),
129            bit_reversed: (0..fft_len)
130                .map(|i| {
131                    if fft_len == 1 {
132                        0
133                    } else {
134                        i.reverse_bits() >> (usize::BITS - fft_len.trailing_zeros())
135                    }
136                })
137                .collect(),
138        }
139    }
140
141    fn imdct_half(&self, dst: &mut [f32], src: &[f32], scale: f64) {
142        let fft_len = self.phase.len();
143        let len = 2 * fft_len;
144        let mut work = vec![Complex::default(); fft_len];
145
146        // The required half IMDCT is the reversed DCT-IV:
147        // D[k] = sum_j x[j] cos(pi/N * (j + 1/2) * (k + 1/2)).
148        // Pack z[j] = x[2j] + i*x[N-1-2j], j = 0..N/2, and rotate
149        // by exp(-i*pi/N * (j + 1/8)). After an N/2-point forward FFT
150        // and the same rotation at index k, the real part is D[2k]
151        // and the imaginary part is -D[N-1-2k]. Both components are
152        // useful; no zero padding or redundant FFT outputs are needed.
153        for j in 0..fft_len {
154            work[self.bit_reversed[j]] = Complex {
155                re: src[2 * j] as f64,
156                im: src[len - 1 - 2 * j] as f64,
157            }
158            .mul(self.phase[j]);
159        }
160
161        let mut width = 2;
162        while width <= fft_len {
163            let half = width / 2;
164            let stride = fft_len / width;
165            for block in work.chunks_exact_mut(width) {
166                for j in 0..half {
167                    let even = block[j];
168                    let odd = block[j + half].mul(self.twiddles[j * stride]);
169                    block[j] = Complex {
170                        re: even.re + odd.re,
171                        im: even.im + odd.im,
172                    };
173                    block[j + half] = Complex {
174                        re: even.re - odd.re,
175                        im: even.im - odd.im,
176                    };
177                }
178            }
179            width *= 2;
180        }
181
182        for k in 0..fft_len {
183            let value = work[k].mul(self.phase[k]);
184            dst[len - 1 - 2 * k] = (value.re * scale) as f32;
185            dst[2 * k] = (-value.im * scale) as f32;
186        }
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::MdctNaive;
193
194    #[test]
195    fn packed_imdct_preserves_each_coefficient_and_scale() {
196        // Exercise both members of every complex pair, including the N=2
197        // case whose one-point FFT has no butterfly stages.
198        for len in [2, 4, 8, 16, 32] {
199            for coefficient in 0..len {
200                let mut input = vec![0.0; len];
201                input[coefficient] = 1.0;
202                for scale in [1.0, -0.25, 0.0] {
203                    let transform = MdctNaive::new(len, scale);
204                    let mut expected = vec![0.0; len];
205                    transform.imdct_half_naive(&mut expected, &input);
206                    let mut actual = vec![123.0; len + 3];
207                    transform.imdct_half(&mut actual, &input);
208                    for i in 0..len {
209                        assert!(
210                            (actual[i] - expected[i]).abs() < 1e-7,
211                            "N={len} coefficient={coefficient} scale={scale} i={i}"
212                        );
213                    }
214                    assert_eq!(&actual[len..], &[123.0; 3]);
215                }
216            }
217        }
218    }
219
220    #[test]
221    fn fft_matches_reference_half_and_full_imdct() {
222        // Include every transform size used by WMAv1/v2 and WMA Pro.
223        for bits in 1..=13 {
224            let len = 1 << bits;
225            let mut seed = 1234567u32;
226            let noise: Vec<f32> = (0..len)
227                .map(|_| {
228                    seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
229                    (seed as i32) as f32 / i32::MAX as f32
230                })
231                .collect();
232            let mut impulse = vec![0.0; len];
233            impulse[len / 3] = 1.0;
234            for input in [noise, impulse, vec![0.0; len]] {
235                for scale in [1.0, 1.0 / 32768.0] {
236                    let transform = MdctNaive::new(len, scale);
237                    let mut expected = vec![0.0; len];
238                    transform.imdct_half_naive(&mut expected, &input);
239                    let mut half = vec![0.0; len];
240                    transform.imdct_half(&mut half, &input);
241                    let mut full = vec![0.0; 2 * len];
242                    transform.imdct_full(&mut full, &input);
243                    for (i, (&actual, &reference)) in half.iter().zip(&expected).enumerate() {
244                        let tolerance = (2e-6 * reference.abs()).max(2e-7 * scale as f32);
245                        assert!(
246                            (actual - reference).abs() <= tolerance,
247                            "N={len} i={i} scale={scale}: {actual} != {reference}"
248                        );
249                        assert_eq!(full[len / 2 + i], actual);
250                    }
251                    for i in 0..len / 2 {
252                        assert_eq!(full[i], -half[len / 2 - i - 1]);
253                        assert_eq!(full[2 * len - i - 1], half[len / 2 + i]);
254                    }
255                }
256            }
257        }
258    }
259}
wmv_decoder/wma/mdct.rs

wmv_decoder/wma/
mdct.rs
