[clangd] Collecting main file macro expansion locations in ParsedAST.

[jvikstrom]

Summary

TokenBuffer does not collect macro expansions inside macro arguments which is needed for semantic higlighting. Therefore collects macro expansions in the main file in a PPCallback when building the ParsedAST instead.

[ilya-biryukov]

Looks great, just requests for comments and more test-cases from my side

clang-tools-extra/clangd/ClangdUnit.h

148

NIT: clarify that this is the start location of the identifier that was macro-expanded.
"expansion location" is an overloaded term

149

Could you please document whether these include:

1. locations outside the main file
2. locations from inside other macro expansions, i.e. those that have isMacroID() == true.

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

248

Could you add a few more interesting cases?

1. Macros outside the main file and the preamble:

// foo.inc
int a = ID(1);

// foo.cpp
#define ID(X) X

int b;
#include "foo.inc"

2. macro expansions from token concatenations

#define FOO(X) X##1()
#define MACRO1() 123

int a = FOO(MACRO);

3. Macro names inside other macros:

#define FOO BAR
#define BAR 1


int a = FOO; // should BAR at line 1 be highlighted?

4. #include not part of the preamble:

#define FOO 1

// Preamble ends here.
int a = 10;
#include "some_file_with_macros.h" // <-- should not get any macros from here

264

NIT: maybe use push_back and for-each loop, the resulting code should be simpler and readability is more important for the test code than performance:

std::vector<Position> Positions;
for (SourceLocation Loc : AST.getMainFileMacros())
  Positions.push_back(sourceLocToPosition(Loc));

[jvikstrom]

Clarified comments. Added tests. Not getting expansions inside other macro expansions.

[jvikstrom]

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

248

Does clangd handle .inc files differently from .h? Because if it doesn't shouldn't ` case 1 cover case 4 as well ?

[ilya-biryukov]

clang-tools-extra/clangd/ClangdUnit.cpp

107

Maybe make this part of CollectMainFileMacros? It looks like a natural fit there and we won't need another instance of PPCallbacks.

120

I believe getSpellingLoc is a no-op if isMacroID() is false. Maybe simplify to:

if (!L.isMacroID() && isInsideMainFile(L, SM))
  ...

clang-tools-extra/clangd/ClangdUnit.h

119

NIT: return llvm::ArrayRef<SourceLocation> instead of const vector

119

NIT: maybe shorten the name to getMainFileExpansions?

150

NIT: a comment like this or a similar one is probably also useful in the public accessor.

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

248

Right, sorry, was rushing to finish a comment. .inc and .h are handled exactly the same.
There is only a difference if they are in the preamble or not.

260

Could we also test?

#define PREPEND(X) Foo##X

It'll probably be the same, but still interesting to see whether it's any differnet.

[jvikstrom]

clang-tools-extra/clangd/ClangdUnit.cpp

107

But CollectMainFileMacros is only used for buildPreamble.
I think this one needs to run in ParsedAST::build?

Is it safe to add a CollectMainFileMacros in ParsedAST::build`?

[ilya-biryukov]

clang-tools-extra/clangd/ClangdUnit.cpp

107

Ah, good point, LG then.
Could you put a small comment beside the class that explicitly mentions this?
It's obvious when reading through the whole code, but not obvious when peeking at different parts of it from a distance.

Something like

// CollectMainFileMacros and CollectMainFileMacroExpansions are two different classes because
// the latter is only used when building preamble and the former only when building the AST for the main file.

[jvikstrom]

Added test for prepending concatenations. Also added made tests pass.

[jvikstrom]

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

260

So it turns out that the tests weren't actually passing before, must have accidentally forgot to save the file or something before compiling. Anyways tokens from concatenations are not included right now (which for highlighting is probably what we want, we don't highlight types/names that are from macro concatenations either)

But the reason they are not are because:

• They do not pass the isInsideMainFile check. Their file id is set to something that prints <scratch space> when dumped
• Even if they did pass the check the SourceLocation does not seem to be correct. They return the same SourceLocation as the parent CONCAT or PREPEND

Don't know how to fix any of this, and don't know if we actually want to collect these expansions either.

[ilya-biryukov]

LGTM, thanks!

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

260

It's definitely ok to not highlight those arguments that get concatenated.
If we wanted to fix this, I think we'd need to detect tokens coming from concatenations separately and finds ways to map them back (not sure if it's possible at all). But let's not bother with this for now.

268

I've realized there's one last interesting case that we are missing:

#define assert(COND) if (!(COND)) { printf("%s", #COND); exit(0); }

void test(int a) {
  assert(0 <= a); // a should be highlighted
}

Could you please add this to the tests? If this works - great; if not - feel free to add a FIXME and land the patch as is, we can fix this later.

