Merge pull request #234 from mozilla/create_sample_table-prealloc

Preallocate storage in create_sample_table

Author: baumanj

mp4parse/src/lib.rs

@@ -58,7 +58,7 @@ impl ToU64 for usize {
 
 /// A trait to indicate a type can be infallibly converted to `usize`.
 /// This should only be implemented for infallible conversions, so only unsigned types are valid.
-trait ToUsize {
+pub trait ToUsize {
     fn to_usize(self) -> usize;
 }
 

mp4parse_capi/src/lib.rs

@@ -43,6 +43,7 @@ use byteorder::WriteBytesExt;
 use num::{CheckedAdd, CheckedSub};
 use num::{PrimInt, Zero};
 use std::convert::TryFrom;
+use std::convert::TryInto;
 
 use std::io::Read;
 use std::ops::Neg;
@@ -60,6 +61,7 @@ use mp4parse::MediaContext;
 use mp4parse::MediaScaledTime;
 use mp4parse::MediaTimeScale;
 use mp4parse::SampleEntry;
+use mp4parse::ToUsize;
 use mp4parse::Track;
 use mp4parse::TrackScaledTime;
 use mp4parse::TrackTimeScale;
@@ -1385,6 +1387,21 @@ impl<'a> TimeToSampleIterator<'a> {
 // For example:
 // (1, 5), (5, 10), (9, 2) => (1, 5), (2, 5), (3, 5), (4, 5), (5, 10), (6, 10),
 // (7, 10), (8, 10), (9, 2)
+fn sample_to_chunk_iter<'a>(
+    stsc_samples: &'a TryVec<mp4parse::SampleToChunk>,
+    stco_offsets: &'a TryVec<u64>,
+) -> SampleToChunkIterator<'a> {
+    SampleToChunkIterator {
+        chunks: (0..0),
+        sample_count: 0,
+        stsc_peek_iter: stsc_samples.as_slice().iter().peekable(),
+        remain_chunk_count: stco_offsets
+            .len()
+            .try_into()
+            .expect("stco.entry_count is u32"),
+    }
+}
+
 struct SampleToChunkIterator<'a> {
     chunks: std::ops::Range<u32>,
     sample_count: u32,
@@ -1399,7 +1416,12 @@ impl<'a> Iterator for SampleToChunkIterator<'a> {
         let has_chunk = self.chunks.next().or_else(|| {
             self.chunks = self.locate();
             self.remain_chunk_count
-                .checked_sub(self.chunks.len() as u32)
+                .checked_sub(
+                    self.chunks
+                        .len()
+                        .try_into()
+                        .expect("len() of a Range<u32> must fit in u32"),
+                )
                 .and_then(|res| {
                     self.remain_chunk_count = res;
                     self.chunks.next()
@@ -1460,19 +1482,20 @@ fn create_sample_table(
         _ => false,
     };
 
-    let mut sample_table = TryVec::new();
     let mut sample_size_iter = stsz.sample_sizes.iter();
 
     // Get 'stsc' iterator for (chunk_id, chunk_sample_count) and calculate the sample
     // offset address.
-    let stsc_iter = SampleToChunkIterator {
-        chunks: (0..0),
-        sample_count: 0,
-        stsc_peek_iter: stsc.samples.as_slice().iter().peekable(),
-        remain_chunk_count: stco.offsets.len() as u32,
-    };
 
-    for i in stsc_iter {
+    // With large numbers of samples, the cost of many allocations dominates,
+    // so it's worth iterating twice to allocate sample_table just once.
+    let total_sample_count = sample_to_chunk_iter(&stsc.samples, &stco.offsets)
+        .by_ref()
+        .map(|(_, sample_counts)| sample_counts.to_usize())
+        .sum();
+    let mut sample_table = TryVec::with_capacity(total_sample_count).ok()?;
+
+    for i in sample_to_chunk_iter(&stsc.samples, &stco.offsets) {
         let chunk_id = i.0 as usize;
         let sample_counts = i.1;
         let mut cur_position = match stco.offsets.get(chunk_id) {
@@ -1565,7 +1588,8 @@ fn create_sample_table(
     // calculate to correct the composition end time.
     if !sample_table.is_empty() {
         // Create an index table refers to sample_table and sorted by start_composisiton time.
-        let mut sort_table = TryVec::new();
+        let mut sort_table = TryVec::with_capacity(sample_table.len()).ok()?;
+
         for i in 0..sample_table.len() {
             sort_table.push(i).ok()?;
         }