[clangd] Fix go-to-type target location

[sammccall]

Summary

Previously it'd ~always jump to the lexically first declaration, which was
often an unhelpful forward declaration.

• consult the index for definition and preferred declaration locations (query code shared with go-to-definition)
• when unwrapping to LSP, prefer definition to declaration This matches "go-to-definition", which is the most common go-to operation

[usaxena95]

clang-tools-extra/clangd/XRefs.cpp

347

nit: ResultIndex looks like an implementation detail of this function and the callers do not need it. SymbolID is already present as LocatedSymbol::ID so we can compute this map in this function itself.

[sammccall]

clang-tools-extra/clangd/XRefs.cpp

347

Nice catch! We were in fact computing all the symbol ids twice, too..

clang-tools-extra/clangd/ClangdServer.cpp

void ClangdServer::foldingRanges(llvm::StringRef File,
   };
   // We want to make sure folding ranges are always available for all the open
   // files, hence prefer runQuick to not wait for operations on other files.
   WorkScheduler->runQuick("FoldingRanges", File, std::move(Action));
 }
 
 void ClangdServer::findType(llvm::StringRef File, Position Pos,
                             Callback<std::vector<LocatedSymbol>> CB) {
-  auto Action =
-      [Pos, CB = std::move(CB)](llvm::Expected<InputsAndAST> InpAST) mutable {
-        if (!InpAST)
-          return CB(InpAST.takeError());
-        CB(clangd::findType(InpAST->AST, Pos));
-      };
+  auto Action = [Pos, CB = std::move(CB),
+                 this](llvm::Expected<InputsAndAST> InpAST) mutable {
+    CB(clangd::findType(InpAST->AST, Pos, Index));
   WorkScheduler->runWithAST("FindType", File, std::move(Action));
 }
 
 void ClangdServer::findImplementations(
     PathRef File, Position Pos, Callback<std::vector<LocatedSymbol>> CB) {
   auto Action = [Pos, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)

clang-tools-extra/clangd/XRefs.h

 ///   - subclasses for a base class;
 std::vector<LocatedSymbol> findImplementations(ParsedAST &AST, Position Pos,
                                                const SymbolIndex *Index);
 
 /// Returns symbols for types referenced at \p Pos.
 ///
 /// For example, given `b^ar()` wher bar return Foo, this function returns the
 /// definition of class Foo.
-std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos);
+std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos,
+                                    const SymbolIndex *Index);
 
 /// Returns references of the symbol at a specified \p Pos.
 /// \p Limit limits the number of results returned (0 means no limit).
 ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
                                 const SymbolIndex *Index = nullptr,
                                 bool AddContext = false);
 
 /// Get info about symbols at \p Pos.

clang-tools-extra/clangd/XRefs.cpp

if (!DefLoc) {
       elog("Failed to convert location: {0}", DefLoc.takeError());
       return;
     }
     Results.back().Definition = *DefLoc;
   });
   return Results;
 }
 
