1use bytes::{BufMut, Bytes, BytesMut};
2use flowly::Fourcc;
3use std::collections::BTreeSet;
4
5use crate::ctts::CttsEntry;
6use crate::error::Error;
7use crate::stsc::StscEntry;
8use crate::stts::SttsEntry;
9use crate::{BoxType, TrackType};
10
11#[derive(Clone)]
12pub struct Mp4SampleOffset {
13 pub offset: u64,
14 pub size: u32,
15 pub duration: u32,
16 pub start_time: u64,
17 pub rendering_offset: i32,
18 pub is_sync: bool,
19 pub chunk_id: u32,
20}
21
22#[derive(Clone)]
23pub struct Mp4Track {
24 pub track_id: u32,
25 pub duration: u64,
26 pub samples: Vec<Mp4SampleOffset>,
27 pub tkhd: crate::TkhdBox,
28 pub mdia: crate::MdiaBox,
29}
30
31impl Mp4Track {
32 pub fn new(trak: crate::TrakBox, offsets: &mut BTreeSet<u64>) -> Result<Mp4Track, Error> {
33 let default_sample_duration = 1024;
34 let mut total_duration = 0;
35 let mut samples = Vec::with_capacity(trak.mdia.minf.stbl.stsz.sample_count as _);
36 let stco = &trak.mdia.minf.stbl.stco;
37 let co64 = &trak.mdia.minf.stbl.co64;
38
39 let mb_iter1 = stco.clone().map(IntoIterator::into_iter);
40 let mb_iter2 = co64.clone().map(IntoIterator::into_iter);
41
42 if let Some(stco) = co64.as_ref().map(IntoIterator::into_iter) {
43 offsets.extend(stco);
44 }
45
46 if let Some(stco) = stco.as_ref().map(IntoIterator::into_iter) {
47 offsets.extend(stco);
48 }
49
50 let chunk_iter = chunk_iter(
51 trak.mdia.minf.stbl.stsc.entries.clone().into_iter(),
52 mb_iter1
53 .into_iter()
54 .flatten()
55 .chain(mb_iter2.into_iter().flatten()),
56 );
57
58 let mut sample_chunk_iter = run_len_iter(chunk_iter);
59
60 let sync_iter_peek = trak
61 .mdia
62 .minf
63 .stbl
64 .stss
65 .as_ref()
66 .map(|x| x.entries.iter().copied().peekable());
67
68 let mut sync_iter =
69 (1..=trak.mdia.minf.stbl.stsz.sample_count).scan(sync_iter_peek, |iter, idx| {
70 let iter = iter.as_mut()?;
71
72 Some(if idx == iter.peek().copied().unwrap_or(u32::MAX) {
73 iter.next();
74 true
75 } else {
76 false
77 })
78 });
79
80 let mut ts_deltas =
81 run_len_iter(trak.mdia.minf.stbl.stts.entries.clone().into_iter().chain(
82 std::iter::once(SttsEntry {
83 sample_count: u32::MAX,
84 sample_delta: default_sample_duration,
85 }),
86 ))
87 .scan(0u64, |s, delta| {
88 let out = *s;
89 *s += delta as u64;
90 Some((out, delta))
91 });
92
93 let mut rend_offset_iter = run_len_iter(
94 trak.mdia
95 .minf
96 .stbl
97 .ctts
98 .clone()
99 .into_iter()
100 .flat_map(|x| x.entries.into_iter()),
101 );
102
103 let mut sample_offset = 0;
104 let mut curr_chunk_index = 0;
105 let mut prev_size = 0;
106
107 for sample_idx in 0..trak.mdia.minf.stbl.stsz.sample_count as usize {
108 let (start_time, duration) = ts_deltas.next().unwrap();
109 let chunk = sample_chunk_iter.next().unwrap();
110 let size = *trak
111 .mdia
112 .minf
113 .stbl
114 .stsz
115 .sample_sizes
116 .get(sample_idx)
117 .unwrap_or(&trak.mdia.minf.stbl.stsz.sample_size);
118
119 if curr_chunk_index != chunk.index {
120 curr_chunk_index = chunk.index;
121 sample_offset = 0;
122 } else {
123 sample_offset += prev_size;
124 }
125
126 prev_size = size;
127 total_duration = start_time + duration as u64;
128 samples.push(Mp4SampleOffset {
129 chunk_id: chunk.index,
130 offset: chunk.offset + sample_offset as u64,
131 size,
132 duration,
133 start_time,
134 rendering_offset: rend_offset_iter.next().unwrap_or(0),
135 is_sync: sync_iter.next().unwrap_or(true),
136 })
137 }
138
139 Ok(Self {
140 track_id: trak.tkhd.track_id,
141 tkhd: trak.tkhd,
142 mdia: trak.mdia,
143 samples,
144 duration: total_duration,
145 })
146 }
147
148 #[inline]
149 pub fn track_type(&self) -> TrackType {
150 TrackType::from(&self.mdia.hdlr.handler_type)
151 }
152
153 #[inline]
154 pub fn codec(&self) -> Fourcc {
155 if self.mdia.minf.stbl.stsd.avc1.is_some() {
156 Fourcc::VIDEO_AVC
157 } else if self.mdia.minf.stbl.stsd.hev1.is_some() {
158 Fourcc::VIDEO_HEVC
159 } else if self.mdia.minf.stbl.stsd.vp09.is_some() {
160 Fourcc::VIDEO_VP9
161 } else if self.mdia.minf.stbl.stsd.mp4a.is_some() {
162 Fourcc::AUDIO_AAC
163 } else if self.mdia.minf.stbl.stsd.tx3g.is_some() {
164 Fourcc::from_static("TTXT")
165 } else {
166 Default::default()
167 }
168 }
169
170 pub(crate) fn add_traf(
171 &mut self,
172 base_moof_offset: u64,
173 chunk_index: u32,
174 traf: crate::TrafBox,
175 offsets: &mut BTreeSet<u64>,
176 ) {
177 let base_data_offset = traf.tfhd.base_data_offset.unwrap_or(base_moof_offset);
178 offsets.insert(base_data_offset);
179
180 let default_sample_size = traf.tfhd.default_sample_size.unwrap_or(0);
181 let default_sample_duration = traf.tfhd.default_sample_duration.unwrap_or(0);
182 let base_start_time = traf
183 .tfdt
184 .map(|x| x.base_media_decode_time)
185 .or_else(|| {
186 self.samples
187 .last()
188 .map(|x| x.start_time + x.duration as u64)
189 })
190 .unwrap_or(0);
191
192 let Some(trun) = traf.trun else {
193 return;
194 };
195
196 let mut sample_offset = 0u64;
197 let mut start_time_offset = 0u64;
198 for sample_idx in 0..trun.sample_count as usize {
199 let size = trun
200 .sample_sizes
201 .get(sample_idx)
202 .copied()
203 .unwrap_or(default_sample_size);
204
205 let duration = trun
206 .sample_durations
207 .get(sample_idx)
208 .copied()
209 .unwrap_or(default_sample_duration);
210
211 let rendering_offset = trun.sample_cts.get(sample_idx).copied().unwrap_or(0) as i32;
212
213 self.samples.push(Mp4SampleOffset {
214 chunk_id: chunk_index,
215 offset: (base_data_offset as i64
216 + trun.data_offset.map(|x| x as i64).unwrap_or(0)
217 + sample_offset as i64) as u64,
218 size,
219 duration,
220 start_time: base_start_time + start_time_offset,
221 rendering_offset,
222 is_sync: sample_idx == 0,
223 });
224
225 sample_offset += size as u64;
226 start_time_offset += duration as u64;
227 }
228 }
229
230 pub fn sequence_parameter_set(&self) -> Result<&[u8], Error> {
231 if let Some(ref avc1) = self.mdia.minf.stbl.stsd.avc1 {
232 match avc1.avcc.sequence_parameter_sets.first() {
233 Some(nal) => Ok(nal.bytes.as_ref()),
234 None => Err(Error::EntryInStblNotFound(
235 self.track_id,
236 BoxType::AvcCBox,
237 0,
238 )),
239 }
240 } else {
241 Err(Error::BoxInStblNotFound(self.track_id, BoxType::Avc1Box))
242 }
243 }
244
245 pub fn picture_parameter_set(&self) -> Result<&[u8], Error> {
246 if let Some(ref avc1) = self.mdia.minf.stbl.stsd.avc1 {
247 match avc1.avcc.picture_parameter_sets.first() {
248 Some(nal) => Ok(nal.bytes.as_ref()),
249 None => Err(Error::EntryInStblNotFound(
250 self.track_id,
251 BoxType::AvcCBox,
252 0,
253 )),
254 }
255 } else {
256 Err(Error::BoxInStblNotFound(self.track_id, BoxType::Avc1Box))
257 }
258 }
259
260 pub fn decode_params(&self) -> Option<Bytes> {
261 match self.codec() {
262 Fourcc::VIDEO_AVC => {
263 let mut buf = BytesMut::new();
264
265 let sps = self.sequence_parameter_set().unwrap();
266 buf.put_u32(sps.len() as u32 + 4);
267 buf.put_slice(&[0, 0, 0, 1]);
268 buf.put_slice(sps);
269
270 let pps = self.picture_parameter_set().unwrap();
271 buf.put_u32(pps.len() as u32 + 4);
272 buf.put_slice(&[0, 0, 0, 1]);
273 buf.put_slice(pps);
274
275 Some(buf.freeze())
276 }
277
278 Fourcc::VIDEO_HEVC => {
279 let mut buf = BytesMut::new();
280 let x = self.mdia.minf.stbl.stsd.hev1.as_ref().unwrap();
281 for arr in &x.hvcc.arrays {
282 for nalu in &arr.nalus {
283 buf.put_u32(nalu.data.len() as u32 + 4);
284 buf.put_slice(&[0, 0, 0, 1]);
285 buf.put_slice(&nalu.data);
286 }
287 }
288 Some(buf.freeze())
289 }
290
291 _ => None,
292 }
293 }
294
295 #[inline]
296 pub fn timescale(&self) -> u32 {
297 self.mdia.mdhd.timescale
298 }
299}
300
301trait RunLenghtItem {
302 type Value: Clone;
303
304 fn count(&self) -> usize;
305 fn value(&self) -> Self::Value;
306}
307
308impl<T: Clone> RunLenghtItem for (usize, T) {
309 type Value = T;
310
311 fn count(&self) -> usize {
312 self.0
313 }
314 fn value(&self) -> Self::Value {
315 self.1.clone()
316 }
317}
318
319impl RunLenghtItem for CttsEntry {
320 type Value = i32;
321
322 fn count(&self) -> usize {
323 self.sample_count as _
324 }
325
326 fn value(&self) -> Self::Value {
327 self.sample_offset
328 }
329}
330
331impl RunLenghtItem for SttsEntry {
332 type Value = u32;
333
334 fn count(&self) -> usize {
335 self.sample_count as _
336 }
337
338 fn value(&self) -> Self::Value {
339 self.sample_delta
340 }
341}
342
343#[derive(Debug, Clone, Copy, PartialEq, Eq)]
344pub struct Chunk {
345 pub index: u32,
346 pub offset: u64,
347 pub samples_per_chunk: u32,
348 pub sample_description_index: u32,
349}
350
351impl RunLenghtItem for Chunk {
352 type Value = Chunk;
353
354 fn count(&self) -> usize {
355 self.samples_per_chunk as _
356 }
357
358 fn value(&self) -> Self::Value {
359 *self
360 }
361}
362
363fn chunk_iter(
364 mut stsc: impl Iterator<Item = StscEntry>,
365 stco: impl Iterator<Item = u64>,
366) -> impl Iterator<Item = Chunk> {
367 let mut prev = stsc.next().unwrap_or(StscEntry {
368 first_chunk: 1,
369 samples_per_chunk: u32::MAX,
370 sample_description_index: 1,
371 first_sample: 1,
372 });
373 let mut curr = stsc.next();
374
375 stco.enumerate().map(move |(idx, offset)| {
376 if let Some(c) = &curr {
377 if idx + 1 >= c.first_chunk as usize {
378 prev = *c;
379 curr = stsc.next();
380 }
381 }
382
383 Chunk {
384 index: idx as _,
385 offset,
386 samples_per_chunk: prev.samples_per_chunk,
387 sample_description_index: prev.sample_description_index,
388 }
389 })
390}
391
392fn run_len_iter<E: RunLenghtItem, I: IntoIterator<Item = E>>(
393 iter: I,
394) -> impl Iterator<Item = E::Value> {
395 let mut iter = iter.into_iter();
396 let mut value = None::<E::Value>;
397 let mut repeat = 0;
398 std::iter::from_fn(move || loop {
399 if let Some(val) = &value {
400 if repeat > 0 {
401 repeat -= 1;
402 return Some(val.clone());
403 } else {
404 value = None;
405 }
406 }
407
408 let x = iter.next()?;
409 value = Some(x.value());
410 repeat = x.count();
411 })
412}