[clangd] Store xref for Macros in ParsedAST.

This patch adds the cross references for Macros in the MainFile.
We add references for the main file to the ParsedAST. We query the
references from it using the SymbolID.
Xref outside main file will be added to the index in a separate patch.

clang-tools-extra/clangd/CollectMacros.h

@@ -9,10 +9,13 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_COLLECTEDMACROS_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_COLLECTEDMACROS_H
 
+#include "AST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
+#include "index/SymbolID.h"
 #include "clang/Basic/IdentifierTable.h"
 #include "clang/Lex/PPCallbacks.h"
+#include "llvm/ADT/DenseMap.h"
 #include <string>
 
 namespace clang {
@@ -22,7 +25,11 @@ struct MainFileMacros {
   llvm::StringSet<> Names;
   // Instead of storing SourceLocation, we have to store the token range because
   // SourceManager from preamble is not available when we build the AST.
-  std::vector<Range> Ranges;
+  llvm::DenseMap<SymbolID, std::vector<Range>> MacroRefs;
+  // Somtimes it is not possible to compute the SymbolID for the Macro, e.g. a
+  // reference to an undefined macro. Store them separately, e.g. for semantic
+  // highlighting.
+  std::vector<Range> UnknownMacros;
 };
 
 /// Collects macro references (e.g. definitions, expansions) in the main file.
@@ -80,8 +87,12 @@ private:
       return;
 
     if (auto Range = getTokenRange(SM, LangOpts, Loc)) {
-      Out.Names.insert(MacroNameTok.getIdentifierInfo()->getName());
-      Out.Ranges.push_back(*Range);
+      auto Name = MacroNameTok.getIdentifierInfo()->getName();
+      Out.Names.insert(Name);
+      if (auto SID = getSymbolID(Name, MI, SM))
+        Out.MacroRefs[*SID].push_back(*Range);
+      else
+        Out.UnknownMacros.push_back(*Range);
     }
   }
   const SourceManager &SM;

clang-tools-extra/clangd/SemanticHighlighting.cpp

@@ -311,9 +311,14 @@ std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
     if (auto Kind = kindForReference(R))
       Builder.addToken(R.NameLoc, *Kind);
   });
-  // Add highlightings for macro expansions.
-  for (const auto &M : AST.getMacros().Ranges)
+  // Add highlightings for macro references.
+  for (const auto &SIDToRefs : AST.getMacros().MacroRefs) {
+    for (const auto &M : SIDToRefs.second)
+      Builder.addToken({HighlightingKind::Macro, M});
+  }
+  for (const auto &M : AST.getMacros().UnknownMacros)
     Builder.addToken({HighlightingKind::Macro, M});
+
   return std::move(Builder).collect();
 }
 

clang-tools-extra/clangd/XRefs.cpp

