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

[Aaron Ballman via cfe-commits]

https://github.com/AaronBallman 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 1/2] [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) {

>From f3f08877f08c77f50f622fb166029ad1f3b1cb7c Mon Sep 17 00:00:00 2001
From: wangmi <wangmi at feysh.com>
Date: Mon, 2 Dec 2024 14:33:29 +0000
Subject: [PATCH 2/2] use DenseSet instead of std::unordered_set and update
 ReleaseNotes

---
 clang/docs/ReleaseNotes.rst                   |  2 ++
 clang/lib/ASTMatchers/ASTMatchersInternal.cpp | 29 ++++---------------
 2 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst
index 999c88455b64a5..94605d6a0a7013 100644
--- a/clang/docs/ReleaseNotes.rst
+++ b/clang/docs/ReleaseNotes.rst
@@ -890,6 +890,8 @@ AST Matchers
 - Ensure ``hasName`` matches template specializations across inline namespaces,
   making `matchesNodeFullSlow` and `matchesNodeFullFast` consistent.
 
+- Improved the performance of the ``getExpansionLocOfMacro`` by tracking already processed macros during recursion.
+
 clang-format
 ------------
 
diff --git a/clang/lib/ASTMatchers/ASTMatchersInternal.cpp b/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
index fe57a2064b0f96..8d4147657e7f42 100644
--- a/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
+++ b/clang/lib/ASTMatchers/ASTMatchersInternal.cpp
@@ -21,6 +21,7 @@
 #include "clang/Basic/LLVM.h"
 #include "clang/Lex/Lexer.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
@@ -36,7 +37,6 @@
 #include <cstddef>
 #include <optional>
 #include <string>
-#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -698,31 +698,14 @@ static bool isTokenAtLoc(const SourceManager &SM, const LangOptions &LangOpts,
   return !Invalid && Text == TokenText;
 }
 
-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) {
+    llvm::DenseSet<SourceLocation> &CheckedLocations) {
   auto &SM = Context.getSourceManager();
   const LangOptions &LangOpts = Context.getLangOpts();
   while (Loc.isMacroID()) {
-    if (CheckedLocations.count(Loc)) {
+    if (CheckedLocations.count(Loc))
       return std::nullopt;
-    }
     CheckedLocations.insert(Loc);
     SrcMgr::ExpansionInfo Expansion =
         SM.getSLocEntry(SM.getFileID(Loc)).getExpansion();
@@ -737,9 +720,8 @@ static std::optional<SourceLocation> getExpansionLocOfMacroRecursive(
       }
     }
     Loc = Expansion.getExpansionLocStart();
-    if (isTokenAtLoc(SM, LangOpts, MacroName, Loc)) {
+    if (isTokenAtLoc(SM, LangOpts, MacroName, Loc))
       return Loc;
-    }
   }
   return std::nullopt;
 }
@@ -747,8 +729,7 @@ static std::optional<SourceLocation> getExpansionLocOfMacroRecursive(
 std::optional<SourceLocation>
 getExpansionLocOfMacro(StringRef MacroName, SourceLocation Loc,
                        const ASTContext &Context) {
-  std::unordered_set<SourceLocation, SourceLocationHash, SourceLocationEqual>
-      CheckedLocations;
+  llvm::DenseSet<SourceLocation> CheckedLocations;
   return getExpansionLocOfMacroRecursive(MacroName, Loc, Context,
                                          CheckedLocations);
 }