+// Given LocatedSymbol results derived from the AST, query the index to obtain
+// definitions and preferred declarations.
+void enhanceLocatedSymbolsFromIndex(llvm::MutableArrayRef<LocatedSymbol> Result,
+                                    const SymbolIndex *Index,
+                                    llvm::StringRef MainFilePath) {
+  LookupRequest QueryRequest;
+  llvm::DenseMap<SymbolID, unsigned> ResultIndex;
+  for (unsigned I = 0; I < Result.size(); ++I) {
+    if (auto ID = Result[I].ID) {
+      ResultIndex.try_emplace(ID, I);
+      QueryRequest.IDs.insert(ID);
+    }
+  }
+  if (!Index || QueryRequest.IDs.empty())
+    return;
+  std::string Scratch;
+  Index->lookup(QueryRequest, [&](const Symbol &Sym) {
+    auto &R = Result[ResultIndex.lookup(Sym.ID)];
+
+    if (R.Definition) { // from AST
+      // Special case: if the AST yielded a definition, then it may not be
+      // the right *declaration*. Prefer the one from the index.
+      if (auto Loc = toLSPLocation(Sym.CanonicalDeclaration, MainFilePath))
+        R.PreferredDeclaration = *Loc;
+
+      // We might still prefer the definition from the index, e.g. for
+      // generated symbols.
+      if (auto Loc = toLSPLocation(
+              getPreferredLocation(*R.Definition, Sym.Definition, Scratch),
+              MainFilePath))
+        R.Definition = *Loc;
+    } else {
+      R.Definition = toLSPLocation(Sym.Definition, MainFilePath);
+
+      // Use merge logic to choose AST or index declaration.
+      if (auto Loc = toLSPLocation(
+              getPreferredLocation(R.PreferredDeclaration,
+                                   Sym.CanonicalDeclaration, Scratch),
+              MainFilePath))
+        R.PreferredDeclaration = *Loc;
+    }
+  });
+}
+
 // Decls are more complicated.
 // The AST contains at least a declaration, maybe a definition.
 // These are up-to-date, and so generally preferred over index results.
 // We perform a single batch index lookup to find additional definitions.
 std::vector<LocatedSymbol>
 locateASTReferent(SourceLocation CurLoc, const syntax::Token *TouchedIdentifier,
                   ParsedAST &AST, llvm::StringRef MainFilePath,
                   const SymbolIndex *Index, ASTNodeKind &NodeKind) {
   const SourceManager &SM = AST.getSourceManager();
   // Results follow the order of Symbols.Decls.
   std::vector<LocatedSymbol> Result;
-  // Keep track of SymbolID -> index mapping, to fill in index data later.
-  llvm::DenseMap<SymbolID, size_t> ResultIndex;
 
   static constexpr trace::Metric LocateASTReferentMetric(
       "locate_ast_referent", trace::Metric::Counter, "case");
   auto AddResultDecl = [&](const NamedDecl *D) {
     D = getPreferredDecl(D);
     auto Loc =
         makeLocation(AST.getASTContext(), nameLocation(*D, SM), MainFilePath);
     if (!Loc)
       return;
 
     Result.emplace_back();
     Result.back().Name = printName(AST.getASTContext(), *D);
     Result.back().PreferredDeclaration = *Loc;
     Result.back().ID = getSymbolID(D);
     if (const NamedDecl *Def = getDefinition(D))
       Result.back().Definition = makeLocation(
           AST.getASTContext(), nameLocation(*Def, SM), MainFilePath);
-
-    // Record SymbolID for index lookup later.
-    if (auto ID = getSymbolID(D))
-      ResultIndex[ID] = Result.size() - 1;
   };
 
   // Emit all symbol locations (declaration or definition) from AST.
   DeclRelationSet Relations =
       DeclRelation::TemplatePattern | DeclRelation::Alias;
   auto Candidates =
       getDeclAtPositionWithRelations(AST, CurLoc, Relations, &NodeKind);
   llvm::DenseSet<SymbolID> VirtualMethods;
if (const auto *CD = dyn_cast<ObjCCategoryDecl>(D))
           LocateASTReferentMetric.record(1, "objc-category-to-class");
           AddResultDecl(ID);
         }
 
     LocateASTReferentMetric.record(1, "regular");
     // Otherwise the target declaration is the right one.
     AddResultDecl(D);
   }
