[MC] Emit a jump table size section (#101962)

This patch will make LLVM emit a new section .llvm_jump_table_sizes 
containing tuples of (jump table address, entry count) in object files.
This section is useful for tools that need to statically reconstruct the
control flow of executables.

At the moment this is only enabled by default for the PS5 target.

10 files changed, 298 insertions, 0 deletions

diff --git a/clang/lib/Driver/ToolChains/PS4CPU.cpp b/clang/lib/Driver/ToolChains/PS4CPU.cpp
index 54ec59e..48d8241 100644
--- a/clang/lib/Driver/ToolChains/PS4CPU.cpp
+++ b/clang/lib/Driver/ToolChains/PS4CPU.cpp
@@ -272,6 +272,8 @@ void tools::PS5cpu::Linker::ConstructJob(Compilation &C, const JobAction &JA,
     CmdArgs.push_back(D.getLTOMode() == LTOK_Thin ? "--lto=thin"
                                                   : "--lto=full");
 
+  AddLTOFlag("-emit-jump-table-sizes-section");
+
   if (UseJMC)
     AddLTOFlag("-enable-jmc-instrument");
 
@@ -484,6 +486,12 @@ void toolchains::PS4PS5Base::addClangTargetOptions(
     else
       CC1Args.push_back("-fvisibility-externs-nodllstorageclass=keep");
   }
+
+  // Enable jump table sizes section for PS5.
+  if (getTriple().isPS5()) {
+    CC1Args.push_back("-mllvm");
+    CC1Args.push_back("-emit-jump-table-sizes-section");
+  }
 }
 
 // PS4 toolchain.
diff --git a/clang/test/Driver/ps4-ps5-toolchain.c b/clang/test/Driver/ps4-ps5-toolchain.c
index 444e9df..c9987c2 100644
--- a/clang/test/Driver/ps4-ps5-toolchain.c
+++ b/clang/test/Driver/ps4-ps5-toolchain.c
@@ -11,3 +11,8 @@
 // RUN: %clang %s -### -target x86_64-sie-ps5 -flto 2>&1 | FileCheck %s --check-prefix=LTO
 // LTO-NOT: error:
 // LTO-NOT: unable to pass LLVM bit-code
+
+// Verify that the jump table sizes section is enabled.
+// RUN: %clang %s -target x86_64-sie-ps5 -### 2>&1 | FileCheck -check-prefix=JUMPTABLESIZES %s
+// JUMPTABLESIZES: "-mllvm" "-emit-jump-table-sizes-section"
+// JUMPTABLESIZES: "-plugin-opt=-emit-jump-table-sizes-section"
diff --git a/llvm/docs/Extensions.rst b/llvm/docs/Extensions.rst
index 74ca8cb..abc34bc3 100644
--- a/llvm/docs/Extensions.rst
+++ b/llvm/docs/Extensions.rst
@@ -554,6 +554,12 @@ time. This section is generated when the compiler enables fat LTO. This section
 has the ``SHF_EXCLUDE`` flag so that it is stripped from the final executable
 or shared library.
 
+``SHT_LLVM_JT_SIZES`` Section (Jump table addresses and sizes)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+This section stores pairs of (jump table address, number of entries).
+This information is useful for tools that need to statically reconstruct
+the control flow of executables.
+
 CodeView-Dependent
 ------------------
 
diff --git a/llvm/include/llvm/BinaryFormat/ELF.h b/llvm/include/llvm/BinaryFormat/ELF.h
index d61d096..2eff87c 100644
--- a/llvm/include/llvm/BinaryFormat/ELF.h
+++ b/llvm/include/llvm/BinaryFormat/ELF.h
@@ -1129,6 +1129,7 @@ enum : unsigned {
   SHT_LLVM_BB_ADDR_MAP = 0x6fff4c0a,        // LLVM Basic Block Address Map.
   SHT_LLVM_OFFLOADING = 0x6fff4c0b,         // LLVM device offloading data.
   SHT_LLVM_LTO = 0x6fff4c0c,                // .llvm.lto for fat LTO.
+  SHT_LLVM_JT_SIZES = 0x6fff4c0d,           // LLVM jump tables sizes.
   // Android's experimental support for SHT_RELR sections.
   // https://android.googlesource.com/platform/bionic/+/b7feec74547f84559a1467aca02708ff61346d2a/libc/include/elf.h#512
   SHT_ANDROID_RELR = 0x6fffff00,   // Relocation entries; only offsets.
diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index 36d1b47..c9a88d7 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -894,6 +894,10 @@ private:
 
   void emitJumpTableEntry(const MachineJumpTableInfo *MJTI,
                           const MachineBasicBlock *MBB, unsigned uid) const;
+
+  void emitJumpTableSizesSection(const MachineJumpTableInfo *MJTI,
+                                 const Function &F) const;
+
   void emitLLVMUsedList(const ConstantArray *InitList);
   /// Emit llvm.ident metadata in an '.ident' directive.
   void emitModuleIdents(Module &M);
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 88e9b9d2..db7adfd 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -99,6 +99,7 @@
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCSymbolCOFF.h"
 #include "llvm/MC/MCSymbolELF.h"
 #include "llvm/MC/MCTargetOptions.h"
 #include "llvm/MC/MCValue.h"
@@ -107,6 +108,7 @@
 #include "llvm/Pass.h"
 #include "llvm/Remarks/RemarkStreamer.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
@@ -155,6 +157,11 @@ static cl::bits<PGOMapFeaturesEnum> PgoAnalysisMapFeatures(
         "Enable extended information within the SHT_LLVM_BB_ADDR_MAP that is "
         "extracted from PGO related analysis."));
 
+static cl::opt<bool> EmitJumpTableSizesSection(
+    "emit-jump-table-sizes-section",
+    cl::desc("Emit a section containing jump table addresses and sizes"),
+    cl::Hidden, cl::init(false));
+
 STATISTIC(EmittedInsts, "Number of machine instrs printed");
 
 char AsmPrinter::ID = 0;
@@ -2786,10 +2793,62 @@ void AsmPrinter::emitJumpTableInfo() {
     for (const MachineBasicBlock *MBB : JTBBs)
       emitJumpTableEntry(MJTI, MBB, JTI);
   }
+
+  if (EmitJumpTableSizesSection)
+    emitJumpTableSizesSection(MJTI, F);
+
   if (!JTInDiffSection)
     OutStreamer->emitDataRegion(MCDR_DataRegionEnd);
 }
 
