From 34cd06a9b3bddaa7a989c606bbf1327ee651711c Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl@google.com>
Date: Wed, 7 Oct 2020 13:21:20 -0700
Subject: [PATCH] [BasicBlockSections] Make sure that the labels for
 address-taken blocks are emitted after switching the seciton.

Currently, AsmPrinter code is organized in a way in which the labels of address-taken blocks are emitted in the previous section, which makes the relocation incorrect.
This patch reorganizes the code to switch to the basic block section before handling address-taken blocks.

Reviewed By: snehasish, MaskRay

Differential Revision: https://reviews.llvm.org/D88517
---
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    | 36 ++++++++++---------
 ...basic-block-sections-blockaddress-taken.ll | 35 ++++++++++++++++++
 2 files changed, 55 insertions(+), 16 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/basic-block-sections-blockaddress-taken.ll

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 613e7ebff2dfd..f45f8b7cb9603 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -3047,6 +3047,16 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
   if (Alignment != Align(1))
     emitAlignment(Alignment);
 
+  // Switch to a new section if this basic block must begin a section. The
+  // entry block is always placed in the function section and is handled
+  // separately.
+  if (MBB.isBeginSection() && !MBB.pred_empty()) {
+    OutStreamer->SwitchSection(
+        getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
+                                                            MBB, TM));
+    CurrentSectionBeginSym = MBB.getSymbol();
+  }
+
   // If the block has its address taken, emit any labels that were used to
   // reference the block.  It is possible that there is more than one label
   // here, because multiple LLVM BB's may have been RAUW'd to this block after
@@ -3077,6 +3087,7 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
     emitBasicBlockLoopComments(MBB, MLI, *this);
   }
 
+  // Print the main label for the block.
   if (MBB.pred_empty() ||
       (!MF->hasBBLabels() && isBlockOnlyReachableByFallthrough(&MBB) &&
        !MBB.isEHFuncletEntry() && !MBB.hasLabelMustBeEmitted())) {
@@ -3086,24 +3097,17 @@ void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
                                   false);
     }
   } else {
-    if (isVerbose() && MBB.hasLabelMustBeEmitted()) {
+    if (isVerbose() && MBB.hasLabelMustBeEmitted())
       OutStreamer->AddComment("Label of block must be emitted");
-    }
-    auto *BBSymbol = MBB.getSymbol();
-    // Switch to a new section if this basic block must begin a section.
-    if (MBB.isBeginSection()) {
-      OutStreamer->SwitchSection(
-          getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
-                                                              MBB, TM));
-      CurrentSectionBeginSym = BBSymbol;
-    }
-    OutStreamer->emitLabel(BBSymbol);
-    // With BB sections, each basic block must handle CFI information on its own
-    // if it begins a section.
-    if (MBB.isBeginSection())
-      for (const HandlerInfo &HI : Handlers)
-        HI.Handler->beginBasicBlock(MBB);
+    OutStreamer->emitLabel(MBB.getSymbol());
   }
+
+  // With BB sections, each basic block must handle CFI information on its own
+  // if it begins a section (Entry block is handled separately by
+  // AsmPrinterHandler::beginFunction).
+  if (MBB.isBeginSection() && !MBB.pred_empty())
+    for (const HandlerInfo &HI : Handlers)
+      HI.Handler->beginBasicBlock(MBB);
 }
 
 void AsmPrinter::emitBasicBlockEnd(const MachineBasicBlock &MBB) {
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-blockaddress-taken.ll b/llvm/test/CodeGen/X86/basic-block-sections-blockaddress-taken.ll
new file mode 100644
index 0000000000000..aaae1cf2a942c
--- /dev/null
+++ b/llvm/test/CodeGen/X86/basic-block-sections-blockaddress-taken.ll
@@ -0,0 +1,35 @@
+;; This test verifies that basic-block-sections works with address-taken basic blocks.
+; RUN: llc < %s -mtriple=x86_64 -basic-block-sections=all | FileCheck %s
+
+define void @foo(i1 zeroext %0) nounwind {
+entry:
+  %1 = select i1 %0, i8* blockaddress(@foo, %bb1), i8* blockaddress(@foo, %bb2) ; <i8*> [#uses=1]
+  indirectbr i8* %1, [label %bb1, label %bb2]
+
+; CHECK:         .text
+; CHECK-LABEL: foo:
+; CHECK:         movl    $.Ltmp0, %eax
+; CHECK-NEXT:    movl    $.Ltmp1, %ecx
+; CHECK-NEXT:    cmovneq %rax, %rcx
+; CHECK-NEXT:    jmpq    *%rcx
+
+bb1:                                                ; preds = %entry
+  %2 = call i32 @bar()
+  ret void
+; CHECK:         .section .text,"ax",@progbits,unique,1
+; CHECK-NEXT:  .Ltmp0:
+; CHECK-NEXT:  foo.1
+; CHECK-NEXT:    callq   bar
+;
+
+bb2:                                                ; preds = %entry
+  %3 = call i32 @baz()
+  ret void
+; CHECK:         .section .text,"ax",@progbits,unique,2
+; CHECK-NEXT:  .Ltmp1:
+; CHECK-NEXT:  foo.2
+; CHECK-NEXT:    callq   baz
+}
+
+declare i32 @bar()
+declare i32 @baz()