[llvm] [BOLT][AArch64] Create entry points for addresses referenced by dynamic relocations and allow getNewFunctionOrDataAddress to map addrs inside functions. (PR #120267)

[Rin Dobrescu via llvm-commits]

https://github.com/Rin18 created https://github.com/llvm/llvm-project/pull/120267

By adding addresses referenced by dynamic relocations as entry points, this patch fixes an issue where bolt fails on code using computing goto's. This also fixes a mapping issue with the bugfix from this PR: https://github.com/llvm/llvm-project/pull/117766.

>From f6435e5499dd0384820381bf85995b823789c279 Mon Sep 17 00:00:00 2001
From: Rin Dobrescu <rin.dobrescu at arm.com>
Date: Tue, 17 Dec 2024 16:43:33 +0000
Subject: [PATCH] [BOLT][AArch64] Create entry points for addresses referenced
 by dynamic relocations and allow getNewFunctionOrDataAddress to map addrs
 inside functions.

By adding addresses referenced by dynamic relocations as entry points,
this patch fixes an issue where bolt fails on code using computing
goto's. This also fixes a mapping issue with the bugfix from this
PR: https://github.com/llvm/llvm-project/pull/117766.
---
 bolt/lib/Rewrite/RewriteInstance.cpp | 10 ++++++-
 bolt/test/AArch64/computed-goto.s    | 39 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 bolt/test/AArch64/computed-goto.s

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 4329235d470497..55fcd6b6e782c4 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -2439,6 +2439,14 @@ void RewriteInstance::readDynamicRelocations(const SectionRef &Section,
     if (Symbol)
       SymbolIndex[Symbol] = getRelocationSymbol(InputFile, Rel);
 
+    const uint64_t SymAddress = SymbolAddress + Addend;
+    BinaryFunction *Func = BC->getBinaryFunctionContainingAddress(SymAddress);
+    if(Func){
+      const uint64_t FunctionOffset = SymAddress - Func->getAddress();
+      if(FunctionOffset)
+        Func->addEntryPointAtOffset(FunctionOffset);
+    }
+
     BC->addDynamicRelocation(Rel.getOffset(), Symbol, RType, Addend);
   }
 }
@@ -5599,7 +5607,7 @@ uint64_t RewriteInstance::getNewFunctionOrDataAddress(uint64_t OldAddress) {
         for (const BinaryBasicBlock &BB : *BF)
           if (BB.isEntryPoint() &&
               (BF->getAddress() + BB.getOffset()) == OldAddress)
-            return BF->getOutputAddress() + BB.getOffset();
+            return BB.getOutputStartAddress();
       }
       BC->errs() << "BOLT-ERROR: unable to get new address corresponding to "
                     "input address 0x"
diff --git a/bolt/test/AArch64/computed-goto.s b/bolt/test/AArch64/computed-goto.s
new file mode 100644
index 00000000000000..043f9a8e37e6b0
--- /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