mp4_edit/atom/leaf/
stco_co64.rs

1use bon::bon;
2use derive_more::{Deref, DerefMut};
3use std::fmt;
4
5use crate::{
6    atom::{
7        util::{DebugList, DebugUpperHex},
8        FourCC,
9    },
10    parser::ParseAtomData,
11    writer::SerializeAtom,
12    ParseError,
13};
14
15#[cfg(feature = "experimental-trim")]
16use {crate::atom::stsz::RemovedSampleSizes, anyhow::anyhow, std::ops::Range};
17
18pub const STCO: FourCC = FourCC::new(b"stco");
19pub const CO64: FourCC = FourCC::new(b"co64");
20
21#[derive(Default, Clone, Deref, DerefMut, PartialEq, Eq)]
22pub struct ChunkOffsets(Vec<u64>);
23
24impl ChunkOffsets {
25    pub fn into_inner(self) -> Vec<u64> {
26        self.0
27    }
28
29    pub fn inner(&self) -> &[u64] {
30        &self.0
31    }
32}
33
34impl From<Vec<u64>> for ChunkOffsets {
35    fn from(value: Vec<u64>) -> Self {
36        Self(value)
37    }
38}
39
40impl FromIterator<u64> for ChunkOffsets {
41    fn from_iter<T: IntoIterator<Item = u64>>(iter: T) -> Self {
42        Self(Vec::from_iter(iter))
43    }
44}
45
46impl fmt::Debug for ChunkOffsets {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        fmt::Debug::fmt(&DebugList::new(self.0.iter().map(DebugUpperHex), 10), f)
49    }
50}
51
52/// Chunk Offset Atom - contains file offsets of chunks
53#[derive(Default, Debug, Clone)]
54pub struct ChunkOffsetAtom {
55    /// Version of the stco atom format (0)
56    pub version: u8,
57    /// Flags for the stco atom (usually all zeros)
58    pub flags: [u8; 3],
59    /// List of chunk offsets
60    pub chunk_offsets: ChunkOffsets,
61    /// Whether this uses 64-bit offsets (co64) or 32-bit (stco)
62    pub is_64bit: bool,
63}
64
65#[cfg(feature = "experimental-trim")]
66#[derive(Debug)]
67pub(crate) enum ChunkOffsetOperationUnresolved {
68    Remove(Range<usize>),
69    Insert {
70        /// unadjusted chunk index used to get reference offset
71        chunk_index_unadjusted: usize,
72        /// chunk index that takes into account all previous operations
73        chunk_index: usize,
74        /// sample indices used to calculate delta from old offset
75        sample_indices: Range<usize>,
76    },
77    ShiftRight {
78        /// unadjusted chunk index used to get reference offset
79        chunk_index_unadjusted: usize,
80        /// chunk index that takes into account all previous operations
81        chunk_index: usize,
82        /// sample indices used to calculate delta from old offset
83        sample_indices: Range<usize>,
84    },
85}
86
87#[cfg(feature = "experimental-trim")]
88impl ChunkOffsetOperationUnresolved {
89    pub fn resolve(
90        self,
91        chunk_offsets: &ChunkOffsets,
92        removed_sample_sizes: &RemovedSampleSizes,
93    ) -> anyhow::Result<ChunkOffsetOperation> {
94        let derive_new_offset = |chunk_index: usize, sample_indices: Range<usize>| {
95            let prev_offset = *chunk_offsets
96                .get(chunk_index)
97                .ok_or_else(|| anyhow!("chunk index {chunk_index} not found"))?;
98
99            let delta = removed_sample_sizes
100                .get_sizes(sample_indices.clone())
101                .ok_or_else(|| anyhow!("sample indices {sample_indices:?} not found"))?
102                .iter()
103                .map(|s| *s as u64)
104                .sum::<u64>();
105
106            Ok::<u64, anyhow::Error>(prev_offset + delta)
107        };
108
109        Ok(match self {
110            Self::Remove(chunk_offsets) => ChunkOffsetOperation::Remove(chunk_offsets),
111            Self::Insert {
112                chunk_index_unadjusted,
113                chunk_index,
114                sample_indices,
115            } => {
116                let offset = derive_new_offset(chunk_index_unadjusted - 1, sample_indices)?;
117                ChunkOffsetOperation::Insert(chunk_index, offset)
118            }
119            Self::ShiftRight {
120                chunk_index_unadjusted,
121                chunk_index,
122                sample_indices,
123            } => {
124                let new_offset = derive_new_offset(chunk_index_unadjusted, sample_indices)?;
125                ChunkOffsetOperation::Replace(chunk_index, new_offset)
126            }
127        })
128    }
129}
130
131#[cfg(feature = "experimental-trim")]
132#[derive(Debug)]
133pub(crate) enum ChunkOffsetOperation {
134    Remove(Range<usize>),
135    Insert(usize, u64),
136    Replace(usize, u64),
137}
138
139#[bon]
140impl ChunkOffsetAtom {
141    #[builder]
142    pub fn new(
143        #[builder(default = 0)] version: u8,
144        #[builder(default = [0u8; 3])] flags: [u8; 3],
145        #[builder(with = FromIterator::from_iter)] chunk_offsets: Vec<u64>,
146        #[builder(default = false)] is_64bit: bool,
147    ) -> Self {
148        Self {
149            version,
150            flags,
151            chunk_offsets: chunk_offsets.into(),
152            is_64bit,
153        }
154    }
155
156    /// Returns the total number of chunks
157    pub fn chunk_count(&self) -> usize {
158        self.chunk_offsets.len()
159    }
160
161    /// Applies a list of operations
162    #[cfg(feature = "experimental-trim")]
163    pub(crate) fn apply_operations(&mut self, ops: Vec<ChunkOffsetOperation>) {
164        for op in ops {
165            match op {
166                ChunkOffsetOperation::Remove(chunk_indices_to_remove) => {
167                    self.chunk_offsets.drain(chunk_indices_to_remove);
168                }
169                ChunkOffsetOperation::Insert(chunk_index, offset) => {
170                    self.chunk_offsets.insert(chunk_index, offset);
171                }
172                ChunkOffsetOperation::Replace(chunk_index, new_offset) => {
173                    let chunk = self
174                        .chunk_offsets
175                        .get_mut(chunk_index)
176                        .expect("chunk offset must exist");
177                    *chunk = new_offset;
178                }
179            }
180        }
181    }
182}
183
184impl<S: chunk_offset_atom_builder::State> ChunkOffsetAtomBuilder<S> {
185    pub fn chunk_offset(
186        self,
187        chunk_offset: impl Into<u64>,
188    ) -> ChunkOffsetAtomBuilder<chunk_offset_atom_builder::SetChunkOffsets<S>>
189    where
190        S::ChunkOffsets: chunk_offset_atom_builder::IsUnset,
191    {
192        self.chunk_offsets(vec![chunk_offset.into()])
193    }
194}
195
196impl ParseAtomData for ChunkOffsetAtom {
197    fn parse_atom_data(atom_type: FourCC, input: &[u8]) -> Result<Self, ParseError> {
198        crate::atom::util::parser::assert_atom_type!(atom_type, STCO, CO64);
199        use crate::atom::util::parser::stream;
200        use winnow::Parser;
201        Ok(match atom_type {
202            STCO => parser::parse_stco_data.parse(stream(input))?,
203            CO64 => parser::parse_co64_data.parse(stream(input))?,
204            _ => unreachable!(),
205        })
206    }
207}
208
209impl SerializeAtom for ChunkOffsetAtom {
210    fn atom_type(&self) -> FourCC {
211        // Use the appropriate atom type based on is_64bit
212        if self.is_64bit {
213            CO64
214        } else {
215            STCO
216        }
217    }
218
219    fn into_body_bytes(self) -> Vec<u8> {
220        serializer::serialize_stco_co64_data(self)
221    }
222}
223
224mod serializer {
225    use crate::atom::{util::serializer::be_u32, ChunkOffsetAtom};
226
227    pub fn serialize_stco_co64_data(atom: ChunkOffsetAtom) -> Vec<u8> {
228        let mut data = Vec::new();
229
230        data.push(atom.version);
231        data.extend(atom.flags);
232        data.extend(be_u32(
233            atom.chunk_offsets
234                .len()
235                .try_into()
236                .expect("chunk offsets length must fit in u32"),
237        ));
238
239        atom.chunk_offsets.0.into_iter().for_each(|offset| {
240            if atom.is_64bit {
241                data.extend(offset.to_be_bytes());
242            } else {
243                data.extend(be_u32(
244                    offset.try_into().expect("chunk offset must fit in u32"),
245                ))
246            }
247        });
248
249        data
250    }
251}
252
253mod parser {
254    use winnow::{
255        binary::{be_u32, be_u64},
256        combinator::{empty, repeat, seq, trace},
257        error::{ContextError, ErrMode, StrContext},
258        ModalResult, Parser,
259    };
260
261    use super::{ChunkOffsetAtom, ChunkOffsets};
262    use crate::atom::util::parser::{byte_array, version, Stream};
263
264    pub fn parse_stco_data(input: &mut Stream<'_>) -> ModalResult<ChunkOffsetAtom> {
265        parse_stco_co64_data_inner(false).parse_next(input)
266    }
267
268    pub fn parse_co64_data(input: &mut Stream<'_>) -> ModalResult<ChunkOffsetAtom> {
269        parse_stco_co64_data_inner(true).parse_next(input)
270    }
271
272    fn parse_stco_co64_data_inner<'i>(
273        is_64bit: bool,
274    ) -> impl Parser<Stream<'i>, ChunkOffsetAtom, ErrMode<ContextError>> {
275        trace(
276            if is_64bit { "co64" } else { "stco" },
277            move |input: &mut Stream<'_>| {
278                seq!(ChunkOffsetAtom {
279                    version: version,
280                    flags: byte_array.context(StrContext::Label("flags")),
281                    chunk_offsets: chunk_offsets(is_64bit)
282                        .map(ChunkOffsets)
283                        .context(StrContext::Label("chunk_offsets")),
284                    is_64bit: empty.value(is_64bit),
285                })
286                .parse_next(input)
287            },
288        )
289    }
290
291    fn chunk_offsets<'i>(
292        is_64bit: bool,
293    ) -> impl Parser<Stream<'i>, Vec<u64>, ErrMode<ContextError>> {
294        trace("chunk_offsets", move |input: &mut Stream<'_>| {
295            let entry_count = be_u32.parse_next(input)?;
296            repeat(entry_count as usize, chunk_offset(is_64bit)).parse_next(input)
297        })
298    }
299
300    fn chunk_offset<'i>(is_64bit: bool) -> impl Parser<Stream<'i>, u64, ErrMode<ContextError>> {
301        trace("chunk_offset", move |input: &mut Stream<'_>| {
302            if is_64bit {
303                be_u64.parse_next(input)
304            } else {
305                be_u32.map(|v| v as u64).parse_next(input)
306            }
307        })
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314    use crate::atom::test_utils::test_atom_roundtrip;
315
316    /// Test round-trip for all available stco/co64 test data files
317    #[test]
318    fn test_stco_co64_roundtrip() {
319        test_atom_roundtrip::<ChunkOffsetAtom>(STCO);
320        test_atom_roundtrip::<ChunkOffsetAtom>(CO64);
321    }
322}
mp4_edit/atom/leaf/stco_co64.rs

mp4_edit/atom/leaf/
stco_co64.rs
