diff --git a/clang-tools-extra/clangd/IncludeCleaner.h b/clang-tools-extra/clangd/IncludeCleaner.h --- a/clang-tools-extra/clangd/IncludeCleaner.h +++ b/clang-tools-extra/clangd/IncludeCleaner.h @@ -33,11 +33,12 @@ using ReferencedLocations = llvm::DenseSet; /// 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". +/// - A 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/IncludeCleaner.cpp b/clang-tools-extra/clangd/IncludeCleaner.cpp --- a/clang-tools-extra/clangd/IncludeCleaner.cpp +++ b/clang-tools-extra/clangd/IncludeCleaner.cpp @@ -11,8 +11,11 @@ #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/Basic/TokenKinds.h" +#include "clang/Tooling/Syntax/Tokens.h" #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/Path.h" @@ -168,13 +171,40 @@ return Result; } +// Finds locations of all macros referenced from within the main file. This +// includes references that were not yet expanded, like `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/SourceCode.cpp b/clang-tools-extra/clangd/SourceCode.cpp --- a/clang-tools-extra/clangd/SourceCode.cpp +++ b/clang-tools-extra/clangd/SourceCode.cpp @@ -975,6 +975,8 @@ llvm::Optional 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 --- a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp +++ b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp @@ -116,6 +116,35 @@ "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 {};",