[BOLT] Allow getNewFunctionOrDataAddress to map addrs inside functions

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

PR llvm#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.

Author: peterwaller-arm

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

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,@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,@object
+        .section        .data.rel.ro,"aw",@progbits
+        .p2align        3, 0x0
+.L__const.main.ptrs:
+        .xword  .Label0
+        .xword  .Label1