[XRay] Fix assertion failure on empty machine basic blocks (PR 31424)

The original version of the code in XRayInstrumentation.cpp assumed that
functions may not have empty machine basic blocks (or that the first one
couldn't be). This change addresses that by special-casing that specific
situation.

We provide two .mir test-cases to make sure we're handling this
appropriately.

Fixes llvm.org/PR31424.

Reviewers: chandlerc

Subscribers: varno, llvm-commits

Differential Revision: https://reviews.llvm.org/D27913

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@290091 91177308-0d34-0410-b5e6-96231b3b80d8

lib/CodeGen/XRayInstrumentation.cpp

@@ -129,7 +129,15 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
       return false; // Function is too small.
   }
 
-  auto &FirstMBB = *MF.begin();
+  // We look for the first non-empty MachineBasicBlock, so that we can insert
+  // the function instrumentation in the appropriate place.
+  auto MBI =
+      find_if(MF, [&](const MachineBasicBlock &MBB) { return !MBB.empty(); });
+  if (MBI == MF.end())
+    return false; // The function is empty.
+
+  auto *TII = MF.getSubtarget().getInstrInfo();
+  auto &FirstMBB = *MBI;
   auto &FirstMI = *FirstMBB.begin();
 
   if (!MF.getSubtarget().isXRaySupported()) {
@@ -142,7 +150,6 @@ bool XRayInstrumentation::runOnMachineFunction(MachineFunction &MF) {
 
   // First, insert an PATCHABLE_FUNCTION_ENTER as the first instruction of the
   // MachineFunction.
-  auto *TII = MF.getSubtarget().getInstrInfo();
   BuildMI(FirstMBB, FirstMI, FirstMI.getDebugLoc(),
           TII->get(TargetOpcode::PATCHABLE_FUNCTION_ENTER));
 

test/CodeGen/X86/xray-empty-firstmbb.mir

@@ -0,0 +1,23 @@
+# RUN: llc -run-pass xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+#
+# Make sure we can handle empty first basic blocks.
+
+--- |
+
+  define i32 @foo() noinline uwtable "xray-instruction-threshold"="1" {
+  entry:
+    unreachable
+  }
+
+...
+
+---
+name: foo
+tracksRegLiveness: true
+liveins:
+  - { reg: '%edi'}
+body:            |
+  bb.0.entry:
+    liveins: %edi
+    ; CHECK-NOT: PATCHABLE_FUNCTION_ENTER
+...

test/CodeGen/X86/xray-empty-function.mir

@@ -0,0 +1,13 @@
+# RUN: llc -run-pass xray-instrumentation -mtriple=x86_64-unknown-linux-gnu -o - %s | FileCheck %s
+#
+# Make sure we can handle empty functions.
+---
+name: empty
+tracksRegLiveness: true
+liveins:
+  - { reg: '%edi'}
+body:            |
+  bb.0:
+    ; CHECK-NOT: PATCHABLE_FUNCTION_ENTER
+...
+