MergeFunctions: Impose a total order on the replacement of functions

We don't want to replace function A by Function B in one module and Function B
by Function A in another module.

If these functions are marked with linkonce_odr we would end up with a function
stub calling B in one module and a function stub calling A in another module. If
the linker decides to pick these two we will have two stubs calling each other.

rdar://21265586

llvm-svn: 239367

3 files changed, 78 insertions, 5 deletions

diff --git a/llvm/lib/Transforms/IPO/MergeFunctions.cpp b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
index 1327645183e..6ef1ac77da3 100644
--- a/llvm/lib/Transforms/IPO/MergeFunctions.cpp
+++ b/llvm/lib/Transforms/IPO/MergeFunctions.cpp
@@ -389,11 +389,21 @@ private:
 };
 
 class FunctionNode {
-  AssertingVH<Function> F;
+  mutable AssertingVH<Function> F;
 
 public:
   FunctionNode(Function *F) : F(F) {}
   Function *getFunc() const { return F; }
+
+  /// Replace the reference to the function F by the function G, assuming their
+  /// implementations are equal.
+  void replaceBy(Function *G) const {
+    assert(!(*this < FunctionNode(G)) && !(FunctionNode(G) < *this) &&
+           "The two functions must be equal");
+
+    F = G;
+  }
+
   void release() { F = 0; }
   bool operator<(const FunctionNode &RHS) const {
     return (FunctionComparator(F, RHS.getFunc()).compare()) == -1;
@@ -1122,6 +1132,9 @@ private:
   /// Replace G with an alias to F. Deletes G.
   void writeAlias(Function *F, Function *G);
 
+  /// Replace function F with function G in the function tree.
+  void replaceFunctionInTree(FnTreeType::iterator &IterToF, Function *G);
+
   /// The set of all distinct functions. Use the insert() and remove() methods
   /// to modify it.
   FnTreeType FnTree;
@@ -1414,6 +1427,20 @@ void MergeFunctions::mergeTwoFunctions(Function *F, Function *G) {
   ++NumFunctionsMerged;
 }
 
+/// Replace function F for function G in the map.
+void MergeFunctions::replaceFunctionInTree(FnTreeType::iterator &IterToF,
+                                           Function *G) {
+  Function *F = IterToF->getFunc();
+
+  // A total order is already guaranteed otherwise because we process strong
+  // functions before weak functions.
+  assert((F->mayBeOverridden() && G->mayBeOverridden()) ||
+         (!F->mayBeOverridden() && !G->mayBeOverridden()) &&
+         "Only change functions if both are strong or both are weak");
+
+  IterToF->replaceBy(G);
+}
+
 // Insert a ComparableFunction into the FnTree, or merge it away if equal to one
 // that was already inserted.
 bool MergeFunctions::insert(Function *NewFunction) {
@@ -1439,6 +1466,22 @@ bool MergeFunctions::insert(Function *NewFunction) {
     }
   }
 
+  // Impose a total order (by name) on the replacement of functions. This is
+  // important when operating on more than one module independently to prevent
+  // cycles of thunks calling each other when the modules are linked together.
+  //
+  // When one function is weak and the other is strong there is an order imposed
+  // already. We process strong functions before weak functions.
+  if ((OldF.getFunc()->mayBeOverridden() && NewFunction->mayBeOverridden()) ||
+      (!OldF.getFunc()->mayBeOverridden() && !NewFunction->mayBeOverridden()))
+    if (OldF.getFunc()->getName() > NewFunction->getName()) {
+      // Swap the two functions.
+      Function *F = OldF.getFunc();
+      replaceFunctionInTree(Result.first, NewFunction);
+      NewFunction = F;
+      assert(OldF.getFunc() != F && "Must have swapped the functions.");
+    }
+
   // Never thunk a strong function to a weak function.
   assert(!OldF.getFunc()->mayBeOverridden() || NewFunction->mayBeOverridden());
 
diff --git a/llvm/test/Transforms/MergeFunc/call-and-invoke-with-ranges.ll b/llvm/test/Transforms/MergeFunc/call-and-invoke-with-ranges.ll
index b2083cb5c70..99eba5e2809 100644
--- a/llvm/test/Transforms/MergeFunc/call-and-invoke-with-ranges.ll
+++ b/llvm/test/Transforms/MergeFunc/call-and-invoke-with-ranges.ll
@@ -63,16 +63,16 @@ lpad:
   resume { i8*, i32 } zeroinitializer
 }
 
-define i8 @call_same_range() {
-; CHECK-LABEL: @call_same_range
+define i8 @call_with_same_range() {
+; CHECK-LABEL: @call_with_same_range
 ; CHECK: tail call i8 @call_with_range
   bitcast i8 0 to i8
   %out = call i8 @dummy(), !range !0
   ret i8 %out
 }
 
-define i8 @invoke_same_range() {
-; CHECK-LABEL: @invoke_same_range()
+define i8 @invoke_with_same_range() {
+; CHECK-LABEL: @invoke_with_same_range()
 ; CHECK: tail call i8 @invoke_with_range()
   %out = invoke i8 @dummy() to label %next unwind label %lpad, !range !0
 
diff --git a/llvm/test/Transforms/MergeFunc/linkonce_odr.ll b/llvm/test/Transforms/MergeFunc/linkonce_odr.ll
new file mode 100644
index 00000000000..1ad0d727f83
--- /dev/null
+++ b/llvm/test/Transforms/MergeFunc/linkonce_odr.ll
@@ -0,0 +1,30 @@
+; RUN: opt -S -mergefunc < %s | FileCheck %s
+
+; Replacments should be totally ordered on the function name.
+; If we don't do this we  can end up with one module defining a thunk for @funA
+; and another module defining a thunk for @funB.
+;
+; The problem with this is that the linker could then choose these two stubs
+; each of the two modules and we end up with two stubs calling each other.
+
+; CHECK-LABEL: define linkonce_odr i32 @funA
+; CHECK-NEXT:    add
+; CHECK:         ret
+
+; CHECK-LABEL: define linkonce_odr i32 @funB
+; CHECK-NEXT:    tail call i32 @funA(i32 %0, i32 %1)
+; CHECK-NEXT:    ret
+
+define linkonce_odr i32 @funB(i32 %x, i32 %y) {
+  %sum = add i32 %x, %y
+  %sum2 = add i32 %x, %sum
+  %sum3 = add i32 %x, %sum2
+  ret i32 %sum3
+}
+
+define linkonce_odr i32 @funA(i32 %x, i32 %y) {
+  %sum = add i32 %x, %y
+  %sum2 = add i32 %x, %sum
+  %sum3 = add i32 %x, %sum2
+  ret i32 %sum3
+}