diff --git a/src/hash_table.rs b/src/hash_table.rs index e818bef..ba5960b 100644 --- a/src/hash_table.rs +++ b/src/hash_table.rs @@ -63,7 +63,6 @@ fn read_value( let mut last_hash = 0; while input.len() > end_len { - let mut counts = vec![]; let (bytes, hash) = le_u64(input)?; last_hash = hash; if bytes.len() <= end_len { @@ -83,6 +82,7 @@ fn read_value( )]; return Err(nom::Err::Failure(VerboseError { errors })); } + let mut counts = Vec::with_capacity(counts_len as usize); for _ in 0..counts_len { let (bytes, count) = le_u64(input)?; input = bytes; @@ -136,8 +136,8 @@ fn read_value( } impl HashTable { - fn new() -> Self { - Self(IndexMap::new()) + fn with_capacity(capacity: usize) -> Self { + Self(IndexMap::with_capacity(capacity)) } /// buckets is the data the hash table buckets start at - the start of the `HashTable` in memory. @@ -155,7 +155,7 @@ impl HashTable { debug!("Number of entries: {}", num_entries); let mut payload = input; - let mut result = Self::new(); + let mut result = Self::with_capacity(num_entries as usize); //TODO is this change right? while num_entries > 0 { let (bytes, entries) = result.parse_bucket(version, payload, num_entries)?; diff --git a/src/instrumentation_profile/indexed_profile.rs b/src/instrumentation_profile/indexed_profile.rs index f05cc6b..70dcb18 100644 --- a/src/instrumentation_profile/indexed_profile.rs +++ b/src/instrumentation_profile/indexed_profile.rs @@ -4,7 +4,7 @@ use crate::summary::*; use anyhow::bail; use nom::{ error::{ContextError, ErrorKind, ParseError}, - number::{complete::*, Endianness}, + number::complete::*, }; use rustc_hash::FxHashMap; use std::convert::TryFrom; @@ -191,13 +191,14 @@ impl InstrProfReader for IndexedInstrProf { )?; debug!("Function hash table: {:?}", table); input = bytes; + profile.reserve_records(table.0.len()); + profile.symtab.names.reserve(table.0.len()); for ((hash, name), v) in &table.0 { let name = name.to_string(); + let name_hash = compute_hash(&name); profile .symtab - .add_func_name(name.clone(), Some(Endianness::Little)); - - let name_hash = compute_hash(&name); + .add_func_name_with_hash(name.clone(), name_hash); let record = NamedInstrProfRecord { name: Some(name), name_hash: Some(name_hash), diff --git a/src/instrumentation_profile/raw_profile.rs b/src/instrumentation_profile/raw_profile.rs index ab056d7..5ab7d38 100644 --- a/src/instrumentation_profile/raw_profile.rs +++ b/src/instrumentation_profile/raw_profile.rs @@ -248,17 +248,41 @@ where } else { let mut counts = Vec::::with_capacity(data.num_counters as usize); bytes = &bytes[(counter_offset as usize)..]; - for _ in 0..(data.num_counters as usize) { - let counter = if header.has_byte_coverage() { + if header.has_byte_coverage() { + if data.num_counters as usize > bytes.len() { + let pos = &bytes[bytes.len()..]; + return Err(Err::Failure(VerboseError::from_error_kind( + pos, + ErrorKind::Eof, + ))); + } + for _ in 0..(data.num_counters as usize) { let counter = bytes[0]; bytes = &bytes[1..]; - (counter == 0) as u64 - } else { - let (b, counter) = nom_u64(header.endianness)(bytes)?; - bytes = b; - counter - }; - counts.push(counter); + counts.push((counter == 0) as u64); + } + } else { + let counters_bytes = data.num_counters as usize * size_of::(); + if counters_bytes > bytes.len() { + let pos = &bytes[bytes.len()..]; + return Err(Err::Failure(VerboseError::from_error_kind( + pos, + ErrorKind::Eof, + ))); + } + for _ in 0..(data.num_counters as usize) { + let (count_bytes, remaining) = bytes.split_at(size_of::()); + bytes = remaining; + let count_bytes = count_bytes + .try_into() + .expect("split_at returned exactly eight counter bytes"); + let counter = match header.endianness { + Endianness::Little => u64::from_le_bytes(count_bytes), + Endianness::Big => u64::from_be_bytes(count_bytes), + Endianness::Native => u64::from_ne_bytes(count_bytes), + }; + counts.push(counter); + } } let record = InstrProfRecord { counts, @@ -294,8 +318,8 @@ where fn parse_bytes(mut input: &[u8]) -> ParseResult<'_, InstrumentationProfile> { if !input.is_empty() { - let mut result = InstrumentationProfile::default(); let (bytes, header) = Self::parse_header(input)?; + let mut result = InstrumentationProfile::with_capacity(header.data_len as usize); // LLVM 11 and 12 are version 5. LLVM 13 is version 7 let version_num = header.version(); result.version = Some(version_num); @@ -313,7 +337,7 @@ where ))); } input = &bytes[(header.binary_ids_len as usize)..]; - let mut data_section = vec![]; + let mut data_section = Vec::with_capacity(header.data_len as usize); for _ in 0..header.data_len { let (bytes, data) = ProfileData::::parse(input, &header)?; debug!("Parsed data section {:?}", data); @@ -334,7 +358,7 @@ where } }; input = bytes; - let mut counters = vec![]; + let mut counters = Vec::with_capacity(data_section.len()); let mut counters_delta = header.counters_delta; // Okay so the counters section looks a bit hairy. So as a brief explanation. @@ -368,7 +392,7 @@ where let (bytes, _) = take(counters_end)(input)?; input = bytes; let end_length = input.len() - header.names_len as usize; - let mut symtab = Symtab::default(); + let mut symtab = Symtab::with_capacity(data_section.len()); while input.len() > end_length { let (new_bytes, names) = parse_string_ref(input)?; debug!( diff --git a/src/instrumentation_profile/types.rs b/src/instrumentation_profile/types.rs index 68e6515..2bcd9d3 100644 --- a/src/instrumentation_profile/types.rs +++ b/src/instrumentation_profile/types.rs @@ -34,6 +34,14 @@ pub struct Symtab { pub names: FxHashMap, } +impl Symtab { + pub(crate) fn with_capacity(capacity: usize) -> Self { + let mut names = FxHashMap::default(); + names.reserve(capacity); + Self { names } + } +} + pub fn compute_hash(data: impl AsRef<[u8]>) -> u64 { let hash = md5::compute(data).0[..8].try_into().unwrap_or_default(); u64::from_le_bytes(hash) @@ -124,6 +132,22 @@ impl InstrumentationProfile { } } + pub(crate) fn with_capacity(capacity: usize) -> Self { + let mut record_name_lookup = FxHashMap::default(); + record_name_lookup.reserve(capacity); + Self { + records: Vec::with_capacity(capacity), + record_name_lookup, + symtab: Symtab::with_capacity(capacity), + ..Default::default() + } + } + + pub(crate) fn reserve_records(&mut self, additional: usize) { + self.records.reserve(additional); + self.record_name_lookup.reserve(additional); + } + pub fn version(&self) -> Option { self.version } diff --git a/src/lib.rs b/src/lib.rs index aa3f02b..fb9c0d7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,18 +25,14 @@ pub fn merge_profiles(files: &[T]) -> std::io::Result where T: AsRef, { - if files.is_empty() { - Ok(InstrumentationProfile::default()) - } else { - let mut profiles = vec![]; - for input in files { - let profile = parse(input)?; - profiles.push(profile); - } - let mut base = profiles.remove(0); - for profile in &profiles { - base.merge(profile); - } - Ok(base) + let Some((first, rest)) = files.split_first() else { + return Ok(InstrumentationProfile::default()); + }; + + let mut base = parse(first)?; + for input in rest { + let profile = parse(input)?; + base.merge(&profile); } + Ok(base) }