+void AsmPrinter::emitJumpTableSizesSection(const MachineJumpTableInfo *MJTI,
+                                           const Function &F) const {
+  const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
+
+  if (JT.empty())
+    return;
+
+  StringRef GroupName = F.hasComdat() ? F.getComdat()->getName() : "";
+  MCSection *JumpTableSizesSection = nullptr;
+  StringRef sectionName = ".llvm_jump_table_sizes";
+
+  bool isElf = TM.getTargetTriple().isOSBinFormatELF();
+  bool isCoff = TM.getTargetTriple().isOSBinFormatCOFF();
+
+  if (!isCoff && !isElf)
+    return;
+
+  if (isElf) {
+    MCSymbolELF *LinkedToSym = dyn_cast<MCSymbolELF>(CurrentFnSym);
+    int Flags = F.hasComdat() ? ELF::SHF_GROUP : 0;
+
+    JumpTableSizesSection = OutContext.getELFSection(
+        sectionName, ELF::SHT_LLVM_JT_SIZES, Flags, 0, GroupName, F.hasComdat(),
+        MCSection::NonUniqueID, LinkedToSym);
+  } else if (isCoff) {
+    if (F.hasComdat()) {
+      JumpTableSizesSection = OutContext.getCOFFSection(
+          sectionName,
+          COFF::IMAGE_SCN_CNT_INITIALIZED_DATA | COFF::IMAGE_SCN_MEM_READ |
+              COFF::IMAGE_SCN_LNK_COMDAT | COFF::IMAGE_SCN_MEM_DISCARDABLE,
+          F.getComdat()->getName(), COFF::IMAGE_COMDAT_SELECT_ASSOCIATIVE);
+    } else {
+      JumpTableSizesSection = OutContext.getCOFFSection(
+          sectionName, COFF::IMAGE_SCN_CNT_INITIALIZED_DATA |
+                           COFF::IMAGE_SCN_MEM_READ |
+                           COFF::IMAGE_SCN_MEM_DISCARDABLE);
+    }
+  }
+
+  OutStreamer->switchSection(JumpTableSizesSection);
+
+  for (unsigned JTI = 0, E = JT.size(); JTI != E; ++JTI) {
+    const std::vector<MachineBasicBlock *> &JTBBs = JT[JTI].MBBs;
+    OutStreamer->emitSymbolValue(GetJTISymbol(JTI), TM.getProgramPointerSize());
+    OutStreamer->emitIntValue(JTBBs.size(), TM.getProgramPointerSize());
+  }
+}
+
 /// EmitJumpTableEntry - Emit a jump table entry for the specified MBB to the
 /// current stream.
 void AsmPrinter::emitJumpTableEntry(const MachineJumpTableInfo *MJTI,
diff --git a/llvm/lib/MC/MCParser/ELFAsmParser.cpp b/llvm/lib/MC/MCParser/ELFAsmParser.cpp
index e8a22d3..c453644 100644
--- a/llvm/lib/MC/MCParser/ELFAsmParser.cpp
+++ b/llvm/lib/MC/MCParser/ELFAsmParser.cpp
@@ -677,6 +677,8 @@ EndStmt:
       Type = ELF::SHT_LLVM_OFFLOADING;
     else if (TypeName == "llvm_lto")
       Type = ELF::SHT_LLVM_LTO;
+    else if (TypeName == "llvm_jt_sizes")
+      Type = ELF::SHT_LLVM_JT_SIZES;
     else if (TypeName.getAsInteger(0, Type))
       return TokError("unknown section type");
   }
