MergeFunc: Transfer the callee's attributes when replacing a direct caller

We insert a bitcast which obfuscates the getCalledFunction for the utility
function which looks up attributes from the called function. Loosing ABI
changing parameter attributes is a bad thing.

rdar://21516488

llvm-svn: 242807

2 files changed, 59 insertions, 0 deletions

diff --git a/llvm/lib/Transforms/IPO/MergeFunctions.cpp b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
index 109cac7fece..7ece66302f9 100644
--- a/llvm/lib/Transforms/IPO/MergeFunctions.cpp
+++ b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
@@ -1282,6 +1282,25 @@ void MergeFunctions::replaceDirectCallers(Function *Old, Function *New) {
     ++UI;
     CallSite CS(U->getUser());
     if (CS && CS.isCallee(U)) {
+      // Transfer the called function's attributes to the call site. Due to the
+      // bitcast we will 'loose' ABI changing attributes because the 'called
+      // function' is no longer a Function* but the bitcast. Code that looks up
+      // the attributes from the called function will fail.
+      auto &Context = New->getContext();
+      auto NewFuncAttrs = New->getAttributes();
+      auto CallSiteAttrs = CS.getAttributes();
+
+      CallSiteAttrs = CallSiteAttrs.addAttributes(
+          Context, AttributeSet::ReturnIndex, NewFuncAttrs.getRetAttributes());
+
+      for (unsigned argIdx = 0; argIdx < CS.arg_size(); argIdx++) {
+        AttributeSet Attrs = NewFuncAttrs.getParamAttributes(argIdx);
+        if (Attrs.getNumSlots())
+          CallSiteAttrs = CallSiteAttrs.addAttributes(Context, argIdx, Attrs);
+      }
+
+      CS.setAttributes(CallSiteAttrs);
+
       remove(CS.getInstruction()->getParent()->getParent());
       U->set(BitcastNew);
     }
diff --git a/llvm/test/Transforms/MergeFunc/apply_function_attributes.ll b/llvm/test/Transforms/MergeFunc/apply_function_attributes.ll
new file mode 100644
index 00000000000..fa7f052f7ff
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/apply_function_attributes.ll
@@ -0,0 +1,40 @@
+; RUN: opt -S -mergefunc < %s | FileCheck %s
+
+%Opaque_type = type opaque
+%S2i = type <{ i64, i64 }>
+%D2i = type <{ i64, i64 }>
+%Di = type <{ i32 }>
+%Si = type <{ i32 }>
+
+define void @B(%Opaque_type* sret %a, %S2i* %b, i32* %xp, i32* %yp) {
+  %x = load i32, i32* %xp
+  %y = load i32, i32* %yp
+  %sum = add i32 %x, %y
+  %sum2 = add i32 %sum, %y
+  %sum3 = add i32 %sum2, %y
+  ret void
+}
+
+define void @C(%Opaque_type* sret %a, %S2i* %b, i32* %xp, i32* %yp) {
+  %x = load i32, i32* %xp
+  %y = load i32, i32* %yp
+  %sum = add i32 %x, %y
+  %sum2 = add i32 %sum, %y
+  %sum3 = add i32 %sum2, %y
+  ret void
+}
+
+define void @A(%Opaque_type* sret %a, %D2i* %b, i32* %xp, i32* %yp) {
+  %x = load i32, i32* %xp
+  %y = load i32, i32* %yp
+  %sum = add i32 %x, %y
+  %sum2 = add i32 %sum, %y
+  %sum3 = add i32 %sum2, %y
+  ret void
+}
+
+; Make sure we transfer the parameter attributes to the call site.
+; CHECK-LABEL: define void @C(%Opaque_type* sret
+; CHECK:  tail call void bitcast (void (%Opaque_type*, %D2i*, i32*, i32*)* @A to void (%Opaque_type*, %S2i*, i32*, i32*)*)(%Opaque_type* sret %0, %S2i* %1, i32* %2, i32* %3)
+; CHECK:  ret void
+