mp4_edit/atom/leaf/
stsz.rs

1use bon::bon;
2use derive_more::{Deref, DerefMut};
3use either::Either;
4use rangemap::{RangeMap, RangeSet};
5use std::{
6    fmt::{self},
7    ops::Range,
8};
9
10use crate::{
11    atom::{util::DebugList, FourCC},
12    parser::ParseAtomData,
13    writer::SerializeAtom,
14    ParseError,
15};
16
17pub const STSZ: FourCC = FourCC::new(b"stsz");
18
19#[derive(Clone, Default, Deref, DerefMut)]
20pub struct SampleEntrySizes(Vec<u32>);
21
22impl SampleEntrySizes {
23    pub fn inner(&self) -> &[u32] {
24        &self.0
25    }
26}
27
28impl From<Vec<u32>> for SampleEntrySizes {
29    fn from(value: Vec<u32>) -> Self {
30        SampleEntrySizes(value)
31    }
32}
33
34impl fmt::Debug for SampleEntrySizes {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        fmt::Debug::fmt(&DebugList::new(self.0.iter(), 10), f)
37    }
38}
39
40impl SampleEntrySizes {
41    /// Create a new SampleEntrySizes from a vector of sample sizes
42    pub fn new(sizes: Vec<u32>) -> Self {
43        Self(sizes)
44    }
45
46    /// Create a new SampleEntrySizes from a vector of sample sizes
47    pub fn from_vec(sizes: Vec<u32>) -> Self {
48        Self(sizes)
49    }
50
51    /// Convert to the inner `Vec<u32>`
52    pub fn to_vec(&self) -> Vec<u32> {
53        self.0.clone()
54    }
55}
56
57/// Sample Size Atom (stsz) - ISO/IEC 14496-12
58/// This atom contains the sample count and a table giving the size in bytes of each sample.
59/// Samples within the media may have different sizes, up to the limit of a 32-bit integer.
60#[derive(Default, Debug, Clone)]
61pub struct SampleSizeAtom {
62    pub version: u8,
63    pub flags: [u8; 3],
64    /// If this field is set to some value other than 0, then it gives the (constant) size
65    /// of every sample in the track. If this field is set to 0, then the samples have
66    /// different sizes, and those sizes are stored in the sample size table.
67    pub sample_size: u32,
68    /// Number of samples in the track
69    pub sample_count: u32,
70    /// If `sample_size` is 0, this contains the size of each sample, indexed by sample number.
71    /// If `sample_size` is non-zero, this table is empty.
72    pub entry_sizes: SampleEntrySizes,
73}
74
75pub(crate) struct RemovedSampleSizes {
76    /// key=removed range start, value=(range.start, removed sizes)
77    removed_sizes: RangeMap<usize, (usize, Vec<u32>)>,
78}
79
80impl RemovedSampleSizes {
81    fn new() -> Self {
82        Self {
83            removed_sizes: RangeMap::new(),
84        }
85    }
86
87    fn insert(&mut self, indices: Range<usize>, sizes: impl Iterator<Item = u32>) {
88        let start_index = indices.start;
89        self.removed_sizes
90            .insert(indices, (start_index, sizes.collect::<Vec<_>>()));
91    }
92
93    /// Get removed sample sizes for the given sample indices
94    pub(crate) fn get_sizes(&self, sample_indices: Range<usize>) -> Option<&[u32]> {
95        // get the list of sizes that the range belongs to
96        let (first_index, sizes) = self.removed_sizes.get(&sample_indices.start)?;
97        // slice to the requested range (invariant: RangeMap Range<K>.len() == V.len())
98        let sizes = &sizes.as_slice()
99            [(sample_indices.start - first_index)..(sample_indices.end - first_index)];
100        Some(sizes)
101    }
102}
103
104impl SampleSizeAtom {
105    #[cfg(feature = "experimental-trim")]
106    pub(crate) fn remove_sample_indices(
107        &mut self,
108        indices_to_remove: &RangeSet<usize>,
109    ) -> RemovedSampleSizes {
110        let num_samples_removed = indices_to_remove
111            .iter()
112            .map(|r| r.end - r.start)
113            .sum::<usize>() as u32;
114
115        fn adjust_range(n_removed: usize, range: &Range<usize>) -> Range<usize> {
116            let start = range.start - n_removed;
117            let end = range.end - n_removed;
118            start..end
119        }
120
121        let mut removed_sizes = RemovedSampleSizes::new();
122
123        if !self.entry_sizes.is_empty() && !indices_to_remove.is_empty() {
124            let mut n_removed = 0;
125            for range in indices_to_remove.iter() {
126                let adjusted_range = adjust_range(n_removed, range);
127                n_removed += adjusted_range.len();
128                removed_sizes.insert(
129                    range.clone(),
130                    self.entry_sizes.drain(adjusted_range.clone()),
131                );
132            }
133        }
134
135        self.sample_count = self.sample_count.saturating_sub(num_samples_removed);
136
137        removed_sizes
138    }
139
140    /// Returns `sample_count` if it's set, otherwise `entry_sizes.len()`
141    pub fn sample_count(&self) -> usize {
142        if self.sample_count > 0 {
143            self.sample_count as usize
144        } else {
145            self.entry_sizes.len()
146        }
147    }
148}
149
150#[bon]
151impl SampleSizeAtom {
152    #[builder]
153    pub fn new(
154        #[builder(setters(vis = "", name = "sample_size_internal"))] sample_size: u32,
155        #[builder(default = 0)] sample_count: u32,
156        /// either set `sample_size` and `sample_count` or `entry_sizes`
157        #[builder(with = FromIterator::from_iter, setters(vis = "", name = "entry_sizes_internal"))]
158        entry_sizes: Vec<u32>,
159    ) -> Self {
160        let entry_sizes: SampleEntrySizes = entry_sizes.into();
161        let sample_count = if sample_count == 0 {
162            u32::try_from(entry_sizes.len()).expect("entry_sizes.len() should fit in a u32")
163        } else {
164            sample_count
165        };
166        Self {
167            version: 0,
168            flags: [0u8; 3],
169            sample_size,
170            sample_count,
171            entry_sizes,
172        }
173    }
174
175    /// Returns an iterator over _all_ sample sizes.
176    ///
177    /// If `sample_size != 0` this will repeat that value
178    /// `sample_count` times; otherwise it will yield
179    /// the values from `entry_sizes`.
180    pub fn sample_sizes(&self) -> impl Iterator<Item = &u32> + '_ {
181        if self.sample_size != 0 {
182            Either::Left(std::iter::repeat_n(
183                &self.sample_size,
184                self.sample_count as usize,
185            ))
186        } else {
187            Either::Right(self.entry_sizes.iter())
188        }
189    }
190}
191
192#[bon]
193impl<S: sample_size_atom_builder::State> SampleSizeAtomBuilder<S> {
194    pub fn sample_size(
195        self,
196        sample_size: u32,
197    ) -> SampleSizeAtomBuilder<
198        sample_size_atom_builder::SetSampleSize<sample_size_atom_builder::SetEntrySizes<S>>,
199    >
200    where
201        S::EntrySizes: sample_size_atom_builder::IsUnset,
202        S::SampleSize: sample_size_atom_builder::IsUnset,
203    {
204        self.entry_sizes_internal(vec![])
205            .sample_size_internal(sample_size)
206    }
207
208    #[builder(finish_fn(name = "build"))]
209    pub fn entry_sizes(
210        self,
211        #[builder(start_fn)] entry_sizes: impl IntoIterator<Item = u32>,
212    ) -> SampleSizeAtom
213    where
214        S::EntrySizes: sample_size_atom_builder::IsUnset,
215        S::SampleSize: sample_size_atom_builder::IsUnset,
216        S::SampleCount: sample_size_atom_builder::IsUnset,
217    {
218        self.entry_sizes_internal(entry_sizes)
219            .sample_size_internal(0)
220            .sample_count(0)
221            .build()
222    }
223}
224
225impl ParseAtomData for SampleSizeAtom {
226    fn parse_atom_data(atom_type: FourCC, input: &[u8]) -> Result<Self, ParseError> {
227        crate::atom::util::parser::assert_atom_type!(atom_type, STSZ);
228        use crate::atom::util::parser::stream;
229        use winnow::Parser;
230        Ok(parser::parse_stsz_data.parse(stream(input))?)
231    }
232}
233
234impl fmt::Display for SampleSizeAtom {
235    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236        write!(f, "SampleSize(count: {}, ", self.sample_count)?;
237
238        if self.sample_size != 0 {
239            write!(f, "constant_size: {})", self.sample_size)
240        } else {
241            write!(f, "variable_sizes: {} entries)", self.entry_sizes.len())
242        }
243    }
244}
245
246impl SerializeAtom for SampleSizeAtom {
247    fn atom_type(&self) -> FourCC {
248        STSZ
249    }
250
251    fn into_body_bytes(self) -> Vec<u8> {
252        serializer::serialize_stsz_data(self)
253    }
254}
255
256mod serializer {
257    use super::SampleSizeAtom;
258
259    pub fn serialize_stsz_data(stsz: SampleSizeAtom) -> Vec<u8> {
260        let mut data = Vec::new();
261
262        data.push(stsz.version);
263        data.extend(stsz.flags);
264        data.extend(stsz.sample_size.to_be_bytes());
265        data.extend(stsz.sample_count.to_be_bytes());
266
267        // If sample_size is 0, write the sample size table
268        if stsz.sample_size == 0 {
269            for size in stsz.entry_sizes.0.into_iter() {
270                data.extend(size.to_be_bytes());
271            }
272        }
273
274        data
275    }
276}
277
278mod parser {
279    use winnow::{
280        binary::be_u32,
281        combinator::{repeat, seq, trace},
282        error::StrContext,
283        ModalResult, Parser,
284    };
285
286    use super::{SampleEntrySizes, SampleSizeAtom};
287    use crate::atom::util::parser::{flags3, version, Stream};
288
289    pub fn parse_stsz_data(input: &mut Stream<'_>) -> ModalResult<SampleSizeAtom> {
290        trace(
291            "stsz",
292            seq!(SampleSizeAtom {
293                version: version,
294                flags: flags3,
295                sample_size: be_u32.context(StrContext::Label("sample_size")),
296                sample_count: be_u32.context(StrContext::Label("sample_count")),
297                entry_sizes: repeat(0.., be_u32.context(StrContext::Label("entry_size")))
298                    .map(SampleEntrySizes)
299                    .context(StrContext::Label("entry_sizes")),
300            })
301            .context(StrContext::Label("stsz")),
302        )
303        .parse_next(input)
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use crate::atom::test_utils::test_atom_roundtrip;
311
312    /// Test round-trip for all available stsz test data files
313    #[test]
314    fn test_stsz_roundtrip() {
315        test_atom_roundtrip::<SampleSizeAtom>(STSZ);
316    }
317}
mp4_edit/atom/leaf/stsz.rs

mp4_edit/atom/leaf/
stsz.rs
