[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 ↗

(On Diff #217789)

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

149 ↗

(On Diff #217789)

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 ↗

(On Diff #217789)

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 ↗

(On Diff #217789)

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 ↗

(On Diff #217789)

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 ↗

(On Diff #217811)

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

120 ↗

(On Diff #217811)

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 ↗

(On Diff #217811)

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

119 ↗

(On Diff #217811)

NIT: maybe shorten the name to getMainFileExpansions?

150 ↗

(On Diff #217811)

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

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

260 ↗

(On Diff #217811)

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.

248 ↗

(On Diff #217789)

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.

[jvikstrom]

clang-tools-extra/clangd/ClangdUnit.cpp

107 ↗

(On Diff #217811)

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 ↗

(On Diff #217811)

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 ↗

(On Diff #217811)

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

268 ↗

(On Diff #217920)

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.

260 ↗

(On Diff #217811)

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.

clang-tools-extra/trunk/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;
 
+  /// The start locations of all macro expansions spelled inside the main file.
+  /// Does not include expansions from inside other macro expansions.
+  llvm::ArrayRef<SourceLocation> getMainFileExpansions() 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/trunk/clangd/ClangdUnit.cpp

 #include "Logger.h"
 #include "SourceCode.h"
 #include "Trace.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/LangOptions.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
 #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;
 };
 
+// CollectMainFileMacroExpansions and CollectMainFileMacros are two different
+// classes as CollectMainFileMacroExpansions is only used when building the AST
+// for the main file. CollectMainFileMacros is only used when building the
+// preamble.
+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 (!L.isMacroID() && isInsideMainFile(L, SM))
+      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<clang::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<clang::CompilerInvocation> CI,
   if (Preamble)
     Diags.insert(Diags.end(), Preamble->Diags.begin(), Preamble->Diags.end());
   // Finally, add diagnostics coming from the AST.
   {
     std::vector<Diag> D = ASTDiags.take(CTContext.getPointer());
     Diags.insert(Diags.end(), D.begin(), D.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;
 }
 
+llvm::ArrayRef<SourceLocation> ParsedAST::getMainFileExpansions() 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/trunk/clangd/unittests/ClangdUnitTests.cpp

TEST(ClangdUnitTest, CanBuildInvocationWithUnknownArgs) {
   EXPECT_NE(buildCompilerInvocation(Inputs, IgnoreDiags), nullptr);
 
   // Unknown forwarded to -cc1 should not a failure either.
   Inputs.CompileCommand.CommandLine = {
       "clang", "-Xclang", "-fsome-unknown-flag", testPath("foo.cpp")};
   EXPECT_NE(buildCompilerInvocation(Inputs, IgnoreDiags), nullptr);
 }
 
+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 not included.
+    #define CONCAT(X) X##A()
+    #define PREPEND(X) MACRO##X()
+    #define MACROA() 123
+    int B = ^CONCAT(MACRO);
+    int D = ^PREPEND(A)
+
+    // Macros included not from preamble not included.
+    #include "foo.inc"
+
+    #define assert(COND) if (!(COND)) { printf("%s", #COND); exit(0); }
+
+    void test() {
+      // Includes macro expansions in arguments that are expressions
+      ^assert(0 <= ^BAR);
+    }
+  )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.getMainFileExpansions();
+  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