1 files changed, 27 insertions, 12 deletions

diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index f8fb2c29570..c75184010ee 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -115,9 +115,7 @@ bool shouldFilterDecl(const NamedDecl *ND, ASTContext *ASTCtx,
   //   * enum constants in unscoped enum decl (e.g. "red" in "enum {red};")
   auto InTopLevelScope = hasDeclContext(
       anyOf(namespaceDecl(), translationUnitDecl(), linkageSpecDecl()));
-  if (match(decl(allOf(Opts.IndexMainFiles
-                           ? decl()
-                           : decl(unless(isExpansionInMainFile())),
+  if (match(decl(allOf(unless(isExpansionInMainFile()),
                        anyOf(InTopLevelScope,
                              hasDeclContext(enumDecl(InTopLevelScope,
                                                      unless(isScoped())))))),
@@ -177,11 +175,9 @@ getIncludeHeader(const SourceManager &SM, SourceLocation Loc,
 // For symbols defined inside macros:
 //   * use expansion location, if the symbol is formed via macro concatenation.
 //   * use spelling location, otherwise.
-llvm::Optional<SymbolLocation>
-getSymbolLocation(const NamedDecl &D, SourceManager &SM,
-                  const SymbolCollector::Options &Opts,
-                  const clang::LangOptions& LangOpts,
-                  std::string &FileURIStorage) {
+llvm::Optional<SymbolLocation> getSymbolLocation(
+    const NamedDecl &D, SourceManager &SM, const SymbolCollector::Options &Opts,
+    const clang::LangOptions &LangOpts, std::string &FileURIStorage) {
   SourceLocation SpellingLoc = SM.getSpellingLoc(D.getLocation());
   if (D.getLocation().isMacroID()) {
     std::string PrintLoc = SpellingLoc.printToString(SM);
@@ -209,6 +205,19 @@ getSymbolLocation(const NamedDecl &D, SourceManager &SM,
   return std::move(Result);
 }
 
+// Checks whether \p ND is a definition of a TagDecl (class/struct/enum/union)
+// in a header file, in which case clangd would prefer to use ND as a canonical
+// declaration.
+// FIXME: handle symbol types that are not TagDecl (e.g. functions), if using
+// the the first seen declaration as canonical declaration is not a good enough
+// heuristic.
+bool isPreferredDeclaration(const NamedDecl &ND, index::SymbolRoleSet Roles) {
+  using namespace clang::ast_matchers;
+  return (Roles & static_cast<unsigned>(index::SymbolRole::Definition)) &&
+         llvm::isa<TagDecl>(&ND) &&
+         match(decl(isExpansionInMainFile()), ND, ND.getASTContext()).empty();
+}
+
 } // namespace
 
 SymbolCollector::SymbolCollector(Options Opts) : Opts(std::move(Opts)) {}
@@ -241,12 +250,20 @@ bool SymbolCollector::handleDeclOccurence(
     if (index::generateUSRForDecl(ND, USR))
       return true;
 
+    const NamedDecl &OriginalDecl = *cast<NamedDecl>(ASTNode.OrigD);
     auto ID = SymbolID(USR);
-    const Symbol* BasicSymbol = Symbols.find(ID);
+    const Symbol *BasicSymbol = Symbols.find(ID);
     if (!BasicSymbol) // Regardless of role, ND is the canonical declaration.
       BasicSymbol = addDeclaration(*ND, std::move(ID));
+    else if (isPreferredDeclaration(OriginalDecl, Roles))
+      // If OriginalDecl is preferred, replace the existing canonical
+      // declaration (e.g. a class forward declaration). There should be at most
+      // one duplicate as we expect to see only one preferred declaration per
+      // TU, because in practice they are definitions.
+      BasicSymbol = addDeclaration(OriginalDecl, std::move(ID));
+
     if (Roles & static_cast<unsigned>(index::SymbolRole::Definition))
-      addDefinition(*cast<NamedDecl>(ASTNode.OrigD), *BasicSymbol);
+      addDefinition(OriginalDecl, *BasicSymbol);
   }
   return true;
 }
@@ -271,8 +288,6 @@ const Symbol *SymbolCollector::addDeclaration(const NamedDecl &ND,
   std::tie(S.Scope, S.Name) = splitQualifiedName(QName);
   S.SymInfo = index::getSymbolInfo(&ND);
   std::string FileURI;
-  // FIXME: we may want a different "canonical" heuristic than clang chooses.
-  // Clang seems to choose the first, which may not have the most information.
   if (auto DeclLoc =
           getSymbolLocation(ND, SM, Opts, ASTCtx->getLangOpts(), FileURI))
     S.CanonicalDeclaration = *DeclLoc;