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