clang-tools-extra/clangd/ClangdUnit.h

public:
   const std::vector<Diag> &getDiagnostics() const;
 
   /// Returns the esitmated size of the AST and the accessory structures, in
   /// bytes. Does not include the size of the preamble.
   std::size_t getUsedBytes() const;
   const IncludeStructure &getIncludeStructure() const;
   const CanonicalIncludes &getCanonicalIncludes() const;
 
+  const std::vector<SourceLocation> &getMainFileMacroExpansionLocations() const;
   /// Tokens recorded while parsing the main file.
   /// (!) does not have tokens from the preamble.
   const syntax::TokenBuffer &getTokens() const { return Tokens; }
 
 private:
   ParsedAST(std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action, syntax::TokenBuffer Tokens,
+            std::vector<SourceLocation> MainFileMacroExpLocs,
             std::vector<Decl *> LocalTopLevelDecls, std::vector<Diag> Diags,
             IncludeStructure Includes, CanonicalIncludes CanonIncludes);
 
   // In-memory preambles must outlive the AST, it is important that this member
   // goes before Clang and Action.
   std::shared_ptr<const PreambleData> Preamble;
   // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
   // on it) and CompilerInstance used to run it. That way we don't have to do
   // complex memory management of all Clang structures on our own. (They are
   // stored in CompilerInstance and cleaned up by
   // FrontendAction.EndSourceFile).
   std::unique_ptr<CompilerInstance> Clang;
   std::unique_ptr<FrontendAction> Action;
   /// Tokens recorded after the preamble finished.
   ///   - Includes all spelled tokens for the main file.
   ///   - Includes expanded tokens produced **after** preabmle.
   ///   - Does not have spelled or expanded tokens for files from preamble.
   syntax::TokenBuffer Tokens;
 
+  /// The start locations of all macro expansions spelled inside the main file.
+  /// Does not include expansions from inside other macro expansions.
+  std::vector<SourceLocation> MainFileMacroExpLocs;
   // Data, stored after parsing.
   std::vector<Diag> Diags;
   // Top-level decls inside the current file. Not that this does not include
   // top-level decls from the preamble.
   std::vector<Decl *> LocalTopLevelDecls;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
 };

clang-tools-extra/clangd/ClangdUnit.cpp

 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/MacroInfo.h"
+#include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Serialization/PCHContainerOperations.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
protected:
   CreateASTConsumer(CompilerInstance &CI, llvm::StringRef InFile) override {
     return std::make_unique<DeclTrackingASTConsumer>(/*ref*/ TopLevelDecls);
   }
 
 private:
   std::vector<Decl *> TopLevelDecls;
 };
 
