From 58b6d7c7d5dbdd1513dca9822bce32bd43f32f3a Mon Sep 17 00:00:00 2001
From: Luke Street <luke.street@encounterpc.com>
Date: Sat, 6 Jan 2024 21:57:42 -0700
Subject: [PATCH] Avoid duplicate unit names when auto-splitting

Fixes #20
---
 Cargo.lock        |  2 +-
 Cargo.toml        |  2 +-
 src/util/split.rs | 82 ++++++++++++++++++++++++++++++++++++-----------
 3 files changed, 65 insertions(+), 21 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 76702cd..9d4089a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -295,7 +295,7 @@ dependencies = [
 
 [[package]]
 name = "decomp-toolkit"
-version = "0.6.5"
+version = "0.7.0"
 dependencies = [
  "anyhow",
  "ar",
diff --git a/Cargo.toml b/Cargo.toml
index 0a1dc62..60c20af 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,7 +3,7 @@ name = "decomp-toolkit"
 description = "Yet another GameCube/Wii decompilation toolkit."
 authors = ["Luke Street <luke@street.dev>"]
 license = "MIT OR Apache-2.0"
-version = "0.6.5"
+version = "0.7.0"
 edition = "2021"
 publish = false
 repository = "https://github.com/encounter/decomp-toolkit"
diff --git a/src/util/split.rs b/src/util/split.rs
index af3518b..4c1a0e2 100644
--- a/src/util/split.rs
+++ b/src/util/split.rs
@@ -65,18 +65,10 @@ fn split_ctors_dtors(obj: &mut ObjInfo, start: SectionAddress, end: SectionAddre
         }
 
         if ctors_split.is_none() || function_split.is_none() {
-            let unit = expected_unit.unwrap_or_else(|| {
-                let section_name = function_symbol
-                    .section
-                    .and_then(|idx| obj.sections.get(idx).map(|s| s.name.clone()))
-                    .unwrap_or_else(|| "unknown".to_string());
-                let name =
-                    sanitize_with_options(&function_symbol.name, &sanitise_file_name::Options {
-                        length_limit: 20,
-                        ..Default::default()
-                    });
-                format!("auto_{}_{}", name, section_name.trim_start_matches('.'))
-            });
+            let unit = match expected_unit {
+                Some(unit) => unit,
+                None => auto_unit_name(obj, function_symbol, &new_splits)?,
+            };
             log::debug!("Adding splits to unit {}", unit);
 
             if ctors_split.is_none() {
@@ -251,13 +243,10 @@ fn split_extabindex(obj: &mut ObjInfo, start: SectionAddress) -> Result<()> {
         }
 
         if extabindex_split.is_none() || extab_split.is_none() || function_split.is_none() {
-            let unit = expected_unit.unwrap_or_else(|| {
-                let section_name = function_symbol
-                    .section
-                    .and_then(|idx| obj.sections.get(idx).map(|s| s.name.clone()))
-                    .unwrap_or_else(|| "unknown".to_string());
-                format!("auto_{}_{}", function_symbol.name, section_name.trim_start_matches('.'))
-            });
+            let unit = match expected_unit {
+                Some(unit) => unit,
+                None => auto_unit_name(obj, function_symbol, &new_splits)?,
+            };
             log::debug!("Adding splits to unit {}", unit);
 
             if extabindex_split.is_none() {
@@ -384,6 +373,14 @@ fn create_gap_splits(obj: &mut ObjInfo) -> Result<()> {
                             symbol.name,
                             symbol.address
                         );
+                        if symbol.address & 3 != 0 {
+                            bail!(
+                                "Need to split at {:#010X} for duplicate symbol name {}, \
+                                but it is not 4-byte aligned. Please split manually.",
+                                symbol.address,
+                                symbol.name,
+                            );
+                        }
                         new_split_end.address = symbol.address as u32;
                         break;
                     }
@@ -1340,3 +1337,50 @@ pub fn end_for_section(obj: &ObjInfo, section_index: usize) -> Result<SectionAdd
     }
     Ok(SectionAddress::new(section_index, section_end))
 }
+
+/// Generates a unit name for an autogenerated split.
+/// The name is based on the symbol name and section name.
+/// If the name is not unique, a number is appended to the end.
+fn auto_unit_name(
+    obj: &ObjInfo,
+    symbol: &ObjSymbol,
+    new_splits: &BTreeMap<SectionAddress, ObjSplit>,
+) -> Result<String> {
+    let section_name = symbol
+        .section
+        .and_then(|idx| obj.sections.get(idx).map(|s| s.name.clone()))
+        .unwrap_or_else(|| "unknown".to_string());
+    let name = sanitize_with_options(&symbol.name, &sanitise_file_name::Options {
+        length_limit: 20,
+        ..Default::default()
+    });
+    let mut unit_name = format!("auto_{}_{}", name, section_name.trim_start_matches('.'));
+    // Ensure the name is unique
+    if unit_exists(&unit_name, obj, new_splits) {
+        let mut i = 1;
+        loop {
+            let new_unit_name = format!("{}_{}", unit_name, i);
+            if !unit_exists(&new_unit_name, obj, new_splits) {
+                unit_name = new_unit_name;
+                break;
+            }
+            i += 1;
+        }
+    }
+    Ok(unit_name)
+}
+
+/// Check if a unit name exists in the object or pending splits.
+fn unit_exists(
+    unit_name: &str,
+    obj: &ObjInfo,
+    new_splits: &BTreeMap<SectionAddress, ObjSplit>,
+) -> bool {
+    obj.link_order.iter().any(|u| u.name.eq_ignore_ascii_case(unit_name))
+        || obj
+            .sections
+            .iter()
+            .flat_map(|(_, s)| s.splits.iter())
+            .any(|(_, split)| split.unit.eq_ignore_ascii_case(unit_name))
+        || new_splits.values().any(|split| split.unit.eq_ignore_ascii_case(unit_name))
+}