[clangd] Add main file macros into the main-file index.

This patch is a try to fix `WorkspaceSymbols.Macros` test after D93796.
If a macro definition is in the preamble section, then it appears to be in the preamble (static) index and not in the main-file (dynamic) index.
Thus, a such macro could not be found at a symbol search according to the logic that we skip symbols from the static index if the location of these symbols is inside the dynamic index files.
To fix this behavior this patch adds main file macros into the main-file (dynamic) index.

Reviewed By: sammccall

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

8 files changed, 48 insertions, 33 deletions

diff --git a/clang-tools-extra/clangd/CollectMacros.cpp b/clang-tools-extra/clangd/CollectMacros.cpp
index bd6165e..0e89b35 100644
--- a/clang-tools-extra/clangd/CollectMacros.cpp
+++ b/clang-tools-extra/clangd/CollectMacros.cpp
@@ -13,8 +13,8 @@
 namespace clang {
 namespace clangd {
 
-void CollectMainFileMacros::add(const Token &MacroNameTok,
-                                const MacroInfo *MI) {
+void CollectMainFileMacros::add(const Token &MacroNameTok, const MacroInfo *MI,
+                                bool IsDefinition) {
   if (!InMainFile)
     return;
   auto Loc = MacroNameTok.getLocation();
@@ -26,9 +26,9 @@ void CollectMainFileMacros::add(const Token &MacroNameTok,
   auto Range = halfOpenToRange(
       SM, CharSourceRange::getCharRange(Loc, MacroNameTok.getEndLoc()));
   if (auto SID = getSymbolID(Name, MI, SM))
-    Out.MacroRefs[SID].push_back(Range);
+    Out.MacroRefs[SID].push_back({Range, IsDefinition});
   else
-    Out.UnknownMacros.push_back(Range);
+    Out.UnknownMacros.push_back({Range, IsDefinition});
 }
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/CollectMacros.h b/clang-tools-extra/clangd/CollectMacros.h
index eecea04..3240111 100644
--- a/clang-tools-extra/clangd/CollectMacros.h
+++ b/clang-tools-extra/clangd/CollectMacros.h
@@ -21,15 +21,20 @@
 namespace clang {
 namespace clangd {
 
-struct MainFileMacros {
-  llvm::StringSet<> Names;
+struct MacroOccurrence {
   // Instead of storing SourceLocation, we have to store the token range because
   // SourceManager from preamble is not available when we build the AST.
-  llvm::DenseMap<SymbolID, std::vector<Range>> MacroRefs;
+  Range Rng;
+  bool IsDefinition;
+};
+
+struct MainFileMacros {
+  llvm::StringSet<> Names;
+  llvm::DenseMap<SymbolID, std::vector<MacroOccurrence>> 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;
+  std::vector<MacroOccurrence> UnknownMacros;
   // Ranges skipped by the preprocessor due to being inactive.
   std::vector<Range> SkippedRanges;
 };
@@ -49,7 +54,7 @@ public:
   }
 
   void MacroDefined(const Token &MacroName, const MacroDirective *MD) override {
-    add(MacroName, MD->getMacroInfo());
+    add(MacroName, MD->getMacroInfo(), /*IsDefinition=*/true);
   }
 
   void MacroExpands(const Token &MacroName, const MacroDefinition &MD,
@@ -87,7 +92,8 @@ public:
   }
 
 private:
-  void add(const Token &MacroNameTok, const MacroInfo *MI);
+  void add(const Token &MacroNameTok, const MacroInfo *MI,
+           bool IsDefinition = false);
   const SourceManager &SM;
   bool InMainFile = true;
   MainFileMacros &Out;
diff --git a/clang-tools-extra/clangd/SemanticHighlighting.cpp b/clang-tools-extra/clangd/SemanticHighlighting.cpp
index be0a543..5e70baf 100644
--- a/clang-tools-extra/clangd/SemanticHighlighting.cpp
+++ b/clang-tools-extra/clangd/SemanticHighlighting.cpp
@@ -397,10 +397,10 @@ std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
   // Add highlightings for macro references.
   for (const auto &SIDToRefs : AST.getMacros().MacroRefs) {
     for (const auto &M : SIDToRefs.second)
-      Builder.addToken({HighlightingKind::Macro, M});
+      Builder.addToken({HighlightingKind::Macro, M.Rng});
   }
   for (const auto &M : AST.getMacros().UnknownMacros)
-    Builder.addToken({HighlightingKind::Macro, M});
+    Builder.addToken({HighlightingKind::Macro, M.Rng});
 
   return std::move(Builder).collect(AST);
 }
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 8bb74ed..8027e05 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -1311,7 +1311,7 @@ ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
       if (Refs != IDToRefs.end()) {
         for (const auto &Ref : Refs->second) {
           Location Result;
-          Result.range = Ref;
+          Result.range = Ref.Rng;
           Result.uri = URIMainFile;
           Results.References.push_back(std::move(Result));
         }
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index 94ab54b..20f2eac 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -386,16 +386,31 @@ void SymbolCollector::handleMacros(const MainFileMacros &MacroRefsToIndex) {
   const auto MainFileURI = toURI(SM, MainFileEntry->getName(), Opts);
   // Add macro references.
   for (const auto &IDToRefs : MacroRefsToIndex.MacroRefs) {
-    for (const auto &Range : IDToRefs.second) {
+    for (const auto &MacroRef : IDToRefs.second) {
+      const auto &Range = MacroRef.Rng;
+      bool IsDefinition = MacroRef.IsDefinition;
       Ref R;
       R.Location.Start.setLine(Range.start.line);
       R.Location.Start.setColumn(Range.start.character);
       R.Location.End.setLine(Range.end.line);
       R.Location.End.setColumn(Range.end.character);
       R.Location.FileURI = MainFileURI.c_str();
-      // FIXME: Add correct RefKind information to MainFileMacros.
-      R.Kind = RefKind::Reference;
+      R.Kind = IsDefinition ? RefKind::Definition : RefKind::Reference;
       Refs.insert(IDToRefs.first, R);
+      if (IsDefinition) {
+        Symbol S;
+        S.ID = IDToRefs.first;
+        auto StartLoc = cantFail(sourceLocationInMainFile(SM, Range.start));
+        auto EndLoc = cantFail(sourceLocationInMainFile(SM, Range.end));
+        S.Name = toSourceCode(SM, SourceRange(StartLoc, EndLoc));
+        S.SymInfo.Kind = index::SymbolKind::Macro;
+        S.SymInfo.SubKind = index::SymbolSubKind::None;
+        S.SymInfo.Properties = index::SymbolPropertySet();
+        S.SymInfo.Lang = index::SymbolLanguage::C;
+        S.Origin = Opts.Origin;
+        S.CanonicalDeclaration = R.Location;
+        Symbols.insert(S);
+      }
     }
   }
 }
diff --git a/clang-tools-extra/clangd/unittests/CollectMacrosTests.cpp b/clang-tools-extra/clangd/unittests/CollectMacrosTests.cpp
index 4576fa9..9912786 100644
--- a/clang-tools-extra/clangd/unittests/CollectMacrosTests.cpp
+++ b/clang-tools-extra/clangd/unittests/CollectMacrosTests.cpp
@@ -95,14 +95,18 @@ TEST(CollectMainFileMacros, SelectedMacros) {
       assert(Macro);
       auto SID = getSymbolID(Macro->Name, Macro->Info, SM);
 
-      EXPECT_THAT(ExpectedRefs,
-                  UnorderedElementsAreArray(ActualMacroRefs.MacroRefs[SID]))
+      std::vector<Range> Ranges;
+      for (const auto &Ref : ActualMacroRefs.MacroRefs[SID])
+        Ranges.push_back(Ref.Rng);
+      EXPECT_THAT(ExpectedRefs, UnorderedElementsAreArray(Ranges))
           << "Annotation=" << I << ", MacroName=" << Macro->Name
           << ", Test = " << Test;
     }
     // Unknown macros.
-    EXPECT_THAT(AST.getMacros().UnknownMacros,
-                UnorderedElementsAreArray(T.ranges("Unknown")))
+    std::vector<Range> Ranges;
+    for (const auto &Ref : AST.getMacros().UnknownMacros)
+      Ranges.push_back(Ref.Rng);
+    EXPECT_THAT(Ranges, UnorderedElementsAreArray(T.ranges("Unknown")))
         << "Unknown macros doesn't match in " << Test;
   }
 }
diff --git a/clang-tools-extra/clangd/unittests/FindSymbolsTests.cpp b/clang-tools-extra/clangd/unittests/FindSymbolsTests.cpp
index bdd3ddc..4365828 100644
--- a/clang-tools-extra/clangd/unittests/FindSymbolsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/FindSymbolsTests.cpp
@@ -52,17 +52,7 @@ std::vector<SymbolInformation> getSymbols(TestTU &TU, llvm::StringRef Query,
   return *SymbolInfos;
 }
 
-// FIXME: We update two indexes during main file processing:
-//        - preamble index (static)
-//        - main-file index (dynamic)
-//        The macro in this test appears to be in the preamble index and not
-//        in the main-file index. According to our logic of indexes merging, we
-//        do not take this macro from the static (preamble) index, because it
-//        location within the file from the dynamic (main-file) index.
-//
-//        Possible solution is to exclude main-file symbols from the preamble
-//        index, after that we can enable this test again.
-TEST(WorkspaceSymbols, DISABLED_Macros) {
+TEST(WorkspaceSymbols, Macros) {
   TestTU TU;
   TU.Code = R"cpp(
        #define MACRO X
diff --git a/clang-tools-extra/clangd/unittests/ParsedASTTests.cpp b/clang-tools-extra/clangd/unittests/ParsedASTTests.cpp
index b3321f5..98a2483 100644
--- a/clang-tools-extra/clangd/unittests/ParsedASTTests.cpp
+++ b/clang-tools-extra/clangd/unittests/ParsedASTTests.cpp
@@ -341,10 +341,10 @@ TEST(ParsedASTTest, CollectsMainFileMacroExpansions) {
   std::vector<Position> MacroExpansionPositions;
   for (const auto &SIDToRefs : AST.getMacros().MacroRefs) {
     for (const auto &R : SIDToRefs.second)
-      MacroExpansionPositions.push_back(R.start);
+      MacroExpansionPositions.push_back(R.Rng.start);
   }
   for (const auto &R : AST.getMacros().UnknownMacros)
-    MacroExpansionPositions.push_back(R.start);
+    MacroExpansionPositions.push_back(R.Rng.start);
   EXPECT_THAT(MacroExpansionPositions,
               testing::UnorderedElementsAreArray(TestCase.points()));
 }