+class CollectMainFileMacroExpansions : public PPCallbacks {
+  const SourceManager &SM;
+  std::vector<SourceLocation> &MainFileMacroLocs;
+
+public:
+  CollectMainFileMacroExpansions(const SourceManager &SM,
+                                 std::vector<SourceLocation> &MainFileMacroLocs)
+      : SM(SM), MainFileMacroLocs(MainFileMacroLocs) {}
+
+  virtual void MacroExpands(const Token &MacroNameTok,
+                            const MacroDefinition &MD, SourceRange Range,
+                            const MacroArgs *Args) {
+    SourceLocation L = MacroNameTok.getLocation();
+    if (isInsideMainFile(SM.getSpellingLoc(L), SM) && !L.isMacroID())
+      MainFileMacroLocs.push_back(L);
+  }
+};
+
 class CollectMainFileMacros : public PPCallbacks {
 public:
   explicit CollectMainFileMacros(const SourceManager &SM,
                                  std::vector<std::string> *Out)
       : SM(SM), Out(Out) {}
 
   void FileChanged(SourceLocation Loc, FileChangeReason,
                    SrcMgr::CharacteristicKind, FileID Prev) {
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   // Replay the preamble includes so that clang-tidy checks can see them.
   if (Preamble)
     ReplayPreamble::attach(Includes, *Clang);
   // Important: collectIncludeStructure is registered *after* ReplayPreamble!
   // Otherwise we would collect the replayed includes again...
   // (We can't *just* use the replayed includes, they don't have Resolved path).
   Clang->getPreprocessor().addPPCallbacks(
       collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
+  // Collect the macro expansions in the main file.
+  std::vector<SourceLocation> MainFileMacroExpLocs;
+  Clang->getPreprocessor().addPPCallbacks(
+      std::make_unique<CollectMainFileMacroExpansions>(
+          Clang->getSourceManager(), MainFileMacroExpLocs));
 
   // Copy over the includes from the preamble, then combine with the
   // non-preamble includes below.
   CanonicalIncludes CanonIncludes;
   if (Preamble)
     CanonIncludes = Preamble->CanonIncludes;
   else
     addSystemHeadersMapping(&CanonIncludes, Clang->getLangOpts());
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   // So just inform the preprocessor of EOF, while keeping everything alive.
   Clang->getPreprocessor().EndSourceFile();
 
   std::vector<Diag> Diags = ASTDiags.take(CTContext.getPointer());
   // Add diagnostics from the preamble, if any.
   if (Preamble)
     Diags.insert(Diags.begin(), Preamble->Diags.begin(), Preamble->Diags.end());
   return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
-                   std::move(Tokens), std::move(ParsedDecls), std::move(Diags),
+                   std::move(Tokens), std::move(MainFileMacroExpLocs),
+                   std::move(ParsedDecls), std::move(Diags),
                    std::move(Includes), std::move(CanonIncludes));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
 ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
 ParsedAST::~ParsedAST() {

 const Preprocessor &ParsedAST::getPreprocessor() const {
   return Clang->getPreprocessor();
 }
 
 llvm::ArrayRef<Decl *> ParsedAST::getLocalTopLevelDecls() {
   return LocalTopLevelDecls;
 }
 
+const std::vector<SourceLocation> &
+ParsedAST::getMainFileMacroExpansionLocations() const {
+  return MainFileMacroExpLocs;
+}
+
 const std::vector<Diag> &ParsedAST::getDiagnostics() const { return Diags; }
 
 std::size_t ParsedAST::getUsedBytes() const {
   auto &AST = getASTContext();
   // FIXME(ibiryukov): we do not account for the dynamically allocated part of
   // Message and Fixes inside each diagnostic.
   std::size_t Total =
       clangd::getUsedBytes(LocalTopLevelDecls) + clangd::getUsedBytes(Diags);
: Preamble(std::move(Preamble)), Diags(std::move(Diags)),
       Includes(std::move(Includes)), MainFileMacros(std::move(MainFileMacros)),
       StatCache(std::move(StatCache)), CanonIncludes(std::move(CanonIncludes)) {
 }
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      syntax::TokenBuffer Tokens,
+                     std::vector<SourceLocation> MainFileMacroExpLocs,
                      std::vector<Decl *> LocalTopLevelDecls,
                      std::vector<Diag> Diags, IncludeStructure Includes,
                      CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Tokens(std::move(Tokens)),
+      MainFileMacroExpLocs(std::move(MainFileMacroExpLocs)),
       Diags(std::move(Diags)),
       LocalTopLevelDecls(std::move(LocalTopLevelDecls)),
       Includes(std::move(Includes)), CanonIncludes(std::move(CanonIncludes)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
 std::shared_ptr<const PreambleData>

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

TEST(ClangdUnitTest, NoCrashOnTokensWithTidyCheck) {
   // Check first token after the preamble.
   EXPECT_EQ(T.expandedTokens().front().text(SM), "inline");
   // Last token is always 'eof'.
   EXPECT_EQ(T.expandedTokens().back().kind(), tok::eof);
   // Check the token before 'eof'.
   EXPECT_EQ(T.expandedTokens().drop_back().back().text(SM), "}");
 }
 
+TEST(ClangdUnitTest, CollectsMainFileMacroExpansions) {
+  Annotations TestCase(R"cpp(
+    #define MACRO_ARGS(X, Y) X Y
+    ^ID(int A);
+    // Macro arguments included.
+    ^MACRO_ARGS(^MACRO_ARGS(^MACRO_EXP(int), A), ^ID(= 2));
+
+    // Macro names inside other macros not included.
+    #define FOO BAR
+    #define BAR 1
+    int A = ^FOO;
+
+    // Macros from token concatenations included.
+    #define CONCAT(X) X##1()
+    #define MACRO1() 123
+    int B = ^CONCAT(^MACRO);
+
+    // Macros included not from preamble not included.
+    #include "foo.inc"
+  )cpp");
+  auto TU = TestTU::withCode(TestCase.code());
+  TU.HeaderCode = R"cpp(
+    #define ID(X) X
+    #define MACRO_EXP(X) ID(X)
+    MACRO_EXP(int B);
+  )cpp";
+  TU.AdditionalFiles["foo.inc"] = R"cpp(
+    int C = ID(1);
+    #define DEF 1
+    int D = DEF;
+  )cpp";
+  ParsedAST AST = TU.build();
+  const std::vector<SourceLocation> &MacroExpansionLocations =
+      AST.getMainFileMacroExpansionLocations();
+  std::vector<Position> MacroExpansionPositions;
+  for (const auto &L : MacroExpansionLocations)
+    MacroExpansionPositions.push_back(
+        sourceLocToPosition(AST.getSourceManager(), L));
+  EXPECT_EQ(MacroExpansionPositions, TestCase.points());
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang