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

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.

Reviewers: hokein, ilya-biryukov

Subscribers: MaskRay, jkorous, arphaman, kadircet, cfe-commits

Tags: #clang

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

llvm-svn: 370452

clang-tools-extra/clangd/ClangdUnit.cpp

@@ -22,6 +22,7 @@
 #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"
@@ -32,6 +33,7 @@
 #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"
@@ -103,6 +105,28 @@ 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,
@@ -417,6 +441,11 @@ ParsedAST::build(std::unique_ptr<clang::CompilerInvocation> CI,
   // (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.
@@ -470,7 +499,8 @@ ParsedAST::build(std::unique_ptr<clang::CompilerInvocation> CI,
     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));
 }
 
@@ -509,6 +539,10 @@ 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 {
@@ -565,11 +599,13 @@ 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)) {

clang-tools-extra/clangd/ClangdUnit.h

@@ -118,6 +118,9 @@ public:
   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; }
@@ -126,6 +129,7 @@ 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);
 
@@ -145,6 +149,9 @@ private:
   ///   - 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

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

@@ -262,6 +262,56 @@ TEST(ClangdUnitTest, CanBuildInvocationWithUnknownArgs) {
   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