Skip to main content

openipc_video/codecs/
access_unit.rs

1use bytes::Bytes;
2
3use crate::{CodecConfig, VideoCodec, VideoError};
4
5use super::{annex_b, h264, h265, H264Config, H265Config};
6
7/// Result of inspecting an access unit for decoder parameter sets.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum ConfigUpdate {
10    /// Required parameter sets have not all been observed.
11    Incomplete,
12    /// Complete configuration is unchanged.
13    Unchanged,
14    /// A complete new configuration was observed.
15    Changed(CodecConfig),
16}
17
18/// Tracks H.264 and H.265 parameter sets across access units.
19#[derive(Debug, Clone, Default)]
20pub struct CodecConfigTracker {
21    h264_sps: Option<Bytes>,
22    h264_pps: Option<Bytes>,
23    h265_vps: Option<Bytes>,
24    h265_sps: Option<Bytes>,
25    h265_pps: Option<Bytes>,
26    current: Option<CodecConfig>,
27}
28
29impl CodecConfigTracker {
30    /// Inspect one access unit for both parameter-set changes and keyframes.
31    ///
32    /// Decoder backends should prefer this method over calling [`Self::observe`]
33    /// and [`Self::is_keyframe`] separately because it scans Annex-B only once.
34    pub fn inspect(
35        &mut self,
36        codec: VideoCodec,
37        data: &[u8],
38    ) -> Result<(ConfigUpdate, bool), VideoError> {
39        let mut parameter_sets_changed = false;
40        let mut keyframe = false;
41        for unit in annex_b::nal_units_iter(data)? {
42            match codec {
43                VideoCodec::H264 => {
44                    keyframe |= h264::is_keyframe(unit.data);
45                    match h264::nal_type(unit.data) {
46                        Some(7) if self.h264_sps.as_deref() != Some(unit.data) => {
47                            self.h264_sps = Some(Bytes::copy_from_slice(unit.data));
48                            parameter_sets_changed = true;
49                        }
50                        Some(8) if self.h264_pps.as_deref() != Some(unit.data) => {
51                            self.h264_pps = Some(Bytes::copy_from_slice(unit.data));
52                            parameter_sets_changed = true;
53                        }
54                        _ => {}
55                    }
56                }
57                VideoCodec::H265 => {
58                    keyframe |= h265::is_keyframe(unit.data);
59                    match h265::nal_type(unit.data) {
60                        Some(32) if self.h265_vps.as_deref() != Some(unit.data) => {
61                            self.h265_vps = Some(Bytes::copy_from_slice(unit.data));
62                            parameter_sets_changed = true;
63                        }
64                        Some(33) if self.h265_sps.as_deref() != Some(unit.data) => {
65                            self.h265_sps = Some(Bytes::copy_from_slice(unit.data));
66                            parameter_sets_changed = true;
67                        }
68                        Some(34) if self.h265_pps.as_deref() != Some(unit.data) => {
69                            self.h265_pps = Some(Bytes::copy_from_slice(unit.data));
70                            parameter_sets_changed = true;
71                        }
72                        _ => {}
73                    }
74                }
75            }
76        }
77
78        if !parameter_sets_changed
79            && self
80                .current
81                .as_ref()
82                .is_some_and(|config| config.codec() == codec)
83        {
84            return Ok((ConfigUpdate::Unchanged, keyframe));
85        }
86
87        let Some(config) = self.complete_config(codec)? else {
88            return Ok((ConfigUpdate::Incomplete, keyframe));
89        };
90        let update = if self.current.as_ref() == Some(&config) {
91            ConfigUpdate::Unchanged
92        } else {
93            log::debug!(
94                target: "openipc_video::codec",
95                "complete codec configuration observed codec={codec}"
96            );
97            self.current = Some(config.clone());
98            ConfigUpdate::Changed(config)
99        };
100        Ok((update, keyframe))
101    }
102
103    /// Inspect an Annex-B access unit and update cached parameter sets.
104    pub fn observe(&mut self, codec: VideoCodec, data: &[u8]) -> Result<ConfigUpdate, VideoError> {
105        self.inspect(codec, data).map(|(update, _)| update)
106    }
107
108    /// Return the latest complete configuration for `codec`.
109    pub fn config(&self, codec: VideoCodec) -> Option<&CodecConfig> {
110        self.current
111            .as_ref()
112            .filter(|config| config.codec() == codec)
113    }
114
115    /// Clear all cached parameter sets and active configuration.
116    pub fn reset(&mut self) {
117        *self = Self::default();
118    }
119
120    /// Return true when the access unit contains a random-access NAL unit.
121    pub fn is_keyframe(codec: VideoCodec, data: &[u8]) -> Result<bool, VideoError> {
122        Ok(annex_b::nal_units_iter(data)?.any(|unit| match codec {
123            VideoCodec::H264 => h264::is_keyframe(unit.data),
124            VideoCodec::H265 => h265::is_keyframe(unit.data),
125        }))
126    }
127
128    fn complete_config(&self, codec: VideoCodec) -> Result<Option<CodecConfig>, VideoError> {
129        match codec {
130            VideoCodec::H264 => match (&self.h264_sps, &self.h264_pps) {
131                (Some(sps), Some(pps)) => Ok(Some(CodecConfig::H264(H264Config::new(
132                    sps.clone(),
133                    pps.clone(),
134                )?))),
135                _ => Ok(None),
136            },
137            VideoCodec::H265 => {
138                match (&self.h265_vps, &self.h265_sps, &self.h265_pps) {
139                    (Some(vps), Some(sps), Some(pps)) => Ok(Some(CodecConfig::H265(
140                        H265Config::new(vps.clone(), sps.clone(), pps.clone())?,
141                    ))),
142                    _ => Ok(None),
143                }
144            }
145        }
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::{CodecConfigTracker, ConfigUpdate};
152    use crate::{CodecConfig, VideoCodec};
153
154    #[test]
155    fn tracks_h264_parameter_sets_across_access_units() {
156        let mut tracker = CodecConfigTracker::default();
157        assert!(matches!(
158            tracker.observe(VideoCodec::H264, &[0, 0, 1, 0x67, 1]),
159            Ok(ConfigUpdate::Incomplete)
160        ));
161        let update = tracker
162            .observe(VideoCodec::H264, &[0, 0, 1, 0x68, 2])
163            .unwrap();
164        assert!(matches!(
165            update,
166            ConfigUpdate::Changed(CodecConfig::H264(_))
167        ));
168    }
169
170    #[test]
171    fn tracks_h265_parameter_sets() {
172        let data = [
173            0,
174            0,
175            1,
176            32 << 1,
177            1,
178            0,
179            0,
180            1,
181            33 << 1,
182            2,
183            0,
184            0,
185            1,
186            34 << 1,
187            3,
188        ];
189        let mut tracker = CodecConfigTracker::default();
190        let update = tracker.observe(VideoCodec::H265, &data).unwrap();
191        assert!(matches!(
192            update,
193            ConfigUpdate::Changed(CodecConfig::H265(_))
194        ));
195    }
196
197    #[test]
198    fn detects_h264_and_h265_keyframes() {
199        assert!(CodecConfigTracker::is_keyframe(VideoCodec::H264, &[0, 0, 1, 0x65, 1]).unwrap());
200        assert!(CodecConfigTracker::is_keyframe(VideoCodec::H265, &[0, 0, 1, 19 << 1, 1]).unwrap());
201    }
202
203    #[test]
204    fn inspect_reports_configuration_and_keyframe_in_one_pass() {
205        let mut tracker = CodecConfigTracker::default();
206        let data = [0, 0, 1, 0x67, 1, 0, 0, 1, 0x68, 2, 0, 0, 1, 0x65, 3];
207        let (update, keyframe) = tracker.inspect(VideoCodec::H264, &data).unwrap();
208        assert!(matches!(
209            update,
210            ConfigUpdate::Changed(CodecConfig::H264(_))
211        ));
212        assert!(keyframe);
213    }
214}