@@ -917,8 +917,7 @@ ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
                      MainFileRefs.end());
   for (const auto &Ref : MainFileRefs) {
     if (auto Range =
-            getTokenRange(AST.getASTContext().getSourceManager(),
-                          AST.getASTContext().getLangOpts(), Ref.Loc)) {
+            getTokenRange(SM, AST.getASTContext().getLangOpts(), Ref.Loc)) {
       Location Result;
       Result.range = *Range;
       Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath);

clang-tools-extra/clangd/unittests/CMakeLists.txt

@@ -30,6 +30,7 @@ add_unittest(ClangdUnitTests ClangdTests
   ClangdTests.cpp
   CodeCompleteTests.cpp
   CodeCompletionStringsTests.cpp
+  CollectMacrosTests.cpp
   ContextTests.cpp
   DexTests.cpp
   DiagnosticsTests.cpp

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

@@ -0,0 +1,109 @@
+//===-- CollectMacrosTests.cpp ----------------------------------*- 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 "Annotations.h"
+#include "CollectMacros.h"
+#include "Matchers.h"
+#include "SourceCode.h"
+#include "TestTU.h"
+#include "index/SymbolID.h"
+#include "clang/Basic/SourceLocation.h"
+#include "llvm/Support/ScopedPrinter.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+
+using testing::UnorderedElementsAreArray;
+
+TEST(CollectMainFileMacros, SelectedMacros) {
+  // References of the same symbol must have the ranges with the same
+  // name(integer). If there are N different symbols then they must be named
+  // from 1 to N. Macros for which SymbolID cannot be computed must be named
+  // "Unknown".
+  const char *Tests[] = {
+      R"cpp(// Macros: Cursor on definition.
+        #define $1[[FOO]](x,y) (x + y)
+        int main() { int x = $1[[FOO]]($1[[FOO]](3, 4), $1[[FOO]](5, 6)); }
+      )cpp",
+      R"cpp(
+        #define $1[[M]](X) X;
+        #define $2[[abc]] 123
+        int s = $1[[M]]($2[[abc]]);
+      )cpp",
+      // FIXME: Locating macro in duplicate definitions doesn't work. Enable
+      // this once LocateMacro is fixed.
+      // R"cpp(// Multiple definitions.
+      //   #define $1[[abc]] 1
+      //   int func1() { int a = $1[[abc]];}
+      //   #undef $1[[abc]]
+
+      //   #define $2[[abc]] 2
+      //   int func2() { int a = $2[[abc]];}
+      //   #undef $2[[abc]]
+      // )cpp",
+      R"cpp(
+        #ifdef $Unknown[[UNDEFINED]]
+        #endif
+      )cpp",
+      R"cpp(
+        #ifndef $Unknown[[abc]]
+        #define $1[[abc]]
+        #ifdef $1[[abc]]
+        #endif
+        #endif
+      )cpp",
+      R"cpp(
+        // Macros from token concatenations not included.
+        #define $1[[CONCAT]](X) X##A()
+        #define $2[[PREPEND]](X) MACRO##X()
+        #define $3[[MACROA]]() 123
+        int B = $1[[CONCAT]](MACRO);
+        int D = $2[[PREPEND]](A)
+      )cpp",
+      R"cpp(
+        // FIXME: Macro names in a definition are not detected.
+        #define $1[[MACRO_ARGS2]](X, Y) X Y
+        #define $2[[FOO]] BAR
+        #define $3[[BAR]] 1
+        int A = $2[[FOO]];
+      )cpp"};
+  for (const char *Test : Tests) {
+    Annotations T(Test);
+    auto AST = TestTU::withCode(T.code()).build();
+    auto ActualMacroRefs = AST.getMacros();
+    auto &SM = AST.getSourceManager();
+    auto &PP = AST.getPreprocessor();
+
+    // Known macros.
+    for (int I = 1;; I++) {
+      const auto ExpectedRefs = T.ranges(llvm::to_string(I));
+      if (ExpectedRefs.empty())
+        break;
+
+      auto Loc = getBeginningOfIdentifier(ExpectedRefs.begin()->start, SM,
+                                          AST.getASTContext().getLangOpts());
+      auto Macro = locateMacroAt(Loc, PP);
+      assert(Macro);
+      auto SID = getSymbolID(Macro->Name, Macro->Info, SM);
+
+      EXPECT_THAT(ExpectedRefs,
+                  UnorderedElementsAreArray(ActualMacroRefs.MacroRefs[*SID]))
+          << "Annotation=" << I << ", MacroName=" << Macro->Name
+          << ", Test = " << Test;
+    }
+    // Unkown macros.
+    EXPECT_THAT(AST.getMacros().UnknownMacros,
+                UnorderedElementsAreArray(T.ranges("Unknown")))
+        << "Unknown macros doesn't match in " << Test;
+  }
+}
+} // namespace
+} // namespace clangd
+} // namespace clang

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

@@ -260,6 +260,15 @@ TEST(ParsedASTTest, CollectsMainFileMacroExpansions) {
       // Includes macro expansions in arguments that are expressions
       ^assert(0 <= ^BAR);
     }
+
+    #ifdef ^UNDEFINED
+    #endif
+
+    #define ^MULTIPLE_DEFINITION 1
+    #undef ^MULTIPLE_DEFINITION
+
+    #define ^MULTIPLE_DEFINITION 2
+    #undef ^MULTIPLE_DEFINITION
   )cpp");
   auto TU = TestTU::withCode(TestCase.code());
   TU.HeaderCode = R"cpp(
@@ -274,7 +283,11 @@ TEST(ParsedASTTest, CollectsMainFileMacroExpansions) {
   )cpp";
   ParsedAST AST = TU.build();
   std::vector<Position> MacroExpansionPositions;
-  for (const auto &R : AST.getMacros().Ranges)
+  for (const auto &SIDToRefs : AST.getMacros().MacroRefs) {
+    for (const auto &R : SIDToRefs.second)
+      MacroExpansionPositions.push_back(R.start);
+  }
+  for (const auto &R : AST.getMacros().UnknownMacros)
     MacroExpansionPositions.push_back(R.start);
   EXPECT_THAT(MacroExpansionPositions,
               testing::UnorderedElementsAreArray(TestCase.points()));