[clang-tools-extra] e3c6090 - [clangd] IncludeCleaner: Support macros

[Kirill Bobyrev via cfe-commits]

Author: Kirill Bobyrev
Date: 2021-10-27T10:30:04+02:00
New Revision: e3c6090e597685c23e9d2aac27d45c5a35b2f9bd

URL: https://github.com/llvm/llvm-project/commit/e3c6090e597685c23e9d2aac27d45c5a35b2f9bd
DIFF: https://github.com/llvm/llvm-project/commit/e3c6090e597685c23e9d2aac27d45c5a35b2f9bd.diff

LOG: [clangd] IncludeCleaner: Support macros

Collect the macro definition locations for all the macros used in the main
file.

Reviewed By: sammccall

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

Added: 
    

Modified: 
    clang-tools-extra/clangd/IncludeCleaner.cpp
    clang-tools-extra/clangd/IncludeCleaner.h
    clang-tools-extra/clangd/SourceCode.cpp
    clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/IncludeCleaner.cpp b/clang-tools-extra/clangd/IncludeCleaner.cpp
index af82c96e494a..95b4b2419cbe 100644
--- a/clang-tools-extra/clangd/IncludeCleaner.cpp
+++ b/clang-tools-extra/clangd/IncludeCleaner.cpp
@@ -11,8 +11,10 @@
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "support/Logger.h"
+#include "support/Trace.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 
@@ -168,13 +170,39 @@ clangd::Range getDiagnosticRange(llvm::StringRef Code, unsigned HashOffset) {
   return Result;
 }
 
+// Finds locations of macros referenced from within the main file. That includes
+// references that were not yet expanded, e.g `BAR` in `#define FOO BAR`.
+void findReferencedMacros(ParsedAST &AST, ReferencedLocations &Result) {
+  trace::Span Tracer("IncludeCleaner::findReferencedMacros");
+  auto &SM = AST.getSourceManager();
+  auto &PP = AST.getPreprocessor();
+  // FIXME(kirillbobyrev): The macros from the main file are collected in
+  // ParsedAST's MainFileMacros. However, we can't use it here because it
+  // doesn't handle macro references that were not expanded, e.g. in macro
+  // definitions or preprocessor-disabled sections.
+  //
+  // Extending MainFileMacros to collect missing references and switching to
+  // this mechanism (as opposed to iterating through all tokens) will improve
+  // the performance of findReferencedMacros and also improve other features
+  // relying on MainFileMacros.
+  for (const syntax::Token &Tok :
+       AST.getTokens().spelledTokens(SM.getMainFileID())) {
+    auto Macro = locateMacroAt(Tok, PP);
+    if (!Macro)
+      continue;
+    auto Loc = Macro->Info->getDefinitionLoc();
+    if (Loc.isValid())
+      Result.insert(Loc);
+  }
+}
+
 } // namespace
 
 ReferencedLocations findReferencedLocations(ParsedAST &AST) {
   ReferencedLocations Result;
   ReferencedLocationCrawler Crawler(Result);
   Crawler.TraverseAST(AST.getASTContext());
-  // FIXME(kirillbobyrev): Handle macros.
+  findReferencedMacros(AST, Result);
   return Result;
 }
 

diff  --git a/clang-tools-extra/clangd/IncludeCleaner.h b/clang-tools-extra/clangd/IncludeCleaner.h
index 19d471e2e761..c7dd7efd4ed6 100644
--- a/clang-tools-extra/clangd/IncludeCleaner.h
+++ b/clang-tools-extra/clangd/IncludeCleaner.h
@@ -33,11 +33,12 @@ namespace clangd {
 using ReferencedLocations = llvm::DenseSet<SourceLocation>;
 /// Finds locations of all symbols used in the main file.
 ///
-/// Uses RecursiveASTVisitor to go through main file AST and computes all the
-/// locations used symbols are coming from. Returned locations may be macro
-/// expansions, and are not resolved to their spelling/expansion location. These
-/// locations are later used to determine which headers should be marked as
-/// "used" and "directly used".
+/// - RecursiveASTVisitor finds references to symbols and records their
+///   associated locations. These may be macro expansions, and are not resolved
+///   to their spelling or expansion location. These locations are later used to
+///   determine which headers should be marked as "used" and "directly used".
+/// - We also examine all identifier tokens in the file in case they reference
+///   macros.
 ///
 /// We use this to compute unused headers, so we:
 ///

diff  --git a/clang-tools-extra/clangd/SourceCode.cpp b/clang-tools-extra/clangd/SourceCode.cpp
index 014bdfc63281..16072f3803c5 100644
--- a/clang-tools-extra/clangd/SourceCode.cpp
+++ b/clang-tools-extra/clangd/SourceCode.cpp
@@ -975,6 +975,8 @@ llvm::Optional<SpelledWord> SpelledWord::touching(SourceLocation SpelledLoc,
 
 llvm::Optional<DefinedMacro> locateMacroAt(const syntax::Token &SpelledTok,
                                            Preprocessor &PP) {
+  if (SpelledTok.kind() != tok::identifier)
+    return None;
   SourceLocation Loc = SpelledTok.location();
   assert(Loc.isFileID());
   const auto &SM = PP.getSourceManager();

diff  --git a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
index 978d3794fa0e..38cd8da23066 100644
--- a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
@@ -116,6 +116,35 @@ TEST(IncludeCleaner, ReferencedLocations) {
           "struct ^X { enum ^Language { ^CXX = 42, Python = 9000}; };",
           "int Lang = X::CXX;",
       },
+      // Macros
+      {
+          "#define ^CONSTANT 42",
+          "int Foo = CONSTANT;",
+      },
+      {
+          "#define ^FOO x",
+          "#define BAR FOO",
+      },
+      {
+          "#define INNER 42\n"
+          "#define ^OUTER INNER",
+          "int answer = OUTER;",
+      },
+      {
+          "#define ^ANSWER 42\n"
+          "#define ^SQUARE(X) X * X",
+          "int sq = SQUARE(ANSWER);",
+      },
+      {
+          "#define ^FOO\n"
+          "#define ^BAR",
+          "#if 0\n"
+          "#if FOO\n"
+          "BAR\n"
+          "#endif\n"
+          "#endif",
+      },
+      // Misc
       {
           "enum class ^Color : int;",
           "enum class Color : int {};",