-
-  // Now query the index for all Symbol IDs we found in the AST.
-  if (Index && !ResultIndex.empty()) {
-    LookupRequest QueryRequest;
-    for (auto It : ResultIndex)
-      QueryRequest.IDs.insert(It.first);
-    std::string Scratch;
-    Index->lookup(QueryRequest, [&](const Symbol &Sym) {
-      auto &R = Result[ResultIndex.lookup(Sym.ID)];
-
-      if (R.Definition) { // from AST
-        // Special case: if the AST yielded a definition, then it may not be
-        // the right *declaration*. Prefer the one from the index.
-        if (auto Loc = toLSPLocation(Sym.CanonicalDeclaration, MainFilePath))
-          R.PreferredDeclaration = *Loc;
-
-        // We might still prefer the definition from the index, e.g. for
-        // generated symbols.
-        if (auto Loc = toLSPLocation(
-                getPreferredLocation(*R.Definition, Sym.Definition, Scratch),
-                MainFilePath))
-          R.Definition = *Loc;
-      } else {
-        R.Definition = toLSPLocation(Sym.Definition, MainFilePath);
-
-        // Use merge logic to choose AST or index declaration.
-        if (auto Loc = toLSPLocation(
-                getPreferredLocation(R.PreferredDeclaration,
-                                     Sym.CanonicalDeclaration, Scratch),
-                MainFilePath))
-          R.PreferredDeclaration = *Loc;
-      }
-    });
-  }
+  enhanceLocatedSymbolsFromIndex(Result, Index, MainFilePath);
 
   auto Overrides = findImplementors(VirtualMethods, RelationKind::OverriddenBy,
                                     Index, MainFilePath);
   Result.insert(Result.end(), Overrides.begin(), Overrides.end());
   return Result;
 }
 
 std::vector<LocatedSymbol> locateSymbolForType(const ParsedAST &AST,
-                                               const QualType &Type) {
+                                               const QualType &Type,
+                                               const SymbolIndex *Index) {
   const auto &SM = AST.getSourceManager();
   auto MainFilePath = AST.tuPath();
 
   // FIXME: this sends unique_ptr<Foo> to unique_ptr<T>.
   // Likely it would be better to send it to Foo (heuristically) or to both.
   auto Decls = targetDecl(DynTypedNode::create(Type.getNonReferenceType()),
                           DeclRelation::TemplatePattern | DeclRelation::Alias,
                           AST.getHeuristicResolver());
for (const NamedDecl *D : Decls) {
     Results.emplace_back();
     Results.back().Name = printName(ASTContext, *D);
     Results.back().PreferredDeclaration = *Loc;
     Results.back().ID = getSymbolID(D);
     if (const NamedDecl *Def = getDefinition(D))
       Results.back().Definition =
           makeLocation(ASTContext, nameLocation(*Def, SM), MainFilePath);
   }
+  enhanceLocatedSymbolsFromIndex(Results, Index, MainFilePath);
 
   return Results;
 }
 
 bool tokenSpelledAt(SourceLocation SpellingLoc, const syntax::TokenBuffer &TB) {
   auto ExpandedTokens = TB.expandedTokens(
       TB.sourceManager().getMacroArgExpandedLocation(SpellingLoc));
   return !ExpandedTokens.empty();
if (Tok.kind() == tok::identifier) {
       TouchedIdentifier = &Tok;
       break;
     }
 
     if (Tok.kind() == tok::kw_auto || Tok.kind() == tok::kw_decltype) {
       // go-to-definition on auto should find the definition of the deduced
       // type, if possible
       if (auto Deduced = getDeducedType(AST.getASTContext(), Tok.location())) {
-        auto LocSym = locateSymbolForType(AST, *Deduced);
+        auto LocSym = locateSymbolForType(AST, *Deduced, Index);
         if (!LocSym.empty())
           return LocSym;
       }
     }
   }
 
   ASTNodeKind NodeKind;
   auto ASTResults = locateASTReferent(*CurLoc, TouchedIdentifier, AST,
if (H)
     }
 
   return Out.push_back(T);
 }
 
 // Convenience overload, to allow calling this without the out-parameter
 static llvm::SmallVector<QualType> unwrapFindType(
     QualType T, const HeuristicResolver* H) {
-    llvm::SmallVector<QualType> Result;
-    unwrapFindType(T, H, Result);
-    return Result;
 }
 
