Update NRVO logic to support early return (Attempt 2)

Summary:
This is the second attempt of r333500 (Update NRVO logic to support early return).
The previous one was reverted for a miscompilation for an incorrect NRVO set up on templates such as:
```
struct Foo {};

template <typename T>
T bar() {
  T t;
  if (false)
    return T();
  return t;
}
```

Where, `t` is marked as non-NRVO variable before its instantiation. However, while its instantiation, it's left an NRVO candidate, turned into an NRVO variable later.

Reviewers: rsmith

Reviewed By: rsmith

Subscribers: cfe-commits

Differential Revision: https://reviews.llvm.org/D47586

llvm-svn: 335019
diff --git a/clang/lib/Sema/Scope.cpp b/clang/lib/Sema/Scope.cpp
index eae5a32..5a46ba2 100644
--- a/clang/lib/Sema/Scope.cpp
+++ b/clang/lib/Sema/Scope.cpp
@@ -92,7 +92,6 @@
   UsingDirectives.clear();
   Entity = nullptr;
   ErrorTrap.reset();
-  NRVO.setPointerAndInt(nullptr, 0);
 }
 
 bool Scope::containedInPrototypeScope() const {
@@ -119,19 +118,15 @@
   Flags |= FlagsToSet;
 }
 
-void Scope::mergeNRVOIntoParent() {
-  if (VarDecl *Candidate = NRVO.getPointer()) {
-    if (isDeclScope(Candidate))
-      Candidate->setNRVOVariable(true);
+void Scope::setNRVOCandidate(VarDecl *Candidate) {
+  for (Decl *D : DeclsInScope) {
+    VarDecl *VD = dyn_cast<VarDecl>(D);
+    if (VD && VD != Candidate && VD->isNRVOCandidate())
+      VD->setNRVOVariable(false);
   }
 
-  if (getEntity())
-    return;
-
-  if (NRVO.getInt())
-    getParent()->setNoNRVO();
-  else if (NRVO.getPointer())
-    getParent()->addNRVOCandidate(NRVO.getPointer());
+  if (Scope *parent = getParent())
+    parent->setNRVOCandidate(Candidate);
 }
 
 LLVM_DUMP_METHOD void Scope::dump() const { dumpImpl(llvm::errs()); }
@@ -191,9 +186,4 @@
   OS << "MSCurManglingNumber: " << getMSCurManglingNumber() << '\n';
   if (const DeclContext *DC = getEntity())
     OS << "Entity : (clang::DeclContext*)" << DC << '\n';
-
-  if (NRVO.getInt())
-    OS << "NRVO not allowed\n";
-  else if (NRVO.getPointer())
-    OS << "NRVO candidate : (clang::VarDecl*)" << NRVO.getPointer() << '\n';
 }
diff --git a/clang/lib/Sema/SemaDecl.cpp b/clang/lib/Sema/SemaDecl.cpp
index f6faf38..4c98ffe 100644
--- a/clang/lib/Sema/SemaDecl.cpp
+++ b/clang/lib/Sema/SemaDecl.cpp
@@ -1798,8 +1798,6 @@
 }
 
 void Sema::ActOnPopScope(SourceLocation Loc, Scope *S) {
-  S->mergeNRVOIntoParent();
-
   if (S->decl_empty()) return;
   assert((S->getFlags() & (Scope::DeclScope | Scope::TemplateParamScope)) &&
          "Scope shouldn't contain decls!");
@@ -12563,21 +12561,24 @@
 /// optimization.
 ///
 /// Each of the variables that is subject to the named return value
-/// optimization will be marked as NRVO variables in the AST, and any
+/// optimization will be marked as NRVO variable candidates in the AST, and any
 /// return statement that has a marked NRVO variable as its NRVO candidate can
 /// use the named return value optimization.
 ///
-/// This function applies a very simplistic algorithm for NRVO: if every return
-/// statement in the scope of a variable has the same NRVO candidate, that
-/// candidate is an NRVO variable.
+/// This function applies a very simplistic algorithm for NRVO: if every
+/// reachable return statement in the scope of a variable has the same NRVO
+/// candidate, that candidate is an NRVO variable.
 void Sema::computeNRVO(Stmt *Body, FunctionScopeInfo *Scope) {
-  ReturnStmt **Returns = Scope->Returns.data();
+  for (ReturnStmt *Return : Scope->Returns) {
+    const VarDecl *Candidate = Return->getNRVOCandidate();
+    if (!Candidate)
+      continue;
 
-  for (unsigned I = 0, E = Scope->Returns.size(); I != E; ++I) {
-    if (const VarDecl *NRVOCandidate = Returns[I]->getNRVOCandidate()) {
-      if (!NRVOCandidate->isNRVOVariable())
-        Returns[I]->setNRVOCandidate(nullptr);
-    }
+    if (Candidate->isNRVOCandidate())
+      const_cast<VarDecl*>(Candidate)->setNRVOVariable(true);
+
+    if (!Candidate->isNRVOVariable())
+      Return->setNRVOCandidate(nullptr);
   }
 }
 
@@ -12712,12 +12713,8 @@
       else if (CXXDestructorDecl *Destructor = dyn_cast<CXXDestructorDecl>(FD))
         MarkVTableUsed(FD->getLocation(), Destructor->getParent());
 
-      // Try to apply the named return value optimization. We have to check
-      // if we can do this here because lambdas keep return statements around
-      // to deduce an implicit return type.
-      if (FD->getReturnType()->isRecordType() &&
-          (!getLangOpts().CPlusPlus || !FD->isDependentContext()))
-        computeNRVO(Body, getCurFunction());
+      // Try to apply the named return value optimization.
+      computeNRVO(Body, getCurFunction());
     }
 
     // GNU warning -Wmissing-prototypes:
diff --git a/clang/lib/Sema/SemaExpr.cpp b/clang/lib/Sema/SemaExpr.cpp
index aeedd6b..220d690 100644
--- a/clang/lib/Sema/SemaExpr.cpp
+++ b/clang/lib/Sema/SemaExpr.cpp
@@ -13385,13 +13385,9 @@
   if (Body && getCurFunction()->HasPotentialAvailabilityViolations)
     DiagnoseUnguardedAvailabilityViolations(BSI->TheDecl);
 
-  // Try to apply the named return value optimization. We have to check again
-  // if we can do this, though, because blocks keep return statements around
-  // to deduce an implicit return type.
-  if (getLangOpts().CPlusPlus && RetTy->isRecordType() &&
-      !BSI->TheDecl->isDependentContext())
-    computeNRVO(Body, BSI);
-  
+  // Try to apply the named return value optimization.
+  computeNRVO(Body, BSI);
+
   BlockExpr *Result = new (Context) BlockExpr(BSI->TheDecl, BlockTy);
   AnalysisBasedWarnings::Policy WP = AnalysisWarnings.getDefaultPolicy();
   PopFunctionScopeInfo(&WP, Result->getBlockDecl(), Result);
diff --git a/clang/lib/Sema/SemaStmt.cpp b/clang/lib/Sema/SemaStmt.cpp
index ccf25ee..04f5114 100644
--- a/clang/lib/Sema/SemaStmt.cpp
+++ b/clang/lib/Sema/SemaStmt.cpp
@@ -3455,12 +3455,9 @@
                            ExpressionEvaluationContext::DiscardedStatement)
     return R;
 
