Skip to main content

pito_codecs/mkv/
parse.rs

1use super::ebml::{Children, float, string, uint};
2use super::ids::*;
3use super::track::{
4    AudioInfo, Colour, ContentEncoding, Mastering, Projection, Track, VideoInfo, gcd,
5};
6use crate::media::TimeBase;
7use crate::mp4::TrackKind;
8
9#[derive(Clone, Debug, PartialEq)]
10pub struct Info {
11    pub timestamp_scale: u64,
12    pub duration: Option<f64>,
13    pub title: Option<String>,
14    pub muxing_app: String,
15    pub writing_app: String,
16}
17
18impl Default for Info {
19    fn default() -> Self {
20        Self {
21            timestamp_scale: 1_000_000,
22            duration: None,
23            title: None,
24            muxing_app: String::new(),
25            writing_app: String::new(),
26        }
27    }
28}
29
30#[derive(Clone, Copy, Debug, PartialEq, Eq)]
31pub struct SeekEntry {
32    pub id: u32,
33    pub position: u64,
34}
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
37pub struct CuePoint {
38    pub time: u64,
39    pub track: u64,
40    pub cluster: u64,
41}
42
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub struct Cluster {
45    pub offset: u64,
46    pub data_start: u64,
47    pub end: u64,
48    pub timestamp: u64,
49}
50
51pub(crate) fn time_base(scale: u64) -> TimeBase {
52    let g = gcd(scale, 1_000_000_000);
53    let (num, den) = (scale / g, 1_000_000_000 / g);
54    if num <= u32::MAX as u64 && den <= u32::MAX as u64 {
55        TimeBase::new(num as u32, den as u32)
56    } else {
57        TimeBase::new(1, 1000)
58    }
59}
60
61pub(crate) fn info(d: &[u8]) -> Info {
62    let mut info = Info::default();
63    for e in Children::new(d) {
64        match e.id {
65            TIMESTAMP_SCALE => {
66                let v = uint(e.data);
67                if v > 0 && v <= u32::MAX as u64 {
68                    info.timestamp_scale = v;
69                }
70            }
71            DURATION => {
72                if let Some(v) = float(e.data) {
73                    info.duration = Some(v).filter(|v| v.is_finite() && *v >= 0.0);
74                }
75            }
76            TITLE => info.title = Some(string(e.data)),
77            MUXING_APP => info.muxing_app = string(e.data),
78            WRITING_APP => info.writing_app = string(e.data),
79            _ => {}
80        }
81    }
82    info
83}
84
85pub(crate) fn seek_head(d: &[u8]) -> Vec<SeekEntry> {
86    let mut out = Vec::new();
87    for s in Children::new(d).filter(|e| e.id == SEEK) {
88        let (mut id, mut pos) = (None, None);
89        for e in Children::new(s.data) {
90            match e.id {
91                SEEK_ID => id = Some(uint(e.data) as u32),
92                SEEK_POSITION => pos = Some(uint(e.data)),
93                _ => {}
94            }
95        }
96        if let (Some(id), Some(position)) = (id, pos) {
97            out.push(SeekEntry { id, position });
98        }
99    }
100    out
101}
102
103pub(crate) fn cues(d: &[u8], out: &mut Vec<CuePoint>) {
104    for cp in Children::new(d).filter(|e| e.id == CUE_POINT) {
105        let mut time = None;
106        let start = out.len();
107        for e in Children::new(cp.data) {
108            match e.id {
109                CUE_TIME => time = Some(uint(e.data)),
110                CUE_TRACK_POSITIONS => {
111                    let (mut track, mut cluster) = (0, u64::MAX);
112                    for c in Children::new(e.data) {
113                        match c.id {
114                            CUE_TRACK => track = uint(c.data),
115                            CUE_CLUSTER_POSITION => cluster = uint(c.data),
116                            _ => {}
117                        }
118                    }
119                    if track != 0 && cluster != u64::MAX {
120                        out.push(CuePoint {
121                            time: 0,
122                            track,
123                            cluster,
124                        });
125                    }
126                }
127                _ => {}
128            }
129        }
130        match time {
131            Some(t) => out[start..].iter_mut().for_each(|c| c.time = t),
132            None => out.truncate(start),
133        }
134    }
135}
136
137fn colour(d: &[u8]) -> Colour {
138    let mut c = Colour::default();
139    for e in Children::new(d) {
140        let v = Some(uint(e.data));
141        match e.id {
142            MATRIX_COEFFICIENTS => c.matrix = v,
143            BITS_PER_CHANNEL => c.bits_per_channel = v,
144            CHROMA_SITING_HORZ => c.chroma_siting.0 = v,
145            CHROMA_SITING_VERT => c.chroma_siting.1 = v,
146            RANGE => c.range = v,
147            TRANSFER_CHARACTERISTICS => c.transfer = v,
148            PRIMARIES => c.primaries = v,
149            MAX_CLL => c.max_cll = v,
150            MAX_FALL => c.max_fall = v,
151            MASTERING_METADATA => {
152                let mut m = Mastering::default();
153                for x in Children::new(e.data) {
154                    let Some(f) = float(x.data) else {
155                        continue;
156                    };
157                    match x.id {
158                        PRIMARY_R_X => m.primaries[0].0 = f,
159                        PRIMARY_R_Y => m.primaries[0].1 = f,
160                        PRIMARY_G_X => m.primaries[1].0 = f,
161                        PRIMARY_G_Y => m.primaries[1].1 = f,
162                        PRIMARY_B_X => m.primaries[2].0 = f,
163                        PRIMARY_B_Y => m.primaries[2].1 = f,
164                        WHITE_POINT_X => m.white_point.0 = f,
165                        WHITE_POINT_Y => m.white_point.1 = f,
166                        LUMINANCE_MAX => m.luminance_max = f,
167                        LUMINANCE_MIN => m.luminance_min = f,
168                        _ => {}
169                    }
170                }
171                c.mastering = Some(m);
172            }
173            _ => {}
174        }
175    }
176    c
177}
178
179fn video(d: &[u8]) -> VideoInfo {
180    let mut v = VideoInfo::default();
181    for e in Children::new(d) {
182        let u = uint(e.data);
183        match e.id {
184            PIXEL_WIDTH => v.pixel_width = u as u32,
185            PIXEL_HEIGHT => v.pixel_height = u as u32,
186            DISPLAY_WIDTH => v.display_width = Some(u as u32),
187            DISPLAY_HEIGHT => v.display_height = Some(u as u32),
188            DISPLAY_UNIT => v.display_unit = u,
189            PIXEL_CROP_TOP => v.crop.0 = u as u32,
190            PIXEL_CROP_BOTTOM => v.crop.1 = u as u32,
191            PIXEL_CROP_LEFT => v.crop.2 = u as u32,
192            PIXEL_CROP_RIGHT => v.crop.3 = u as u32,
193            FLAG_INTERLACED => v.interlaced = u,
194            COLOUR => v.colour = Some(colour(e.data)),
195            PROJECTION => {
196                let mut p = Projection::default();
197                for x in Children::new(e.data) {
198                    match x.id {
199                        PROJECTION_TYPE => p.kind = uint(x.data),
200                        PROJECTION_POSE_YAW => p.yaw = float(x.data).unwrap_or(p.yaw),
201                        PROJECTION_POSE_PITCH => p.pitch = float(x.data).unwrap_or(p.pitch),
202                        PROJECTION_POSE_ROLL => p.roll = float(x.data).unwrap_or(p.roll),
203                        _ => {}
204                    }
205                }
206                v.projection = Some(p);
207            }
208            _ => {}
209        }
210    }
211    v
212}
213
214fn audio(d: &[u8]) -> AudioInfo {
215    let mut a = AudioInfo::default();
216    for e in Children::new(d) {
217        match e.id {
218            SAMPLING_FREQUENCY => {
219                a.sampling_frequency = float(e.data).unwrap_or(a.sampling_frequency)
220            }
221            OUTPUT_SAMPLING_FREQUENCY => {
222                a.output_sampling_frequency = float(e.data).or(a.output_sampling_frequency)
223            }
224            CHANNELS => a.channels = uint(e.data),
225            BIT_DEPTH => a.bit_depth = Some(uint(e.data)),
226            _ => {}
227        }
228    }
229    a
230}
231
232fn encodings(d: &[u8]) -> Vec<ContentEncoding> {
233    let mut list: Vec<(u64, ContentEncoding)> = Vec::new();
234    for ce in Children::new(d).filter(|e| e.id == CONTENT_ENCODING) {
235        let (mut order, mut scope, mut kind) = (0u64, 1u64, 0u64);
236        let mut comp: Option<(u64, Vec<u8>)> = None;
237        for e in Children::new(ce.data) {
238            match e.id {
239                CONTENT_ENCODING_ORDER => order = uint(e.data),
240                CONTENT_ENCODING_SCOPE => scope = uint(e.data),
241                CONTENT_ENCODING_TYPE => kind = uint(e.data),
242                CONTENT_COMPRESSION => {
243                    let (mut algo, mut settings) = (0u64, Vec::new());
244                    for x in Children::new(e.data) {
245                        match x.id {
246                            CONTENT_COMP_ALGO => algo = uint(x.data),
247                            CONTENT_COMP_SETTINGS => settings = x.data.to_vec(),
248                            _ => {}
249                        }
250                    }
251                    comp = Some((algo, settings));
252                }
253                _ => {}
254            }
255        }
256        let enc = if kind == 1 {
257            ContentEncoding::Encryption { scope }
258        } else {
259            match comp.unwrap_or((0, Vec::new())) {
260                (3, bytes) => ContentEncoding::HeaderStripping { bytes, scope },
261                (algo, _) => ContentEncoding::Compression { algo, scope },
262            }
263        };
264        list.push((order, enc));
265    }
266    list.sort_by_key(|e| std::cmp::Reverse(e.0));
267    list.into_iter().map(|e| e.1).collect()
268}
269
270pub(crate) fn tracks(d: &[u8], timestamp_scale: u64) -> Vec<Track> {
271    let time_base = time_base(timestamp_scale);
272    let mut out = Vec::new();
273    for te in Children::new(d).filter(|e| e.id == TRACK_ENTRY) {
274        let mut t = Track {
275            number: 0,
276            uid: 0,
277            kind: TrackKind::Other,
278            enabled: true,
279            default: true,
280            name: String::new(),
281            language: "eng".into(),
282            codec_id: String::new(),
283            codec_private: Vec::new(),
284            default_duration: None,
285            codec_delay: 0,
286            seek_pre_roll: 0,
287            width: 0,
288            height: 0,
289            video: None,
290            audio: None,
291            encodings: Vec::new(),
292            time_base,
293            samples: Vec::new(),
294            sync: Vec::new(),
295            pts_order: Vec::new(),
296        };
297        for e in Children::new(te.data) {
298            match e.id {
299                TRACK_NUMBER => t.number = uint(e.data),
300                TRACK_UID => t.uid = uint(e.data),
301                TRACK_TYPE => {
302                    t.kind = match uint(e.data) {
303                        1 => TrackKind::Video,
304                        2 => TrackKind::Audio,
305                        0x11 => TrackKind::Subtitle,
306                        _ => TrackKind::Other,
307                    }
308                }
309                FLAG_ENABLED => t.enabled = uint(e.data) != 0,
310                FLAG_DEFAULT => t.default = uint(e.data) != 0,
311                DEFAULT_DURATION => t.default_duration = Some(uint(e.data)).filter(|&v| v > 0),
312                NAME => t.name = string(e.data),
313                LANGUAGE => t.language = string(e.data),
314                CODEC_ID => t.codec_id = string(e.data),
315                CODEC_PRIVATE => t.codec_private = e.data.to_vec(),
316                CODEC_DELAY => t.codec_delay = uint(e.data),
317                SEEK_PRE_ROLL => t.seek_pre_roll = uint(e.data),
318                VIDEO => t.video = Some(video(e.data)),
319                AUDIO => t.audio = Some(audio(e.data)),
320                CONTENT_ENCODINGS => t.encodings = encodings(e.data),
321                _ => {}
322            }
323        }
324        if t.number == 0 {
325            continue;
326        }
327        if t.kind == TrackKind::Audio && t.audio.is_none() {
328            t.audio = Some(AudioInfo::default());
329        }
330        if let Some(v) = &t.video {
331            t.width = v.pixel_width;
332            t.height = v.pixel_height;
333        }
334        for enc in &t.encodings {
335            if let ContentEncoding::HeaderStripping { bytes, scope } = enc
336                && scope & 2 != 0
337            {
338                let mut p = bytes.clone();
339                p.extend_from_slice(&t.codec_private);
340                t.codec_private = p;
341            }
342        }
343        out.push(t);
344    }
345    out
346}