[llvm] [BOLT] Allow getNewFunctionOrDataAddress to map addrs inside functions (PR #117766)

[Peter Waller via llvm-commits]

https://github.com/peterwaller-arm created https://github.com/llvm/llvm-project/pull/117766

Add logic to map addresses referring to non-entry basic blocks.

PR #101466 extended this function to enable it to map addresses for
the entry points of multi-entry functions, but this still left
references to individual basic blocks unmappable.


>From 1557cfd2ced9fdc03a547704ef7e1fd47e37803f Mon Sep 17 00:00:00 2001
From: Peter Waller <peter.waller at arm.com>
Date: Mon, 25 Nov 2024 13:26:19 +0000
Subject: [PATCH] [BOLT] Allow getNewFunctionOrDataAddress to map addrs inside
 functions

Add logic to map addresses referring to non-entry basic blocks.

PR #101466 extended this function to enable it to map addresses for
the entry points of multi-entry functions, but this still left
references to individual basic blocks unmappable.
---
 bolt/lib/Rewrite/RewriteInstance.cpp | 12 +++------
 bolt/test/AArch64/computed-goto.s    | 39 ++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 8 deletions(-)
 create mode 100644 bolt/test/AArch64/computed-goto.s

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 7059a3dd231099..be1c905f6000de 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -5569,14 +5569,10 @@ uint64_t RewriteInstance::getNewFunctionOrDataAddress(uint64_t OldAddress) {
   if (const BinaryFunction *BF =
           BC->getBinaryFunctionContainingAddress(OldAddress)) {
     if (BF->isEmitted()) {
-      // If OldAddress is the another entry point of
-      // the function, then BOLT could get the new address.
-      if (BF->isMultiEntry()) {
-        for (const BinaryBasicBlock &BB : *BF)
-          if (BB.isEntryPoint() &&
-              (BF->getAddress() + BB.getOffset()) == OldAddress)
-            return BF->getOutputAddress() + BB.getOffset();
-      }
+      for (const BinaryBasicBlock &BB : *BF)
+        if ((BF->getAddress() + BB.getOffset()) == OldAddress)
+          return BB.getOutputStartAddress();
+
       BC->errs() << "BOLT-ERROR: unable to get new address corresponding to "
                     "input address 0x"
                  << Twine::utohexstr(OldAddress) << " in function " << *BF
diff --git a/bolt/test/AArch64/computed-goto.s b/bolt/test/AArch64/computed-goto.s
new file mode 100644
index 00000000000000..ad850f4e7f8d9e
--- /dev/null
+++ b/bolt/test/AArch64/computed-goto.s
@@ -0,0 +1,39 @@
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.bolt 2>&1 | FileCheck %s
+
+# Before bolt could handle mapping addresses within moved functions, it
+# would bail out with an error of the form:
+# BOLT-ERROR: unable to get new address corresponding to input address 0x10390 in function main. Consider adding this function to --skip-funcs=...
+# These addresses arise if computed GOTO is in use.
+# Check that bolt does not emit any error.
+
+# CHECK-NOT: BOLT-ERROR
+
+.globl  main
+.p2align        2
+.type   main, at function
+main:
+.cfi_startproc
+        adrp    x8, .L__const.main.ptrs+8
+        add     x8, x8, :lo12:.L__const.main.ptrs+8
+        ldr     x9, [x8], #8
+        br      x9
+
+.Label0: // Block address taken
+        ldr     x9, [x8], #8
+        br      x9
+
+.Label1: // Block address taken
+        mov     w0, #42
+        ret
+
+.Lfunc_end0:
+.size   main, .Lfunc_end0-main
+.cfi_endproc
+        .type   .L__const.main.ptrs, at object
+        .section        .data.rel.ro,"aw", at progbits
+        .p2align        3, 0x0
+.L__const.main.ptrs:
+        .xword  .Label0
+        .xword  .Label1