[xray] Handle conditional calls when lowering patchable tail calls.

xray instruments tail call function exits by inserting a nop sled before
the tail call. When tracing is enabled, the nop sled is replaced with a
call to `__xray_FunctionTailExit()`. This currently does not work for
conditional tail calls, as the instrumentation assumes that the tail
call will be unconditional. This causes two issues:
 - `__xray_FunctionTailExit()` is inappropately called even when the
   tail call is not taken.
 - `__xray_FunctionTailExit()`'s prologue/epilogue adjusts the stack
   pointer with add/sub instructions. This clobbers condition flags,
   which can flip the condition used for the tail call, leading to
   incorrect program behavior.

Fix this by rewriting conditional calls when lowering patchable tail
calls.

With this change, a conditional patchable tail call like:
```
  je target
```

Will be lowered to:
```
  jne .fallthrough
  .p2align 1, ..
.Lxray_sled_N:
  SLED_CODE
  jmp target
.fallthrough:
```

Author: rickyz

llvm/lib/Target/X86/X86MCInstLower.cpp

@@ -22,6 +22,7 @@
 #include "X86RegisterInfo.h"
 #include "X86ShuffleDecodeConstantPool.h"
 #include "X86Subtarget.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
@@ -1361,6 +1362,35 @@ void X86AsmPrinter::LowerPATCHABLE_RET(const MachineInstr &MI,
 
 void X86AsmPrinter::LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI,
                                              X86MCInstLower &MCIL) {
+  MCInst TC;
+  TC.setOpcode(convertTailJumpOpcode(MI.getOperand(0).getImm()));
+  // Drop the tail jump opcode.
+  auto TCOperands = drop_begin(MI.operands());
+  bool IsConditional = TC.getOpcode() == X86::JCC_1;
+  MCSymbol *FallthroughLabel;
+  if (IsConditional) {
+    // Rewrite:
+    //   je target
+    //
+    // To:
+    //   jne .fallthrough
+    //   .p2align 1, ...
+    // .Lxray_sled_N:
+    //   SLED_CODE
+    //   jmp target
+    // .fallthrough:
+    FallthroughLabel = OutContext.createTempSymbol();
+    EmitToStreamer(
+        *OutStreamer,
+        MCInstBuilder(X86::JCC_1)
+            .addExpr(MCSymbolRefExpr::create(FallthroughLabel, OutContext))
+            .addImm(X86::GetOppositeBranchCondition(
+                static_cast<X86::CondCode>(MI.getOperand(2).getImm()))));
+    TC.setOpcode(X86::JMP_1);
+    // Drop the condition code.
+    TCOperands = drop_end(TCOperands);
+  }
+
   NoAutoPaddingScope NoPadScope(*OutStreamer);
 
   // Like PATCHABLE_RET, we have the actual instruction in the operands to this
@@ -1382,18 +1412,16 @@ void X86AsmPrinter::LowerPATCHABLE_TAIL_CALL(const MachineInstr &MI,
   OutStreamer->emitLabel(Target);
   recordSled(CurSled, MI, SledKind::TAIL_CALL, 2);
 
-  unsigned OpCode = MI.getOperand(0).getImm();
-  OpCode = convertTailJumpOpcode(OpCode);
-  MCInst TC;
-  TC.setOpcode(OpCode);
-
   // Before emitting the instruction, add a comment to indicate that this is
   // indeed a tail call.
   OutStreamer->AddComment("TAILCALL");
-  for (auto &MO : drop_begin(MI.operands()))
+  for (auto &MO : TCOperands)
     if (auto MaybeOperand = MCIL.LowerMachineOperand(&MI, MO))
       TC.addOperand(*MaybeOperand);
   OutStreamer->emitInstruction(TC, getSubtargetInfo());
+
+  if (IsConditional)
+    OutStreamer->emitLabel(FallthroughLabel);
 }
 
 // Returns instruction preceding MBBI in MachineFunction.

llvm/test/CodeGen/X86/xray-conditional-tail-call.ll

@@ -0,0 +1,35 @@
+; RUN: llc -mtriple=x86_64 < %s | FileCheck %s
+
+declare void @tail_call_target()
+
+define void @conditional_tail_call(i32 %cond) "function-instrument"="xray-always" nounwind {
+  ; CHECK-LABEL: conditional_tail_call:
+  ; CHECK-NEXT:  .Lfunc_begin0:
+  ; CHECK-NEXT:  # %bb.0:
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_sled_0:
+  ; CHECK-NEXT:    .ascii "\353\t"
+  ; CHECK-NEXT:    nopw 512(%rax,%rax)
+  ; CHECK-NEXT:    testl %edi, %edi
+  ; CHECK-NEXT:    je .Ltmp0
+  ; CHECK-NEXT:    .p2align 1, 0x90
+  ; CHECK-NEXT:  .Lxray_sled_1:
+  ; CHECK-NEXT:    .ascii "\353\t"
+  ; CHECK-NEXT:    nopw  512(%rax,%rax)
+  ; CHECK-NEXT:  .Ltmp1:
+  ; CHECK-NEXT:    jmp tail_call_target@PLT # TAILCALL
+  ; CHECK-NEXT:  .Ltmp0:
+  ; CHECK-NEXT:  # %bb.1:
+  ; CHECK-NEXT:   .p2align  1, 0x90
+  ; CHECK-NEXT:  .Lxray_sled_2:
+  ; CHECK-NEXT:    retq
+  ; CHECK-NEXT:    nopw %cs:512(%rax,%rax)
+  ; CHECK-NEXT:  .Lfunc_end0:
+  %cmp = icmp ne i32 %cond, 0
+  br i1 %cmp, label %docall, label %ret
+docall:
+  tail call void @tail_call_target()
+  ret void
+ret:
+  ret void
+}