[clang-tidy] ElseAfterReturn check wont suggest fixes if preprocessor branches are involved

Consider this code:
```
if (Cond) {
#ifdef X_SUPPORTED
  X();
#else
  return;
#endif
} else {
  Y();
}
Z();```

In this example, if `X_SUPPORTED` is not defined, currently we'll get a warning from the else-after-return check. However If we apply that fix, and then the code is recompiled with `X_SUPPORTED` defined, we have inadvertently changed the behaviour of the if statement due to the else being removed. Code flow when `Cond` is `true` will be:
```
X();
Y();
Z();```
 where as before the fix it was:
 ```
 X();
 Z();```

 This patch adds checks that guard against `#endif` directives appearing between the control flow interrupter and the else and not applying the fix if they are detected.

Reviewed By: aaron.ballman

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

Author: njames93

clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.cpp

@@ -10,6 +10,7 @@
 #include "clang/AST/ASTContext.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/Lex/Lexer.h"
+#include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/FixIt.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -19,10 +20,30 @@ namespace clang {
 namespace tidy {
 namespace readability {
 
-static const char ReturnStr[] = "return";
-static const char ContinueStr[] = "continue";
-static const char BreakStr[] = "break";
-static const char ThrowStr[] = "throw";
+namespace {
+
+class PPConditionalCollector : public PPCallbacks {
+public:
+  PPConditionalCollector(
+      ElseAfterReturnCheck::ConditionalBranchMap &Collections,
+      const SourceManager &SM)
+      : Collections(Collections), SM(SM) {}
+  void Endif(SourceLocation Loc, SourceLocation IfLoc) override {
+    if (!SM.isWrittenInSameFile(Loc, IfLoc))
+      return;
+    SmallVectorImpl<SourceRange> &Collection = Collections[SM.getFileID(Loc)];
+    assert(Collection.empty() || Collection.back().getEnd() < Loc);
+    Collection.emplace_back(IfLoc, Loc);
+  }
+
+private:
+  ElseAfterReturnCheck::ConditionalBranchMap &Collections;
+  const SourceManager &SM;
+};
+
+} // namespace
+
+static const char InterruptingStr[] = "interrupting";
 static const char WarningMessage[] = "do not use 'else' after '%0'";
 static const char WarnOnUnfixableStr[] = "WarnOnUnfixable";
 static const char WarnOnConditionVariablesStr[] = "WarnOnConditionVariables";
@@ -140,11 +161,18 @@ void ElseAfterReturnCheck::storeOptions(ClangTidyOptions::OptionMap &Opts) {
   Options.store(Opts, WarnOnConditionVariablesStr, WarnOnConditionVariables);
 }
 
+void ElseAfterReturnCheck::registerPPCallbacks(const SourceManager &SM,
+                                               Preprocessor *PP,
+                                               Preprocessor *ModuleExpanderPP) {
+  PP->addPPCallbacks(
+      std::make_unique<PPConditionalCollector>(this->PPConditionals, SM));
+}
+
 void ElseAfterReturnCheck::registerMatchers(MatchFinder *Finder) {
-  const auto InterruptsControlFlow =
-      stmt(anyOf(returnStmt().bind(ReturnStr), continueStmt().bind(ContinueStr),
-                 breakStmt().bind(BreakStr),
-                 expr(ignoringImplicit(cxxThrowExpr().bind(ThrowStr)))));
+  const auto InterruptsControlFlow = stmt(anyOf(
+      returnStmt().bind(InterruptingStr), continueStmt().bind(InterruptingStr),
+      breakStmt().bind(InterruptingStr),
+      expr(ignoringImplicit(cxxThrowExpr().bind(InterruptingStr)))));
   Finder->addMatcher(
       compoundStmt(
           forEach(ifStmt(unless(isConstexpr()),
@@ -157,21 +185,72 @@ void ElseAfterReturnCheck::registerMatchers(MatchFinder *Finder) {
       this);
 }
 
+static bool hasPreprocessorBranchEndBetweenLocations(
+    const ElseAfterReturnCheck::ConditionalBranchMap &ConditionalBranchMap,
+    const SourceManager &SM, SourceLocation StartLoc, SourceLocation EndLoc) {
+
+  SourceLocation ExpandedStartLoc = SM.getExpansionLoc(StartLoc);
+  SourceLocation ExpandedEndLoc = SM.getExpansionLoc(EndLoc);
+  if (!SM.isWrittenInSameFile(ExpandedStartLoc, ExpandedEndLoc))
+    return false;
+
+  // StartLoc and EndLoc expand to the same macro.
+  if (ExpandedStartLoc == ExpandedEndLoc)
+    return false;
+
+  assert(ExpandedStartLoc < ExpandedEndLoc);
+
+  auto Iter = ConditionalBranchMap.find(SM.getFileID(ExpandedEndLoc));
+
+  if (Iter == ConditionalBranchMap.end() || Iter->getSecond().empty())
+    return false;
+
+  const SmallVectorImpl<SourceRange> &ConditionalBranches = Iter->getSecond();
+
+  assert(llvm::is_sorted(ConditionalBranches,
+                         [](const SourceRange &LHS, const SourceRange &RHS) {
+                           return LHS.getEnd() < RHS.getEnd();
+                         }));
+
+  // First conditional block that ends after ExpandedStartLoc.
+  const auto *Begin =
+      llvm::lower_bound(ConditionalBranches, ExpandedStartLoc,
+                        [](const SourceRange &LHS, const SourceLocation &RHS) {
+                          return LHS.getEnd() < RHS;
+                        });
+  const auto *End = ConditionalBranches.end();
+  for (; Begin != End && Begin->getEnd() < ExpandedEndLoc; ++Begin)
+    if (Begin->getBegin() < ExpandedStartLoc)
+      return true;
+  return false;
+}
+
+static StringRef getControlFlowString(const Stmt &Stmt) {
+  if (isa<ReturnStmt>(Stmt))
+    return "return";
+  if (isa<ContinueStmt>(Stmt))
+    return "continue";
+  if (isa<BreakStmt>(Stmt))
+    return "break";
+  if (isa<CXXThrowExpr>(Stmt))
+    return "throw";
+  llvm_unreachable("Unknown control flow interruptor");
+}
+
 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;
+  const auto *Interrupt = Result.Nodes.getNodeAs<Stmt>(InterruptingStr);
   SourceLocation ElseLoc = If->getElseLoc();
 
-  auto ControlFlowInterruptor = [&]() -> llvm::StringRef {
-    for (llvm::StringRef BindingName :
-         {ReturnStr, ContinueStr, BreakStr, ThrowStr})
-      if (Result.Nodes.getNodeAs<Stmt>(BindingName))
-        return BindingName;
-    return {};
-  }();
+  if (hasPreprocessorBranchEndBetweenLocations(
+          PPConditionals, *Result.SourceManager, Interrupt->getBeginLoc(),
+          ElseLoc))
+    return;
+
+  bool IsLastInScope = OuterScope->body_back() == If;
+  StringRef ControlFlowInterruptor = getControlFlowString(*Interrupt);
 
   if (!IsLastInScope && containsDeclInScope(Else)) {
     if (WarnOnUnfixable) {

clang-tools-extra/clang-tidy/readability/ElseAfterReturnCheck.h

@@ -10,6 +10,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANG_TIDY_READABILITY_ELSEAFTERRETURNCHECK_H
 
 #include "../ClangTidyCheck.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace clang {
 namespace tidy {
@@ -23,12 +24,18 @@ class ElseAfterReturnCheck : public ClangTidyCheck {
   ElseAfterReturnCheck(StringRef Name, ClangTidyContext *Context);
 
   void storeOptions(ClangTidyOptions::OptionMap &Opts) override;
+  void registerPPCallbacks(const SourceManager &SM, Preprocessor *PP,
+                           Preprocessor *ModuleExpanderPP) override;
   void registerMatchers(ast_matchers::MatchFinder *Finder) override;
   void check(const ast_matchers::MatchFinder::MatchResult &Result) override;
 
+  using ConditionalBranchMap =
+      llvm::DenseMap<FileID, SmallVector<SourceRange, 1>>;
+
 private:
   const bool WarnOnUnfixable;
   const bool WarnOnConditionVariables;
+  ConditionalBranchMap PPConditionals;
 };
 
 } // namespace readability

clang-tools-extra/test/clang-tidy/checkers/readability-else-after-return-pp-no-crash.cpp

@@ -0,0 +1,22 @@
+// RUN: clang-tidy %s -checks=-*,readability-else-after-return
+
+// We aren't concerned about the output here, just want to ensure clang-tidy doesn't crash.
+void foo() {
+#if 1
+  if (true) {
+    return;
+#else
+  {
+#endif
+  } else {
+    return;
+  }
+
+  if (true) {
+#if 1
+    return;
+  } else {
+#endif
+    return;
+  }
+}

clang-tools-extra/test/clang-tidy/checkers/readability-else-after-return.cpp

@@ -226,3 +226,89 @@ void test_B44745() {
   }
   return;
 }
+
+void testPPConditionals() {
+
+  // These cases the return isn't inside the conditional so diagnose as normal.
+  if (true) {
+    return;
+#if 1
+#endif
+  } else {
+    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: do not use 'else' after 'return'
+    return;
+  }
+  if (true) {
+#if 1
+#endif
+    return;
+  } else {
+    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: do not use 'else' after 'return'
+    return;
+  }
+
+  // No return here in the AST, no special handling needed.
+  if (true) {
+#if 0
+    return;
+#endif
+  } else {
+    return;
+  }
+
+  // Return here is inside a preprocessor conditional block, ignore this case.
+  if (true) {
+#if 1
+    return;
+#endif
+  } else {
+    return;
+  }
+
+  // These cases, same as above but with an #else block.
+  if (true) {
+#if 1
+    return;
+#else
+#endif
+  } else {
+    return;
+  }
+  if (true) {
+#if 0
+#else
+    return;
+#endif
+  } else {
+    return;
+  }
+
+// Ensure it can handle macros.
+#define RETURN return
+  if (true) {
+#if 1
+    RETURN;
+#endif
+  } else {
+    return;
+  }
+#define ELSE else
+  if (true) {
+#if 1
+    return;
+#endif
+  }
+  ELSE {
+    return;
+  }
+
+  // Whole statement is in a conditional block so diagnose as normal.
+#if 1
+  if (true) {
+    return;
+  } else {
+    // CHECK-MESSAGES: :[[@LINE-1]]:5: warning: do not use 'else' after 'return'
+    return;
+  }
+#endif
+}