simplify code

Author: mingmingl-llvm

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -453,7 +453,7 @@ class AsmPrinter : public MachineFunctionPass {
   /// function to the current output stream.
   virtual void emitJumpTableInfo();
 
-  virtual void emitJumpTables(const std::vector<unsigned> &JumpTableIndices,
+  virtual void emitJumpTables(ArrayRef<unsigned> JumpTableIndices,
                               MCSection *JumpTableSection, bool JTInDiffSection,
                               const MachineJumpTableInfo &MJTI);
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -2881,42 +2881,41 @@ void AsmPrinter::emitJumpTableInfo() {
           MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64,
       F);
 
+  std::vector<unsigned> JumpTableIndices;
+  for (unsigned JTI = 0, JTSize = JT.size(); JTI != JTSize; ++JTI)
+    JumpTableIndices.push_back(JTI);
   if (!EmitStaticDataHotnessSuffix) {
-    std::vector<unsigned> JumpTableIndices;
-    for (unsigned JTI = 0, e = JT.size(); JTI != e; ++JTI)
-      JumpTableIndices.push_back(JTI);
-    // Prior code in this function verifies JT is not empty. Use JT[0] directly.
     emitJumpTables(JumpTableIndices, TLOF.getSectionForJumpTable(F, TM),
                    JTInDiffSection, *MJTI);
     return;
   }
 
-  // Iterate all jump tables, put them into two vectors, hot and lukewarm
-  std::vector<unsigned> HotOrWarmJumpTableIndices, ColdJumpTableIndices;
-
-  for (unsigned JTI = 0, e = JT.size(); JTI != e; ++JTI) {
-    if (JT[JTI].Hotness == llvm::DataHotness::Cold)
-      ColdJumpTableIndices.push_back(JTI);
-    else
-      HotOrWarmJumpTableIndices.push_back(JTI);
+  // Iterate all jump tables, put hot jump table indices towards the beginning
+  // of the vector, and cold jump table indices towards the end.
+  unsigned ColdJumpTableStart = JT.size();
+  for (unsigned JTI = 0; JTI != ColdJumpTableStart;) {
+    if (JT[JTI].Hotness == llvm::DataHotness::Cold) {
+      ColdJumpTableStart -= 1;
+      std::swap(JumpTableIndices[JTI], JumpTableIndices[ColdJumpTableStart]);
+    } else
+      JTI += 1;
   }
 
-  if (!HotOrWarmJumpTableIndices.empty())
-    emitJumpTables(HotOrWarmJumpTableIndices,
-                   TLOF.getSectionForJumpTable(
-                       F, TM, &JT[*HotOrWarmJumpTableIndices.begin()]),
-                   JTInDiffSection, *MJTI);
+  if (ColdJumpTableStart != 0)
+    emitJumpTables(
+        ArrayRef<unsigned>(JumpTableIndices).slice(0, ColdJumpTableStart),
+        TLOF.getSectionForJumpTable(F, TM, &JT[0]), JTInDiffSection, *MJTI);
 
-  if (!ColdJumpTableIndices.empty())
+  if (ColdJumpTableStart != JT.size())
     emitJumpTables(
-        ColdJumpTableIndices,
-        TLOF.getSectionForJumpTable(F, TM, &JT[*ColdJumpTableIndices.begin()]),
+        ArrayRef<unsigned>(JumpTableIndices).slice(ColdJumpTableStart),
+        TLOF.getSectionForJumpTable(F, TM, &JT[ColdJumpTableStart]),
         JTInDiffSection, *MJTI);
 
   return;
 }
 
-void AsmPrinter::emitJumpTables(const std::vector<unsigned> &JumpTableIndices,
+void AsmPrinter::emitJumpTables(ArrayRef<unsigned> JumpTableIndices,
                                 MCSection *JumpTableSection,
                                 bool JTInDiffSection,
                                 const MachineJumpTableInfo &MJTI) {

llvm/test/CodeGen/X86/jump-table-partition.ll

@@ -2,7 +2,7 @@
 ; requires: asserts
 
 ; RUN: llc -stop-after=block-placement %s -o - | llc --run-pass=static-data-splitter -stats -x mir -o - 2>&1 | FileCheck %s --check-prefix=STAT
-; RUN: llc -enable-split-machine-functions -split-static-data --function-sections -data-sections %s -o - 2>&1 | FileCheck %s --check-prefix=SECTION
+; RUN: llc -enable-split-machine-functions -split-static-data -emit-static-data-hotness-suffix=true -function-sections -data-sections %s -o - 2>&1 | FileCheck %s --check-prefix=SECTION
 
 ; `func_with_hot_jumptable` contains a hot jump table and `func_with_cold_jumptable` contains a cold one. 
 ; `func_without_entry_count` simulates the functions without profile information (e.g., not instrumented or not profiled),
@@ -13,7 +13,6 @@
 ; STAT-DAG: 1 static-data-splitter - Number of hot jump tables seen
 ; STAT-DAG: 1 static-data-splitter - Number of jump tables with unknown hotness
 
-; TODO: Construct a test case where a hot function has a cold jump table.
 ; SECTION: .section .rodata.hot.func_with_hot_jumptable
 ; SECTION: .section .rodata.unlikely.func_with_cold_jumptable
 ; SECTION: .section .rodata.hot.func_without_entry_count
@@ -28,6 +27,7 @@ target triple = "x86_64-unknown-linux-gnu"
 @str.10 = private constant [7 x i8] c"case 1\00"
 @str.11 = private constant [8 x i8] c"default\00"
 
+; TODO: Add a cold jump table in this function.
 define i32 @func_with_hot_jumptable(i32 %num) !prof !13 {
 entry:
   switch i32 %num, label %sw.default [