ts_fix/
pes.rs

1//! PES access-unit reconstruction from TS payloads — framing only, no codec
2//! bitstream parsing (ISO/IEC 13818-1 §2.4.3.6).
3//!
4//! The public entry point is [`reconstruct_access_units`]. Each PUSI-delimited
5//! PES packet on the requested PIDs becomes one [`AccessUnit`] carrying the
6//! reassembled PES bytes and any PTS/DTS from the PES header.
7//!
8//! # Spec
9//!
10//! ISO/IEC 13818-1 (= ITU-T H.222.0) — §2.4.3.6 (PES packet), §2.4.3.7 (PES
11//! header / PTS/DTS).
12
13use alloc::vec::Vec;
14
15/// A reassembled PES access unit — the complete PES packet bytes with optional
16/// timing from the PES header.
17///
18/// The `data` field holds the **full** PES packet (from `0x00 0x00 0x01`
19/// `packet_start_code_prefix` through the last `PES_packet_data_byte`). No codec
20/// parsing is performed; the bytes are opaque.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct AccessUnit {
23    /// PID this access unit was carried on.
24    pub pid: u16,
25    /// Presentation time stamp from the PES header, if present (33-bit, 90 kHz).
26    pub pts: Option<u64>,
27    /// Decoding time stamp from the PES header, if present (33-bit, 90 kHz).
28    pub dts: Option<u64>,
29    /// The reassembled PES packet bytes (`00 00 01 stream_id ...`).
30    pub data: Vec<u8>,
31}
32
33/// Reconstruct PES access units on the given PIDs from a contiguous TS buffer.
34///
35/// Iterates 188-byte TS packets, reassembles PES payloads per PID via
36/// [`mpeg_pes::PesAssembler`], and parses each completed PES packet to extract
37/// PTS/DTS. Access units are returned in arrival (completion) order; AUs from
38/// different PIDs may interleave.
39///
40/// # Panics
41///
42/// Panics if `ts` is not a multiple of 188 bytes (caller must pre-chop to
43/// packet boundaries).
44pub fn reconstruct_access_units(ts: &[u8], pids: &[u16]) -> Vec<AccessUnit> {
45    const TS_PACKET_SIZE: usize = 188;
46
47    assert_eq!(
48        ts.len() % TS_PACKET_SIZE,
49        0,
50        "ts buffer length {} is not a multiple of 188",
51        ts.len()
52    );
53
54    // Build a set for O(1) PID membership checks.
55    let pid_set = {
56        let mut set = alloc::collections::BTreeSet::new();
57        for &pid in pids {
58            set.insert(pid);
59        }
60        set
61    };
62
63    // Per-PID assemblers.
64    let mut assemblers: alloc::collections::BTreeMap<u16, mpeg_pes::PesAssembler> =
65        alloc::collections::BTreeMap::new();
66
67    // Completed AUs, in order.
68    let mut result: Vec<AccessUnit> = Vec::new();
69
70    for chunk in ts.chunks(TS_PACKET_SIZE) {
71        let raw: [u8; TS_PACKET_SIZE] = match chunk.try_into() {
72            Ok(a) => a,
73            Err(_) => continue, // not 188 bytes; should not happen due to the assert above
74        };
75
76        let pkt = match mpeg_ts::OwnedTsPacket::parse(raw) {
77            Ok(p) => p,
78            Err(_) => continue,
79        };
80
81        if !pid_set.contains(&pkt.pid) {
82            continue;
83        }
84
85        let payload = match pkt.payload() {
86            Some(p) => p,
87            None => continue,
88        };
89
90        let asm = assemblers.entry(pkt.pid).or_default();
91
92        if let Some(completed) = asm.feed(pkt.pusi, payload) {
93            let au = parse_au(pkt.pid, completed);
94            result.push(au);
95        }
96    }
97
98    // Flush any remaining PES on each PID.
99    for (&pid, asm) in assemblers.iter_mut() {
100        if let Some(completed) = asm.flush() {
101            let au = parse_au(pid, completed);
102            result.push(au);
103        }
104    }
105
106    result
107}
108
109/// Parse a completed PES packet `Vec<u8>` into an `AccessUnit`.
110fn parse_au(pid: u16, data: Vec<u8>) -> AccessUnit {
111    let (pts, dts) = match mpeg_pes::PesPacket::parse(&data) {
112        Ok(pkt) => {
113            let pts = pkt.header.as_ref().and_then(|h| h.pts).map(|p| p.ticks());
114            let dts = pkt.header.as_ref().and_then(|h| h.dts).map(|d| d.ticks());
115            (pts, dts)
116        }
117        Err(_) => (None, None),
118    };
119
120    AccessUnit {
121        pid,
122        pts,
123        dts,
124        data,
125    }
126}
ts_fix/pes.rs

ts_fix/
pes.rs
