rtc_media/io/h26x_reader/
sample_reader.rs

1use bytes::{Bytes, BytesMut};
2use shared::error::{Error, Result};
3use std::io::Read;
4
5use super::{H26xNAL, H26xReader, H264NalUnitType, H265NalUnitType};
6
7const ANNEXB_START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10/// One access unit read from an Annex B stream — the NAL units making up a single frame.
11pub struct H26xSample {
12    /// The sample's bytes, start codes included.
13    pub data: Bytes,
14    /// Whether this sample advances presentation time, i.e. completes a frame.
15    pub timed: bool,
16}
17
18/// Reads whole samples from an H.264 or H.265 Annex B stream.
19pub struct H26xSampleReader<R: Read> {
20    reader: H26xReader<R>,
21    is_hevc: bool,
22    pending_hevc_nals: Vec<BytesMut>,
23}
24
25impl<R: Read> H26xSampleReader<R> {
26    /// Wraps `reader`, buffering up to `capacity` bytes.
27    ///
28    /// Set `is_hevc` for H.265; the two codecs differ in how NAL headers are parsed.
29    pub fn new(reader: R, capacity: usize, is_hevc: bool) -> Self {
30        Self {
31            reader: H26xReader::new(reader, capacity, is_hevc),
32            is_hevc,
33            pending_hevc_nals: vec![],
34        }
35    }
36
37    /// Reads the next sample.
38    ///
39    /// # Errors
40    ///
41    /// Fails on an I/O error, or at end of stream.
42    pub fn next_sample(&mut self) -> Result<H26xSample> {
43        loop {
44            let nal = match self.reader.next_nal() {
45                Ok(nal) => nal,
46                Err(Error::ErrIoEOF) if self.is_hevc && !self.pending_hevc_nals.is_empty() => {
47                    let data = build_hevc_access_unit(&mut self.pending_hevc_nals, None);
48                    return Ok(H26xSample { data, timed: false });
49                }
50                Err(err) => return Err(err),
51            };
52
53            let timed = !should_skip_timing(&nal);
54            if self.is_hevc && should_buffer_with_next_hevc_vcl(&nal) {
55                self.pending_hevc_nals.push(nal.data().clone());
56                continue;
57            }
58
59            let data = if self.is_hevc && !self.pending_hevc_nals.is_empty() {
60                build_hevc_access_unit(&mut self.pending_hevc_nals, Some(nal.data()))
61            } else {
62                nal.data().clone().freeze()
63            };
64
65            return Ok(H26xSample { data, timed });
66        }
67    }
68}
69
70fn should_skip_timing(nal: &H26xNAL) -> bool {
71    match nal {
72        H26xNAL::H264(nal) => {
73            matches!(
74                nal.unit_type,
75                H264NalUnitType::SPS
76                    | H264NalUnitType::PPS
77                    | H264NalUnitType::SEI
78                    | H264NalUnitType::AUD
79            )
80        }
81        H26xNAL::H265(nal) => {
82            matches!(
83                nal.unit_type,
84                H265NalUnitType::VPS
85                    | H265NalUnitType::SPS
86                    | H265NalUnitType::PPS
87                    | H265NalUnitType::PrefixSEI
88                    | H265NalUnitType::SuffixSEI
89                    | H265NalUnitType::AUD
90            )
91        }
92    }
93}
94
95fn should_buffer_with_next_hevc_vcl(nal: &H26xNAL) -> bool {
96    matches!(
97        nal,
98        H26xNAL::H265(nal)
99            if matches!(
100                nal.unit_type,
101                H265NalUnitType::VPS
102                    | H265NalUnitType::SPS
103                    | H265NalUnitType::PPS
104                    | H265NalUnitType::PrefixSEI
105                    | H265NalUnitType::SuffixSEI
106                    | H265NalUnitType::AUD
107            )
108    )
109}
110
111fn build_hevc_access_unit(
112    buffered_nals: &mut Vec<BytesMut>,
113    current_nal: Option<&BytesMut>,
114) -> Bytes {
115    let total_len = buffered_nals
116        .iter()
117        .map(|nal| ANNEXB_START_CODE.len() + nal.len())
118        .sum::<usize>()
119        + current_nal.map_or(0, |nal| ANNEXB_START_CODE.len() + nal.len());
120    let mut access_unit = BytesMut::with_capacity(total_len);
121
122    for nal in buffered_nals.drain(..) {
123        access_unit.extend_from_slice(&ANNEXB_START_CODE);
124        access_unit.extend_from_slice(&nal);
125    }
126    if let Some(current_nal) = current_nal {
127        access_unit.extend_from_slice(&ANNEXB_START_CODE);
128        access_unit.extend_from_slice(current_nal);
129    }
130
131    access_unit.freeze()
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use std::io::Cursor;
138
139    #[test]
140    fn h265_sample_reader_groups_parameter_sets_with_following_vcl() -> Result<()> {
141        let stream = vec![
142            0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x01, //
143            0x00, 0x00, 0x00, 0x01, 0x42, 0x01, 0x02, //
144            0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0x03, //
145            0x00, 0x00, 0x00, 0x01, 0x28, 0x01, 0xaa, //
146        ];
147        let mut reader = H26xSampleReader::new(Cursor::new(stream), 1024, true);
148
149        let sample = reader.next_sample()?;
150
151        assert!(sample.timed);
152        assert_eq!(
153            sample.data,
154            Bytes::from_static(&[
155                0x00, 0x00, 0x00, 0x01, 0x40, 0x01, 0x01, //
156                0x00, 0x00, 0x00, 0x01, 0x42, 0x01, 0x02, //
157                0x00, 0x00, 0x00, 0x01, 0x44, 0x01, 0x03, //
158                0x00, 0x00, 0x00, 0x01, 0x28, 0x01, 0xaa, //
159            ])
160        );
161        assert!(matches!(reader.next_sample(), Err(Error::ErrIoEOF)));
162
163        Ok(())
164    }
165
166    #[test]
167    fn h264_sample_reader_keeps_parameter_sets_separate_and_untimed() -> Result<()> {
168        let stream = vec![
169            0x00, 0x00, 0x00, 0x01, 0x67, 0x42, 0x00, 0x1f, //
170            0x00, 0x00, 0x00, 0x01, 0x68, 0xce, 0x06, 0xe2, //
171            0x00, 0x00, 0x00, 0x01, 0x65, 0x88, 0x84, 0x21, //
172        ];
173        let mut reader = H26xSampleReader::new(Cursor::new(stream), 1024, false);
174
175        let sps = reader.next_sample()?;
176        assert!(!sps.timed);
177        assert_eq!(sps.data, Bytes::from_static(&[0x67, 0x42, 0x00, 0x1f]));
178
179        let pps = reader.next_sample()?;
180        assert!(!pps.timed);
181        assert_eq!(pps.data, Bytes::from_static(&[0x68, 0xce, 0x06, 0xe2]));
182
183        let idr = reader.next_sample()?;
184        assert!(idr.timed);
185        assert_eq!(idr.data, Bytes::from_static(&[0x65, 0x88, 0x84, 0x21]));
186
187        Ok(())
188    }
189}
rtc_media/io/h26x_reader/sample_reader.rs

rtc_media/io/h26x_reader/
sample_reader.rs
