Don't use a SmallSet for returned attribute inference

Suggested post-commit by David Majnemer on IRC (following-up on a pre-commit
review comment).

llvm-svn: 275033

1 files changed, 19 insertions, 11 deletions

diff --git a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
index 3f87a1b233f..9c9e4822199 100644
--- a/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionAttrs.cpp
@@ -502,18 +502,26 @@ static bool addArgumentReturnedAttrs(const SCCNodeSet &SCCNodes) {
     if (F->getReturnType()->isVoidTy())
       continue;
 
-    SmallPtrSet<Value *, 2> RetArgs;
-    for (BasicBlock &BB : *F)
-      if (auto *Ret = dyn_cast<ReturnInst>(BB.getTerminator())) {
-        // Note that stripPointerCasts should look through functions with
-        // returned arguments.
-        Value *RetVal = Ret->getReturnValue()->stripPointerCasts();
-        if (RetVal->getType() == F->getReturnType() && isa<Argument>(RetVal))
-          RetArgs.insert(RetVal);
-      }
+    auto FindRetArg = [&]() -> Value * {
+      Value *RetArg = nullptr;
+      for (BasicBlock &BB : *F)
+        if (auto *Ret = dyn_cast<ReturnInst>(BB.getTerminator())) {
+          // Note that stripPointerCasts should look through functions with
+          // returned arguments.
+          Value *RetVal = Ret->getReturnValue()->stripPointerCasts();
+          if (RetVal->getType() == F->getReturnType() && isa<Argument>(RetVal)) {
+            if (!RetArg)
+              RetArg = RetVal;
+            else if (RetArg != RetVal)
+              return nullptr;
+          }
+        }
+
+      return RetArg;
+    };
 
-    if (RetArgs.size() == 1) {
-      auto *A = cast<Argument>(*RetArgs.begin());
+    if (Value *RetArg = FindRetArg()) {
+      auto *A = cast<Argument>(RetArg);
       A->addAttr(AttributeSet::get(F->getContext(), A->getArgNo() + 1, B));
       ++NumReturned;
       Changed = true;