[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.

Author: jdoerfert

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();
     }
 

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
+}