diff --git a/llvm/lib/MC/MCSectionELF.cpp b/llvm/lib/MC/MCSectionELF.cpp
index 5cd6590..25e62b7 100644
--- a/llvm/lib/MC/MCSectionELF.cpp
+++ b/llvm/lib/MC/MCSectionELF.cpp
@@ -172,6 +172,8 @@ void MCSectionELF::printSwitchToSection(const MCAsmInfo &MAI, const Triple &T,
     OS << "llvm_offloading";
   else if (Type == ELF::SHT_LLVM_LTO)
     OS << "llvm_lto";
+  else if (Type == ELF::SHT_LLVM_JT_SIZES)
+    OS << "llvm_jt_sizes";
   else
     OS << "0x" << Twine::utohexstr(Type);
 
diff --git a/llvm/lib/Object/ELF.cpp b/llvm/lib/Object/ELF.cpp
index e47a40b8..c66736f 100644
--- a/llvm/lib/Object/ELF.cpp
+++ b/llvm/lib/Object/ELF.cpp
@@ -319,6 +319,7 @@ StringRef llvm::object::getELFSectionTypeName(uint32_t Machine, unsigned Type) {
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_BB_ADDR_MAP);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_OFFLOADING);
     STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_LTO);
+    STRINGIFY_ENUM_CASE(ELF, SHT_LLVM_JT_SIZES)
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_ATTRIBUTES);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_HASH);
     STRINGIFY_ENUM_CASE(ELF, SHT_GNU_verdef);
