[clang] Fix redundant macro expansion checks in ASTMatchers (PR #117143)

[via cfe-commits]

https://github.com/mandymimi updated https://github.com/llvm/llvm-project/pull/117143

>From b411d9185a6f761901bcc7c8ed147a8efd3b5301 Mon Sep 17 00:00:00 2001
From: wangmi <wangmi at feysh.com>
Date: Tue, 26 Nov 2024 07:11:58 +0000
Subject: [PATCH] [ASTMatcher] Fix redundant macro expansion checks in
 getExpansionLocOfMacro When dealing with macro parameters,
 getExpansionLocOfMacro ecursively calls itself to check the expansion of
 macro arguments. This recursive logic redundantly checks previous macro
 expansions, leading to significant performance degradation when macros are
 heavily nested.

This fix tracks already processed macros during recursion.
---
 clang/lib/ASTMatchers/ASTMatchersInternal.cpp | 49 ++++++++++++++++---
 1 file changed, 42 insertions(+), 7 deletions(-)

diff --git a/clang/lib/ASTMatchers/ASTMatchersInternal.cpp b/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
index 46dd44e6f2b24f..fe57a2064b0f96 100644
--- a/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
+++ b/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
@@ -36,6 +36,7 @@
 #include <cstddef>
 #include <optional>
 #include <string>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -697,27 +698,61 @@ static bool isTokenAtLoc(const SourceManager &SM, const LangOptions &LangOpts,
   return !Invalid && Text == TokenText;
 }
 
-std::optional<SourceLocation>
-getExpansionLocOfMacro(StringRef MacroName, SourceLocation Loc,
-                       const ASTContext &Context) {
+namespace {
+struct SourceLocationHash {
+  std::size_t operator()(const SourceLocation &Loc) const {
+    return Loc.getHashValue();
+  }
+};
+
+struct SourceLocationEqual {
+  bool operator()(const SourceLocation &LHS, const SourceLocation &RHS) const {
+    return LHS == RHS;
+  }
+};
+
+} // namespace
+
+static std::optional<SourceLocation> getExpansionLocOfMacroRecursive(
+    StringRef MacroName, SourceLocation Loc, const ASTContext &Context,
+    std::unordered_set<SourceLocation, SourceLocationHash, SourceLocationEqual>
+        &CheckedLocations) {
   auto &SM = Context.getSourceManager();
   const LangOptions &LangOpts = Context.getLangOpts();
   while (Loc.isMacroID()) {
+    if (CheckedLocations.count(Loc)) {
+      return std::nullopt;
+    }
+    CheckedLocations.insert(Loc);
     SrcMgr::ExpansionInfo Expansion =
         SM.getSLocEntry(SM.getFileID(Loc)).getExpansion();
-    if (Expansion.isMacroArgExpansion())
+    if (Expansion.isMacroArgExpansion()) {
       // Check macro argument for an expansion of the given macro. For example,
       // `F(G(3))`, where `MacroName` is `G`.
-      if (std::optional<SourceLocation> ArgLoc = getExpansionLocOfMacro(
-              MacroName, Expansion.getSpellingLoc(), Context))
+      if (std::optional<SourceLocation> ArgLoc =
+              getExpansionLocOfMacroRecursive(MacroName,
+                                              Expansion.getSpellingLoc(),
+                                              Context, CheckedLocations)) {
         return ArgLoc;
+      }
+    }
     Loc = Expansion.getExpansionLocStart();
-    if (isTokenAtLoc(SM, LangOpts, MacroName, Loc))
+    if (isTokenAtLoc(SM, LangOpts, MacroName, Loc)) {
       return Loc;
+    }
   }
   return std::nullopt;
 }
 
+std::optional<SourceLocation>
+getExpansionLocOfMacro(StringRef MacroName, SourceLocation Loc,
+                       const ASTContext &Context) {
+  std::unordered_set<SourceLocation, SourceLocationHash, SourceLocationEqual>
+      CheckedLocations;
+  return getExpansionLocOfMacroRecursive(MacroName, Loc, Context,
+                                         CheckedLocations);
+}
+
 std::shared_ptr<llvm::Regex> createAndVerifyRegex(StringRef Regex,
                                                   llvm::Regex::RegexFlags Flags,
                                                   StringRef MatcherID) {