From 1e4ee588fbb5047a82dd33c4357fcfdebaf38bab Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Fri, 21 Jul 2023 16:21:44 -0700
Subject: [PATCH] [BOLT] Accept function start as valid jump table entry

Jump tables may contain a function start address. One real-world example
is when a target basic block contains a recursive tail call that is
later optimized/folded into a jump table target.

While analyzing a jump table, we treat start address similar to an
address past the end of the containing function (a result of
__builtin_unreachable), i.e. we require another "regular" entry for the
heuristic to proceed.

Reviewed By: Amir

Differential Revision: https://reviews.llvm.org/D156206
---
 bolt/lib/Core/BinaryContext.cpp       | 30 +++++++----
 bolt/test/X86/jump-table-func-entry.s | 72 +++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 11 deletions(-)
 create mode 100644 bolt/test/X86/jump-table-func-entry.s

diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 2d2b35ee2bd9..1e95d5fe38ab 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -503,6 +503,9 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
   // Is one of the targets __builtin_unreachable?
   bool HasUnreachable = false;
 
+  // Does one of the entries match function start address?
+  bool HasStartAsEntry = false;
+
   // Number of targets other than __builtin_unreachable.
   uint64_t NumRealEntries = 0;
 
@@ -567,14 +570,21 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
       continue;
     }
 
+    // Function start is another special case. It is allowed in the jump table,
+    // but we need at least one another regular entry to distinguish the table
+    // from, e.g. a function pointer array.
+    if (Value == BF.getAddress()) {
+      HasStartAsEntry = true;
+      addEntryAddress(Value);
+      continue;
+    }
+
     // Function or one of its fragments.
     const BinaryFunction *TargetBF = getBinaryFunctionContainingAddress(Value);
-
-    bool DoesBelongToFunction = BF.containsAddress(Value) ||
-                                (TargetBF && TargetBF->isParentOrChildOf(BF));
-
-    // We assume that a jump table cannot have function start as an entry.
-    if (!DoesBelongToFunction || Value == BF.getAddress()) {
+    const bool DoesBelongToFunction =
+        BF.containsAddress(Value) ||
+        (TargetBF && TargetBF->isParentOrChildOf(BF));
+    if (!DoesBelongToFunction) {
       LLVM_DEBUG({
         if (!BF.containsAddress(Value)) {
           dbgs() << "FAIL: function doesn't contain this address\n";
@@ -589,8 +599,6 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
             }
           }
         }
-        if (Value == BF.getAddress())
-          dbgs() << "FAIL: jump table cannot have function start as an entry\n";
       });
       break;
     }
@@ -611,9 +619,9 @@ bool BinaryContext::analyzeJumpTable(const uint64_t Address,
   }
 
   // It's a jump table if the number of real entries is more than 1, or there's
-  // one real entry and "unreachable" targets. If there are only multiple
-  // "unreachable" targets, then it's not a jump table.
-  return NumRealEntries + HasUnreachable >= 2;
+  // one real entry and one or more special targets. If there are only multiple
+  // special targets, then it's not a jump table.
+  return NumRealEntries + (HasUnreachable || HasStartAsEntry) >= 2;
 }
 
 void BinaryContext::populateJumpTables() {
diff --git a/bolt/test/X86/jump-table-func-entry.s b/bolt/test/X86/jump-table-func-entry.s
new file mode 100644
index 000000000000..77b444d520a1
--- /dev/null
+++ b/bolt/test/X86/jump-table-func-entry.s
@@ -0,0 +1,72 @@
+# REQUIRES: system-linux
+
+## Check that BOLT correctly processes jump table that contains function start
+## as one of its entries.
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -no-pie -Wl,-q
+
+# RUN: llvm-bolt %t.exe --print-normalized --print-only=foo -o %t.out \
+# RUN:   |& FileCheck %s
+
+
+
+  .text
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+  call foo
+  ret
+  .cfi_endproc
+  .size _start, .-_start
+
+  .globl foo
+  .type foo, %function
+foo:
+	.cfi_startproc
+.LBB00:
+          movq	0x8(%rdi), %rdi
+          movzbl	0x1(%rdi), %eax
+.LBB00_br:
+	        jmpq	*"JUMP_TABLE/foo.0"(,%rax,8)
+# CHECK:  jmpq {{.*}} # JUMPTABLE
+# CHECK-NEXT: Successors: {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}
+
+.Ltmp87085:
+	xorl	%eax, %eax
+	retq
+
+.Ltmp87086:
+	cmpb	$0x0, 0x8(%rdi)
+	setne	%al
+	retq
+
+.Ltmp87088:
+	movb	$0x1, %al
+	retq
+
+.Ltmp87087:
+	movzbl	0x14(%rdi), %eax
+	andb	$0x2, %al
+	shrb	%al
+	retq
+
+	.cfi_endproc
+.size foo, .-foo
+
+# Jump tables
+.section .rodata
+"JUMP_TABLE/foo.0":
+	.quad	.Ltmp87085
+	.quad	.Ltmp87086
+	.quad	.Ltmp87087
+	.quad	.LBB00
+	.quad	.Ltmp87088
+
+# CHECK: Jump table {{.*}} for function foo
+# CHECK-NEXT: 0x{{.*}} :
+# CHECK-NEXT: 0x{{.*}} :
+# CHECK-NEXT: 0x{{.*}} :
+# CHECK-NEXT: 0x{{.*}} :
+# CHECK-NEXT: 0x{{.*}} :