[Clang][ASTMatchers] Make hasConditionVariableStatement support for loop, while loop and switch statements (#154298)

Co-authored-by: Aaron Ballman <aaron@aaronballman.com>

clang/docs/ReleaseNotes.rst

@@ -329,6 +329,9 @@ AST Matchers
 - Add a boolean member ``IgnoreSystemHeaders`` to ``MatchFinderOptions``. This
   allows it to ignore nodes in system headers when traversing the AST.
 
+- ``hasConditionVariableStatement`` now supports ``for`` loop, ``while`` loop
+  and ``switch`` statements.
+
 clang-format
 ------------
 - Add ``SpaceInEmptyBraces`` option and set it to ``Always`` for WebKit style.

clang/include/clang/ASTMatchers/ASTMatchers.h

@@ -5661,8 +5661,8 @@ AST_POLYMORPHIC_MATCHER_P(hasInitStatement,
   return Init != nullptr && InnerMatcher.matches(*Init, Finder, Builder);
 }
 
-/// Matches the condition expression of an if statement, for loop,
-/// switch statement or conditional operator.
+/// Matches the condition expression of an if statement, for loop, while loop,
+/// do-while loop, switch statement or conditional operator.
 ///
 /// Example matches true (matcher = hasCondition(cxxBoolLiteral(equals(true))))
 /// \code
@@ -5747,16 +5747,21 @@ AST_MATCHER_P(Decl, declaresSameEntityAsBoundNode, std::string, ID) {
   });
 }
 
-/// Matches the condition variable statement in an if statement.
+/// Matches the condition variable statement in an if statement, for loop,
+/// while loop or switch statement.
 ///
 /// Given
 /// \code
 ///   if (A* a = GetAPointer()) {}
+///   for (; A* a = GetAPointer(); ) {}
 /// \endcode
 /// hasConditionVariableStatement(...)
-///   matches 'A* a = GetAPointer()'.
-AST_MATCHER_P(IfStmt, hasConditionVariableStatement,
-              internal::Matcher<DeclStmt>, InnerMatcher) {
+///   matches both 'A* a = GetAPointer()'.
+AST_POLYMORPHIC_MATCHER_P(hasConditionVariableStatement,
+                          AST_POLYMORPHIC_SUPPORTED_TYPES(IfStmt, ForStmt,
+                                                          WhileStmt,
+                                                          SwitchStmt),
+                          internal::Matcher<DeclStmt>, InnerMatcher) {
   const DeclStmt* const DeclarationStatement =
     Node.getConditionVariableDeclStmt();
   return DeclarationStatement != nullptr &&

clang/unittests/ASTMatchers/ASTMatchersNodeTest.cpp

@@ -1183,6 +1183,39 @@ TEST_P(ASTMatchersTest, AsmStatement) {
   EXPECT_TRUE(matches("void foo() { __asm(\"mov al, 2\"); }", asmStmt()));
 }
 
+TEST_P(ASTMatchersTest, HasConditionVariableStatement) {
+  if (!GetParam().isCXX()) {
+    // FIXME: Add a test for `hasConditionVariableStatement()` that does not
+    // depend on C++.
+    return;
+  }
+
+  StatementMatcher IfCondition =
+      ifStmt(hasConditionVariableStatement(declStmt()));
+
+  EXPECT_TRUE(matches("void x() { if (int* a = 0) {} }", IfCondition));
+  EXPECT_TRUE(notMatches("void x() { if (true) {} }", IfCondition));
+  EXPECT_TRUE(notMatches("void x() { int x; if ((x = 42)) {} }", IfCondition));
+
+  StatementMatcher SwitchCondition =
+      switchStmt(hasConditionVariableStatement(declStmt()));
+
+  EXPECT_TRUE(matches("void x() { switch (int a = 0) {} }", SwitchCondition));
+  if (GetParam().isCXX17OrLater()) {
+    EXPECT_TRUE(
+        notMatches("void x() { switch (int a = 0; a) {} }", SwitchCondition));
+  }
+
+  StatementMatcher ForCondition =
+      forStmt(hasConditionVariableStatement(declStmt()));
+
+  EXPECT_TRUE(matches("void x() { for (; int a = 0; ) {} }", ForCondition));
+  EXPECT_TRUE(notMatches("void x() { for (int a = 0; ; ) {} }", ForCondition));
+
+  EXPECT_TRUE(matches("void x() { while (int a = 0) {} }",
+                      whileStmt(hasConditionVariableStatement(declStmt()))));
+}
+
 TEST_P(ASTMatchersTest, HasCondition) {
   if (!GetParam().isCXX()) {
     // FIXME: Add a test for `hasCondition()` that does not depend on C++.

clang/unittests/ASTMatchers/ASTMatchersTraversalTest.cpp

@@ -5142,21 +5142,6 @@ TEST(ForEachLambdaCapture, MatchExplicitCapturesOnly) {
       matcher, std::make_unique<VerifyIdIsBoundTo<LambdaCapture>>("LC", 1)));
 }
 
-TEST(HasConditionVariableStatement, DoesNotMatchCondition) {
-  EXPECT_TRUE(notMatches(
-    "void x() { if(true) {} }",
-    ifStmt(hasConditionVariableStatement(declStmt()))));
-  EXPECT_TRUE(notMatches(
-    "void x() { int x; if((x = 42)) {} }",
-    ifStmt(hasConditionVariableStatement(declStmt()))));
-}
-
-TEST(HasConditionVariableStatement, MatchesConditionVariables) {
-  EXPECT_TRUE(matches(
-    "void x() { if(int* a = 0) {} }",
-    ifStmt(hasConditionVariableStatement(declStmt()))));
-}
-
 TEST(ForEach, BindsOneNode) {
   EXPECT_TRUE(matchAndVerifyResultTrue("class C { int x; };",
                                        recordDecl(hasName("C"), forEach(fieldDecl(hasName("x")).bind("x"))),