From 61cd72cd11126be957ba568eba09ac09eb03114d Mon Sep 17 00:00:00 2001
From: Luke Street <luke.street@encounterpc.com>
Date: Mon, 3 Jun 2024 18:08:06 -0600
Subject: [PATCH] Support address in config symbol references

Example: `symbol_name!.data:0x1234`
Allows disambiguating local symbols
with the same name.

Supported in `extract` and
`add_relocations` in config.yml

Resolves #58
---
 src/cmd/dol.rs     |  9 ++++-----
 src/obj/symbols.rs | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/cmd/dol.rs b/src/cmd/dol.rs
index 6391a26..74073a9 100644
--- a/src/cmd/dol.rs
+++ b/src/cmd/dol.rs
@@ -947,11 +947,10 @@ fn split_write_obj(
         let (_, symbol) = module
             .obj
             .symbols
-            .by_name(&extract.symbol)?
+            .by_ref(&module.obj.sections, &extract.symbol)?
             .with_context(|| format!("Failed to locate symbol '{}'", extract.symbol))?;
-        let section_index = symbol
-            .section
-            .with_context(|| format!("Symbol '{}' has no section", extract.symbol))?;
+        let section_index =
+            symbol.section.with_context(|| format!("Symbol '{}' has no section", symbol.name))?;
         let section = &module.obj.sections[section_index];
         let data = section.symbol_data(symbol)?;
 
@@ -1893,7 +1892,7 @@ fn apply_block_relocations(
 fn apply_add_relocations(obj: &mut ObjInfo, relocations: &[AddRelocationConfig]) -> Result<()> {
     for reloc in relocations {
         let SectionAddress { section, address } = reloc.source.resolve(obj)?;
-        let (target_symbol, _) = match obj.symbols.by_name(&reloc.target)? {
+        let (target_symbol, _) = match obj.symbols.by_ref(&obj.sections, &reloc.target)? {
             Some(v) => v,
             None => {
                 // Assume external symbol
diff --git a/src/obj/symbols.rs b/src/obj/symbols.rs
index 731189d..cce02f7 100644
--- a/src/obj/symbols.rs
+++ b/src/obj/symbols.rs
@@ -12,9 +12,9 @@ use serde_repr::{Deserialize_repr, Serialize_repr};
 
 use crate::{
     analysis::cfa::SectionAddress,
-    obj::{ObjKind, ObjRelocKind},
+    obj::{ObjKind, ObjRelocKind, ObjSections},
     util::{
-        config::{is_auto_jump_table, is_auto_label, is_auto_symbol},
+        config::{is_auto_jump_table, is_auto_label, is_auto_symbol, parse_u32},
         nested::NestedVec,
         split::is_linker_generated_label,
     },
@@ -477,6 +477,40 @@ impl ObjSymbols {
         Ok(result)
     }
 
+    /// Locate a symbol by name, with optional reference attributes. Example:
+    /// `symbol_name!.data:0x1234` will find the symbol named `symbol_name`
+    /// in the `.data` section at address `0x1234`.
+    pub fn by_ref<'a>(
+        &'a self,
+        sections: &ObjSections,
+        symbol_ref: &str,
+    ) -> Result<Option<(SymbolIndex, &'a ObjSymbol)>> {
+        if let Some((name, rest)) = symbol_ref.split_once('!') {
+            let (section_index, address) = if let Some((section_name, rest)) = rest.split_once(':')
+            {
+                let section_index = sections
+                    .by_name(section_name)?
+                    .map(|(idx, _)| idx)
+                    .ok_or_else(|| anyhow!("Section not found: {}", section_name))?;
+                (Some(section_index), parse_u32(rest)?)
+            } else {
+                (None, parse_u32(rest)?)
+            };
+            let mut out = None;
+            for (index, symbol) in self.for_name(name) {
+                if (section_index.is_none() || symbol.section == section_index)
+                    && symbol.address == address as u64
+                {
+                    ensure!(out.is_none(), "Multiple symbols matched {}", symbol_ref);
+                    out = Some((index, symbol));
+                }
+            }
+            Ok(out)
+        } else {
+            self.by_name(symbol_ref)
+        }
+    }
+
     pub fn by_kind(
         &self,
         kind: ObjSymbolKind,