From 34c6c5e72f48de65a7e332033af9566576c1895d Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Fri, 24 Jan 2025 14:20:24 -0800
Subject: [PATCH] [BOLT][AArch64] Fix PLT optimization (#124192)

Preserve C++ exception metadata while running PLT optimization on
AArch64.
---
 bolt/include/bolt/Core/MCPlusBuilder.h        |  7 ++++---
 bolt/lib/Passes/PLTCall.cpp                   |  4 ++--
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  5 ++---
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |  2 +-
 bolt/test/AArch64/exceptions-plt.cpp          | 21 +++++++++++++++++++
 bolt/test/runtime/exceptions-plt.cpp          | 16 ++++++++++++++
 6 files changed, 46 insertions(+), 9 deletions(-)
 create mode 100644 bolt/test/AArch64/exceptions-plt.cpp
 create mode 100644 bolt/test/runtime/exceptions-plt.cpp

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 5d77e6faff2f..c1460b2aac8a 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1426,11 +1426,12 @@ public:
   }
 
   /// Creates an indirect call to the function within the \p DirectCall PLT
-  /// stub. The function's memory location is pointed by the \p TargetLocation
+  /// stub. The function's address location is pointed by the \p TargetLocation
   /// symbol.
+  /// Move instruction annotations from \p DirectCall to the indirect call.
   virtual InstructionListType
-  createIndirectPltCall(const MCInst &DirectCall,
-                        const MCSymbol *TargetLocation, MCContext *Ctx) {
+  createIndirectPLTCall(MCInst &&DirectCall, const MCSymbol *TargetLocation,
+                        MCContext *Ctx) {
     llvm_unreachable("not implemented");
     return {};
   }
diff --git a/bolt/lib/Passes/PLTCall.cpp b/bolt/lib/Passes/PLTCall.cpp
index 2ed996fadbb9..31c2d92ebc20 100644
--- a/bolt/lib/Passes/PLTCall.cpp
+++ b/bolt/lib/Passes/PLTCall.cpp
@@ -70,8 +70,8 @@ Error PLTCall::runOnFunctions(BinaryContext &BC) {
         const BinaryFunction *CalleeBF = BC.getFunctionForSymbol(CallSymbol);
         if (!CalleeBF || !CalleeBF->isPLTFunction())
           continue;
-        const InstructionListType NewCode = BC.MIB->createIndirectPltCall(
-            *II, CalleeBF->getPLTSymbol(), BC.Ctx.get());
+        const InstructionListType NewCode = BC.MIB->createIndirectPLTCall(
+            std::move(*II), CalleeBF->getPLTSymbol(), BC.Ctx.get());
         II = BB.replaceInstruction(II, NewCode);
         assert(!NewCode.empty() && "PLT Call replacement must be non-empty");
         std::advance(II, NewCode.size() - 1);
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 0b6f21527f0a..ac709c5dd063 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -1263,7 +1263,7 @@ public:
     return true;
   }
 
-  InstructionListType createIndirectPltCall(const MCInst &DirectCall,
+  InstructionListType createIndirectPLTCall(MCInst &&DirectCall,
                                             const MCSymbol *TargetLocation,
                                             MCContext *Ctx) override {
     const bool IsTailCall = isTailCall(DirectCall);
@@ -1297,8 +1297,7 @@ public:
     MCInst InstCall;
     InstCall.setOpcode(IsTailCall ? AArch64::BR : AArch64::BLR);
     InstCall.addOperand(MCOperand::createReg(AArch64::X17));
-    if (IsTailCall)
-      setTailCall(InstCall);
+    moveAnnotations(std::move(DirectCall), InstCall);
     Code.emplace_back(InstCall);
 
     return Code;
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 63086c06d74f..465533ee71f2 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1605,7 +1605,7 @@ public:
     return true;
   }
 
-  InstructionListType createIndirectPltCall(const MCInst &DirectCall,
+  InstructionListType createIndirectPLTCall(MCInst &&DirectCall,
                                             const MCSymbol *TargetLocation,
                                             MCContext *Ctx) override {
     assert((DirectCall.getOpcode() == X86::CALL64pcrel32 ||
diff --git a/bolt/test/AArch64/exceptions-plt.cpp b/bolt/test/AArch64/exceptions-plt.cpp
new file mode 100644
index 000000000000..576f0fc91a9d
--- /dev/null
+++ b/bolt/test/AArch64/exceptions-plt.cpp
@@ -0,0 +1,21 @@
+// Verify that PLT optimization in BOLT preserves exception-handling info.
+
+// REQUIRES: system-linux
+
+// RUN: %clangxx %cxxflags -O1 -Wl,-q,-znow %s -o %t.exe
+// RUN: llvm-bolt %t.exe -o %t.bolt.exe --plt=all --print-only=.*main.* \
+// RUN:   --print-finalized 2>&1 | FileCheck %s
+
+// CHECK-LABEL: Binary Function
+// CHECK:      adrp {{.*}}__cxa_throw
+// CHECK-NEXT: ldr {{.*}}__cxa_throw
+// CHECK-NEXT: blr x17 {{.*}} handler: {{.*}} PLTCall:
+
+int main() {
+  try {
+    throw new int;
+  } catch (...) {
+    return 0;
+  }
+  return 1;
+}
diff --git a/bolt/test/runtime/exceptions-plt.cpp b/bolt/test/runtime/exceptions-plt.cpp
new file mode 100644
index 000000000000..8a75a3cb384b
--- /dev/null
+++ b/bolt/test/runtime/exceptions-plt.cpp
@@ -0,0 +1,16 @@
+// Verify that PLT optimization in BOLT preserves exception-handling info.
+
+// REQUIRES: system-linux
+
+// RUN: %clangxx %cxxflags -O1 -Wl,-q,-znow %s -o %t.exe
+// RUN: llvm-bolt %t.exe -o %t.bolt.exe --plt=all
+// RUN: %t.bolt.exe
+
+int main() {
+  try {
+    throw new int;
+  } catch (...) {
+    return 0;
+  }
+  return 1;
+}