Extract splits from PDB Section Contributions

Cargo.toml

@@ -71,7 +71,7 @@ tracing-attributes = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 xxhash-rust = { version = "0.8", features = ["xxh3"] }
 zerocopy = { version = "0.8", features = ["derive"] }
-pdb = "0.8.0"
+pdb2 = "0.10.1"
 lzxd = "0.2.6"
 
 [target.'cfg(target_env = "musl")'.dependencies]

src/cmd/xex.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::BTreeMap,
+    collections::{BTreeMap, HashSet},
     fs::{self, DirBuilder, File},
     io::{BufWriter, Write},
     time::UNIX_EPOCH,
@@ -32,7 +32,7 @@ use crate::{
     },
     obj::{
         best_match_for_reloc, ObjInfo, ObjKind, ObjRelocKind, ObjSectionKind, ObjSections,
-        ObjSymbolKind, ObjSymbolScope, SectionIndex, SymbolIndex,
+        ObjSymbolKind, ObjSymbolScope, ObjUnit, SectionIndex, SymbolIndex,
     },
     util::{
         asm::write_asm,
@@ -472,7 +472,32 @@ fn load_analyze_xex(config: &ProjectConfig) -> Result<ExeAnalyzeResult> {
 
     if let Some(pdb_path) = &config.base.pdb {
         let pdb_path: Utf8NativePathBuf = pdb_path.with_encoding();
-        let pdb_syms = try_parse_pdb(&pdb_path, &obj.sections)?;
+        let (pdb_units, pdb_splits, pdb_syms) = try_parse_pdb(&pdb_path, &obj.sections)?;
+
+        // Apply all the splits
+        // FIXME: Don't add splits unconditionally here; it may conflict with
+        // user-provided splits. For now, users can comment out the pdb key
+        // in config.yml after initial analysis
+        for (i, splits_for_section) in pdb_splits.into_iter().enumerate() {
+            for (start, split) in splits_for_section.iter() {
+                obj.sections[i as u32].splits.push(start, split.clone());
+            }
+        }
+
+        // Apply all the units, discarding the ones with no splits
+        let mut nonempty_mods = HashSet::new();
+        for split in obj.sections.all_splits() {
+            nonempty_mods.insert(&split.3.unit);
+        }
+        for unit in pdb_units {
+            if nonempty_mods.contains(&unit) {
+                obj.link_order.push(ObjUnit { name: unit, autogenerated: false, order: None });
+            } else {
+                log::debug!("Module {} is empty", unit);
+            }
+        }
+
+        // Apply all the symbols
         for mut sym in pdb_syms.into_iter() {
             if !is_reg_intrinsic(&sym.name) && sym.name != "__NLG_Return" {
                 match obj.sections.at_address(sym.address as u32).ok() {

src/obj/splits.rs

@@ -9,7 +9,7 @@ use crate::{
 };
 
 /// Marks a split point within a section.
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Eq, PartialEq)]
 pub struct ObjSplit {
     pub unit: String,
     pub end: u32,
@@ -106,10 +106,27 @@ impl ObjSplits {
             .map_err(|_| anyhow!("Multiple splits for unit {}", unit))
     }
 
-    pub fn push(&mut self, address: u32, split: ObjSplit) {
+    /// Get the ObjSplit provided by unit with the specified rename,
+    /// if it exists
+    pub fn for_unit_rename(
+        &mut self,
+        unit: &str,
+        rename: Option<&str>,
+    ) -> Result<Option<(u32, &mut ObjSplit)>> {
+        self.splits
+            .iter_mut()
+            .flat_map(|(addr, v)| v.iter_mut().map(move |u| (*addr, u)))
+            .filter(|(_, split)| split.unit == unit && split.rename.as_deref() == rename)
+            .at_most_one()
+            .map_err(|_| anyhow!("Multiple splits for unit {} with rename {:?}", unit, rename))
+    }
+
+    /// Add the split, returning a mutable reference to it within the vector
+    pub fn push(&mut self, address: u32, split: ObjSplit) -> &mut ObjSplit {
         let out = self.splits.entry(address).or_default();
         out.push(split);
         out.sort_by_key(|s| s.end);
+        out.last_mut().unwrap()
     }
 
     pub fn remove(&mut self, address: u32) -> Option<Vec<ObjSplit>> { self.splits.remove(&address) }