1use revelo_util::Ztring;
12use std::collections::BTreeMap;
13
14#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
24#[repr(u8)]
25pub enum StreamKind {
26 General = 0,
27 Video = 1,
28 Audio = 2,
29 Text = 3,
30 Other = 4,
31 Image = 5,
32 Menu = 6,
33 Exif = 7,
34 Iptc = 8,
35 Xmp = 9,
36 Icc = 10,
37 C2pa = 11,
38 MakerNotes = 12,
39}
40
41impl StreamKind {
42 pub fn name(self) -> &'static str {
43 match self {
44 StreamKind::General => "General",
45 StreamKind::Video => "Video",
46 StreamKind::Audio => "Audio",
47 StreamKind::Text => "Text",
48 StreamKind::Other => "Other",
49 StreamKind::Image => "Image",
50 StreamKind::Menu => "Menu",
51 StreamKind::Exif => "Exif",
52 StreamKind::Iptc => "Iptc",
53 StreamKind::Xmp => "Xmp",
54 StreamKind::Icc => "Icc",
55 StreamKind::C2pa => "C2pa",
56 StreamKind::MakerNotes => "MakerNotes",
57 }
58 }
59}
60
61#[derive(Clone, Debug, Default)]
64pub struct Stream {
65 fields: BTreeMap<String, Ztring>,
66 insertion_order: Vec<String>,
68 extras: Vec<(String, Ztring)>,
73}
74
75impl Stream {
76 pub fn new() -> Self {
77 Stream::default()
78 }
79
80 pub fn set(&mut self, parameter: &str, value: Ztring, replace: bool) {
81 let key = parameter.to_owned();
82 let existed = self.fields.contains_key(&key);
83 if existed && !replace {
84 return;
85 }
86 if !existed {
87 self.insertion_order.push(key.clone());
88 }
89 self.fields.insert(key, value);
90 }
91
92 pub fn set_extra(&mut self, parameter: &str, value: Ztring, replace: bool) {
96 if replace {
97 if let Some(slot) = self.extras.iter_mut().find(|(k, _)| k == parameter) {
98 slot.1 = value;
99 return;
100 }
101 } else if self.extras.iter().any(|(k, _)| k == parameter) {
102 return;
103 }
104 self.extras.push((parameter.to_owned(), value));
105 }
106
107 pub fn get(&self, parameter: &str) -> Option<&Ztring> {
108 self.fields.get(parameter)
109 }
110
111 pub fn contains(&self, parameter: &str) -> bool {
112 self.fields.contains_key(parameter)
113 }
114
115 pub fn count(&self) -> usize {
116 self.fields.len()
117 }
118
119 pub fn iter(&self) -> impl Iterator<Item = (&str, &Ztring)> {
122 self.insertion_order
123 .iter()
124 .filter_map(|k| self.fields.get_key_value(k.as_str()).map(|(k, v)| (k.as_str(), v)))
125 }
126
127 pub fn extras_iter(&self) -> impl Iterator<Item = (&str, &Ztring)> {
129 self.extras.iter().map(|(k, v)| (k.as_str(), v))
130 }
131}
132
133#[derive(Clone, Debug, Default)]
135pub struct StreamCollection {
136 by_kind: BTreeMap<StreamKind, Vec<Stream>>,
137}
138
139impl StreamCollection {
140 pub fn new() -> Self {
141 StreamCollection::default()
142 }
143
144 pub fn stream_prepare(&mut self, kind: StreamKind) -> usize {
147 let v = self.by_kind.entry(kind).or_default();
148 v.push(Stream::new());
149 v.len() - 1
150 }
151
152 pub fn stream_count(&self, kind: StreamKind) -> usize {
153 self.by_kind.get(&kind).map(|v| v.len()).unwrap_or(0)
154 }
155
156 pub fn set_field(
160 &mut self,
161 kind: StreamKind,
162 pos: usize,
163 parameter: &str,
164 value: impl Into<Ztring>,
165 ) {
166 self.fill(kind, pos, parameter, value, false)
167 }
168
169 pub fn force_field(
172 &mut self,
173 kind: StreamKind,
174 pos: usize,
175 parameter: &str,
176 value: impl Into<Ztring>,
177 ) {
178 self.fill(kind, pos, parameter, value, true)
179 }
180
181 fn fill(
183 &mut self,
184 kind: StreamKind,
185 pos: usize,
186 parameter: &str,
187 value: impl Into<Ztring>,
188 replace: bool,
189 ) {
190 let v = self.by_kind.entry(kind).or_default();
191 while v.len() <= pos {
192 v.push(Stream::new());
193 }
194 v[pos].set(parameter, value.into(), replace);
195 }
196
197 pub fn set_extra_field(
200 &mut self,
201 kind: StreamKind,
202 pos: usize,
203 parameter: &str,
204 value: impl Into<Ztring>,
205 ) {
206 self.fill_extra(kind, pos, parameter, value, false)
207 }
208
209 pub fn force_extra_field(
211 &mut self,
212 kind: StreamKind,
213 pos: usize,
214 parameter: &str,
215 value: impl Into<Ztring>,
216 ) {
217 self.fill_extra(kind, pos, parameter, value, true)
218 }
219
220 fn fill_extra(
222 &mut self,
223 kind: StreamKind,
224 pos: usize,
225 parameter: &str,
226 value: impl Into<Ztring>,
227 replace: bool,
228 ) {
229 let v = self.by_kind.entry(kind).or_default();
230 while v.len() <= pos {
231 v.push(Stream::new());
232 }
233 v[pos].set_extra(parameter, value.into(), replace);
234 }
235
236 pub fn retrieve(&self, kind: StreamKind, pos: usize, parameter: &str) -> Option<&Ztring> {
237 self.by_kind.get(&kind)?.get(pos)?.get(parameter)
238 }
239
240 pub fn stream(&self, kind: StreamKind, pos: usize) -> Option<&Stream> {
241 self.by_kind.get(&kind)?.get(pos)
242 }
243
244 pub fn filter_keep(&mut self, keep: &[StreamKind], specific: &[(StreamKind, usize)]) {
249 let keep_set: std::collections::HashSet<StreamKind> = keep.iter().copied().collect();
250 let specific_map: std::collections::HashMap<StreamKind, std::collections::HashSet<usize>> =
251 specific.iter().fold(std::collections::HashMap::new(), |mut map, &(k, pos)| {
252 map.entry(k).or_default().insert(pos);
253 map
254 });
255
256 let mut specific_keep_all = true;
257 if !specific.is_empty() {
258 specific_keep_all = false;
259 }
260
261 self.by_kind.retain(|&kind, streams| {
262 if !keep_set.contains(&kind) {
263 return false;
264 }
265 if specific_keep_all {
266 return true;
267 }
268 if let Some(positions) = specific_map.get(&kind) {
269 let mut i = 0;
271 streams.retain(|_| {
272 let keep = positions.contains(&i);
273 i += 1;
274 keep
275 });
276 } else {
277 }
280 true
281 });
282 }
283
284 pub fn iter(&self) -> impl Iterator<Item = (StreamKind, usize, &Stream)> {
285 self.by_kind.iter().flat_map(|(k, v)| v.iter().enumerate().map(move |(i, s)| (*k, i, s)))
286 }
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292
293 #[test]
294 fn stream_prepare_returns_sequential_indices() {
295 let mut c = StreamCollection::new();
296 assert_eq!(c.stream_prepare(StreamKind::Audio), 0);
297 assert_eq!(c.stream_prepare(StreamKind::Audio), 1);
298 assert_eq!(c.stream_prepare(StreamKind::Video), 0);
299 assert_eq!(c.stream_count(StreamKind::Audio), 2);
300 assert_eq!(c.stream_count(StreamKind::Video), 1);
301 assert_eq!(c.stream_count(StreamKind::Text), 0);
302 }
303
304 #[test]
305 fn fill_and_retrieve_round_trip() {
306 let mut c = StreamCollection::new();
307 c.stream_prepare(StreamKind::Audio);
308 c.set_field(StreamKind::Audio, 0, "Format", "FLAC");
309 c.set_field(StreamKind::Audio, 0, "SamplingRate", "48000");
310 assert_eq!(c.retrieve(StreamKind::Audio, 0, "Format").map(|z| z.as_str()), Some("FLAC"));
311 assert_eq!(
312 c.retrieve(StreamKind::Audio, 0, "SamplingRate").map(|z| z.as_str()),
313 Some("48000")
314 );
315 assert_eq!(c.retrieve(StreamKind::Audio, 0, "Missing"), None);
316 }
317
318 #[test]
319 fn set_field_keeps_first_value() {
320 let mut c = StreamCollection::new();
321 c.set_field(StreamKind::General, 0, "Format", "MP4");
322 c.set_field(StreamKind::General, 0, "Format", "MOV");
323 assert_eq!(c.retrieve(StreamKind::General, 0, "Format").map(|z| z.as_str()), Some("MP4"));
324 }
325
326 #[test]
327 fn force_field_overwrites() {
328 let mut c = StreamCollection::new();
329 c.set_field(StreamKind::General, 0, "Format", "MP4");
330 c.force_field(StreamKind::General, 0, "Format", "MOV");
331 assert_eq!(c.retrieve(StreamKind::General, 0, "Format").map(|z| z.as_str()), Some("MOV"));
332 }
333
334 #[test]
335 fn set_field_auto_creates_stream_if_pos_unset() {
336 let mut c = StreamCollection::new();
337 c.set_field(StreamKind::Audio, 2, "Format", "AAC");
338 assert_eq!(c.stream_count(StreamKind::Audio), 3);
339 assert_eq!(c.retrieve(StreamKind::Audio, 2, "Format").map(|z| z.as_str()), Some("AAC"));
340 assert_eq!(c.retrieve(StreamKind::Audio, 0, "Format"), None);
342 }
343
344 #[test]
345 fn iter_preserves_insertion_order_within_stream() {
346 let mut c = StreamCollection::new();
347 c.set_field(StreamKind::Video, 0, "Format", "AVC");
348 c.set_field(StreamKind::Video, 0, "Width", "1920");
349 c.set_field(StreamKind::Video, 0, "Height", "1080");
350 let s = c.stream(StreamKind::Video, 0).unwrap();
351 let order: Vec<&str> = s.iter().map(|(k, _)| k).collect();
352 assert_eq!(order, vec!["Format", "Width", "Height"]);
353 }
354
355 #[test]
356 fn iter_walks_all_streams_grouped_by_kind() {
357 let mut c = StreamCollection::new();
358 c.set_field(StreamKind::General, 0, "Format", "MP4");
359 c.set_field(StreamKind::Video, 0, "Format", "AVC");
360 c.set_field(StreamKind::Audio, 0, "Format", "AAC");
361 c.set_field(StreamKind::Audio, 1, "Format", "AC3");
362 let pairs: Vec<(StreamKind, usize)> = c.iter().map(|(k, i, _)| (k, i)).collect();
363 assert_eq!(
364 pairs,
365 vec![
366 (StreamKind::General, 0),
367 (StreamKind::Video, 0),
368 (StreamKind::Audio, 0),
369 (StreamKind::Audio, 1),
370 ]
371 );
372 }
373
374 #[test]
375 fn stream_kind_name_matches_cpp_output_strings() {
376 assert_eq!(StreamKind::General.name(), "General");
377 assert_eq!(StreamKind::Video.name(), "Video");
378 assert_eq!(StreamKind::Audio.name(), "Audio");
379 assert_eq!(StreamKind::Text.name(), "Text");
380 assert_eq!(StreamKind::Other.name(), "Other");
381 assert_eq!(StreamKind::Image.name(), "Image");
382 assert_eq!(StreamKind::Menu.name(), "Menu");
383 }
384}