Mark autogenerated splits & rework ObjInfo::add_split

src/analysis/signatures.rs

@@ -202,6 +202,7 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
     {
         apply_signature(obj, entry, &signature)?;
     }
+
     for &(name, sig_str) in SIGNATURES {
         if let Some((_, symbol)) = obj.symbols.by_name(name)? {
             let addr = symbol.address as u32;
@@ -210,6 +211,7 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
             }
         }
     }
+
     if let Some((_, symbol)) = obj.symbols.by_name("__init_user")? {
         // __init_user can be overridden, but we can still look for __init_cpp from it
         let mut analyzer = AnalyzerState::default();
@@ -225,6 +227,7 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
             }
         }
     }
+
     if let Some((_, symbol)) = obj.symbols.by_name("_ctors")? {
         // First entry of ctors is __init_cpp_exceptions
         let section = obj.section_at(symbol.address as u32)?;
@@ -260,30 +263,32 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
                         end: address as u32 + 4,
                         align: None,
                         common: false,
-                    });
+                        autogenerated: true,
+                    })?;
                 }
             }
         }
     }
+
     if let Some((_, symbol)) = obj.symbols.by_name("_dtors")? {
         let section = obj.section_at(symbol.address as u32)?;
         let address = symbol.address;
         let section_address = section.address;
         let section_index = section.index;
         // First entry of dtors is __destroy_global_chain
-        let target = read_u32(&section.data, address as u32, section_address as u32)
+        let dgc_target = read_u32(&section.data, address as u32, section_address as u32)
             .ok_or_else(|| anyhow!("Failed to read _dtors data"))?;
-        let target2 = read_u32(&section.data, address as u32 + 4, section_address as u32)
+        let fce_target = read_u32(&section.data, address as u32 + 4, section_address as u32)
             .ok_or_else(|| anyhow!("Failed to read _dtors data"))?;
-        let mut target_ok = false;
-        let mut target2_ok = false;
-        if target != 0 {
+        let mut found_dgc = false;
+        let mut found_fce = false;
+        if dgc_target != 0 {
             if let Some(signature) = check_signatures_str(
                 obj,
-                target,
+                dgc_target,
                 include_str!("../../assets/signatures/__destroy_global_chain.yml"),
             )? {
-                apply_signature(obj, target, &signature)?;
+                apply_signature(obj, dgc_target, &signature)?;
                 obj.add_symbol(
                     ObjSymbol {
                         name: "__destroy_global_chain_reference".to_string(),
@@ -299,17 +304,22 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
                     },
                     true,
                 )?;
-                target_ok = true;
+                found_dgc = true;
+            } else {
+                log::warn!(
+                    "Failed to match __destroy_global_chain signature ({:#010X})",
+                    dgc_target
+                );
             }
         }
         // Second entry of dtors is __fini_cpp_exceptions
-        if target2 != 0 {
+        if fce_target != 0 {
             if let Some(signature) = check_signatures_str(
                 obj,
-                target2,
+                fce_target,
                 include_str!("../../assets/signatures/__fini_cpp_exceptions.yml"),
             )? {
-                apply_signature(obj, target2, &signature)?;
+                apply_signature(obj, fce_target, &signature)?;
                 obj.add_symbol(
                     ObjSymbol {
                         name: "__fini_cpp_exceptions_reference".to_string(),
@@ -325,19 +335,27 @@ pub fn apply_signatures(obj: &mut ObjInfo) -> Result<()> {
                     },
                     true,
                 )?;
-                target2_ok = true;
+                found_fce = true;
             }
         }
 
-        if target_ok && target2_ok && obj.split_for(address as u32).is_none() {
-            obj.add_split(address as u32, ObjSplit {
-                unit: "__init_cpp_exceptions.cpp".to_string(),
-                end: address as u32 + 8,
-                align: None,
-                common: false,
-            });
+        if found_dgc {
+            let mut end = address as u32 + 4;
+            if found_fce {
+                end += 4;
+            }
+            if obj.split_for(address as u32).is_none() {
+                obj.add_split(address as u32, ObjSplit {
+                    unit: "__init_cpp_exceptions.cpp".to_string(),
+                    end,
+                    align: None,
+                    common: false,
+                    autogenerated: true,
+                })?;
+            }
         }
     }
+
     Ok(())
 }
 

src/obj/mod.rs

@@ -2,8 +2,8 @@ pub mod signatures;
 pub mod split;
 
 use std::{
-    cmp::min,
-    collections::{btree_map, BTreeMap, HashMap},
+    cmp::{max, min},
+    collections::{btree_map, BTreeMap, BTreeSet, HashMap},
     hash::{Hash, Hasher},
     ops::{Range, RangeBounds},
 };
@@ -144,7 +144,10 @@ pub struct ObjSplit {
     pub unit: String,
     pub end: u32,
     pub align: Option<u32>,
+    /// Common BSS
     pub common: bool,
+    /// Generated, replaceable by user
+    pub autogenerated: bool,
 }
 
 type SymbolIndex = usize;
@@ -580,10 +583,186 @@ impl ObjInfo {
         self.splits.range(range).flat_map(|(addr, v)| v.iter().map(move |u| (*addr, u)))
     }
 
-    pub fn add_split(&mut self, address: u32, split: ObjSplit) {
-        log::debug!("Adding split @ {:#010X}: {:?}", address, split);
-        // TODO merge with preceding split if possible
+    pub fn split_for_unit(
+        &self,
+        unit: &str,
+        section: &ObjSection,
+    ) -> Result<Option<(u32, &ObjSplit)>> {
+        let mut result = None::<(u32, &ObjSplit)>;
+        for (addr, split) in self
+            .splits_for_range(section.address as u32..(section.address + section.size) as u32)
+            .filter(|(_, split)| split.unit == unit)
+        {
+            ensure!(
+                result.is_none(),
+                "Multiple splits for unit {} in section {}: {:#010X}, {:#010X}",
+                unit,
+                section.name,
+                result.unwrap().0,
+                addr
+            );
+            result = Some((addr, split));
+        }
+        Ok(result)
+    }
+
+    pub fn add_split(&mut self, address: u32, split: ObjSplit) -> Result<()> {
+        let section = self.section_at(address)?;
+        let section_start = section.address as u32;
+        let section_end = (section.address + section.size) as u32;
+        ensure!(
+            split.end == 0 || split.end <= section_end,
+            "Split {} {:#010X}-{:#010X} is outside section {} {:#010X}-{:#010X}",
+            split.unit,
+            address,
+            split.end,
+            section.name,
+            section_start,
+            section_end
+        );
+
+        if let Some((existing_addr, existing_split)) = self.split_for_unit(&split.unit, section)? {
+            let new_start = min(existing_addr, address);
+            let new_end = max(existing_split.end, split.end);
+
+            // TODO use highest alignment?
+            let new_align = match (split.align, existing_split.align) {
+                (Some(a), Some(b)) if a == b => Some(a),
+                (Some(a), Some(b)) => {
+                    bail!(
+                        "Conflicting alignment for split {} {} {:#010X}-{:#010X}: {:#X} != {:#X}",
+                        split.unit,
+                        section.name,
+                        existing_addr,
+                        existing_split.end,
+                        a,
+                        b
+                    );
+                }
+                (Some(a), _) => Some(a),
+                (_, Some(b)) => Some(b),
+                _ => None,
+            };
+
+            // TODO don't merge if common flag is different?
+            ensure!(
+                split.common == existing_split.common,
+                "Conflicting common flag for split {} {} {:#010X}-{:#010X} ({}) and {:#010X}-{:#010X} ({})",
+                split.unit,
+                section.name,
+                existing_addr,
+                existing_split.end,
+                existing_split.common,
+                address,
+                split.end,
+                split.common
+            );
+
+            // Only set autogenerated flag if both splits are autogenerated
+            let new_autogenerated = split.autogenerated && existing_split.autogenerated;
+
+            // If the new split is contained within the existing split, do nothing
+            if new_start >= existing_addr && new_end <= existing_split.end {
+                log::debug!(
+                    "Split {} {} {:#010X}-{:#010X} already covers {:#010X}-{:#010X}",
+                    split.unit,
+                    section.name,
+                    existing_addr,
+                    existing_split.end,
+                    address,
+                    split.end
+                );
+                return Ok(());
+            }
+
+            log::debug!(
+                "Extending split {} {} {:#010X}-{:#010X} to include {:#010X}-{:#010X}: {:#010X}-{:#010X}",
+                split.unit,
+                section.name,
+                existing_addr,
+                existing_split.end,
+                address,
+                split.end,
+                new_start,
+                new_end
+            );
+
+            // Check if new split overlaps any existing splits
+            let mut to_remove = BTreeSet::new();
+            let mut to_rename = BTreeSet::new();
+            for (existing_addr, existing_split) in self.splits_for_range(new_start..new_end) {
+                // TODO the logic in this method should be reworked, this is a hack
+                if split.autogenerated && !existing_split.autogenerated {
+                    log::debug!(
+                        "-> Found existing split {} {} {:#010X}-{:#010X} (not autogenerated)",
+                        existing_split.unit,
+                        section.name,
+                        existing_addr,
+                        existing_split.end
+                    );
+                    return Ok(());
+                }
+
+                ensure!(
+                    existing_split.autogenerated || existing_split.unit == split.unit,
+                    "New split {} {} {:#010X}-{:#010X} overlaps existing split {} {:#010X}-{:#010X}",
+                    split.unit,
+                    section.name,
+                    new_start,
+                    new_end,
+                    existing_split.unit,
+                    existing_addr,
+                    existing_split.end
+                );
+                log::debug!(
+                    "-> Replacing existing split {} {} {:#010X}-{:#010X}",
+                    existing_split.unit,
+                    section.name,
+                    existing_addr,
+                    existing_split.end
+                );
+                to_remove.insert(existing_addr);
+                if existing_split.unit != split.unit {
+                    to_rename.insert(existing_split.unit.clone());
+                }
+            }
+
+            // Remove overlapping splits
+            for addr in to_remove {
+                self.splits.remove(&addr);
+            }
+            // Rename any units that were overwritten
+            // TODO this should also merge with existing splits
+            for unit in to_rename {
+                for (existing_addr, existing) in self
+                    .splits
+                    .iter_mut()
+                    .flat_map(|(addr, v)| v.iter_mut().map(move |u| (addr, u)))
+                    .filter(|(_, split)| split.unit == unit)
+                {
+                    log::debug!(
+                        "-> Renaming {} {:#010X}-{:#010X} to {}",
+                        existing.unit,
+                        existing_addr,
+                        existing.end,
+                        split.unit
+                    );
+                    existing.unit = split.unit.clone();
+                }
+            }
+            self.add_split(new_start, ObjSplit {
+                unit: split.unit,
+                end: new_end,
+                align: new_align,
+                common: split.common,
+                autogenerated: new_autogenerated,
+            })?;
+            return Ok(());
+        }
+
+        log::debug!("Adding split @ {} {:#010X}: {:?}", section.name, address, split);
         self.splits.entry(address).or_default().push(split);
+        Ok(())
     }
 }