Skip to main content

na_mpeg2_decoder/
av.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4#[cfg(feature = "audio")]
5use crate::audio::{MpaAudioChunk, MpaAudioDecoder};
6use crate::convert::frame_to_rgba_bt601_limited;
7use crate::demux::{Demuxer, Packet, StreamType};
8use crate::error::Result;
9use crate::video::{Decoder as VideoDecoder, Frame};
10
11#[derive(Clone)]
12pub struct MpegRgbaFrame {
13    pub pts_ms: i64,
14    pub width: u32,
15    pub height: u32,
16    pub rgba: Vec<u8>,
17}
18
19// Always present: the struct itself has no symphonia dependencies.
20// In video-only builds (no "audio" feature) this type is defined but the
21// MpegAvEvent::Audio variant is never emitted by MpegAvPipeline.
22#[derive(Clone)]
23pub struct MpegAudioF32 {
24    pub pts_ms: i64,
25    pub sample_rate: u32,
26    pub channels: u16,
27    pub samples: Vec<f32>,
28}
29
30#[derive(Clone)]
31pub enum MpegAvEvent {
32    Video(MpegRgbaFrame),
33    Audio(MpegAudioF32),
34}
35
36#[derive(Default)]
37pub struct MpegAvPipeline {
38    demux: Demuxer,
39    vdec: VideoDecoder,
40    #[cfg(feature = "audio")]
41    adec: MpaAudioDecoder,
42
43    pkts: Vec<Packet>,
44    pub stash: VecDeque<MpegAvEvent>,
45}
46
47impl MpegAvPipeline {
48    pub fn new() -> Self {
49        Self {
50            demux: Demuxer::new_auto(),
51            vdec: VideoDecoder::new(),
52            #[cfg(feature = "audio")]
53            adec: MpaAudioDecoder::new(),
54            pkts: Vec::new(),
55            stash: VecDeque::new(),
56        }
57    }
58
59    #[inline]
60    pub fn demuxer_mut(&mut self) -> &mut Demuxer {
61        &mut self.demux
62    }
63
64    #[inline]
65    pub fn video_decoder_mut(&mut self) -> &mut VideoDecoder {
66        &mut self.vdec
67    }
68
69    #[cfg(feature = "audio")]
70    #[inline]
71    pub fn audio_decoder_mut(&mut self) -> &mut MpaAudioDecoder {
72        &mut self.adec
73    }
74
75    pub fn push_with<F>(&mut self, data: &[u8], pts_90k: Option<i64>, mut on_event: F) -> Result<()>
76    where
77        F: FnMut(MpegAvEvent),
78    {
79        self.pkts.clear();
80        self.demux.push_into(data, pts_90k, &mut self.pkts);
81
82        // Move packets out to avoid borrowing self.pkts while calling &mut self handlers.
83        let mut local_pkts: Vec<Packet> = Vec::new();
84        std::mem::swap(&mut self.pkts, &mut local_pkts);
85
86        for pkt in local_pkts.drain(..) {
87            match pkt.stream_type {
88                StreamType::MpegVideo => self.handle_video_pkt(&pkt, &mut on_event)?,
89                StreamType::MpegAudio => {
90                    #[cfg(feature = "audio")]
91                    self.handle_audio_pkt(&pkt, &mut on_event)?;
92                }
93                StreamType::Unknown => {}
94            }
95        }
96
97        std::mem::swap(&mut self.pkts, &mut local_pkts);
98        self.pkts.clear();
99
100        Ok(())
101    }
102
103    pub fn push(&mut self, data: &[u8], pts_90k: Option<i64>) -> Result<()> {
104        let mut tmp: Vec<MpegAvEvent> = Vec::new();
105        self.push_with(data, pts_90k, |ev| tmp.push(ev))?;
106        for ev in tmp {
107            self.stash.push_back(ev);
108        }
109        Ok(())
110    }
111
112    pub fn flush_with<F>(&mut self, mut on_event: F) -> Result<()>
113    where
114        F: FnMut(MpegAvEvent),
115    {
116        // Video: flush delayed frames.
117        for f in self.vdec.flush_shared()? {
118            self.emit_video_frame(f, &mut on_event)?;
119        }
120        Ok(())
121    }
122
123    pub fn flush(&mut self) -> Result<()> {
124        let mut tmp: Vec<MpegAvEvent> = Vec::new();
125        self.flush_with(|ev| tmp.push(ev))?;
126        for ev in tmp {
127            self.stash.push_back(ev);
128        }
129        Ok(())
130    }
131
132    fn handle_video_pkt<F>(&mut self, pkt: &Packet, on_event: &mut F) -> Result<()>
133    where
134        F: FnMut(MpegAvEvent),
135    {
136        let decoded: Vec<Arc<Frame>> = self.vdec.decode_shared(&pkt.data, pkt.pts_90k)?;
137        for f in decoded {
138            self.emit_video_frame(f, on_event)?;
139        }
140        Ok(())
141    }
142
143    fn emit_video_frame<F>(&mut self, f: Arc<Frame>, on_event: &mut F) -> Result<()>
144    where
145        F: FnMut(MpegAvEvent),
146    {
147        let w = f.width as u32;
148        let h = f.height as u32;
149        let mut rgba = vec![0u8; (w as usize) * (h as usize) * 4];
150        frame_to_rgba_bt601_limited(&f, &mut rgba);
151
152        let pts_ms = pts90k_opt_to_ms(f.pts_90k);
153        on_event(MpegAvEvent::Video(MpegRgbaFrame {
154            pts_ms,
155            width: w,
156            height: h,
157            rgba,
158        }));
159        Ok(())
160    }
161
162    #[cfg(feature = "audio")]
163    fn handle_audio_pkt<F>(&mut self, pkt: &Packet, on_event: &mut F) -> Result<()>
164    where
165        F: FnMut(MpegAvEvent),
166    {
167        let pts_ms_opt = pkt.pts_90k.map(pts90k_to_ms);
168        self.adec
169            .push_with(&pkt.data, pts_ms_opt, |ch: MpaAudioChunk| {
170                on_event(MpegAvEvent::Audio(MpegAudioF32 {
171                    pts_ms: ch.pts_ms,
172                    sample_rate: ch.sample_rate,
173                    channels: ch.channels,
174                    samples: ch.samples,
175                }))
176            })?;
177        Ok(())
178    }
179}
180
181#[inline]
182fn pts90k_to_ms(v: i64) -> i64 {
183    (v * 1000) / 90000
184}
185
186#[inline]
187fn pts90k_opt_to_ms(v: Option<i64>) -> i64 {
188    v.map(pts90k_to_ms).unwrap_or(0)
189}