1use bytes::Bytes;
2
3use crate::{CodecConfig, VideoCodec, VideoError};
4
5use super::{annex_b, h264, h265, H264Config, H265Config};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum ConfigUpdate {
10 Incomplete,
12 Unchanged,
14 Changed(CodecConfig),
16}
17
18#[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 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 pub fn observe(&mut self, codec: VideoCodec, data: &[u8]) -> Result<ConfigUpdate, VideoError> {
105 self.inspect(codec, data).map(|(update, _)| update)
106 }
107
108 pub fn config(&self, codec: VideoCodec) -> Option<&CodecConfig> {
110 self.current
111 .as_ref()
112 .filter(|config| config.codec() == codec)
113 }
114
115 pub fn reset(&mut self) {
117 *self = Self::default();
118 }
119
120 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}