[BOLT] Handle ifuncs trampolines for aarch64

The aarch64 uses the trampolines located in .iplt section, which
contains plt-like trampolines on the value stored in .got. In this case
we don't have JUMP_SLOT relocation, but we have a symbol that belongs to
ifunc trampoline, so use it and set set plt symbol for such functions.

Vladislav Khmelevsky,
Advanced Software Technology Lab, Huawei

Differential Revision: https://reviews.llvm.org/D120850

3 files changed, 49 insertions, 8 deletions

diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index d94fc9c6e709..7f39b70cbdfc 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -498,7 +498,8 @@ private:
   };
 
   /// AArch64 PLT sections.
-  const PLTSectionInfo AArch64_PLTSections[2] = {{".plt"}, {nullptr}};
+  const PLTSectionInfo AArch64_PLTSections[3] = {
+      {".plt"}, {".iplt"}, {nullptr}};
 
   /// Return PLT information for a section with \p SectionName or nullptr
   /// if the section is not PLT.
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 6c67eb3a87a2..1c5d52b227e6 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -1251,19 +1251,30 @@ void RewriteInstance::createPLTBinaryFunction(uint64_t TargetAddress,
   if (!TargetAddress)
     return;
 
+  auto setPLTSymbol = [&](BinaryFunction *BF, StringRef Name) {
+    const unsigned PtrSize = BC->AsmInfo->getCodePointerSize();
+    MCSymbol *TargetSymbol = BC->registerNameAtAddress(
+        Name.str() + "@GOT", TargetAddress, PtrSize, PtrSize);
+    BF->setPLTSymbol(TargetSymbol);
+  };
+
+  BinaryFunction *BF = BC->getBinaryFunctionAtAddress(EntryAddress);
+  if (BF && BC->isAArch64()) {
+    // Handle IFUNC trampoline
+    setPLTSymbol(BF, BF->getOneName());
+    return;
+  }
+
   const Relocation *Rel = BC->getDynamicRelocationAt(TargetAddress);
   if (!Rel || !Rel->Symbol)
     return;
 
-  const unsigned PtrSize = BC->AsmInfo->getCodePointerSize();
   ErrorOr<BinarySection &> Section = BC->getSectionForAddress(EntryAddress);
   assert(Section && "cannot get section for address");
-  BinaryFunction *BF = BC->createBinaryFunction(
-      Rel->Symbol->getName().str() + "@PLT", *Section, EntryAddress, 0,
-      EntrySize, Section->getAlignment());
-  MCSymbol *TargetSymbol = BC->registerNameAtAddress(
-      Rel->Symbol->getName().str() + "@GOT", TargetAddress, PtrSize, PtrSize);
-  BF->setPLTSymbol(TargetSymbol);
+  BF = BC->createBinaryFunction(Rel->Symbol->getName().str() + "@PLT", *Section,
+                                EntryAddress, 0, EntrySize,
+                                Section->getAlignment());
+  setPLTSymbol(BF, Rel->Symbol->getName());
 }
 
 void RewriteInstance::disassemblePLTSectionAArch64(BinarySection &Section) {
diff --git a/bolt/test/runtime/AArch64/iplt.c b/bolt/test/runtime/AArch64/iplt.c
new file mode 100644
index 000000000000..30a169c1265d
--- /dev/null
+++ b/bolt/test/runtime/AArch64/iplt.c
@@ -0,0 +1,29 @@
+// This test checks that the ifuncs works after bolt.
+
+// RUN: %clang %cflags %s -fuse-ld=lld \
+// RUN:    -o %t.exe -Wl,-q
+// RUN: llvm-bolt %t.exe -o %t.bolt.exe -use-old-text=0 -lite=0
+// RUN: %t.bolt.exe  | FileCheck %s
+
+// CHECK: foo
+
+#include <stdio.h>
+#include <string.h>
+
+static void foo() { printf("foo\n"); }
+
+static void *resolver_foo(void) { return foo; }
+
+__attribute__((ifunc("resolver_foo"))) void ifoo();
+
+static void *resolver_memcpy(void) { return memcpy; }
+
+__attribute__((ifunc("resolver_memcpy"))) void *
+imemcpy(void *dest, const void *src, size_t n);
+
+int main() {
+  int a = 0xdeadbeef, b = 0;
+  ifoo();
+  imemcpy(&b, &a, sizeof(b));
+  return a != b;
+}