-
-std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos) {
+std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos,
+                                    const SymbolIndex *Index) {
   const SourceManager &SM = AST.getSourceManager();
   auto Offset = positionToOffset(SM.getBufferData(SM.getMainFileID()), Pos);
   std::vector<LocatedSymbol> Result;
   if (!Offset) {
     elog("failed to convert position {0} for findTypes: {1}", Pos,
          Offset.takeError());
     return Result;
   }
   // The general scheme is: position -> AST node -> type -> declaration.
   auto SymbolsFromNode =
-      [&AST](const SelectionTree::Node *N) -> std::vector<LocatedSymbol> {
+      [&](const SelectionTree::Node *N) -> std::vector<LocatedSymbol> {
     std::vector<LocatedSymbol> LocatedSymbols;
 
     // NOTE: unwrapFindType might return duplicates for something like
     // unique_ptr<unique_ptr<T>>. Let's *not* remove them, because it gives you some
     // information about the type you may have not known before
     // (since unique_ptr<unique_ptr<T>> != unique_ptr<T>).
     for (const QualType& Type : unwrapFindType(typeForNode(N), AST.getHeuristicResolver()))
-        llvm::copy(locateSymbolForType(AST, Type), std::back_inserter(LocatedSymbols));
+      llvm::copy(locateSymbolForType(AST, Type, Index),
+                 std::back_inserter(LocatedSymbols));
 
     return LocatedSymbols;
   };
   SelectionTree::createEach(AST.getASTContext(), AST.getTokens(), *Offset,
                             *Offset, [&](SelectionTree ST) {
                               Result = SymbolsFromNode(ST.commonAncestor());
                               return !Result.empty();
                             });

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

 #include <string>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
-using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::Matcher;
-using ::testing::Not;
 using ::testing::UnorderedElementsAre;
 using ::testing::UnorderedElementsAreArray;
 using ::testing::UnorderedPointwise;
 
 std::string guard(llvm::StringRef Code) {
   return "#pragma once\n" + Code.str();
 }
 
for (const llvm::StringRef Case : {
            "Target ^tarray[3];",
        }) {
     Annotations A(Case);
     TU.Code = A.code().str();
     ParsedAST AST = TU.build();
 
     ASSERT_GT(A.points().size(), 0u) << Case;
     for (auto Pos : A.points())
-      EXPECT_THAT(findType(AST, Pos),
+      EXPECT_THAT(findType(AST, Pos, nullptr),
                   ElementsAre(
                     sym("Target", HeaderA.range("Target"), HeaderA.range("Target"))))
           << Case;
   }
 
   for (const llvm::StringRef Case : {
            "smart_ptr<Target> ^tsmart;",
        }) {
     Annotations A(Case);
     TU.Code = A.code().str();
     ParsedAST AST = TU.build();
 
-    EXPECT_THAT(findType(AST, A.point()),
+    EXPECT_THAT(findType(AST, A.point(), nullptr),
                 UnorderedElementsAre(
                   sym("Target", HeaderA.range("Target"), HeaderA.range("Target")),
                   sym("smart_ptr", HeaderA.range("smart_ptr"), HeaderA.range("smart_ptr"))
                 ))
         << Case;
   }
 }
 
+TEST(FindType, Definition) {
+  Annotations A(R"cpp(
+    class $decl[[X]];
+    X *^x;
+    class $def[[X]] {};
+  )cpp");
+  auto TU = TestTU::withCode(A.code().str());
+  ParsedAST AST = TU.build();
+
+  EXPECT_THAT(findType(AST, A.point(), nullptr),
+              ElementsAre(sym("X", A.range("decl"), A.range("def"))));
+}
+
+TEST(FindType, Index) {
+  Annotations Def(R"cpp(
+    // This definition is only available through the index.
+    class [[X]] {};
+  )cpp");
+  TestTU DefTU = TestTU::withHeaderCode(Def.code());
+  DefTU.HeaderFilename = "def.h";
+  auto DefIdx = DefTU.index();
+
+  Annotations A(R"cpp(
+    class [[X]];
+    X *^x;
+  )cpp");
+  auto TU = TestTU::withCode(A.code().str());
+  ParsedAST AST = TU.build();
+
+  EXPECT_THAT(findType(AST, A.point(), DefIdx.get()),
+              ElementsAre(sym("X", A.range(), Def.range())));
+}
+
 void checkFindRefs(llvm::StringRef Test, bool UseIndex = false) {
   Annotations T(Test);
   auto TU = TestTU::withCode(T.code());
   TU.ExtraArgs.push_back("-std=c++20");
 
   auto AST = TU.build();
   std::vector<Matcher<ReferencesResult::Reference>> ExpectedLocations;
   for (const auto &[R, Context] : T.rangesWithPayload())