diff --git a/mp4parse/src/lib.rs b/mp4parse/src/lib.rs index a8997fd2..854a1364 100644 --- a/mp4parse/src/lib.rs +++ b/mp4parse/src/lib.rs @@ -161,6 +161,31 @@ struct HashMap; #[allow(dead_code)] struct String; +// Arbitrary 1MB limit. +const CAPACITY_HINT_LIMIT: usize = 1024 * 1024; + +/// Returns a `TryVec` with preallocated capacity for up to `expected_items`. +/// `expected_items` is treated as a hint only and the returned `TryVec`'s +/// capacity may be less than `expected_items`. +fn vec_with_capacity_hint(expected_items: usize) -> Result, TryReserveError> { + let reserved_items = expected_items.min(CAPACITY_HINT_LIMIT / std::mem::size_of::()); + TryVec::with_capacity(reserved_items) +} + +/// Returns a `TryHashMap` with preallocated capacity for up to `expected_items`. +/// `expected_items` is treated as a hint only and the returned `TryHashMap`'s +/// capacity may be less than `expected_items`. +fn hashmap_with_capacity_hint( + expected_items: usize, +) -> Result, TryReserveError> +where + K: Eq + std::hash::Hash, +{ + let reserved_items = expected_items + .min(CAPACITY_HINT_LIMIT / (std::mem::size_of::() + std::mem::size_of::())); + TryHashMap::with_capacity(reserved_items) +} + /// The return value to the C API /// Any detail that needs to be communicated to the caller must be encoded here /// since the [`Error`] type's associated data is part of the FFI. @@ -2962,7 +2987,7 @@ fn read_iinf( } else { be_u32(src)?.to_usize() }; - let mut item_infos = TryVec::with_capacity(entry_count)?; + let mut item_infos = vec_with_capacity_hint(entry_count)?; let mut iter = src.box_iter(); while let Some(mut b) = iter.next_box()? { @@ -4071,7 +4096,7 @@ fn read_iloc(src: &mut BMFFBox) -> Result iloc.read_u32(32)?, }; - let mut items = TryHashMap::with_capacity(item_count.to_usize())?; + let mut items = hashmap_with_capacity_hint(item_count.to_usize())?; for _ in 0..item_count { let item_id = ItemId(match version { @@ -4691,7 +4716,7 @@ fn read_tkhd(src: &mut BMFFBox) -> Result { fn read_elst(src: &mut BMFFBox) -> Result { let (version, flags) = read_fullbox_extra(src)?; let edit_count = be_u32(src)?; - let mut edits = TryVec::with_capacity(edit_count.to_usize())?; + let mut edits = vec_with_capacity_hint(edit_count.to_usize())?; for _ in 0..edit_count { let (segment_duration, media_time) = match version { 1 => { @@ -4770,7 +4795,7 @@ fn read_mdhd(src: &mut BMFFBox) -> Result { fn read_stco(src: &mut BMFFBox) -> Result { let (_, _) = read_fullbox_extra(src)?; let offset_count = be_u32(src)?; - let mut offsets = TryVec::with_capacity(offset_count.to_usize())?; + let mut offsets = vec_with_capacity_hint(offset_count.to_usize())?; for _ in 0..offset_count { offsets.push(be_u32(src)?.into())?; } @@ -4786,7 +4811,7 @@ fn read_stco(src: &mut BMFFBox) -> Result { fn read_co64(src: &mut BMFFBox) -> Result { let (_, _) = read_fullbox_extra(src)?; let offset_count = be_u32(src)?; - let mut offsets = TryVec::with_capacity(offset_count.to_usize())?; + let mut offsets = vec_with_capacity_hint(offset_count.to_usize())?; for _ in 0..offset_count { offsets.push(be_u64(src)?)?; } @@ -4802,7 +4827,7 @@ fn read_co64(src: &mut BMFFBox) -> Result { fn read_stss(src: &mut BMFFBox) -> Result { let (_, _) = read_fullbox_extra(src)?; let sample_count = be_u32(src)?; - let mut samples = TryVec::with_capacity(sample_count.to_usize())?; + let mut samples = vec_with_capacity_hint(sample_count.to_usize())?; for _ in 0..sample_count { samples.push(be_u32(src)?)?; } @@ -4818,7 +4843,7 @@ fn read_stss(src: &mut BMFFBox) -> Result { fn read_stsc(src: &mut BMFFBox) -> Result { let (_, _) = read_fullbox_extra(src)?; let sample_count = be_u32(src)?; - let mut samples = TryVec::with_capacity(sample_count.to_usize())?; + let mut samples = vec_with_capacity_hint(sample_count.to_usize())?; for _ in 0..sample_count { let first_chunk = be_u32(src)?; let samples_per_chunk = be_u32(src)?; @@ -4850,7 +4875,7 @@ fn read_ctts(src: &mut BMFFBox) -> Result { return Status::CttsBadSize.into(); } - let mut offsets = TryVec::with_capacity(counts.to_usize())?; + let mut offsets = vec_with_capacity_hint(counts.to_usize())?; for _ in 0..counts { let (sample_count, time_offset) = match version { // According to spec, Version0 shoule be used when version == 0; @@ -4904,7 +4929,7 @@ fn read_stsz(src: &mut BMFFBox) -> Result { fn read_stts(src: &mut BMFFBox) -> Result { let (_, _) = read_fullbox_extra(src)?; let sample_count = be_u32(src)?; - let mut samples = TryVec::with_capacity(sample_count.to_usize())?; + let mut samples = vec_with_capacity_hint(sample_count.to_usize())?; for _ in 0..sample_count { let sample_count = be_u32(src)?; let sample_delta = be_u32(src)?; @@ -5892,7 +5917,7 @@ fn read_stsd(src: &mut BMFFBox, track: &mut Track) -> Result