1use std::collections::VecDeque;
2
3use super::ebml::Size;
4use super::ids::*;
5use super::parse::Cluster;
6use super::reader::{
7 Cache, FrameRef, Hdr, Io, Reader, expand, parse_block_elem, read_frame_io, scan_cluster,
8};
9use crate::media::{Error, Packet, Result};
10use crate::mp4::ByteSource;
11
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub struct Block {
14 pub track: usize,
15 pub pts: i64,
16 pub duration: u32,
17 pub keyframe: bool,
18}
19
20const BISECT: u64 = 8 << 20;
21
22fn cluster_ts<S: ByteSource + ?Sized>(io: &mut Io<'_, S>, h: &Hdr, end: u64) -> Result<i64> {
23 let mut pos = h.data_start();
24 for _ in 0..3 {
25 match io.header(pos, end)? {
26 Some(c) if c.id == TIMESTAMP => return Ok(io.uint(&c, c.end(end))? as i64),
27 Some(c) if matches!(c.id, CRC32 | VOID) && c.end(end) > pos => pos = c.end(end),
28 _ => break,
29 }
30 }
31 Ok(0)
32}
33
34fn cluster_from<S: ByteSource + ?Sized>(
35 io: &mut Io<'_, S>,
36 from: u64,
37 limit: u64,
38) -> Result<Option<(u64, i64)>> {
39 let mut pos = from;
40 for _ in 0..16 {
41 let Some(at) = io.resync(pos, limit)? else {
42 return Ok(None);
43 };
44 let Some(h) = io.header(at, limit)? else {
45 return Ok(None);
46 };
47 if h.id == CLUSTER {
48 let end = h.end(limit);
49 return Ok(Some((at, cluster_ts(io, &h, end)?)));
50 }
51 if h.size == Size::Unknown {
52 return Ok(None);
53 }
54 pos = h.end(limit).max(at + 1);
55 }
56 Ok(None)
57}
58
59fn block(f: &FrameRef) -> Block {
60 Block {
61 track: f.track,
62 pts: f.pts,
63 duration: f.duration,
64 keyframe: f.keyframe,
65 }
66}
67
68#[derive(Clone, Copy, Debug)]
69struct InCluster {
70 ts: i64,
71 end: Option<u64>,
72}
73
74pub(crate) struct Cursor {
75 cache: Cache,
76 pos: u64,
77 cluster: Option<InCluster>,
78 pending: VecDeque<FrameRef>,
79 frames: Vec<(u64, u32)>,
80 skip: Option<(u64, u16)>,
81 done: bool,
82}
83
84pub struct Stream<'r, S: ByteSource> {
85 reader: &'r Reader<S>,
86 cursor: Cursor,
87}
88
89impl<S: ByteSource> Reader<S> {
90 pub fn stream(&self) -> Stream<'_, S> {
91 Stream {
92 reader: self,
93 cursor: Cursor::new(self),
94 }
95 }
96}
97
98impl<S: ByteSource> Stream<'_, S> {
99 pub fn rewind(&mut self) {
100 self.cursor.rewind(self.reader);
101 }
102
103 pub fn next_packet(&mut self) -> Result<Option<(usize, Packet)>> {
104 let mut data = Vec::new();
105 Ok(self.next_into(&mut data)?.map(|b| {
106 (
107 b.track,
108 Packet {
109 data,
110 pts: b.pts,
111 dts: b.pts,
112 keyframe: b.keyframe,
113 },
114 )
115 }))
116 }
117
118 pub fn next_into(&mut self, data: &mut Vec<u8>) -> Result<Option<Block>> {
119 self.cursor.next_of(self.reader, None, data)
120 }
121
122 pub fn seek(&mut self, track: usize, seconds: f64) -> Result<f64> {
123 self.cursor.seek(self.reader, track, seconds)
124 }
125
126 #[cfg_attr(
127 not(all(
128 feature = "h264",
129 feature = "hevc",
130 feature = "aac",
131 feature = "gpu",
132 feature = "filter",
133 feature = "vp9"
134 )),
135 allow(dead_code)
136 )]
137 pub(crate) fn next_block(&mut self, track: Option<usize>) -> Result<Option<Block>> {
138 self.cursor.next_block(self.reader, track)
139 }
140
141 #[cfg_attr(
142 not(all(
143 feature = "h264",
144 feature = "hevc",
145 feature = "aac",
146 feature = "gpu",
147 feature = "filter",
148 feature = "vp9"
149 )),
150 allow(dead_code)
151 )]
152 pub(crate) fn tail(&mut self, span: u64) -> Result<bool> {
153 self.cursor.tail(self.reader, span)
154 }
155}
156
157impl Cursor {
158 pub(crate) fn new<S: ByteSource>(r: &Reader<S>) -> Cursor {
159 Cursor {
160 cache: Cache::default(),
161 pos: r.first_cluster.unwrap_or(r.segment_end),
162 cluster: None,
163 pending: VecDeque::new(),
164 frames: Vec::new(),
165 skip: None,
166 done: false,
167 }
168 }
169
170 pub(crate) fn rewind<S: ByteSource>(&mut self, r: &Reader<S>) {
171 self.pos = r.first_cluster.unwrap_or(r.segment_end);
172 self.cluster = None;
173 self.pending.clear();
174 self.skip = None;
175 self.done = false;
176 }
177
178 #[cfg_attr(
179 not(all(
180 feature = "h264",
181 feature = "hevc",
182 feature = "aac",
183 feature = "gpu",
184 feature = "filter",
185 feature = "vp9"
186 )),
187 allow(dead_code)
188 )]
189 pub(crate) fn tail<S: ByteSource>(&mut self, r: &Reader<S>, span: u64) -> Result<bool> {
190 let first = r.first_cluster.unwrap_or(r.segment_end);
191 let from = r.segment_end.saturating_sub(span);
192 let at = if from <= first {
193 Some(first)
194 } else {
195 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
196 cluster_from(&mut io, from, r.segment_end)?.map(|(at, _)| at)
197 };
198 self.rewind(r);
199 match at {
200 Some(at) => {
201 self.pos = at;
202 Ok(true)
203 }
204 None => {
205 self.done = true;
206 Ok(false)
207 }
208 }
209 }
210
211 pub(crate) fn next_of<S: ByteSource>(
212 &mut self,
213 r: &Reader<S>,
214 track: Option<usize>,
215 data: &mut Vec<u8>,
216 ) -> Result<Option<Block>> {
217 let Some(f) = self.next_ref(r, track)? else {
218 return Ok(None);
219 };
220 let t = r
221 .tracks
222 .get(f.track)
223 .ok_or_else(|| Error::Invalid(format!("no track {}", f.track)))?;
224 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
225 read_frame_io(&mut io, t, f.offset, f.size, data)?;
226 Ok(Some(block(&f)))
227 }
228
229 #[cfg_attr(
230 not(all(
231 feature = "h264",
232 feature = "hevc",
233 feature = "aac",
234 feature = "gpu",
235 feature = "filter",
236 feature = "vp9"
237 )),
238 allow(dead_code)
239 )]
240 pub(crate) fn next_block<S: ByteSource>(
241 &mut self,
242 r: &Reader<S>,
243 track: Option<usize>,
244 ) -> Result<Option<Block>> {
245 Ok(self.next_ref(r, track)?.as_ref().map(block))
246 }
247
248 fn next_ref<S: ByteSource>(
249 &mut self,
250 r: &Reader<S>,
251 track: Option<usize>,
252 ) -> Result<Option<FrameRef>> {
253 loop {
254 if let Some(f) = self.pending.pop_front() {
255 if track.is_some_and(|t| t != f.track) {
256 continue;
257 }
258 return Ok(Some(f));
259 }
260 if self.done || !self.fill(r)? {
261 self.done = true;
262 return Ok(None);
263 }
264 }
265 }
266
267 pub(crate) fn seek<S: ByteSource>(
268 &mut self,
269 r: &Reader<S>,
270 track: usize,
271 seconds: f64,
272 ) -> Result<f64> {
273 let t = r
274 .tracks
275 .get(track)
276 .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
277 let ticks = t.ticks(seconds);
278 if !t.samples.is_empty() {
279 let k = t
280 .keyframe_at(ticks)
281 .ok_or_else(|| Error::Invalid(format!("track {track} has no keyframes")))?;
282 let s = t.samples[k];
283 let c = r
284 .clusters
285 .get(s.cluster as usize)
286 .copied()
287 .ok_or_else(|| Error::Invalid(format!("no cluster for track {track}")))?;
288 self.position_at(c, s.block, s.lace);
289 return Ok(t.seconds(s.pts));
290 }
291 let cue = r
292 .cues
293 .iter()
294 .filter(|c| c.track == t.number)
295 .take_while(|c| c.time as i64 <= ticks)
296 .last()
297 .or_else(|| r.cues.iter().find(|c| c.track == t.number));
298 if let Some(at) = cue.and_then(|c| r.segment_start.checked_add(c.cluster))
299 && let Some(pts) = self.seek_in(r, track, ticks, &[at], true)?
300 {
301 return Ok(t.seconds(pts));
302 }
303 self.seek_scan(r, track, ticks)
304 }
305
306 #[cfg_attr(
307 not(all(
308 feature = "h264",
309 feature = "hevc",
310 feature = "aac",
311 feature = "gpu",
312 feature = "filter",
313 feature = "vp9"
314 )),
315 allow(dead_code)
316 )]
317 pub(crate) fn seek_near<S: ByteSource>(
318 &mut self,
319 r: &Reader<S>,
320 track: usize,
321 seconds: f64,
322 ) -> Result<f64> {
323 let t = r
324 .tracks
325 .get(track)
326 .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
327 if !t.samples.is_empty() {
328 return self.seek(r, track, seconds);
329 }
330 let ticks = t.ticks(seconds);
331 let mut near: Vec<u64> = Vec::with_capacity(3);
332 for c in r.cues.iter().take_while(|c| c.time as i64 <= ticks) {
333 let Some(at) = r.segment_start.checked_add(c.cluster) else {
334 continue;
335 };
336 if near.last() != Some(&at) {
337 near.retain(|&p| p != at);
338 near.push(at);
339 if near.len() > 2 {
340 near.remove(0);
341 }
342 }
343 }
344 near.reverse();
345 if !near.is_empty()
346 && let Some(pts) = self.seek_in(r, track, ticks, &near, true)?
347 {
348 return Ok(t.seconds(pts));
349 }
350 self.seek_scan(r, track, ticks)
351 }
352
353 fn seek_scan<S: ByteSource>(&mut self, r: &Reader<S>, track: usize, ticks: i64) -> Result<f64> {
354 let t = r
355 .tracks
356 .get(track)
357 .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
358 let first = r.first_cluster.unwrap_or(r.segment_end);
359 let near = self.bisect(r, ticks)?;
360 if near > first {
361 let candidates = self.clusters_before(r, ticks, near)?;
362 if let Some(pts) = self.seek_in(r, track, ticks, &candidates, false)? {
363 return Ok(t.seconds(pts));
364 }
365 }
366 let candidates = self.clusters_before(r, ticks, first)?;
367 match self.seek_in(r, track, ticks, &candidates, true)? {
368 Some(pts) => Ok(t.seconds(pts)),
369 None => Err(Error::Invalid(format!("track {track} has no keyframes"))),
370 }
371 }
372
373 fn bisect<S: ByteSource>(&mut self, r: &Reader<S>, ticks: i64) -> Result<u64> {
374 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
375 let mut lo = r.first_cluster.unwrap_or(r.segment_end);
376 let mut hi = r.segment_end;
377 while hi.saturating_sub(lo) > BISECT {
378 let mid = lo + (hi - lo) / 2;
379 match cluster_from(&mut io, mid, hi)? {
380 Some((at, ts)) if ts <= ticks => lo = at,
381 _ => hi = mid,
382 }
383 }
384 Ok(lo)
385 }
386
387 fn clusters_before<S: ByteSource>(
388 &mut self,
389 r: &Reader<S>,
390 ticks: i64,
391 from: u64,
392 ) -> Result<Vec<u64>> {
393 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
394 let mut pos = from;
395 let mut out = Vec::new();
396 while let Some(h) = io.header(pos, r.segment_end)? {
397 if h.id != CLUSTER {
398 if matches!(h.id, EBML | SEGMENT) || h.size == Size::Unknown {
399 break;
400 }
401 pos = h.end(r.segment_end);
402 continue;
403 }
404 let end = match h.size {
405 Size::Known(_) => h.end(r.segment_end),
406 Size::Unknown => io.unknown_end(h.data_start(), r.segment_end)?,
407 };
408 let ts = cluster_ts(&mut io, &h, end)?;
409 if ts > ticks && !out.is_empty() {
410 break;
411 }
412 out.push(h.pos);
413 if end <= pos {
414 break;
415 }
416 pos = end;
417 }
418 out.reverse();
419 Ok(out)
420 }
421
422 fn seek_in<S: ByteSource>(
423 &mut self,
424 r: &Reader<S>,
425 track: usize,
426 ticks: i64,
427 clusters: &[u64],
428 loose: bool,
429 ) -> Result<Option<i64>> {
430 let scale = r.info.timestamp_scale;
431 let t = r
432 .tracks
433 .get(track)
434 .ok_or_else(|| Error::Invalid(format!("no track {track}")))?;
435 let last = clusters.last().copied().filter(|_| loose);
436 let mut warnings = Vec::new();
437 for &at in clusters {
438 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
439 let Some(h) = io.header(at, r.segment_end)? else {
440 continue;
441 };
442 if h.id != CLUSTER {
443 continue;
444 }
445 let mut keys: Vec<FrameRef> = Vec::new();
446 let ci = scan_cluster(
447 &mut io,
448 &h,
449 r.segment_end,
450 &r.tracks,
451 &mut warnings,
452 &mut self.frames,
453 |b, frames| {
454 if b.track == track && b.keyframe {
455 expand(b, frames, t, scale, 0, |f| keys.push(f));
456 }
457 },
458 )?;
459 let best = keys
460 .iter()
461 .filter(|f| f.pts <= ticks)
462 .max_by_key(|f| f.pts)
463 .or_else(|| {
464 (Some(at) == last)
465 .then(|| keys.iter().min_by_key(|f| f.pts))
466 .flatten()
467 })
468 .copied();
469 if let Some(f) = best {
470 self.position_at(ci, f.block, f.lace);
471 return Ok(Some(f.pts));
472 }
473 }
474 Ok(None)
475 }
476
477 fn position_at(&mut self, c: Cluster, block: u64, lace: u16) {
478 self.pending.clear();
479 self.done = false;
480 self.pos = block;
481 self.cluster = Some(InCluster {
482 ts: c.timestamp as i64,
483 end: Some(c.end),
484 });
485 self.skip = (lace > 0).then_some((block, lace));
486 }
487
488 fn fill<S: ByteSource>(&mut self, r: &Reader<S>) -> Result<bool> {
489 let seg_end = r.segment_end;
490 let scale = r.info.timestamp_scale;
491 let mut io = Io::with_len(r.source(), r.len(), &mut self.cache);
492 loop {
493 if let Some(c) = self.cluster {
494 let lim = c.end.unwrap_or(seg_end);
495 if self.pos >= lim {
496 self.cluster = None;
497 continue;
498 }
499 let h = match io.header(self.pos, lim)? {
500 Some(h) if !(h.size == Size::Unknown || h.overruns(lim)) => h,
501 Some(h) if c.end.is_none() && is_top_level(h.id) => h,
502 _ => {
503 self.cluster = None;
504 match io.resync(self.pos + 1, seg_end)? {
505 Some(p) => {
506 self.pos = p;
507 continue;
508 }
509 None => return Ok(false),
510 }
511 }
512 };
513 if c.end.is_none() && is_top_level(h.id) {
514 self.cluster = None;
515 continue;
516 }
517 let he = h.end(lim);
518 self.pos = he;
519 match h.id {
520 TIMESTAMP => {
521 let ts = io.uint(&h, he)? as i64;
522 self.cluster = Some(InCluster { ts, end: c.end });
523 }
524 SIMPLE_BLOCK | BLOCK_GROUP => {
525 if let Some(b) =
526 parse_block_elem(&mut io, &h, he, c.ts, &r.tracks, &mut self.frames)?
527 {
528 let skip = match self.skip.take() {
529 Some((off, n)) if off == h.pos => n,
530 _ => 0,
531 };
532 let pending = &mut self.pending;
533 if let Some(t) = r.tracks.get(b.track) {
534 expand(&b, &self.frames, t, scale, skip, |f| pending.push_back(f));
535 }
536 if !self.pending.is_empty() {
537 return Ok(true);
538 }
539 }
540 }
541 _ => {}
542 }
543 continue;
544 }
545 if self.pos >= seg_end {
546 return Ok(false);
547 }
548 let Some(h) = io.header(self.pos, seg_end)? else {
549 match io.resync(self.pos + 1, seg_end)? {
550 Some(p) => {
551 self.pos = p;
552 continue;
553 }
554 None => return Ok(false),
555 }
556 };
557 match h.id {
558 CLUSTER => {
559 self.cluster = Some(InCluster {
560 ts: 0,
561 end: h.size.known().map(|_| h.end(seg_end)),
562 });
563 self.pos = h.data_start();
564 }
565 EBML | SEGMENT => return Ok(false),
566 _ => {
567 self.pos = if h.size == Size::Unknown {
568 io.unknown_end(h.data_start(), seg_end)?
569 } else {
570 h.end(seg_end)
571 };
572 if self.pos <= h.pos {
573 return Ok(false);
574 }
575 }
576 }
577 }
578 }
579}