na_mpeg2_decoder/audio/
mpa.rs

1use crate::error::{AvError, Result};
2
3use symphonia::core::audio::SampleBuffer;
4use symphonia::core::codecs::{CodecParameters, Decoder, DecoderOptions, CODEC_TYPE_MP1, CODEC_TYPE_MP2, CODEC_TYPE_MP3};
5use symphonia::core::errors::Error as SymphError;
6use symphonia::core::formats::Packet;
7
8#[derive(Clone)]
9pub struct MpaAudioChunk {
10    pub pts_ms: i64,
11    pub sample_rate: u32,
12    pub channels: u16,
13    pub samples: Vec<f32>,
14}
15
16#[derive(Default)]
17pub struct MpaAudioDecoder {
18    buf: Vec<u8>,
19
20    dec: Option<Box<dyn Decoder>>,
21    sample_buf: Option<SampleBuffer<f32>>,
22
23    // Best-effort PTS tracking.
24    next_pts_ms: Option<i64>,
25
26    // Symphonia packet track id (arbitrary but must be consistent).
27    track_id: u32,
28}
29
30impl MpaAudioDecoder {
31    pub fn new() -> Self {
32        Self { buf: Vec::new(), dec: None, sample_buf: None, next_pts_ms: None, track_id: 0 }
33    }
34
35    pub fn push_with<F>(&mut self, data: &[u8], pts_ms: Option<i64>, mut on_chunk: F) -> Result<()>
36    where
37        F: FnMut(MpaAudioChunk),
38    {
39        if let Some(pts) = pts_ms {
40            self.next_pts_ms = Some(pts);
41        }
42
43        self.buf.extend_from_slice(data);
44
45        let mut pos = 0usize;
46        while pos + 4 <= self.buf.len() {
47            let Some(h) = MpaHeader::parse(&self.buf[pos..]) else {
48                pos += 1;
49                continue;
50            };
51
52            if pos + h.frame_len > self.buf.len() {
53                break;
54            }
55
56            // Avoid borrowing self.buf while calling into self (decoder state).
57            let pkt_owned = self.buf[pos..pos + h.frame_len].to_vec();
58            pos += h.frame_len;
59
60            let pts0 = self.next_pts_ms.unwrap_or(0);
61            self.decode_one_packet(&pkt_owned, pts0, h.codec_type, &mut on_chunk)?;
62        }
63
64        if pos > 0 {
65            self.buf.drain(0..pos);
66        }
67
68        Ok(())
69    }
70
71    fn decode_one_packet<F>(&mut self, pkt_bytes: &[u8], pts_ms: i64, codec_type: symphonia::core::codecs::CodecType, on_chunk: &mut F) -> Result<()>
72    where
73        F: FnMut(MpaAudioChunk),
74    {
75        if self.dec.is_none() {
76            let mut cp = CodecParameters::new();
77            cp.for_codec(codec_type);
78
79            let dec = symphonia::default::get_codecs()
80                .make(&cp, &DecoderOptions::default())
81                .map_err(AvError::from)?;
82            self.dec = Some(dec);
83        }
84
85        let pkt = Packet::new_from_boxed_slice(self.track_id, 0, 0, pkt_bytes.to_vec().into_boxed_slice());
86
87        let dec = self.dec.as_mut().expect("decoder must be initialized");
88        match dec.decode(&pkt) {
89            Ok(decoded) => {
90                let spec = *decoded.spec();
91                let duration = decoded.capacity();
92                let duration_u64 = duration as u64;
93
94                let sb = match self.sample_buf.as_mut() {
95                    None => {
96                        self.sample_buf = Some(SampleBuffer::<f32>::new(duration_u64, spec));
97                        self.sample_buf.as_mut().unwrap()
98                    }
99                    Some(sb) => {
100                        if sb.capacity() < duration {
101                            *sb = SampleBuffer::<f32>::new(duration_u64, spec);
102                        }
103                        sb
104                    }
105                };
106
107                sb.copy_interleaved_ref(decoded.clone());
108
109                let channels = spec.channels.count() as u16;
110                let samples = sb.samples().to_vec();
111
112                let sample_rate = spec.rate;
113                on_chunk(MpaAudioChunk { pts_ms, sample_rate, channels, samples });
114
115                // Advance PTS based on decoded frames.
116                let frames = decoded.frames() as i64;
117                if frames > 0 && sample_rate > 0 {
118                    let dur_ms = (frames * 1000) / (sample_rate as i64);
119                    self.next_pts_ms = Some(pts_ms + dur_ms);
120                }
121            }
122            Err(SymphError::DecodeError(_)) => {
123                // Best-effort: ignore bad frames.
124            }
125            Err(e) => return Err(e.into()),
126        }
127
128        Ok(())
129    }
130}
131
132#[derive(Clone, Copy)]
133struct MpaHeader {
134    frame_len: usize,
135    codec_type: symphonia::core::codecs::CodecType,
136}
137
138impl MpaHeader {
139    fn parse(buf: &[u8]) -> Option<Self> {
140        if buf.len() < 4 {
141            return None;
142        }
143        let b0 = buf[0];
144        let b1 = buf[1];
145        let b2 = buf[2];
146
147        // Sync.
148        if b0 != 0xFF || (b1 & 0xE0) != 0xE0 {
149            return None;
150        }
151
152        let version_id = (b1 >> 3) & 0x03;
153        let layer_id = (b1 >> 1) & 0x03;
154        if version_id == 0x01 || layer_id == 0x00 {
155            return None;
156        }
157
158        let bitrate_idx = (b2 >> 4) & 0x0F;
159        let sr_idx = (b2 >> 2) & 0x03;
160        if bitrate_idx == 0 || bitrate_idx == 0x0F || sr_idx == 0x03 {
161            return None;
162        }
163
164        let padding: u32 = ((b2 >> 1) & 0x01) as u32;
165
166        let (sr, is_v1) = match version_id {
167            0x03 => (SAMPLE_RATES_V1[sr_idx as usize], true),
168            0x02 => (SAMPLE_RATES_V2[sr_idx as usize], false),
169            0x00 => (SAMPLE_RATES_V25[sr_idx as usize], false),
170            _ => return None,
171        };
172
173        let (codec_type, bitrate_kbps, frame_len) = match layer_id {
174            0x03 => {
175                // Layer I
176                let br = if is_v1 {
177                    BITRATES_V1_L1[bitrate_idx as usize]
178                } else {
179                    BITRATES_V2_L1[bitrate_idx as usize]
180                };
181                let fl = (((12u64 * (br as u64) * 1000u64) / (sr as u64)) + (padding as u64)) * 4u64;
182                (CODEC_TYPE_MP1, br, fl as usize)
183            }
184            0x02 => {
185                // Layer II
186                let br = if is_v1 {
187                    BITRATES_V1_L2[bitrate_idx as usize]
188                } else {
189                    BITRATES_V2_L2L3[bitrate_idx as usize]
190                };
191                let fl = ((144u64 * (br as u64) * 1000u64) / (sr as u64)) + (padding as u64);
192                (CODEC_TYPE_MP2, br, fl as usize)
193            }
194            0x01 => {
195                // Layer III
196                let br = if is_v1 {
197                    BITRATES_V1_L3[bitrate_idx as usize]
198                } else {
199                    BITRATES_V2_L2L3[bitrate_idx as usize]
200                };
201                let coeff: u64 = if is_v1 { 144 } else { 72 };
202                let fl = ((coeff * (br as u64) * 1000u64) / (sr as u64)) + (padding as u64);
203                (CODEC_TYPE_MP3, br, fl as usize)
204            }
205            _ => return None,
206        };
207
208        if bitrate_kbps == 0 || frame_len < 4 {
209            return None;
210        }
211
212        Some(Self { frame_len, codec_type })
213    }
214}
215
216const SAMPLE_RATES_V1: [u32; 3] = [44100, 48000, 32000];
217const SAMPLE_RATES_V2: [u32; 3] = [22050, 24000, 16000];
218const SAMPLE_RATES_V25: [u32; 3] = [11025, 12000, 8000];
219
220const BITRATES_V1_L1: [u32; 16] = [0, 32, 64, 96, 128, 160, 192, 224, 256, 288, 320, 352, 384, 416, 448, 0];
221const BITRATES_V1_L2: [u32; 16] = [0, 32, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 384, 0];
222const BITRATES_V1_L3: [u32; 16] = [0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0];
223
224const BITRATES_V2_L1: [u32; 16] = [0, 32, 48, 56, 64, 80, 96, 112, 128, 144, 160, 176, 192, 224, 256, 0];
225const BITRATES_V2_L2L3: [u32; 16] = [0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0];