[clangd] IncludeCleaner: Support macros

[kbobyrev]

Summary

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

[sammccall]

clang-tools-extra/clangd/IncludeCleaner.cpp

161

can you add a trace for this and verify that it's not a signficant amount of time on a macro-heavy file (e.g. unittests?)

locateMacroAt wasn't designed to be called in a loop like this, but it seems plausible it'll be fine

161

iterating over all the tokens and examining them individually seems like we're doing too much work.

ParsedAST.MainFileMacros.Names already knows the names of all macros referenced from the main file.
PP.getLocalMacroDirectiveHistory() will give you the last seen definition for a given macro name. If we want, we can also get the previous definitions.

This seems like it should be close enough, and wade through many fewer tokens?

165

locateMacroAt doesn't check that the token is an identifier before trying to fetch its text, maybe do that? either here or there...

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

128

#define INNER 42
#define OUTER INNER

int answer = OUTER;

(whatever the behavior is, let's add a test for it)

128

#define ANSWER 42
#define SQUARE(X) X * X

int sq = SQUARE(ANSWER);

[kbobyrev]

Try another strategy to reduce the number of processed tokens (doesn't work in
all cases).

clang-tools-extra/clangd/IncludeCleaner.cpp

161

Hmm, this doesn't really work for "unmaterialized" macros, i.e. ParsedAST::getMacros().Names doesn't have FOO in this case:

// foo.h
#define FOO

// foo.cpp
#define BAR FOO

Is that a bug in CollectMainFileMacros we want to fix or is this intended behavior and we'd need to iterate through the tokens?

[kbobyrev]

Skip the tokens that are not identifiers or weren't defined as macros, attach
tracer.

[kbobyrev]

Rebase on top of main.

[kbobyrev]

Add an elaborate comment on potential future improvements.

[sammccall]

Very nice. As discussed it'd be nice to enhance MainFileMacros at some point but it doesn't seem urgent

clang-tools-extra/clangd/IncludeCleaner.cpp

173

This describes the implementation, not the purpose.

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`.

180

I'd rephrase a little:

• emphasize the data structure MainFileMacros and the data that's missing, rather than the process how it's collected
• "in the definitions" isn't the essential property, it's rather that they weren't expanded. (e.g. macros uses in PP-disabled sections too). Then macro definitions are an example.

190

this is just a fast-fail for non-identifiers, right? The hadMacroDefinition check is already done in locateMacroAt.

Probably cleaner to move the == tok::identifier check inside locateMacroAt and remove the checks here.

clang-tools-extra/clangd/IncludeCleaner.h

36–41

The start of this sentence is hard to parse: "non-macros" as a noun is vague, and probably needs a comma after. "uses" and "go through" are a bit vague too.

Maybe "A RecursiveASTVisitor finds references to symbols and records their associated locations. These may be macro expansions..."
And "We also examine all identifier tokens in the file in case they reference macros".

Really this is all describing the implementation though, and could be left out of the header entirely...

40

something's happened to the punctuation here

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

127

Maybe PP-disabled too?

"#define ^FOO\n"
"#define ^BAR

"#if 0
#if FOO
BAR
#endif
#endif
"

[kbobyrev]

Address review comments.

[kbobyrev]

Use IncludeCleaner to clean the includes and clean up a little.

clang-tools-extra/clangd/IncludeCleaner.h

 #include <vector>
 
 namespace clang {
 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:
 ///
 /// - cover the whole file in a single traversal for efficiency
 /// - don't attempt to describe where symbols were referenced from in
 ///   ambiguous cases (e.g. implicitly used symbols, multiple declarations)
 /// - err on the side of reporting all possible locations
 ReferencedLocations findReferencedLocations(ParsedAST &AST);

clang-tools-extra/clangd/IncludeCleaner.cpp

 //===--- IncludeCleaner.cpp - Unused/Missing Headers Analysis ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IncludeCleaner.h"
 #include "Config.h"
 #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"
 
 namespace clang {
 namespace clangd {
 namespace {
 
 /// Crawler traverses the AST and feeds in the locations of (sometimes
if (!SM.isWrittenInScratchSpace(Exp.getExpansionLocStart()))
       add(Exp.getExpansionLocStart());
     if (!SM.isWrittenInScratchSpace(Exp.getExpansionLocEnd()))
       add(Exp.getExpansionLocEnd());
   }
 };
 
 // Returns the range starting at '#' and ending at EOL. Escaped newlines are not
 // handled.
 clangd::Range getDiagnosticRange(llvm::StringRef Code, unsigned HashOffset) {
   clangd::Range Result;
   Result.end = Result.start = offsetToPosition(Code, HashOffset);
 
   // Span the warning until the EOL or EOF.
   Result.end.character +=
       lspLength(Code.drop_front(HashOffset).take_until([](char C) {
         return C == '\n' || C == '\r';
       }));
   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;
 }
 
 llvm::DenseSet<FileID>
 findReferencedFiles(const llvm::DenseSet<SourceLocation> &Locs,
                     const SourceManager &SM) {
   std::vector<SourceLocation> Sorted{Locs.begin(), Locs.end()};
   llvm::sort(Sorted); // Group by FileID.

clang-tools-extra/clangd/SourceCode.cpp

llvm::Optional<SpelledWord> SpelledWord::touching(SourceLocation SpelledLoc,
   for (const auto &T : Touching)
     if (T.location() <= Result.Location)
       Result.PartOfSpelledToken = &T;
   return Result;
 }
 
 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();
   IdentifierInfo *IdentifierInfo = PP.getIdentifierInfo(SpelledTok.text(SM));
   if (!IdentifierInfo || !IdentifierInfo->hadMacroDefinition())
     return None;
 
   // We need to take special case to handle #define and #undef.

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

TestCase Cases[] = {
       {
           "enum ^Color { ^Red = 42, Green = 9000};",
           "int MyColor = Red;",
       },
       {
           "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 {};",
       },
       {
           "enum class Color : int {};",
           "enum class Color : int;",
       },