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

[ArcsinX]

Summary

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.

[ArcsinX]

I am not happy with my solution, hope to hear an advice.
Initial discussion about this problem: https://reviews.llvm.org/D93683#2468842

[kadircet]

Ah this is really unfortunate :( In theory we already know which slabs belong to main file while updating the index for the preamble, as we shard them per file. But it is really inconvenient layering wise to transfer that slab from updatePreamble to updateMain in parsing callbacks :/
Another layer to fix the issue would be to somehow make main file parsing, re-parse those macro definitions through replay mechanism or preamble patching, but these might have unwanted consequences as we would see a re-definition of the macro now. Also they are considerably hard :/

So it feels like the solution proposed here isn't so bad after all (unless there are other alternatives i am missing ATM). I would suggest having a separate container for macro definitions though, rather than modifying the existing "references container". We could directly have SymbolSlab/vector<Symbol> for the definitions and merge them with SymbolCollector's existing slab during the post-processing. That way we can also re-use the logic in SymbolCollector::handleMacroOccurence to generate a Symbol from MacroInfo.

On a side note, we can also easily skip the shard belonging to main file inside FileIndex::updatePreamble, which is going to be suppressed by new merged-index logic anyway. Though I think it deserves a separate patch. Also I don't fully grasp the workings of the suppression logic yet, so it might not be wise to do so if there are any cases in which we still merge information from static index.

[sammccall]

Ah, this is a bit ugly... I'm not really happy with our representation of info extracted from the preprocessor, but that's something to tackle another time.

Thanks for fixing this!

clang-tools-extra/clangd/CollectMacros.h

27–33

this pair should be a struct instead - MacroOccurrence?

37

this should probably also be MacroOccurrence?

clang-tools-extra/clangd/index/SymbolCollector.cpp

397–398

this fixme is now addressed

404

nit: prefer SourceLocation StartLoc = cantFail(sourceLocationMainFile(...));

407

nit: we have toSourceCode(SourceRange(StartLoc, EndLoc)) for this

[ArcsinX]

• std::pair<Range, bool> => struct MacroOccurrence
• remove addressed fixme comment
• assert() => cantFail()
• use toSourceCode() to get the macro name

clang-tools-extra/clangd/CollectMacros.h

 #include "clang/Basic/IdentifierTable.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "llvm/ADT/DenseMap.h"
 #include <string>
 
 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;
 };
 
 /// Collects macro references (e.g. definitions, expansions) in the main file.
 /// It is used to:
 ///  - collect macros in the preamble section of the main file (in Preamble.cpp)
 ///  - collect macros after the preamble of the main file (in ParsedAST.cpp)
 class CollectMainFileMacros : public PPCallbacks {
 public:
   explicit CollectMainFileMacros(const SourceManager &SM, MainFileMacros &Out)
       : SM(SM), Out(Out) {}
 
   void FileChanged(SourceLocation Loc, FileChangeReason,
                    SrcMgr::CharacteristicKind, FileID) override {
     InMainFile = isInsideMainFile(Loc, SM);
   }
 
   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,
                     SourceRange Range, const MacroArgs *Args) override {
     add(MacroName, MD.getMacroInfo());
   }
 
   void MacroUndefined(const clang::Token &MacroName,
void SourceRangeSkipped(SourceRange R, SourceLocation EndifLoc) override {
     if (!InMainFile)
       return;
     Position Begin = sourceLocToPosition(SM, R.getBegin());
     Position End = sourceLocToPosition(SM, R.getEnd());
     Out.SkippedRanges.push_back(Range{Begin, End});
   }
 
 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;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_COLLECTEDMACROS_H

clang-tools-extra/clangd/CollectMacros.cpp

 //===--- CollectMacros.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 "CollectMacros.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Lex/Lexer.h"
 
 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();
   if (Loc.isInvalid() || Loc.isMacroID())
     return;
 
   auto Name = MacroNameTok.getIdentifierInfo()->getName();
   Out.Names.insert(Name);
   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

clang-tools-extra/clangd/SemanticHighlighting.cpp

std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
   // Highlight all decls and references coming from the AST.
   findExplicitReferences(C, [&](ReferenceLoc R) {
     if (auto Kind = kindForReference(R))
       Builder.addToken(R.NameLoc, *Kind);
   });
   // 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);
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, HighlightingKind K) {
   switch (K) {
   case HighlightingKind::Variable:
     return OS << "Variable";

clang-tools-extra/clangd/XRefs.cpp

if (Macro) {
     // Handle references to macro.
     if (auto MacroSID = getSymbolID(Macro->Name, Macro->Info, SM)) {
       // Collect macro references from main file.
       const auto &IDToRefs = AST.getMacros().MacroRefs;
       auto Refs = IDToRefs.find(MacroSID);
       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));
         }
       }
       Req.IDs.insert(MacroSID);
     }
   } else {
     // Handle references to Decls.

clang-tools-extra/clangd/index/SymbolCollector.cpp

void SymbolCollector::handleMacros(const MainFileMacros &MacroRefsToIndex) {
   assert(PP.get());
   const auto &SM = PP->getSourceManager();
   const auto *MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
   assert(MainFileEntry);
 
   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);
+      }
     }
   }
 }
 
 bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
                                             const MacroInfo *MI,
                                             index::SymbolRoleSet Roles,
                                             SourceLocation Loc) {

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

for (int I = 1;; I++) {
       auto Loc = sourceLocationInMainFile(SM, ExpectedRefs.begin()->start);
       ASSERT_TRUE(bool(Loc));
       const auto *Id = syntax::spelledIdentifierTouching(*Loc, AST.getTokens());
       ASSERT_TRUE(Id);
       auto Macro = locateMacroAt(*Id, PP);
       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;
   }
 }
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 std::vector<SymbolInformation> getSymbols(TestTU &TU, llvm::StringRef Query,
                                           int Limit = 0) {
   auto SymbolInfos = getWorkspaceSymbols(Query, Limit, TU.index().get(),
                                          testPath(TU.Filename));
   EXPECT_TRUE(bool(SymbolInfos)) << "workspaceSymbols returned an error";
   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
        )cpp";
 
   // LSP's SymbolKind doesn't have a "Macro" kind, and
   // indexSymbolKindToSymbolKind() currently maps macros
   // to SymbolKind::String.

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

TU.AdditionalFiles["foo.inc"] = R"cpp(
     int C = ID(1);
     #define DEF 1
     int D = DEF;
   )cpp";
   ParsedAST AST = TU.build();
   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()));
 }
 
 MATCHER_P(WithFileName, Inc, "") { return arg.FileName == Inc; }
 
 TEST(ParsedASTTest, ReplayPreambleForTidyCheckers) {
   struct Inclusion {