Handle init statements in readability-else-after-return

Adds a new ASTMatcher condition called 'hasInitStatement()' that matches if,
switch and range-for statements with an initializer. Reworked clang-tidy
readability-else-after-return to handle variables in the if condition or init
statements in c++17 ifs. Also checks if removing the else would affect object
lifetimes in the else branch.

Fixes PR44364.

2 files changed, 216 insertions, 30 deletions

diff --git a/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.cpp b/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.cpp
index 520586dfe25..93b197a52a5 100644
--- a/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.cpp
+++ b/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.cpp
@@ -10,6 +10,7 @@
 #include "clang/AST/ASTContext.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/Tooling/FixIt.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace clang::ast_matchers;
 
@@ -17,45 +18,226 @@ namespace clang {
 namespace tidy {
 namespace readability {
 
+namespace {
+static const char ReturnStr[] = "return";
+static const char ContinueStr[] = "continue";
+static const char BreakStr[] = "break";
+static const char ThrowStr[] = "throw";
+static const char WarningMessage[] = "do not use 'else' after '%0'";
+static const char WarnOnUnfixableStr[] = "WarnOnUnfixable";
+
+const DeclRefExpr *findUsage(const Stmt *Node, int64_t DeclIdentifier) {
+  if (const auto *DeclRef = dyn_cast<DeclRefExpr>(Node)) {
+    if (DeclRef->getDecl()->getID() == DeclIdentifier) {
+      return DeclRef;
+    }
+  } else {
+    for (const Stmt *ChildNode : Node->children()) {
+      if (const DeclRefExpr *Result = findUsage(ChildNode, DeclIdentifier)) {
+        return Result;
+      }
+    }
+  }
+  return nullptr;
+}
+
+const DeclRefExpr *
+findUsageRange(const Stmt *Node,
+               const llvm::iterator_range<int64_t *> &DeclIdentifiers) {
+  if (const auto *DeclRef = dyn_cast<DeclRefExpr>(Node)) {
+    if (llvm::is_contained(DeclIdentifiers, DeclRef->getDecl()->getID())) {
+      return DeclRef;
+    }
+  } else {
+    for (const Stmt *ChildNode : Node->children()) {
+      if (const DeclRefExpr *Result =
+              findUsageRange(ChildNode, DeclIdentifiers)) {
+        return Result;
+      }
+    }
+  }
+  return nullptr;
+}
+
+const DeclRefExpr *checkInitDeclUsageInElse(const IfStmt *If) {
+  const auto *InitDeclStmt = dyn_cast_or_null<DeclStmt>(If->getInit());
+  if (!InitDeclStmt)
+    return nullptr;
+  if (InitDeclStmt->isSingleDecl()) {
+    const Decl *InitDecl = InitDeclStmt->getSingleDecl();
+    assert(isa<VarDecl>(InitDecl) && "SingleDecl must be a VarDecl");
+    return findUsage(If->getElse(), InitDecl->getID());
+  }
+  llvm::SmallVector<int64_t, 4> DeclIdentifiers;
+  for (const Decl *ChildDecl : InitDeclStmt->decls()) {
+    assert(isa<VarDecl>(ChildDecl) && "Init Decls must be a VarDecl");
+    DeclIdentifiers.push_back(ChildDecl->getID());
+  }
+  return findUsageRange(If->getElse(), DeclIdentifiers);
+}
+
+const DeclRefExpr *checkConditionVarUsageInElse(const IfStmt *If) {
+  const VarDecl *CondVar = If->getConditionVariable();
+  return CondVar != nullptr ? findUsage(If->getElse(), CondVar->getID())
+                            : nullptr;
+}
+
+bool containsDeclInScope(const Stmt *Node) {
+  if (isa<DeclStmt>(Node)) {
+    return true;
+  }
+  if (const auto *Compound = dyn_cast<CompoundStmt>(Node)) {
+    return llvm::any_of(Compound->body(), [](const Stmt *SubNode) {
+      return isa<DeclStmt>(SubNode);
+    });
+  }
+  return false;
+}
+
+void removeElseAndBrackets(DiagnosticBuilder &Diag, ASTContext &Context,
+                           const Stmt *Else, SourceLocation ElseLoc) {
+  auto Remap = [&](SourceLocation Loc) {
+    return Context.getSourceManager().getExpansionLoc(Loc);
+  };
+  auto TokLen = [&](SourceLocation Loc) {
+    return Lexer::MeasureTokenLength(Loc, Context.getSourceManager(),
+                                     Context.getLangOpts());
+  };
+
+  if (const auto *CS = dyn_cast<CompoundStmt>(Else)) {
+    Diag << tooling::fixit::createRemoval(ElseLoc);
+    SourceLocation LBrace = CS->getLBracLoc();
+    SourceLocation RBrace = CS->getRBracLoc();
+    SourceLocation RangeStart =
+        Remap(LBrace).getLocWithOffset(TokLen(LBrace) + 1);
+    SourceLocation RangeEnd = Remap(RBrace).getLocWithOffset(-1);
+
+    llvm::StringRef Repl = Lexer::getSourceText(
+        CharSourceRange::getTokenRange(RangeStart, RangeEnd),
+        Context.getSourceManager(), Context.getLangOpts());
+    Diag << tooling::fixit::createReplacement(CS->getSourceRange(), Repl);
+  } else {
+    SourceLocation ElseExpandedLoc = Remap(ElseLoc);
+    SourceLocation EndLoc = Remap(Else->getEndLoc());
+
+    llvm::StringRef Repl = Lexer::getSourceText(
+        CharSourceRange::getTokenRange(
+            ElseExpandedLoc.getLocWithOffset(TokLen(ElseLoc) + 1), EndLoc),
+        Context.getSourceManager(), Context.getLangOpts());
+    Diag << tooling::fixit::createReplacement(
+        SourceRange(ElseExpandedLoc, EndLoc), Repl);
+  }
+}
+} // namespace
+
+ElseAfterReturnCheck::ElseAfterReturnCheck(StringRef Name,
+                                           ClangTidyContext *Context)
+    : ClangTidyCheck(Name, Context),
+      WarnOnUnfixable(Options.get(WarnOnUnfixableStr, true)) {}
+
+void ElseAfterReturnCheck::storeOptions(ClangTidyOptions::OptionMap &Opts) {
+  Options.store(Opts, WarnOnUnfixableStr, WarnOnUnfixable);
+}
+
 void ElseAfterReturnCheck::registerMatchers(MatchFinder *Finder) {
   const auto InterruptsControlFlow =
-      stmt(anyOf(returnStmt().bind("return"), continueStmt().bind("continue"),
-                 breakStmt().bind("break"),
-                 expr(ignoringImplicit(cxxThrowExpr().bind("throw")))));
+      stmt(anyOf(returnStmt().bind(ReturnStr), continueStmt().bind(ContinueStr),
+                 breakStmt().bind(BreakStr),
+                 expr(ignoringImplicit(cxxThrowExpr().bind(ThrowStr)))));
   Finder->addMatcher(
-      compoundStmt(forEach(
-          ifStmt(unless(isConstexpr()),
-                 // FIXME: Explore alternatives for the
-                 // `if (T x = ...) {... return; } else { <use x> }`
-                 // pattern:
-                 //   * warn, but don't fix;
-                 //   * fix by pulling out the variable declaration out of
-                 //     the condition.
-                 unless(hasConditionVariableStatement(anything())),
-                 hasThen(stmt(anyOf(InterruptsControlFlow,
-                                    compoundStmt(has(InterruptsControlFlow))))),
-                 hasElse(stmt().bind("else")))
-              .bind("if"))),
+      compoundStmt(
+          forEach(ifStmt(unless(isConstexpr()),
+                         hasThen(stmt(
+                             anyOf(InterruptsControlFlow,
+                                   compoundStmt(has(InterruptsControlFlow))))),
+                         hasElse(stmt().bind("else")))
+                      .bind("if")))
+          .bind("cs"),
       this);
 }
 
 void ElseAfterReturnCheck::check(const MatchFinder::MatchResult &Result) {
   const auto *If = Result.Nodes.getNodeAs<IfStmt>("if");
+  const auto *Else = Result.Nodes.getNodeAs<Stmt>("else");
+  const auto *OuterScope = Result.Nodes.getNodeAs<CompoundStmt>("cs");
+
+  bool IsLastInScope = OuterScope->body_back() == If;
   SourceLocation ElseLoc = If->getElseLoc();
-  std::string ControlFlowInterruptor;
-  for (const auto *BindingName : {"return", "continue", "break", "throw"})
-    if (Result.Nodes.getNodeAs<Stmt>(BindingName))
-      ControlFlowInterruptor = BindingName;
 
-  DiagnosticBuilder Diag = diag(ElseLoc, "do not use 'else' after '%0'")
-                           << ControlFlowInterruptor;
-  Diag << tooling::fixit::createRemoval(ElseLoc);
+  auto ControlFlowInterruptor = [&]() -> llvm::StringRef {
+    for (llvm::StringRef BindingName :
+         {ReturnStr, ContinueStr, BreakStr, ThrowStr})
+      if (Result.Nodes.getNodeAs<Stmt>(BindingName))
+        return BindingName;
+    return {};
+  }();
+
+  if (!IsLastInScope && containsDeclInScope(Else)) {
+    if (WarnOnUnfixable) {
+      // Warn, but don't attempt an autofix.
+      diag(ElseLoc, WarningMessage) << ControlFlowInterruptor;
+    }
+    return;
+  }
 
-  // FIXME: Removing the braces isn't always safe. Do a more careful analysis.
-  // FIXME: Change clang-format to correctly un-indent the code.
-  if (const auto *CS = Result.Nodes.getNodeAs<CompoundStmt>("else"))
-    Diag << tooling::fixit::createRemoval(CS->getLBracLoc())
-         << tooling::fixit::createRemoval(CS->getRBracLoc());
+  if (checkConditionVarUsageInElse(If) != nullptr) {
+    if (IsLastInScope) {
+      // If the if statement is the last statement its enclosing statements
+      // scope, we can pull the decl out of the if statement.
+      DiagnosticBuilder Diag =
+          diag(ElseLoc, WarningMessage, clang::DiagnosticIDs::Level::Remark)
+          << ControlFlowInterruptor;
+      if (checkInitDeclUsageInElse(If) != nullptr) {
+        Diag << tooling::fixit::createReplacement(
+                    SourceRange(If->getIfLoc()),
+                    (tooling::fixit::getText(*If->getInit(), *Result.Context) +
+                     llvm::StringRef("\n"))
+                        .str())
+             << tooling::fixit::createRemoval(If->getInit()->getSourceRange());
+      }
+      const DeclStmt *VDeclStmt = If->getConditionVariableDeclStmt();
+      const VarDecl *VDecl = If->getConditionVariable();
+      std::string Repl =
+          (tooling::fixit::getText(*VDeclStmt, *Result.Context) +
+           llvm::StringRef(";\n") +
+           tooling::fixit::getText(If->getIfLoc(), *Result.Context))
+              .str();
+      Diag << tooling::fixit::createReplacement(SourceRange(If->getIfLoc()),
+                                                Repl)
+           << tooling::fixit::createReplacement(VDeclStmt->getSourceRange(),
+                                                VDecl->getName());
+      removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
+    } else if (WarnOnUnfixable) {
+      // Warn, but don't attempt an autofix.
+      diag(ElseLoc, WarningMessage) << ControlFlowInterruptor;
+    }
+    return;
+  }
+
+  if (checkInitDeclUsageInElse(If) != nullptr) {
+    if (IsLastInScope) {
+      // If the if statement is the last statement its enclosing statements
+      // scope, we can pull the decl out of the if statement.
+      DiagnosticBuilder Diag = diag(ElseLoc, WarningMessage)
+                               << ControlFlowInterruptor;
+      Diag << tooling::fixit::createReplacement(
+                  SourceRange(If->getIfLoc()),
+                  (tooling::fixit::getText(*If->getInit(), *Result.Context) +
+                   "\n" +
+                   tooling::fixit::getText(If->getIfLoc(), *Result.Context))
+                      .str())
+           << tooling::fixit::createRemoval(If->getInit()->getSourceRange());
+      removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
+    } else if (WarnOnUnfixable) {
+      // Warn, but don't attempt an autofix.
+      diag(ElseLoc, WarningMessage) << ControlFlowInterruptor;
+    }
+    return;
+  }
+
+  DiagnosticBuilder Diag = diag(ElseLoc, WarningMessage)
+                           << ControlFlowInterruptor;
+  removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
 }
 
 } // namespace readability
diff --git a/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.h b/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.h
index 0aa47b5169e..990d4b4c3f6 100644
--- a/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.h
+++ b/clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.h
@@ -20,10 +20,14 @@ namespace readability {
 /// http://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return
 class ElseAfterReturnCheck : public ClangTidyCheck {
 public:
-  ElseAfterReturnCheck(StringRef Name, ClangTidyContext *Context)
-      : ClangTidyCheck(Name, Context) {}
+  ElseAfterReturnCheck(StringRef Name, ClangTidyContext *Context);
+
+  void storeOptions(ClangTidyOptions::OptionMap &Opts) override;
   void registerMatchers(ast_matchers::MatchFinder *Finder) override;
   void check(const ast_matchers::MatchFinder::MatchResult &Result) override;
+
+private:
+  const bool WarnOnUnfixable;
 };
 
 } // namespace readability