From 38cccb90660347939dd7bfdd57ba66760cba39c2 Mon Sep 17 00:00:00 2001
From: Sam McCall <sam.mccall@gmail.com>
Date: Tue, 8 Nov 2022 16:33:08 +0100
Subject: [PATCH] [include-cleaner] pass through recorded macro refs in
 walkUsed

Differential Revision: https://reviews.llvm.org/D137644
---
 .../include/clang-include-cleaner/Analysis.h       |  5 +--
 clang-tools-extra/include-cleaner/lib/Analysis.cpp | 14 ++++++--
 .../include-cleaner/unittests/AnalysisTest.cpp     | 37 +++++++++++++++++++++-
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
index a75fd9a..4c5c90e 100644
--- a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
+++ b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h
@@ -35,7 +35,7 @@ using UsedSymbolCB = llvm::function_ref<void(SymbolReference SymRef,
 ///
 /// The AST traversal is rooted at ASTRoots - typically top-level declarations
 /// of a single source file.
-/// FIXME: Handle macro uses.
+/// The references to macros must be recorded separately and provided.
 ///
 /// This is the main entrypoint of the include-cleaner library, and can be used:
 ///  - to diagnose missing includes: a referenced symbol is provided by
@@ -44,7 +44,8 @@ using UsedSymbolCB = llvm::function_ref<void(SymbolReference SymRef,
 ///    the headers for any referenced symbol
 /// FIXME: Take in an include structure to improve location to header mappings
 /// (e.g. IWYU pragmas).
-void walkUsed(llvm::ArrayRef<Decl *> ASTRoots, UsedSymbolCB CB);
+void walkUsed(const SourceManager &, llvm::ArrayRef<Decl *> ASTRoots,
+              llvm::ArrayRef<SymbolReference> MacroRefs, UsedSymbolCB CB);
 
 } // namespace include_cleaner
 } // namespace clang
diff --git a/clang-tools-extra/include-cleaner/lib/Analysis.cpp b/clang-tools-extra/include-cleaner/lib/Analysis.cpp
index 95c6c6a..b4dfd12 100644
--- a/clang-tools-extra/include-cleaner/lib/Analysis.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Analysis.cpp
@@ -26,9 +26,10 @@ toHeader(llvm::ArrayRef<tooling::stdlib::Header> Headers) {
   });
   return Result;
 }
-
 } // namespace
-void walkUsed(llvm::ArrayRef<Decl *> ASTRoots, UsedSymbolCB CB) {
+
+void walkUsed(const SourceManager &SM, llvm::ArrayRef<Decl *> ASTRoots,
+              llvm::ArrayRef<SymbolReference> MacroRefs, UsedSymbolCB CB) {
   tooling::stdlib::Recognizer Recognizer;
   for (auto *Root : ASTRoots) {
     auto &SM = Root->getASTContext().getSourceManager();
@@ -45,7 +46,14 @@ void walkUsed(llvm::ArrayRef<Decl *> ASTRoots, UsedSymbolCB CB) {
         return CB({Loc, Symbol(ND), RT}, {Header(FE)});
     });
   }
-  // FIXME: Handle references of macros.
+  for (const SymbolReference &MacroRef : MacroRefs) {
+    assert(MacroRef.Target.kind() == Symbol::Macro);
+    // FIXME: Handle IWYU pragmas, non self-contained files.
+    // FIXME: Handle macro locations.
+    if (auto *FE = SM.getFileEntryForID(
+            SM.getFileID(MacroRef.Target.macro().Definition)))
+      CB(MacroRef, {Header(FE)});
+  }
 }
 
 } // namespace clang::include_cleaner
diff --git a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
index cb7bc1a..bd79403 100644
--- a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
+++ b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
@@ -50,7 +50,7 @@ TEST(WalkUsed, Basic) {
 
   auto &SM = AST.sourceManager();
   llvm::DenseMap<size_t, std::vector<Header>> OffsetToProviders;
-  walkUsed(TopLevelDecls,
+  walkUsed(SM, TopLevelDecls, /*MacroRefs=*/{},
            [&](SymbolReference SymRef, llvm::ArrayRef<Header> Providers) {
              auto [FID, Offset] = SM.getDecomposedLoc(SymRef.RefLocation);
              EXPECT_EQ(FID, SM.getMainFileID());
@@ -68,5 +68,40 @@ TEST(WalkUsed, Basic) {
           Pair(Code.point("vconstructor"), UnorderedElementsAre(VectorSTL))));
 }
 
+TEST(WalkUsed, MacroRefs) {
+  llvm::Annotations Hdr(R"cpp(
+    #define ^ANSWER 42
+  )cpp");
+  llvm::Annotations Main(R"cpp(
+    #include "hdr.h"
+    int x = ^ANSWER;
+  )cpp");
+
+  SourceManagerForFile SMF("main.cpp", Main.code());
+  auto &SM = SMF.get();
+  const FileEntry *HdrFile =
+      SM.getFileManager().getVirtualFile("hdr.h", Hdr.code().size(), 0);
+  SM.overrideFileContents(HdrFile,
+                          llvm::MemoryBuffer::getMemBuffer(Hdr.code().str()));
+  FileID HdrID = SM.createFileID(HdrFile, SourceLocation(), SrcMgr::C_User);
+
+  IdentifierTable Idents;
+  Symbol Answer =
+      Macro{&Idents.get("ANSWER"), SM.getComposedLoc(HdrID, Hdr.point())};
+  llvm::DenseMap<size_t, std::vector<Header>> OffsetToProviders;
+  walkUsed(SM, /*ASTRoots=*/{}, /*MacroRefs=*/
+           {SymbolReference{SM.getComposedLoc(SM.getMainFileID(), Main.point()),
+                            Answer, RefType::Explicit}},
+           [&](SymbolReference SymRef, llvm::ArrayRef<Header> Providers) {
+             auto [FID, Offset] = SM.getDecomposedLoc(SymRef.RefLocation);
+             EXPECT_EQ(FID, SM.getMainFileID());
+             OffsetToProviders.try_emplace(Offset, Providers.vec());
+           });
+
+  EXPECT_THAT(
+      OffsetToProviders,
+      UnorderedElementsAre(Pair(Main.point(), UnorderedElementsAre(HdrFile))));
+}
+
 } // namespace
 } // namespace clang::include_cleaner
-- 
2.7.4