[llvm] r275033 - Don't use a SmallSet for returned attribute inference

[Hal Finkel via llvm-commits]

Author: hfinkel
Date: Sun Jul 10 20:14:21 2016
New Revision: 275033

URL: http://llvm.org/viewvc/llvm-project?rev=275033&view=rev
Log:
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).

Modified:
    llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp

Modified: llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp
URL: http://llvm.org/viewvc/llvm-project/llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp?rev=275033&r1=275032&r2=275033&view=diff
==============================================================================
--- llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp (original)
+++ llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp Sun Jul 10 20:14:21 2016
@@ -502,18 +502,26 @@ static bool addArgumentReturnedAttrs(con
     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;
+          }
+        }
 
-    if (RetArgs.size() == 1) {
-      auto *A = cast<Argument>(*RetArgs.begin());
+      return RetArg;
+    };
+
+    if (Value *RetArg = FindRetArg()) {
+      auto *A = cast<Argument>(RetArg);
       A->addAttr(AttributeSet::get(F->getContext(), A->getArgNo() + 1, B));
       ++NumReturned;
       Changed = true;