[BOLT] Refactor handling of branch targets. NFCI

[maksfb]

Refactor code that verifies external branch destinations and creates secondary entry points.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

Refactor code that verifies external branch destinations and creates secondary entry points.

---

Full diff: https://github.com/llvm/llvm-project/pull/165828.diff

4 Files Affected:

• (modified) bolt/include/bolt/Core/BinaryContext.h (+10)
• (modified) bolt/lib/Core/BinaryContext.cpp (+21-11)
• (modified) bolt/lib/Core/BinaryFunction.cpp (+5-14)
• (modified) bolt/test/AArch64/constant-island-entry.s (+1-1)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 5cbc28fb38a33..085c0265de3ed 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -932,6 +932,16 @@ class BinaryContext {
   std::pair<const MCSymbol *, uint64_t>
   handleAddressRef(uint64_t Address, BinaryFunction &BF, bool IsPCRel);
 
+  /// When \p Address inside function \p BF is a target of a control transfer
+  /// instruction (branch) from another function, return a corresponding symbol
+  /// that should be used by the branch. For example, main or secondary entry
+  /// point.
+  ///
+  /// If \p Address is an invalid destination, such as a constant island, return
+  /// nullptr and mark \p BF as ignored, since we cannot properly handle a
+  /// branch to a constant island.
+  MCSymbol *handleExternalBranchTarget(uint64_t Address, BinaryFunction &BF);
+
   /// Analyze memory contents at the given \p Address and return the type of
   /// memory contents (such as a possible jump table).
   MemoryContentsType analyzeMemoryAt(uint64_t Address, BinaryFunction &BF);
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index c33540ada8a05..a383ced1712e3 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -518,6 +518,23 @@ BinaryContext::handleAddressRef(uint64_t Address, BinaryFunction &BF,
   return std::make_pair(TargetSymbol, 0);
 }
 
+MCSymbol *BinaryContext::handleExternalBranchTarget(uint64_t Address,
+                                                    BinaryFunction &BF) {
+  if (BF.isInConstantIsland(Address)) {
+    BF.setIgnored();
+    this->outs() << "BOLT-WARNING: ignoring entry point at address 0x"
+                 << Twine::utohexstr(Address)
+                 << " in constant island of function " << BF << '\n';
+    return nullptr;
+  }
+
+  const uint64_t Offset = Address - BF.getAddress();
+  assert(Offset < BF.getSize() &&
+         "Address should be inside the referenced function");
+
+  return Offset ? BF.addEntryPointAtOffset(Offset) : BF.getSymbol();
+}
+
 MemoryContentsType BinaryContext::analyzeMemoryAt(uint64_t Address,
                                                   BinaryFunction &BF) {
   if (!isX86())
@@ -1399,17 +1416,10 @@ void BinaryContext::processInterproceduralReferences() {
             << Function.getPrintName() << " and "
             << TargetFunction->getPrintName() << '\n';
       }
-      if (uint64_t Offset = Address - TargetFunction->getAddress()) {
-        if (!TargetFunction->isInConstantIsland(Address)) {
-          TargetFunction->addEntryPointAtOffset(Offset);
-        } else {
-          TargetFunction->setIgnored();
-          this->outs() << "BOLT-WARNING: Ignoring entry point at address 0x"
-                       << Twine::utohexstr(Address)
-                       << " in constant island of function " << *TargetFunction
-                       << '\n';
-        }
-      }
+
+      // Create an extra entry point if needed. Can also render the target
+      // function ignored if the reference is invalid.
+      handleExternalBranchTarget(Address, *TargetFunction);
 
       continue;
     }
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index fbe186454351c..ddaad6eef6140 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1697,21 +1697,12 @@ bool BinaryFunction::scanExternalRefs() {
       if (!TargetFunction || ignoreFunctionRef(*TargetFunction))
         continue;
 
-      const uint64_t FunctionOffset =
-          TargetAddress - TargetFunction->getAddress();
-      if (!TargetFunction->isInConstantIsland(TargetAddress)) {
-        BranchTargetSymbol =
-            FunctionOffset
-                ? TargetFunction->addEntryPointAtOffset(FunctionOffset)
-                : TargetFunction->getSymbol();
-      } else {
-        TargetFunction->setIgnored();
-        BC.outs() << "BOLT-WARNING: Ignoring entry point at address 0x"
-                  << Twine::utohexstr(Address)
-                  << " in constant island of function " << *TargetFunction
-                  << '\n';
+      // Get a reference symbol for the function when address is a valid code
+      // reference.
+      BranchTargetSymbol =
+          BC.handleExternalBranchTarget(TargetAddress, *TargetFunction);
+      if (!BranchTargetSymbol)
         continue;
-      }
     }
 
     // Can't find more references. Not creating relocations since we are not
diff --git a/bolt/test/AArch64/constant-island-entry.s b/bolt/test/AArch64/constant-island-entry.s
index 7f8449deea130..a82b876fde46d 100644
--- a/bolt/test/AArch64/constant-island-entry.s
+++ b/bolt/test/AArch64/constant-island-entry.s
@@ -10,7 +10,7 @@
 ## Skip caller to check the identical warning is triggered from ScanExternalRefs().
 # RUN: llvm-bolt %t.exe -o %t.bolt -skip-funcs=caller 2>&1 | FileCheck %s
 
-# CHECK: BOLT-WARNING: Ignoring entry point at address 0x{{[0-9a-f]+}} in constant island of function func
+# CHECK: BOLT-WARNING: ignoring entry point at address 0x{{[0-9a-f]+}} in constant island of function func
 
 .globl func
 .type func, %function

[Jinjie-Huang]

Thanks! LGTM