From 5f10dae16e0979ceea71ff2c89f9704b200a8da0 Mon Sep 17 00:00:00 2001
From: Hans Wennborg <hans@hanshq.net>
Date: Fri, 6 Apr 2018 10:14:09 +0000
Subject: [PATCH] EntryExitInstrumenter: Handle musttail calls

Inserting instrumentation between a musttail call and ret instruction
would create invalid IR. Instead, treat musttail calls as function
exits.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@329385 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../Utils/EntryExitInstrumenter.cpp           | 20 +++++++++++----
 .../EntryExitInstrumenter/mcount.ll           | 25 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/lib/Transforms/Utils/EntryExitInstrumenter.cpp b/lib/Transforms/Utils/EntryExitInstrumenter.cpp
index d7c11de21a8..569ea58a304 100644
--- a/lib/Transforms/Utils/EntryExitInstrumenter.cpp
+++ b/lib/Transforms/Utils/EntryExitInstrumenter.cpp
@@ -91,17 +91,27 @@ static bool runOnFunction(Function &F, bool PostInlining) {
 
   if (!ExitFunc.empty()) {
     for (BasicBlock &BB : F) {
-      TerminatorInst *T = BB.getTerminator();
+      Instruction *T = BB.getTerminator();
+      if (!isa<ReturnInst>(T))
+        continue;
+
+      // If T is preceded by a musttail call, that's the real terminator.
+      Instruction *Prev = T->getPrevNode();
+      if (BitCastInst *BCI = dyn_cast_or_null<BitCastInst>(Prev))
+        Prev = BCI->getPrevNode();
+      if (CallInst *CI = dyn_cast_or_null<CallInst>(Prev)) {
+        if (CI->isMustTailCall())
+          T = CI;
+      }
+
       DebugLoc DL;
       if (DebugLoc TerminatorDL = T->getDebugLoc())
         DL = TerminatorDL;
       else if (auto SP = F.getSubprogram())
         DL = DebugLoc::get(0, 0, SP);
 
-      if (isa<ReturnInst>(T)) {
-        insertCall(F, ExitFunc, T, DL);
-        Changed = true;
-      }
+      insertCall(F, ExitFunc, T, DL);
+      Changed = true;
     }
     F.removeAttribute(AttributeList::FunctionIndex, ExitAttr);
   }
diff --git a/test/Transforms/EntryExitInstrumenter/mcount.ll b/test/Transforms/EntryExitInstrumenter/mcount.ll
index 50e6b1321bb..30483606caf 100644
--- a/test/Transforms/EntryExitInstrumenter/mcount.ll
+++ b/test/Transforms/EntryExitInstrumenter/mcount.ll
@@ -44,7 +44,6 @@ entry:
 ; CHECK-NEXT ret void
 }
 
-attributes #0 = { "instrument-function-entry-inlined"="mcount" "instrument-function-entry"="__cyg_profile_func_enter" "instrument-function-exit"="__cyg_profile_func_exit" }
 
 
 ; The mcount function has many different names.
@@ -78,10 +77,33 @@ define void @f7() #7 { entry: ret void }
 ; CHECK: call void @__cyg_profile_func_enter_bare
 
 
+; Treat musttail calls as terminators; inserting between the musttail call and
+; ret is not allowed.
+declare i32* @tailcallee()
+define i32* @tailcaller() #8 {
+  %1 = musttail call i32* @tailcallee()
+  ret i32* %1
+; CHECK-LABEL: define i32* @tailcaller
+; CHECK: call void @__cyg_profile_func_exit
+; CHECK: musttail call i32* @tailcallee
+; CHECK: ret
+}
+define i8* @tailcaller2() #8 {
+  %1 = musttail call i32* @tailcallee()
+  %2 = bitcast i32* %1 to i8*
+  ret i8* %2
+; CHECK-LABEL: define i8* @tailcaller2
+; CHECK: call void @__cyg_profile_func_exit
+; CHECK: musttail call i32* @tailcallee
+; CHECK: bitcast
+; CHECK: ret
+}
+
 ; The attributes are "consumed" when the instrumentation is inserted.
 ; CHECK: attributes
 ; CHECK-NOT: instrument-function
 
+attributes #0 = { "instrument-function-entry-inlined"="mcount" "instrument-function-entry"="__cyg_profile_func_enter" "instrument-function-exit"="__cyg_profile_func_exit" }
 attributes #1 = { "instrument-function-entry-inlined"=".mcount" }
 attributes #2 = { "instrument-function-entry-inlined"="\01__gnu_mcount_nc" }
 attributes #3 = { "instrument-function-entry-inlined"="\01_mcount" }
@@ -89,3 +111,4 @@ attributes #4 = { "instrument-function-entry-inlined"="\01mcount" }
 attributes #5 = { "instrument-function-entry-inlined"="__mcount" }
 attributes #6 = { "instrument-function-entry-inlined"="_mcount" }
 attributes #7 = { "instrument-function-entry-inlined"="__cyg_profile_func_enter_bare" }
+attributes #8 = { "instrument-function-exit"="__cyg_profile_func_exit" }