[clangd] Fix Origin and MainFileOnly-ness for macros

Summary:
This was resulting in macros coming from preambles vanishing when user
have opened the source header. For example:

```
// test.h:
 #define X
```

and

```
// test.cc
 #include "test.h
^
```

If user only opens test.cc, we'll get `X` as a completion candidate,
since it is indexed as part of the preamble. But if the user opens
test.h afterwards we would index it as part of the main file and lose
the symbol (as new index shard for test.h will override the existing one
in dynamic index).

Also we were not setting origins for macros correctly, this patch also
fixes it.

Fixes clangd/clangd#461

Reviewers: hokein

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

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

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

@@ -373,16 +373,12 @@ bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
                                             index::SymbolRoleSet Roles,
                                             SourceLocation Loc) {
   assert(PP.get());
-
-  const auto &SM = PP->getSourceManager();
-  auto DefLoc = MI->getDefinitionLoc();
-  auto SpellingLoc = SM.getSpellingLoc(Loc);
-  bool IsMainFileSymbol = SM.isInMainFile(SM.getExpansionLoc(DefLoc));
-
   // Builtin macros don't have useful locations and aren't needed in completion.
   if (MI->isBuiltinMacro())
     return true;
 
+  const auto &SM = PP->getSourceManager();
+  auto DefLoc = MI->getDefinitionLoc();
   // Also avoid storing predefined macros like __DBL_MIN__.
   if (SM.isWrittenInBuiltinFile(DefLoc))
     return true;
@@ -391,8 +387,13 @@ bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
   if (!ID)
     return true;
 
+  auto SpellingLoc = SM.getSpellingLoc(Loc);
+  bool IsMainFileOnly =
+      SM.isInMainFile(SM.getExpansionLoc(DefLoc)) &&
+      !isHeaderFile(SM.getFileEntryForID(SM.getMainFileID())->getName(),
+                    ASTCtx->getLangOpts());
   // Do not store references to main-file macros.
-  if ((static_cast<unsigned>(Opts.RefFilter) & Roles) && !IsMainFileSymbol &&
+  if ((static_cast<unsigned>(Opts.RefFilter) & Roles) && !IsMainFileOnly &&
       (Opts.RefsInHeaders || SM.getFileID(SpellingLoc) == SM.getMainFileID()))
     MacroRefs[*ID].push_back({Loc, Roles});
 
@@ -401,7 +402,7 @@ bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
     return true;
 
   // Skip main-file macros if we are not collecting them.
-  if (IsMainFileSymbol && !Opts.CollectMainFileSymbols)
+  if (IsMainFileOnly && !Opts.CollectMainFileSymbols)
     return false;
 
   // Mark the macro as referenced if this is a reference coming from the main
@@ -425,11 +426,12 @@ bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
   Symbol S;
   S.ID = std::move(*ID);
   S.Name = Name->getName();
-  if (!IsMainFileSymbol) {
+  if (!IsMainFileOnly) {
     S.Flags |= Symbol::IndexedForCodeCompletion;
     S.Flags |= Symbol::VisibleOutsideFile;
   }
   S.SymInfo = index::getSymbolInfoForMacro(*MI);
+  S.Origin = Opts.Origin;
   std::string FileURI;
   // FIXME: use the result to filter out symbols.
   shouldIndexFile(SM.getFileID(Loc));

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

@@ -1402,6 +1402,9 @@ TEST_F(SymbolCollectorTest, Origin) {
   runSymbolCollector("class Foo {};", /*Main=*/"");
   EXPECT_THAT(Symbols, UnorderedElementsAre(
                            Field(&Symbol::Origin, SymbolOrigin::Static)));
+  runSymbolCollector("#define FOO", /*Main=*/"");
+  EXPECT_THAT(Symbols, UnorderedElementsAre(
+                           Field(&Symbol::Origin, SymbolOrigin::Static)));
 }
 
 TEST_F(SymbolCollectorTest, CollectMacros) {
@@ -1504,6 +1507,13 @@ TEST_F(SymbolCollectorTest, BadUTF8) {
   EXPECT_THAT(Refs, Contains(Pair(findSymbol(Symbols, "PUNCT").ID, _)));
 }
 
+TEST_F(SymbolCollectorTest, MacrosInHeaders) {
+  CollectorOpts.CollectMacro = true;
+  TestFileName = testPath("test.h");
+  runSymbolCollector("", "#define X");
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(AllOf(QName("X"), ForCodeCompletion(true))));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang