[llvm] [BOLT] Store FileSymRefs in a multimap (PR #98992)

[Amir Ayupov via llvm-commits]

https://github.com/aaupov created https://github.com/llvm/llvm-project/pull/98992



Test Plan: TBD


>From c893e49a13f59af26ad027b944baa32fb3213085 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Mon, 15 Jul 2024 22:16:33 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.4
---
 bolt/include/bolt/Rewrite/RewriteInstance.h |  2 +-
 bolt/lib/Rewrite/RewriteInstance.cpp        | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index af1d9b4b70a3d..16a82d5687de9 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -490,7 +490,7 @@ class RewriteInstance {
   std::unordered_map<const MCSymbol *, uint32_t> SymbolIndex;
 
   /// Store all non-zero symbols in this map for a quick address lookup.
-  std::map<uint64_t, llvm::object::SymbolRef> FileSymRefs;
+  std::multimap<uint64_t, llvm::object::SymbolRef> FileSymRefs;
 
   /// FILE symbols used for disambiguating split function parents.
   std::vector<ELFSymbolRef> FileSymbols;
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index ded2f577237fe..205be34c2336a 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -886,7 +886,7 @@ void RewriteInstance::discoverFileObjects() {
     if (SymName == "__hot_start" || SymName == "__hot_end")
       continue;
 
-    FileSymRefs[SymbolAddress] = Symbol;
+    FileSymRefs.emplace(SymbolAddress, Symbol);
 
     // Skip section symbols that will be registered by disassemblePLT().
     if (SymbolType == SymbolRef::ST_Debug) {
@@ -1433,7 +1433,16 @@ void RewriteInstance::registerFragments() {
     const uint64_t Address = BF->getAddress();
 
     // Get fragment's own symbol
-    const auto SymIt = FileSymRefs.find(Address);
+    auto SymIt = FileSymRefs.end();
+    auto EqualAddressSymRange = FileSymRefs.equal_range(Address);
+    while (EqualAddressSymRange.first != EqualAddressSymRange.second) {
+      auto EqualAddressSymIt = EqualAddressSymRange.first;
+      StringRef Name = cantFail(EqualAddressSymIt->second.getName());
+      if (Name.contains(ParentName)) {
+        SymIt = EqualAddressSymIt;
+        break;
+      }
+    }
     if (SymIt == FileSymRefs.end()) {
       BC->errs()
           << "BOLT-ERROR: symbol lookup failed for function at address 0x"