diff --git a/llvm/test/CodeGen/X86/jump-table-size-section.ll b/llvm/test/CodeGen/X86/jump-table-size-section.ll
new file mode 100644
index 0000000..c0b57e9
--- /dev/null
+++ b/llvm/test/CodeGen/X86/jump-table-size-section.ll
@@ -0,0 +1,210 @@
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=PS5-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-sie-ps5 -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %s
+
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=ELF-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-unknown-linux-gnu -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %s
+
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -emit-jump-table-sizes-section -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=COFF-CHECK %s
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOFLAG %s
+; RUN: llc %s -o - -mtriple x86_64-pc-windows-msvc -verify-machineinstrs --relocation-model=pic | FileCheck --check-prefix=NOTABLE %s
+
+; This test verifies the jump table size section. Currently only enabled by default on the PS5 target.
+
+$foo1 = comdat any
+
+; Ensure proper comdat handling.
+define void @foo1(i32 %x, ptr %to) comdat {
+
+; PS5-CHECK-LABEL: foo1
+; PS5-CHECK:       .section        .llvm_jump_table_sizes,"G",@llvm_jt_sizes,foo1,comdat
+; PS5-CHECK-NEXT: .quad   .LJTI0_0
+; PS5-CHECK-NEXT: .quad   6
+
+; ELF-CHECK-LABEL: foo1
+; ELF-CHECK:       .section        .llvm_jump_table_sizes,"G",@llvm_jt_sizes,foo1,comdat
+; ELF-CHECK-NEXT: .quad   .LJTI0_0
+; ELF-CHECK-NEXT: .quad   6
+
+; COFF-CHECK-LABEL: foo1
+; COFF-CHECK:      .section         .llvm_jump_table_sizes,"drD",associative,foo1
+; COFF-CHECK-NEXT: .quad   .LJTI0_0
+; COFF-CHECK-NEXT: .quad   6
+
+; NOFLAG-LABEL: foo1
+; NOFLAG-NOT: .section        .llvm_jump_table_sizes
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+    i32 5, label %bb4
+  ]
+bb0:
+  store i32 0, ptr %to
+  br label %exit
+bb1:
+  store i32 1, ptr %to
+  br label %exit
+bb2:
+  store i32 2, ptr %to
+  br label %exit
+bb3:
+  store i32 3, ptr %to
+  br label %exit
+bb4:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+default:
+  unreachable
+}
+
+define void @foo2(i32 %x, ptr %to) {
+
+; PS5-CHECK-LABEL:    foo2
+; PS5-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; PS5-CHECK-NEXT:     .quad .LJTI1_0
+; PS5-CHECK-NEXT:     .quad   5
+
+; ELF-CHECK-LABEL:    foo2
+; ELF-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; ELF-CHECK-NEXT:     .quad .LJTI1_0
+; ELF-CHECK-NEXT:     .quad   5
+
+; COFF-CHECK-LABEL:   foo2
+; COFF-CHECK:         .section         .llvm_jump_table_sizes,"drD"
+; COFF-CHECK-NEXT:    .quad .LJTI1_0
+; COFF-CHECK-NEXT:    .quad   5
+
+; NOFLAG-LABEL:       foo1
+; NOFLAG-NOT:         .section        .llvm_jump_table_sizes
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+  ]
+bb0:
+  store i32 0, ptr %to
+  br label %exit
+bb1:
+  store i32 1, ptr %to
+  br label %exit
+bb2:
+  store i32 2, ptr %to
+  br label %exit
+bb3:
+  store i32 3, ptr %to
+  br label %exit
+bb4:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+default:
+  unreachable
+}
+
+; Ensure that the section isn't produced if there is no jump table.
+
+define void @foo3(i32 %x, ptr %to) {
+
+; NOTABLE-LABEL:    foo3
+; NOTABLE-NOT:      .section        .llvm_jump_table_sizes
+
+exit:
+  ret void
+}
+
+; Ensure we can deal with nested jump tables.
+
+define void @nested(i32 %x, i32 %y, ptr %to) {
+
+; PS5-CHECK-LABEL:    nested
+; PS5-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; PS5-CHECK-NEXT:     .quad .LJTI3_0
+; PS5-CHECK-NEXT:     .quad   5
+; PS5-CHECK-NEXT:     .quad .LJTI3_1
+; PS5-CHECK-NEXT:     .quad 6
+
+; ELF-CHECK-LABEL:    nested
+; ELF-CHECK:          .section .llvm_jump_table_sizes,"",@llvm_jt_sizes
+; ELF-CHECK-NEXT:     .quad .LJTI3_0
+; ELF-CHECK-NEXT:     .quad   5
+; ELF-CHECK-NEXT:     .quad .LJTI3_1
+; ELF-CHECK-NEXT:     .quad 6
+
+; COFF-CHECK-LABEL:   nested
+; COFF-CHECK:         .section         .llvm_jump_table_sizes,"drD"
+; COFF-CHECK-NEXT:     .quad .LJTI3_0
+; COFF-CHECK-NEXT:     .quad   5
+; COFF-CHECK-NEXT:     .quad .LJTI3_1
+; COFF-CHECK-NEXT:     .quad 6
+
+; NOFLAG-LABEL:       nested
+; NOFLAG-NOT:         .section        .llvm_jump_table_sizes
+
+entry:
+  switch i32 %x, label %default [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+  ]
+bb0:
+  store i32 0, ptr %to
+  br label %exit
+bb1:
+  store i32 1, ptr %to
+  br label %exit
+bb2:
+  store i32 2, ptr %to
+  br label %exit
+bb3:
+  store i32 3, ptr %to
+  br label %exit
+bb4:
+  switch i32 %y, label %default [
+    i32 1, label %bb5
+    i32 2, label %bb6
+    i32 3, label %bb7
+    i32 4, label %bb8
+    i32 5, label %bb9
+    i32 6, label %bb10
+  ]
+  br label %exit2
+bb5:
+  store i32 4, ptr %to
+  br label %exit
+bb6:
+  store i32 4, ptr %to
+  br label %exit
+bb7:
+  store i32 4, ptr %to
+  br label %exit
+bb8:
+  store i32 4, ptr %to
+  br label %exit
+bb9:
+  store i32 4, ptr %to
+  br label %exit
+bb10:
+  store i32 4, ptr %to
+  br label %exit
+exit:
+  ret void
+exit2:
+  ret void
+default:
+  unreachable
+}
\ No newline at end of file