Skip to main content

pito_codecs/mp3/
mod.rs

1#![allow(clippy::needless_range_loop, clippy::excessive_precision)]
2
3mod bits;
4mod hdr;
5mod layer12;
6mod layer3;
7mod simd;
8mod synth;
9mod tables;
10
11#[cfg(test)]
12mod malformed;
13#[cfg(test)]
14mod tests;
15
16use bits::Bits;
17use hdr::{HDR_SIZE, Head};
18use layer3::GrInfo;
19
20use crate::media::{Error, Result};
21
22pub const MAX_SAMPLES_PER_FRAME: usize = 1152 * 2;
23
24const MAX_FREE_FORMAT_FRAME_SIZE: usize = 2304;
25const MAX_FRAME_SYNC_MATCHES: usize = 10;
26const MAX_L3_FRAME_PAYLOAD_BYTES: usize = MAX_FREE_FORMAT_FRAME_SIZE;
27const MAX_BITRESERVOIR_BYTES: usize = 511;
28
29#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
30pub struct FrameInfo {
31    pub frame_bytes: usize,
32    pub frame_offset: usize,
33    pub channels: u16,
34    pub rate: u32,
35    pub layer: u8,
36    pub bitrate_kbps: u32,
37    pub samples: usize,
38}
39
40struct State {
41    mdct_overlap: [[f32; 288]; 2],
42    qmf_state: [f32; 960],
43    reserv: usize,
44    free_format_bytes: usize,
45    header: Head,
46    reserv_buf: [u8; MAX_BITRESERVOIR_BYTES],
47}
48
49struct Scratch {
50    maindata: [u8; MAX_BITRESERVOIR_BYTES + MAX_L3_FRAME_PAYLOAD_BYTES],
51    gr_info: [GrInfo; 4],
52    grbuf: [f32; 1156],
53    scf: [f32; 40],
54    syn: [f32; 2112],
55    ist_pos: [[u8; 39]; 2],
56    sci: layer12::ScaleInfo,
57}
58
59pub struct Decoder {
60    st: Box<State>,
61    sc: Box<Scratch>,
62    pcm: Box<[f32; MAX_SAMPLES_PER_FRAME]>,
63    rate: u32,
64    channels: u16,
65}
66
67impl Default for Decoder {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73fn match_frame(b: &[u8], frame_bytes: usize) -> bool {
74    let Some(first) = Head::at(b, 0) else {
75        return false;
76    };
77    let mut i = 0usize;
78    for nmatch in 0..MAX_FRAME_SYNC_MATCHES {
79        let Some(h) = Head::at(b, i) else {
80            return nmatch > 0;
81        };
82        i += h.frame_bytes(frame_bytes) + h.padding();
83        let Some(next) = Head::at(b, i) else {
84            return nmatch > 0;
85        };
86        if !first.compare(next) {
87            return false;
88        }
89    }
90    true
91}
92
93fn find_frame(mp3: &[u8], free_format_bytes: &mut usize) -> (usize, usize) {
94    let n = mp3.len();
95    for i in 0..n.saturating_sub(HDR_SIZE) {
96        let Some(h) = Head::at(mp3, i) else {
97            break;
98        };
99        if !h.valid() {
100            continue;
101        }
102        let mut frame_bytes = h.frame_bytes(*free_format_bytes);
103        let mut frame_and_padding = frame_bytes + h.padding();
104        let mut k = HDR_SIZE;
105        while frame_bytes == 0 && k < MAX_FREE_FORMAT_FRAME_SIZE && i + 2 * k < n - HDR_SIZE {
106            if let Some(hk) = Head::at(mp3, i + k)
107                && h.compare(hk)
108            {
109                let fb = k.saturating_sub(h.padding());
110                let nextfb = fb + hk.padding();
111                if let Some(hn) = Head::at(mp3, i + k + nextfb)
112                    && h.compare(hn)
113                {
114                    frame_and_padding = k;
115                    frame_bytes = fb;
116                    *free_format_bytes = fb;
117                }
118            }
119            k += 1;
120        }
121        if (frame_bytes != 0 && i + frame_and_padding <= n && match_frame(&mp3[i..], frame_bytes))
122            || (i == 0 && frame_and_padding == n)
123        {
124            return (i, frame_and_padding);
125        }
126        *free_format_bytes = 0;
127    }
128    (n, 0)
129}
130
131impl State {
132    fn reset(&mut self) {
133        self.mdct_overlap = [[0.0; 288]; 2];
134        self.qmf_state = [0.0; 960];
135        self.reserv = 0;
136        self.free_format_bytes = 0;
137        self.header = Head::default();
138        self.reserv_buf = [0; MAX_BITRESERVOIR_BYTES];
139    }
140}
141
142fn decode_at(st: &mut State, sc: &mut Scratch, frame: &[u8], pcm: &mut [f32]) -> usize {
143    let Some(h) = Head::at(frame, 0) else {
144        return 0;
145    };
146    let nch = h.channels();
147    if pcm.len() < h.frame_samples() * nch {
148        return 0;
149    }
150    st.header = h;
151    let Scratch {
152        maindata,
153        gr_info,
154        grbuf,
155        scf,
156        syn,
157        ist_pos,
158        sci,
159    } = sc;
160    let mut bs = Bits::new(&frame[HDR_SIZE..]);
161    if h.crc() {
162        bs.get(16);
163    }
164    if h.layer() == 3 {
165        let main_data_begin = layer3::read_side_info(&mut bs, gr_info, h);
166        if main_data_begin < 0 || bs.pos > bs.limit {
167            st.header = Head::default();
168            return 0;
169        }
170        let main_data_begin = main_data_begin as usize;
171        let frame_bytes = (bs.limit - bs.pos) / 8;
172        let have = st.reserv.min(main_data_begin);
173        let start = bs.pos / 8;
174        if have + frame_bytes > maindata.len() {
175            st.header = Head::default();
176            return 0;
177        }
178        let from = st.reserv.saturating_sub(main_data_begin);
179        maindata[..have].copy_from_slice(&st.reserv_buf[from..from + have]);
180        maindata[have..have + frame_bytes].copy_from_slice(&bs.buf[start..start + frame_bytes]);
181        let success = st.reserv >= main_data_begin;
182        let mut w = layer3::Work {
183            bs: Bits::new(&maindata[..have + frame_bytes]),
184            grbuf,
185            scf,
186            syn,
187            ist_pos,
188        };
189        if success {
190            let granules = if h.mpeg1() { 2 } else { 1 };
191            for igr in 0..granules {
192                w.grbuf.fill(0.0);
193                layer3::decode(&mut w, h, &mut st.mdct_overlap, &gr_info[igr * nch..], nch);
194                synth::granule(
195                    &mut st.qmf_state,
196                    &mut w.grbuf[..],
197                    18,
198                    nch,
199                    &mut pcm[576 * nch * igr..],
200                    &mut w.syn[..],
201                );
202            }
203        }
204        let mut pos = w.bs.pos.div_ceil(8);
205        let limit = w.bs.limit / 8;
206        let mut remains = limit as isize - pos as isize;
207        if remains > MAX_BITRESERVOIR_BYTES as isize {
208            pos += remains as usize - MAX_BITRESERVOIR_BYTES;
209            remains = MAX_BITRESERVOIR_BYTES as isize;
210        }
211        if remains > 0 {
212            let r = remains as usize;
213            st.reserv_buf[..r].copy_from_slice(&maindata[pos..pos + r]);
214            st.reserv = r;
215        } else {
216            st.reserv = 0;
217        }
218        if success { h.frame_samples() } else { 0 }
219    } else {
220        layer12::read_scale_info(h, &mut bs, sci);
221        grbuf.fill(0.0);
222        let group = (h.layer() | 1) as usize;
223        let mut i = 0;
224        let mut out = 0;
225        for igr in 0..3 {
226            i += layer12::dequantize_granule(&mut grbuf[..], i, &mut bs, sci, group);
227            if i == 12 {
228                i = 0;
229                layer12::apply_scf_384(sci, igr, &mut grbuf[..]);
230                synth::granule(
231                    &mut st.qmf_state,
232                    &mut grbuf[..],
233                    12,
234                    nch,
235                    &mut pcm[out..],
236                    &mut syn[..],
237                );
238                grbuf.fill(0.0);
239                out += 384 * nch;
240            }
241            if bs.pos > bs.limit {
242                st.header = Head::default();
243                return 0;
244            }
245        }
246        h.frame_samples()
247    }
248}
249
250impl Decoder {
251    pub fn new() -> Decoder {
252        Decoder {
253            st: Box::new(State {
254                mdct_overlap: [[0.0; 288]; 2],
255                qmf_state: [0.0; 960],
256                reserv: 0,
257                free_format_bytes: 0,
258                header: Head::default(),
259                reserv_buf: [0; MAX_BITRESERVOIR_BYTES],
260            }),
261            sc: Box::new(Scratch {
262                maindata: [0; MAX_BITRESERVOIR_BYTES + MAX_L3_FRAME_PAYLOAD_BYTES],
263                gr_info: [GrInfo::default(); 4],
264                grbuf: [0.0; 1156],
265                scf: [0.0; 40],
266                syn: [0.0; 2112],
267                ist_pos: [[0; 39]; 2],
268                sci: layer12::ScaleInfo::default(),
269            }),
270            pcm: Box::new([0.0; MAX_SAMPLES_PER_FRAME]),
271            rate: 0,
272            channels: 0,
273        }
274    }
275
276    pub fn with_format(rate: u32, channels: u16) -> Decoder {
277        let mut d = Decoder::new();
278        d.rate = rate;
279        d.channels = channels.clamp(1, 2);
280        d
281    }
282
283    pub fn rate(&self) -> u32 {
284        self.rate
285    }
286
287    pub fn channels(&self) -> u16 {
288        self.channels.max(1)
289    }
290
291    pub fn reset(&mut self) {
292        self.st.reset();
293    }
294
295    pub fn decode(&mut self, mp3: &[u8], pcm: &mut [f32]) -> FrameInfo {
296        let n = mp3.len();
297        let mut i = 0;
298        let mut frame_size = 0;
299        if n > HDR_SIZE
300            && self.st.header.0[0] == 0xFF
301            && let Some(h) = Head::at(mp3, 0)
302            && self.st.header.compare(h)
303        {
304            frame_size = h.frame_bytes(self.st.free_format_bytes) + h.padding();
305            if frame_size != n && Head::at(mp3, frame_size).is_none_or(|next| !h.compare(next)) {
306                frame_size = 0;
307            }
308        }
309        if frame_size == 0 {
310            self.st.reset();
311            (i, frame_size) = find_frame(mp3, &mut self.st.free_format_bytes);
312            if frame_size == 0 || i + frame_size > n {
313                return FrameInfo {
314                    frame_bytes: i,
315                    ..FrameInfo::default()
316                };
317            }
318        }
319        let Some(h) = Head::at(mp3, i) else {
320            return FrameInfo {
321                frame_bytes: n,
322                ..FrameInfo::default()
323            };
324        };
325        let samples = decode_at(&mut self.st, &mut self.sc, &mp3[i..i + frame_size], pcm);
326        FrameInfo {
327            frame_bytes: i + frame_size,
328            frame_offset: i,
329            channels: h.channels() as u16,
330            rate: h.sample_rate(),
331            layer: h.layer(),
332            bitrate_kbps: h.bitrate_kbps(),
333            samples,
334        }
335    }
336
337    pub fn decode_packet(&mut self, data: &[u8], out: &mut Vec<f32>) -> Result<()> {
338        out.clear();
339        let mut p = 0;
340        let mut frames = 0;
341        while p + HDR_SIZE <= data.len() {
342            let rest = &data[p..];
343            let Some(h) = Head::at(rest, 0) else {
344                break;
345            };
346            let size = if h.valid() {
347                let fb = h.frame_bytes(0);
348                if fb == 0 {
349                    rest.len()
350                } else {
351                    fb + h.padding()
352                }
353            } else {
354                0
355            };
356            let (at, size) = if size != 0 && size <= rest.len() {
357                (0, size)
358            } else {
359                let mut free = self.st.free_format_bytes;
360                let (i, size) = find_frame(rest, &mut free);
361                if size == 0 || i + size > rest.len() {
362                    break;
363                }
364                (i, size)
365            };
366            let frame = &rest[at..at + size];
367            let Some(h) = Head::at(frame, 0) else {
368                break;
369            };
370            if !self.st.header.compare(h) {
371                self.st.reset();
372            }
373            if self.rate == 0 {
374                self.rate = h.sample_rate();
375            }
376            if self.channels == 0 {
377                self.channels = h.channels() as u16;
378            }
379            let n = decode_at(&mut self.st, &mut self.sc, frame, &mut self.pcm[..]);
380            push(&self.pcm[..], n, h, self.channels() as usize, out);
381            p += at + size;
382            frames += 1;
383        }
384        if frames == 0 {
385            return Err(Error::Invalid(format!(
386                "an MP3 packet of {} bytes without a frame",
387                data.len()
388            )));
389        }
390        Ok(())
391    }
392}
393
394fn push(pcm: &[f32], n: usize, h: Head, want: usize, out: &mut Vec<f32>) {
395    let have = h.channels();
396    if n == 0 {
397        out.resize(out.len() + h.frame_samples() * want, 0.0);
398        return;
399    }
400    let s = &pcm[..(n * have).min(pcm.len())];
401    match (have, want) {
402        (1, 2) => {
403            out.reserve(s.len() * 2);
404            for &x in s {
405                out.push(x);
406                out.push(x);
407            }
408        }
409        (2, 1) => out.extend(s.chunks_exact(2).map(|c| 0.5 * (c[0] + c[1]))),
410        _ => out.extend_from_slice(s),
411    }
412}