-  if (VarDecl *VD =
-      const_cast<VarDecl*>(cast<ReturnStmt>(R.get())->getNRVOCandidate())) {
-    CurScope->addNRVOCandidate(VD);
-  } else {
-    CurScope->setNoNRVO();
-  }
+  VarDecl *VD =
+      const_cast<VarDecl*>(cast<ReturnStmt>(R.get())->getNRVOCandidate());
+  CurScope->setNRVOCandidate(VD);
 
   CheckJumpOutOfSEHFinally(*this, ReturnLoc, *CurScope->getFnParent());
 
diff --git a/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp b/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
index 68857d9..7e38dcc 100644
--- a/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
+++ b/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
@@ -740,12 +740,13 @@
   SemaRef.BuildVariableInstantiation(Var, D, TemplateArgs, LateAttrs, Owner,
                                      StartingScope, InstantiatingVarTemplate);
 
+  bool NRVO = false;
   if (D->isNRVOVariable()) {
     QualType ReturnType = cast<FunctionDecl>(DC)->getReturnType();
     if (SemaRef.isCopyElisionCandidate(ReturnType, Var, Sema::CES_Strict))
-      Var->setNRVOVariable(true);
+      NRVO = true;
   }
-
+  Var->setNRVOVariable(NRVO);
   Var->setImplicit(D->isImplicit());
 
   return Var;