Mark autogenerated splits & rework ObjInfo::add_split

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(())
     }
 }
 

src/obj/split.rs

@@ -8,8 +8,8 @@ use itertools::Itertools;
 use petgraph::{graph::NodeIndex, Graph};
 
 use crate::obj::{
-    ObjArchitecture, ObjInfo, ObjKind, ObjReloc, ObjSection, ObjSectionKind, ObjSplit, ObjSymbol,
-    ObjSymbolFlagSet, ObjSymbolFlags, ObjSymbolKind,
+    ObjArchitecture, ObjInfo, ObjKind, ObjReloc, ObjRelocKind, ObjSection, ObjSectionKind,
+    ObjSplit, ObjSymbol, ObjSymbolFlagSet, ObjSymbolFlags, ObjSymbolKind, SymbolIndex,
 };
 
 /// Create splits for function pointers in the given section.
@@ -68,6 +68,7 @@ fn split_ctors_dtors(obj: &mut ObjInfo, section_start: u32, section_end: u32) ->
                     end: current_address + 4,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
             }
             if function_split.is_none() {
@@ -77,6 +78,7 @@ fn split_ctors_dtors(obj: &mut ObjInfo, section_start: u32, section_end: u32) ->
                     end: function_addr + function_symbol.size as u32,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
             }
         }
@@ -85,7 +87,7 @@ fn split_ctors_dtors(obj: &mut ObjInfo, section_start: u32, section_end: u32) ->
     }
 
     for (addr, split) in new_splits {
-        obj.add_split(addr, split);
+        obj.add_split(addr, split)?;
     }
 
     Ok(())
@@ -203,30 +205,40 @@ fn split_extabindex(obj: &mut ObjInfo, section_index: usize, section_start: u32)
             log::debug!("Adding splits to unit {}", unit);
 
             if extabindex_split.is_none() {
-                log::debug!("Adding split for extabindex entry @ {:#010X}", current_address);
+                let end = current_address + 12;
+                log::debug!(
+                    "Adding split for extabindex entry @ {:#010X}-{:#010X}",
+                    current_address,
+                    end
+                );
                 new_splits.insert(current_address, ObjSplit {
                     unit: unit.clone(),
-                    end: current_address + 12,
+                    end,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
             }
             if extab_split.is_none() {
-                log::debug!("Adding split for extab @ {:#010X}", extab_addr);
+                let end = extab_addr + extab_symbol.size as u32;
+                log::debug!("Adding split for extab @ {:#010X}-{:#010X}", extab_addr, end);
                 new_splits.insert(extab_addr, ObjSplit {
                     unit: unit.clone(),
-                    end: extab_addr + extab_symbol.size as u32,
+                    end,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
             }
             if function_split.is_none() {
-                log::debug!("Adding split for function @ {:#010X}", function_addr);
+                let end = function_addr + function_symbol.size as u32;
+                log::debug!("Adding split for function @ {:#010X}-{:#010X}", function_addr, end);
                 new_splits.insert(function_addr, ObjSplit {
                     unit,
-                    end: function_addr + function_symbol.size as u32,
+                    end,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
             }
         }
@@ -235,7 +247,7 @@ fn split_extabindex(obj: &mut ObjInfo, section_index: usize, section_start: u32)
     }
 
     for (addr, split) in new_splits {
-        obj.add_split(addr, split);
+        obj.add_split(addr, split)?;
     }
 
     Ok(())
@@ -310,6 +322,7 @@ fn create_gap_splits(obj: &mut ObjInfo) -> Result<()> {
                     end: new_split_end,
                     align: None,
                     common: false,
+                    autogenerated: true,
                 });
                 current_address = new_split_end;
                 continue;
@@ -330,7 +343,7 @@ fn create_gap_splits(obj: &mut ObjInfo) -> Result<()> {
 
     // Add new splits
     for (addr, split) in new_splits {
-        obj.add_split(addr, split);
+        obj.add_split(addr, split)?;
     }
 
     Ok(())