[BOLT] Accept function start as valid jump table entry

[maksfb]

Summary

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.

[Amir]

Agree to allow function start as jump table entry conditional on presence of another valid entry.

bolt/lib/Core/BinaryContext.cpp

bool BinaryContext::analyzeJumpTable(const uint64_t Address,
                                      const JumpTable::JumpTableType Type,
                                      const BinaryFunction &BF,
                                      const uint64_t NextJTAddress,
                                      JumpTable::AddressesType *EntriesAsAddress,
                                      bool *HasEntryInFragment) const {
   // 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;
 
   auto addEntryAddress = [&](uint64_t EntryAddress) {
     if (EntriesAsAddress)
       EntriesAsAddress->emplace_back(EntryAddress);
   };
 
for (uint64_t EntryAddress = Address; EntryAddress <= UpperBound - EntrySize;
     // __builtin_unreachable() case.
     if (Value == BF.getAddress() + BF.getSize()) {
       addEntryAddress(Value);
       HasUnreachable = true;
       LLVM_DEBUG(dbgs() << formatv("OK: {0:x} __builtin_unreachable\n", Value));
       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) ||
+    if (!DoesBelongToFunction) {
       LLVM_DEBUG({
         if (!BF.containsAddress(Value)) {
           dbgs() << "FAIL: function doesn't contain this address\n";
           if (TargetBF) {
             dbgs() << "  ! function containing this address: "
                    << TargetBF->getPrintName() << '\n';
             if (TargetBF->isFragment()) {
               dbgs() << "  ! is a fragment";
               for (BinaryFunction *Parent : TargetBF->ParentFragments)
                 dbgs() << ", parent: " << Parent->getPrintName();
               dbgs() << '\n';
             }
           }
         }
-        if (Value == BF.getAddress())
-          dbgs() << "FAIL: jump table cannot have function start as an entry\n";
       });
       break;
     }
 
     // Check there's an instruction at this offset.
     if (TargetBF->getState() == BinaryFunction::State::Disassembled &&
         !TargetBF->getInstructionAtOffset(Value - TargetBF->getAddress())) {
       LLVM_DEBUG(dbgs() << formatv("FAIL: no instruction at {0:x}\n", Value));
       break;
     }
 
     ++NumRealEntries;
     LLVM_DEBUG(dbgs() << formatv("OK: {0:x} real entry\n", Value));
 
     if (TargetBF != &BF && HasEntryInFragment)
       *HasEntryInFragment = true;
     addEntryAddress(Value);
   }
 
   // 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() {
   LLVM_DEBUG(dbgs() << "DataPCRelocations: " << DataPCRelocations.size()
                     << '\n');
   for (auto JTI = JumpTables.begin(), JTE = JumpTables.end(); JTI != JTE;
        ++JTI) {
     JumpTable *JT = JTI->second;

bolt/test/X86/jump-table-func-entry.s

+# 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{{.*}} :