Improve performance

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)?;

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),

src/instrumentation_profile/raw_profile.rs

@@ -248,17 +248,41 @@ where
         } else {
             let mut counts = Vec::<u64>::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::<u64>();
+                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::<u64>());
+                    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::<T>::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!(

src/instrumentation_profile/types.rs

@@ -34,6 +34,14 @@ pub struct Symtab {
     pub names: FxHashMap<u64, String>,
 }
 
+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<u64> {
         self.version
     }

src/lib.rs

@@ -25,18 +25,14 @@ pub fn merge_profiles<T>(files: &[T]) -> std::io::Result<InstrumentationProfile>
 where
     T: AsRef<Path>,
 {
-    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)
 }