From 8d94d3c3b44c3a27a69b153cef9be4b8e481150e Mon Sep 17 00:00:00 2001
From: Johannes Doerfert <johannes@jdoerfert.de>
Date: Mon, 11 May 2020 15:22:14 -0500
Subject: [PATCH] [Attributor][FIX] Disallow function signature rewrite for
 casted calls

We will now ensure ensure the return type of called function is the type
of all call sites we are going to rewrite. This avoids a problem
partially fixed by D79680. The part that was not covered is a use of
this "weird" casted call site (see `@func3` in `misc_crash.ll`).

misc_crash.ll checks are auto-generated now.
---
 llvm/lib/Transforms/IPO/Attributor.cpp        | 12 +++++--
 llvm/test/Transforms/Attributor/misc_crash.ll | 48 ++++++++++++++++++++++++---
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index e118923..c1a089b 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -1341,6 +1341,13 @@ bool Attributor::isValidFunctionSignatureRewrite(
     Argument &Arg, ArrayRef<Type *> ReplacementTypes) {
 
   auto CallSiteCanBeChanged = [](AbstractCallSite ACS) {
+    // Forbid the call site to cast the function return type. If we need to
+    // rewrite these functions we need to re-create a cast for the new call site
+    // (if the old had uses).
+    if (!ACS.getCalledFunction() ||
+        ACS.getInstruction()->getType() !=
+            ACS.getCalledFunction()->getReturnType())
+      return false;
     // Forbid must-tail calls for now.
     return !ACS.isCallbackCall() && !ACS.getInstruction()->isMustTailCall();
   };
@@ -1594,10 +1601,11 @@ ChangeStatus Attributor::rewriteFunctionSignatures(
     for (auto &CallSitePair : CallSitePairs) {
       CallBase &OldCB = *CallSitePair.first;
       CallBase &NewCB = *CallSitePair.second;
+      assert(OldCB.getType() == NewCB.getType() &&
+             "Cannot handle call sites with different types!");
       ModifiedFns.insert(OldCB.getFunction());
       CGUpdater.replaceCallSite(OldCB, NewCB);
-      if (!OldCB.use_empty())
-        OldCB.replaceAllUsesWith(&NewCB);
+      OldCB.replaceAllUsesWith(&NewCB);
       OldCB.eraseFromParent();
     }
 
diff --git a/llvm/test/Transforms/Attributor/misc_crash.ll b/llvm/test/Transforms/Attributor/misc_crash.ll
index 476ebbe..989b2bd 100644
--- a/llvm/test/Transforms/Attributor/misc_crash.ll
+++ b/llvm/test/Transforms/Attributor/misc_crash.ll
@@ -1,39 +1,77 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
 ; RUN: opt -attributor -S %s | FileCheck %s
 ; RUN: opt -passes=attributor -S %s | FileCheck %s
 
 @var1 = internal global [1 x i32] undef
 @var2 = internal global i32 0
 
-; CHECK-LABEL: define i32 addrspace(1)* @foo(i32 addrspace(4)* nofree readnone %arg)
 define i32 addrspace(1)* @foo(i32 addrspace(4)* %arg) {
+; CHECK-LABEL: define {{[^@]+}}@foo
+; CHECK-SAME: (i32 addrspace(4)* nofree readnone [[ARG:%.*]]) #0
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast i32 addrspace(4)* [[ARG]] to i32 addrspace(1)*
+; CHECK-NEXT:    ret i32 addrspace(1)* [[TMP0]]
+;
 entry:
   %0 = addrspacecast i32 addrspace(4)* %arg to i32 addrspace(1)*
   ret i32 addrspace(1)* %0
 }
 
 define i32* @func1() {
+; CHECK-LABEL: define {{[^@]+}}@func1() #0
+; CHECK-NEXT:    [[PTR:%.*]] = call i32* @func1a() #0
+; CHECK-NEXT:    ret i32* [[PTR]]
+;
   %ptr = call i32* @func1a([1 x i32]* @var1)
   ret i32* %ptr
 }
 
+; UTC_ARGS: --disable
 ; CHECK-LABEL: define internal nonnull align 4 dereferenceable(4) i32* @func1a()
 ; CHECK-NEXT: ret i32* getelementptr inbounds ([1 x i32], [1 x i32]* @var1, i32 0, i32 0)
 define internal i32* @func1a([1 x i32]* %arg) {
   %ptr = getelementptr inbounds [1 x i32], [1 x i32]* %arg, i64 0, i64 0
   ret i32* %ptr
 }
+; UTC_ARGS: --enable
 
 define internal void @func2a(i32* %0) {
+; CHECK-LABEL: define {{[^@]+}}@func2a
+; CHECK-SAME: (i32* nocapture nofree nonnull writeonly align 4 dereferenceable(4) [[TMP0:%.*]]) #1
+; CHECK-NEXT:    store i32 0, i32* @var2, align 4
+; CHECK-NEXT:    ret void
+;
   store i32 0, i32* %0
   ret void
 }
 
-; CHECK-LABEL: define i32 @func2()
-; CHECK-NEXT:   tail call void @func2a()
-; CHECK-NEXT:   %1 = load i32, i32* @var2, align 4
-; CHECK-NEXT:   ret i32 %1
 define i32 @func2() {
+; CHECK-LABEL: define {{[^@]+}}@func2()
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* nonnull align 4 dereferenceable(4) @var2)
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* @var2, align 4
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
   %1 = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* @var2)
   %2 = load i32, i32* @var2
   ret i32 %2
 }
+
+define i32 @func3(i1 %false) {
+; CHECK-LABEL: define {{[^@]+}}@func3
+; CHECK-SAME: (i1 [[FALSE:%.*]])
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* nonnull align 4 dereferenceable(4) @var2)
+; CHECK-NEXT:    br i1 [[FALSE]], label [[USE_BB:%.*]], label [[RET_BB:%.*]]
+; CHECK:       use_bb:
+; CHECK-NEXT:    ret i32 [[TMP1]]
+; CHECK:       ret_bb:
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* @var2, align 4
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* @var2)
+  br i1 %false, label %use_bb, label %ret_bb
+use_bb:
+  ret i32 %1
+ret_bb:
+  %2 = load i32, i32* @var2
+  ret i32 %2
+}
-- 
2.7.4