na_mpeg2_decoder/demux/
mod.rs

1//! Minimal MPEG demux helpers.
2//!
3//! This module is intentionally small: it only extracts *video elementary stream*
4//! payload bytes (and optional PTS) from common MPEG containers.
5//!
6//! Supported:
7//! - Raw ES (start-code byte stream)
8//! - MPEG-TS (188-byte packets; video PID auto-sniffed from PES stream_id 0xE0..0xEF)
9//! - MPEG-PS (pack/system headers + PES; extracts video PES payload)
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum StreamType {
13    MpegVideo,
14    MpegAudio,
15    Unknown,
16}
17
18impl Default for StreamType {
19    #[inline]
20    fn default() -> Self {
21        StreamType::Unknown
22    }
23}
24
25#[derive(Clone, Debug)]
26pub struct Packet {
27    pub stream_type: StreamType,
28    pub pts_90k: Option<i64>,
29    pub data: Vec<u8>,
30}
31
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33enum ContainerKind {
34    Auto,
35    Es,
36    MpegTs,
37    MpegPs,
38}
39
40impl Default for ContainerKind {
41    #[inline]
42    fn default() -> Self {
43        ContainerKind::Auto
44    }
45}
46
47#[derive(Debug, Default)]
48pub struct Demuxer {
49    kind: ContainerKind,
50    stream_type: StreamType,
51
52    buf: Vec<u8>,
53
54    // TS state
55    ts_video_pid: Option<u16>,
56    ts_audio_pid: Option<u16>,
57}
58
59impl Demuxer {
60    /// Create an auto-detecting demuxer.
61    pub fn new_auto() -> Self {
62        Self {
63            kind: ContainerKind::Auto,
64            stream_type: StreamType::MpegVideo,
65            buf: Vec::new(),
66            ts_video_pid: None,
67            ts_audio_pid: None,
68        }
69    }
70
71    /// Create a demuxer with explicit container kind.
72    pub fn new(kind: StreamType) -> Self {
73        Self {
74            kind: ContainerKind::Es,
75            stream_type: kind,
76            buf: Vec::new(),
77            ts_video_pid: None,
78            ts_audio_pid: None,
79        }
80    }
81
82    /// Feed bytes and return extracted video ES chunks.
83    pub fn push(&mut self, data: &[u8], pts_90k: Option<i64>) -> Vec<Packet> {
84        let mut out = Vec::new();
85        self.push_into(data, pts_90k, &mut out);
86        out
87    }
88
89    /// Feed bytes and append extracted video ES chunks into `out`.
90    ///
91    /// This is the preferred API for memory-sensitive callers because it allows
92    /// reusing `out` capacity across calls.
93    pub fn push_into(&mut self, data: &[u8], pts_90k: Option<i64>, out: &mut Vec<Packet>) {
94        self.buf.extend_from_slice(data);
95
96        if self.kind == ContainerKind::Auto {
97            self.kind = detect_kind(&self.buf);
98        }
99
100        match self.kind {
101            ContainerKind::Es => {
102                if !self.buf.is_empty() {
103                    out.push(Packet {
104                        stream_type: self.stream_type,
105                        pts_90k,
106                        data: std::mem::take(&mut self.buf),
107                    });
108                }
109            }
110            ContainerKind::MpegTs => self.push_ts_into(out),
111            ContainerKind::MpegPs => self.push_ps_into(out),
112            ContainerKind::Auto => {}
113        }
114    }
115
116    fn push_ts_into(&mut self, out: &mut Vec<Packet>) {
117        const TS_SIZE: usize = 188;
118
119        // Try to resync to TS packet boundary.
120        while self.buf.len() >= TS_SIZE {
121            if self.buf[0] != 0x47 {
122                if let Some(pos) = self.buf.iter().position(|&b| b == 0x47) {
123                    self.buf.drain(0..pos);
124                } else {
125                    self.buf.clear();
126                    break;
127                }
128                if self.buf.len() < TS_SIZE {
129                    break;
130                }
131            }
132
133            // Parse directly from the buffered TS packet without allocating.
134            // Important: always consume exactly one TS packet per iteration.
135            if self.buf[0] != 0x47 {
136                self.buf.drain(0..1);
137                continue;
138            }
139
140            let pkt = &self.buf[..TS_SIZE];
141
142            let pusi = (pkt[1] & 0x40) != 0;
143            let pid: u16 = (((pkt[1] & 0x1F) as u16) << 8) | (pkt[2] as u16);
144            let afc = (pkt[3] >> 4) & 0x3;
145
146            let mut idx = 4usize;
147            if afc == 2 || afc == 3 {
148                if idx >= TS_SIZE {
149                    self.buf.drain(0..TS_SIZE);
150                    continue;
151                }
152                let afl = pkt[idx] as usize;
153                idx += 1 + afl;
154                if idx > TS_SIZE {
155                    self.buf.drain(0..TS_SIZE);
156                    continue;
157                }
158            }
159
160            if afc == 0 || afc == 2 {
161                self.buf.drain(0..TS_SIZE);
162                continue; // no payload
163            }
164            if idx >= TS_SIZE {
165                self.buf.drain(0..TS_SIZE);
166                continue;
167            }
168            let payload = &pkt[idx..TS_SIZE];
169            if payload.is_empty() {
170                self.buf.drain(0..TS_SIZE);
171                continue;
172            }
173
174            // Auto-sniff video/audio PID from PES headers.
175            if pusi && (self.ts_video_pid.is_none() || self.ts_audio_pid.is_none()) {
176                if payload.len() >= 4 && payload[0] == 0 && payload[1] == 0 && payload[2] == 1 {
177                    let sid = payload[3];
178                    if self.ts_video_pid.is_none() && (0xE0..=0xEF).contains(&sid) {
179                        self.ts_video_pid = Some(pid);
180                    }
181                    if self.ts_audio_pid.is_none() && (0xC0..=0xDF).contains(&sid) {
182                        self.ts_audio_pid = Some(pid);
183                    }
184                }
185            }
186
187            let mut st: Option<StreamType> = None;
188            if let Some(vpid) = self.ts_video_pid {
189                if pid == vpid {
190                    st = Some(StreamType::MpegVideo);
191                }
192            }
193            if let Some(apid) = self.ts_audio_pid {
194                if pid == apid {
195                    st = Some(StreamType::MpegAudio);
196                }
197            }
198            let Some(stream_type) = st else {
199                self.buf.drain(0..TS_SIZE);
200                continue;
201            };
202
203            if pusi {
204                // PES start.
205                if let Some((pts, off)) = parse_pes_header(payload) {
206                    if off <= payload.len() {
207                        let es = &payload[off..];
208                        if !es.is_empty() {
209                            out.push(Packet {
210                                stream_type,
211                                pts_90k: pts,
212                                data: es.to_vec(),
213                            });
214                        }
215                    }
216                } else {
217                    // No PES header; forward payload.
218                    out.push(Packet {
219                        stream_type,
220                        pts_90k: None,
221                        data: payload.to_vec(),
222                    });
223                }
224            } else {
225                // PES continuation: payload is pure ES bytes.
226                out.push(Packet {
227                    stream_type,
228                    pts_90k: None,
229                    data: payload.to_vec(),
230                });
231            }
232
233            // Drop processed TS packet bytes.
234            self.buf.drain(0..TS_SIZE);
235        }
236    }
237
238    fn push_ps_into(&mut self, out: &mut Vec<Packet>) {
239        // Scan for PES start codes; keep the last partial chunk.
240        let mut pos = 0usize;
241        while let Some((sc_pos, sid)) = find_start_code(&self.buf, pos) {
242            if sc_pos + 4 > self.buf.len() {
243                break;
244            }
245            let stream_type = if (0xE0..=0xEF).contains(&sid) {
246                StreamType::MpegVideo
247            } else if (0xC0..=0xDF).contains(&sid) {
248                StreamType::MpegAudio
249            } else {
250                pos = sc_pos + 4;
251                continue;
252            };
253            if sc_pos + 6 > self.buf.len() {
254                break;
255            }
256
257            let pes_len = u16::from_be_bytes([self.buf[sc_pos + 4], self.buf[sc_pos + 5]]) as usize;
258            let pes_end = if pes_len != 0 {
259                sc_pos + 6 + pes_len
260            } else {
261                // Unbounded PES: end at the next *system-layer* start code.
262                // Important: video ES itself contains 00 00 01 xx start codes
263                // (e.g., 0x00, 0xB3, 0xB5, 0x01..0xAF). We must not cut on those.
264                let mut search = sc_pos + 6;
265                let mut end_opt: Option<usize> = None;
266                while let Some((next_sc, next_id)) = find_start_code(&self.buf, search) {
267                    // System / PES start codes are >= 0xB9 in program streams.
268                    if next_id >= 0xB9 {
269                        end_opt = Some(next_sc);
270                        break;
271                    }
272                    search = next_sc + 4;
273                }
274                let Some(end_pos) = end_opt else {
275                    break;
276                };
277                end_pos
278            };
279            if pes_end > self.buf.len() {
280                break;
281            }
282
283            let pes = &self.buf[sc_pos..pes_end];
284            if let Some((pts, off)) = parse_pes_header(pes) {
285                if off < pes.len() {
286                    out.push(Packet {
287                        stream_type,
288                        pts_90k: pts,
289                        data: pes[off..].to_vec(),
290                    });
291                }
292            } else {
293                // Could not parse; forward raw PES bytes after start code.
294                out.push(Packet {
295                    stream_type,
296                    pts_90k: None,
297                    data: pes[4..].to_vec(),
298                });
299            }
300
301            pos = pes_end;
302        }
303
304        // Keep tail for next push.
305        if pos > 0 {
306            self.buf.drain(0..pos);
307        }
308    }
309}
310
311fn detect_kind(buf: &[u8]) -> ContainerKind {
312    // TS: sync byte 0x47 with 188-byte periodicity.
313    if buf.len() >= 188 * 3 {
314        if buf[0] == 0x47 && buf[188] == 0x47 && buf[376] == 0x47 {
315            return ContainerKind::MpegTs;
316        }
317    }
318    // PS: pack start code 00 00 01 BA.
319    if buf
320        .windows(4)
321        .take(4096)
322        .any(|w| w == [0x00, 0x00, 0x01, 0xBA])
323    {
324        return ContainerKind::MpegPs;
325    }
326    ContainerKind::Es
327}
328
329fn find_start_code(buf: &[u8], from: usize) -> Option<(usize, u8)> {
330    if buf.len() < 4 {
331        return None;
332    }
333    let mut i = from;
334    while i + 3 < buf.len() {
335        if buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1 {
336            return Some((i, buf[i + 3]));
337        }
338        i += 1;
339    }
340    None
341}
342
343/// Parse PES header and return (PTS, payload_offset).
344/// The returned offset is relative to the provided `buf`.
345fn parse_pes_header(buf: &[u8]) -> Option<(Option<i64>, usize)> {
346    if buf.len() < 9 {
347        return None;
348    }
349    if !(buf[0] == 0 && buf[1] == 0 && buf[2] == 1) {
350        return None;
351    }
352    let _sid = buf[3];
353    // buf[4..6] is PES_packet_length.
354
355    // Prefer MPEG-2 PES header syntax: '10' in buf[6] bits 7..6.
356    if (buf[6] & 0xC0) == 0x80 {
357        let flags = buf[7];
358        let hdr_len = buf[8] as usize;
359        let hdr_start = 9usize;
360        let payload_off = hdr_start + hdr_len;
361        if payload_off > buf.len() {
362            return None;
363        }
364        let pts_dts = (flags >> 6) & 0x3;
365        let mut pts: Option<i64> = None;
366        if (pts_dts == 2 || pts_dts == 3) && hdr_len >= 5 && hdr_start + 5 <= buf.len() {
367            pts = Some(parse_pts_90k(&buf[hdr_start..hdr_start + 5]));
368        }
369        return Some((pts, payload_off));
370    }
371
372    // MPEG-1 PES: skip stuffing and parse optional PTS.
373    // Reference: ISO/IEC 11172-1.
374    let mut idx = 6usize;
375    while idx < buf.len() && buf[idx] == 0xFF {
376        idx += 1;
377    }
378    if idx + 1 < buf.len() && (buf[idx] & 0xC0) == 0x40 {
379        idx += 2; // STD_buffer_scale/size
380    }
381    if idx >= buf.len() {
382        return None;
383    }
384    let mut pts: Option<i64> = None;
385    if (buf[idx] & 0xF0) == 0x20 {
386        // PTS only
387        if idx + 5 <= buf.len() {
388            pts = Some(parse_pts_90k(&buf[idx..idx + 5]));
389            idx += 5;
390        }
391    } else if (buf[idx] & 0xF0) == 0x30 {
392        // PTS + DTS, ignore DTS
393        if idx + 10 <= buf.len() {
394            pts = Some(parse_pts_90k(&buf[idx..idx + 5]));
395            idx += 10;
396        }
397    } else if buf[idx] == 0x0F {
398        idx += 1; // no pts
399    }
400    Some((pts, idx))
401}
402
403fn parse_pts_90k(p: &[u8]) -> i64 {
404    // p must be 5 bytes.
405    if p.len() < 5 {
406        return 0;
407    }
408    let pts = (((p[0] & 0x0E) as i64) << 29)
409        | ((p[1] as i64) << 22)
410        | (((p[2] & 0xFE) as i64) << 14)
411        | ((p[3] as i64) << 7)
412        | (((p[4] & 0xFE) as i64) >> 1);
413    pts
414}
na_mpeg2_decoder/demux/mod.rs

na_mpeg2_decoder/demux/
mod.rs
