From 28e7e71251dc4b79c29aa0d4904cb424f9081455 Mon Sep 17 00:00:00 2001
From: rfwang07 <wangrufeng5@huawei.com>
Date: Fri, 21 Jun 2024 11:23:42 +0800
Subject: [PATCH] [BOLT][AArch64] Don't change layout in PatchEntries

---
 bolt/lib/Passes/PatchEntries.cpp      | 11 ++++++++
 bolt/test/AArch64/patch-entries.s     | 36 ++++++++++++++++++++++++
 bolt/unittests/Core/BinaryContext.cpp | 40 +++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 bolt/test/AArch64/patch-entries.s

diff --git a/bolt/lib/Passes/PatchEntries.cpp b/bolt/lib/Passes/PatchEntries.cpp
index 02a044d8b..ee7512d89 100644
--- a/bolt/lib/Passes/PatchEntries.cpp
+++ b/bolt/lib/Passes/PatchEntries.cpp
@@ -98,6 +98,17 @@ void PatchEntries::runOnFunctions(BinaryContext &BC) {
     });
 
     if (!Success) {
+      // We can't change output layout for AArch64 due to LongJmp pass
+      if (BC.isAArch64()) {
+        if (opts::ForcePatch) {
+          errs() << "BOLT-ERROR: unable to patch entries in " << Function
+                 << "\n";
+          exit(1);
+        }
+
+        continue;
+      }
+
       // If the original function entries cannot be patched, then we cannot
       // safely emit new function body.
       errs() << "BOLT-WARNING: failed to patch entries in " << Function
diff --git a/bolt/test/AArch64/patch-entries.s b/bolt/test/AArch64/patch-entries.s
new file mode 100644
index 000000000..cf6f72a0b
--- /dev/null
+++ b/bolt/test/AArch64/patch-entries.s
@@ -0,0 +1,36 @@
+# This test checks patch entries functionality
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
+# RUN:   %s -o %t.o
+# RUN: %clang %cflags -pie %t.o -o %t.exe -nostdlib -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.bolt --use-old-text=0 --lite=0 --skip-funcs=_start
+# RUN: llvm-objdump -dz %t.bolt | FileCheck %s
+
+# CHECK: <pathedEntries.org.0>:
+# CHECK-NEXT: adrp x16, 0x[[#%x,ADRP:]]
+# CHECK-NEXT: add x16, x16, #0x[[#%x,ADD:]]
+# CHECK-NEXT: br x16
+
+# CHECK: [[#ADRP + ADD]] <pathedEntries>:
+# CHECK-NEXT: [[#ADRP + ADD]]: {{.*}} ret
+
+.text
+.balign 4
+.global pathedEntries
+.type pathedEntries, %function
+pathedEntries:
+  .rept 32
+  nop
+  .endr
+  ret
+.size pathedEntries, .-pathedEntries
+
+.global _start
+.type _start, %function
+_start:
+  bl pathedEntries
+  .inst 0xdeadbeef
+  ret
+.size _start, .-_start
diff --git a/bolt/unittests/Core/BinaryContext.cpp b/bolt/unittests/Core/BinaryContext.cpp
index 5a80cb4a2..7ac1c1435 100644
--- a/bolt/unittests/Core/BinaryContext.cpp
+++ b/bolt/unittests/Core/BinaryContext.cpp
@@ -62,6 +62,46 @@ INSTANTIATE_TEST_SUITE_P(X86, BinaryContextTester,
 INSTANTIATE_TEST_SUITE_P(AArch64, BinaryContextTester,
                          ::testing::Values(Triple::aarch64));
 
+TEST_P(BinaryContextTester, FlushPendingRelocCALL26) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  // This test checks that encodeValueAArch64 used by flushPendingRelocations
+  // returns correctly encoded values for CALL26 relocation for both backward
+  // and forward branches.
+  //
+  // The offsets layout is:
+  // 4:  func1
+  // 8:  bl func1
+  // 12: bl func2
+  // 16: func2
+
+  char Data[20] = {};
+  BinarySection &BS = BC->registerOrUpdateSection(
+      ".text", ELF::SHT_PROGBITS, ELF::SHF_EXECINSTR | ELF::SHF_ALLOC,
+      (uint8_t *)Data, sizeof(Data), 4);
+  MCSymbol *RelSymbol1 = BC->getOrCreateGlobalSymbol(4, "Func1");
+  ASSERT_TRUE(RelSymbol1);
+  BS.addRelocation(8, RelSymbol1, ELF::R_AARCH64_CALL26, 0, 0, true);
+  MCSymbol *RelSymbol2 = BC->getOrCreateGlobalSymbol(16, "Func2");
+  ASSERT_TRUE(RelSymbol2);
+  BS.addRelocation(12, RelSymbol2, ELF::R_AARCH64_CALL26, 0, 0, true);
+
+  std::error_code EC;
+  SmallVector<char> Vect(sizeof(Data));
+  raw_svector_ostream OS(Vect);
+
+  BS.flushPendingRelocations(OS, [&](const MCSymbol *S) {
+    return S == RelSymbol1 ? 4 : S == RelSymbol2 ? 16 : 0;
+  });
+
+  const uint8_t Func1Call[4] = {255, 255, 255, 151};
+  const uint8_t Func2Call[4] = {1, 0, 0, 148};
+
+  EXPECT_FALSE(memcmp(Func1Call, &Vect[8], 4)) << "Wrong backward call value\n";
+  EXPECT_FALSE(memcmp(Func2Call, &Vect[12], 4)) << "Wrong forward call value\n";
+}
+
 #endif
 
 TEST_P(BinaryContextTester, BaseAddress) {
-- 
2.39.2 (Apple Git-143)