Skip to main content

revelo_core/
stream.rs

1//! Stream model — transliteration of MediaInfoLib's per-stream `Fill` /
2//! `Retrieve` state.
3//!
4//! On the C++ side this is `MediaInfo_Internal::Stream` indexed by
5//! `(stream_t, size_t)` (kind + position-within-kind) and stores parsed
6//! fields as [`Ztring`]. Output formatters (in the `revelo-export` crate)
7//! walk this state to produce XML, JSON, and text results, so this is the
8//! canonical place every parser writes into under the direction of a
9//! [`FileAnalyze`](super::FileAnalyze).
10
11use revelo_util::Ztring;
12use std::collections::BTreeMap;
13
14/// Discriminants 0..=6 match MediaInfo's `stream_t` from
15/// `MediaInfo_Const.h`, so external consumers binding through the C ABI
16/// get matching values for those kinds. Values 7+ are revelo extensions
17/// that have no MediaInfo equivalent — they model the distinct embedded
18/// metadata standards exiftool exposes as family-0 groups (EXIF, IPTC,
19/// XMP, ICC_Profile, C2PA, MakerNotes) rather than folding them all into
20/// General. Note: upstream MediaInfo terminates the enum at
21/// `Stream_Max = 7`, so any C consumer iterating `0..Stream_Max` will not
22/// see these kinds (and value 7 collides with its `Stream_Max` sentinel).
23#[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/// One stream's fields. BTreeMap keeps iteration order stable, which the
62/// output formatters depend on for deterministic XML/JSON.
63#[derive(Clone, Debug, Default)]
64pub struct Stream {
65    fields: BTreeMap<String, Ztring>,
66    /// Insertion order, for formatters that want the order parsers wrote in.
67    insertion_order: Vec<String>,
68    /// "Extra" fields — emitted in their own `<extra>...</extra>` block
69    /// at the end of the stream, distinct from standard fields. Order
70    /// preserved as inserted. Mirrors MediaInfoLib's
71    /// `Stream::Extra` / `Fill_Measure` two-tier output model.
72    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    /// Append (or, if `replace`, overwrite) an `<extra>`-bucket entry.
93    /// Extras don't appear in `iter()` — formatters call `extras_iter()`
94    /// separately so they end up in their own XML/JSON block.
95    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    /// Iterate fields in insertion order — matches the C++ behavior of
120    /// emitting in the order parsers filled.
121    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    /// Iterate `<extra>`-bucket fields in insertion order.
128    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/// Container of per-kind, indexed-by-position streams.
134#[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    /// Allocate a new stream of `kind`, return its `StreamPos`.
145    /// Matches `File__Analyze::Stream_Prepare`.
146    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    /// Set a field on a stream. If the field already exists, it is NOT
157    /// overwritten (first-write-wins). Auto-creates the stream if it
158    /// doesn't exist yet.
159    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    /// Set a field on a stream, ALWAYS overwriting any existing value.
170    /// Auto-creates the stream if it doesn't exist yet.
171    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    /// Inner helper: `Fill(StreamKind, StreamPos, Parameter, Value, Replace)`.
182    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    /// Set a field in the stream's `<extra>` bucket instead of the
198    /// standard field list. First-write-wins.
199    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    /// Like [`set_extra_field`], but ALWAYS overwrites.
210    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    /// Inner helper for extra-bucket fields.
221    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    /// Remove all streams whose [`StreamKind`] is not in `keep`, and
245    /// remove all streams at a kind whose positional index is not in
246    /// `positions` (when that map entry exists). Used to implement
247    /// `--video-only`, `--audio-only`, and `--stream` filtering.
248    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                // Keep only the positions that were specified.
270                let mut i = 0;
271                streams.retain(|_| {
272                    let keep = positions.contains(&i);
273                    i += 1;
274                    keep
275                });
276            } else {
277                // This kind is in `keep` but has no specific positions
278                // selected — keep all streams of this kind.
279            }
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        // The auto-created earlier streams are empty
341        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}