[clangd] Add extension for adding context (enclosing function or class) in references results

[tom-anders]

Summary

Relevant issue: https://github.com/clangd/clangd/issues/177

[sammccall]

clang-tools-extra/clangd/Protocol.h

218

Location is a common vocabulary type.
If we're going to add an extension that relates to a particular method, we should define a new type that's wire-compatible if that's at all feasible. Could be a subclass of Location or just a clone.

(If this is just a shortcut taken while drafting this patch, please ignore me!)

clang-tools-extra/clangd/XRefs.cpp

1479

Being able to get all the containers in one ~constant-time index query is great! May still not want to have this on-by-default when the client doesn't advertise support, as it's a fixed nontrivial latency hit for remote index.

Also we should skip this query when the set of IDs is empty. That can go here or (maybe better) in the remote index client.

[tom-anders]

clang-tools-extra/clangd/XRefs.cpp

1479

Being able to get all the containers in one ~constant-time index query is great! May still not want to have this on-by-default when the client doesn't advertise support, as it's a fixed nontrivial latency hit for remote index.

Makes sense, we could either add a ClientCapability for this (Looks like we already have some capabilities related to extensions, so this wouldn't be a novelty), or it could just be a command line flag. Any preferences?

Also we should skip this query when the set of IDs is empty. That can go here or (maybe better) in the remote index client.

Agreed, probably more foolproof to let the index do this (otherwise the next person will forget about that as well probably)

[tom-anders]

Add tests, add client capability, add new ReferenceLocation type

[tom-anders]

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

1939 ↗

(On Diff #476438)

@nridge See https://reviews.llvm.org/D137909 for this new syntax

[sammccall]

clang-tools-extra/clangd/Protocol.h

235

This doesn't seem as clear as it could be: special how, additional to what? Maybe:

/// Extends Locations returned by textDocument/references with extra info.
/// This is a clangd extenison: LSP uses `Location`.

240

"context" is vague, and is generally used in LSP to mean "details about how a request was triggered in the client".

Unless we have a reason to keep this abstract, I'd use containerName as it's used elsewhere in LSP for this purpose.
(I kind of hate the name, but at least people will be familiar with it)

248

are these operators actually needed? especially operator< we usually get away without

clang-tools-extra/clangd/XRefs.cpp

1307

a few things to consider here:

• for e.g. lambdas, do you want foo::bar::(anonymous class)::operator(), or foo::bar, or something else?
• there are a few other function-like decls, see Decl::isFunctionOrMethod().
• only functions? suppose you have struct S { T member; }; and you're looking up xrefs for T. Don't we want to return S as the context/container name
• in general, I think want to get the same container for the index vs AST path here. This suggests we should be sharing getRefContainer from SymbolCollector.cpp rather than implementing something ifferent

1310

again for stringifying, I think we want the index/non-index cases to be the same, which suggests the logic to generate Scope+Name (i.e. clangd::printQualifiedName)+Signature and then concatenating them.

Generating the signature is a bit complicated, so maybe leave it out (from both AST + index codepath) from initial patch?

1441

To me, this doesn't seem like a sufficient reason to have irregular behavior for different refs.

1445

nit: unconditional drop_back(2) seems brittle, consume_back instead?

[sammccall]

Since this is a protocol extension, you might want to add an end-to-end test: variant of clang-tools-extra/clangd/test/xrefs.test (xrefs-container.test? to avoid complicating all the existing tests).

And one this lands, one of us should update https://github.com/llvm/clangd-www/blob/main/extensions.md

[tom-anders]

Address comments

[tom-anders]

In D137894#3940600, @sammccall wrote:

Since this is a protocol extension, you might want to add an end-to-end test: variant of clang-tools-extra/clangd/test/xrefs.test (xrefs-container.test? to avoid complicating all the existing tests).

I'll look into adding this, but I'm not familiar with with these kind of tests yet, so I'll might need some time to figure it out - I already uploaded an updated version of the patch anyway to get some more feedback in the meantime.

And one this lands, one of us should update https://github.com/llvm/clangd-www/blob/main/extensions.md

Yeah I could do that, I think I can even commit directly to that repo? Or should I open a review on github/phabricator for this anyway?

clang-tools-extra/clangd/Protocol.h

248

Correct, they are indeed not needed as of now.

clang-tools-extra/clangd/XRefs.cpp

1307

Changed this to use SymbolCollector::getRefContainer, this should address most of the things you mentioned. Now AST and Index results should be consistent here

1310

There's CodeCompletionString::getSignature, we could use that one probably?

But anyway, the Signature field for index container results is actually empty right now for some reason (not sure if this a bug or a feature), so I removed Signature from index container results as well for now. Let's look into this in a follow-up patch.

1441

Ok, changed it to Scope + Name here as well. If users complain about this, we can still change it later.

[tom-anders]

Add end-to-end test

[nridge]

clang-tools-extra/clangd/XRefs.cpp

1412

Unlike the QueryIndex case, I think multiple references having the same container does not apply here (different overrides will be contained in different classes), but it's worth adding a comment to clarify that.

1456

We can have multiple references with the same container (e.g. multiple references in the same function), so I think we need DenseMap<SymbolID, std::vector<size_t>> here.

1473

For good measure, perhaps condition the container-related logic here on AddContext?

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

2188 ↗

(On Diff #477566)

Did you mean to add a call to B->func() as well here?

2238 ↗

(On Diff #477566)

Did you mean to test the payload here?

2310 ↗

(On Diff #477566)

nit: the payloads aren't actually used here

2357 ↗

(On Diff #477566)

nit: we're not actually using the context

2376 ↗

(On Diff #477566)

likewise

[tom-anders]

Rebase, fix review comments

clang-tools-extra/clangd/XRefs.cpp

1456

Good catch! Turns out this fixes the FIXME I added in FindReferences.RefsToBaseMethod

1473

Good point, otherwise we're just wasting cycles and memory here if AddContext is false anyway

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

2238 ↗

(On Diff #477566)

Hmm I think I instead missed that this test does not use checkFindRefs. Removed the payload here, this path is already covered by other tests anyway.

[tom-anders]

s/llvm::Optional/std::optional/ for containerName field

[nridge]

LGTM, and it looks like Sam's comments have been substantially addressed as well, so I'm going to go ahead and approve this, hope that's cool.

[tom-anders]

Thanks both of you for reviewing!

I added documentation to the website here: https://github.com/llvm/clangd-www/pull/77

clang-tools-extra/clangd/Protocol.h

 llvm::json::Value toJSON(const Range &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Range &);
 
 struct Location {
   /// The text document's URI.
   URIForFile uri;
   Range range;
 
+  /// clangd extension: If this a reference, name of the function in which the
+  /// reference occurs
+  std::string container;
+
   friend bool operator==(const Location &LHS, const Location &RHS) {
-    return LHS.uri == RHS.uri && LHS.range == RHS.range;
+    return LHS.uri == RHS.uri && LHS.range == RHS.range && LHS.container == RHS.container;
   }
 
   friend bool operator!=(const Location &LHS, const Location &RHS) {
     return !(LHS == RHS);
   }
 
   friend bool operator<(const Location &LHS, const Location &RHS) {
-    return std::tie(LHS.uri, LHS.range) < std::tie(RHS.uri, RHS.range);
+    return std::tie(LHS.uri, LHS.range, LHS.container) < std::tie(RHS.uri, RHS.range, RHS.container);
   }
 };
 llvm::json::Value toJSON(const Location &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Location &);
 
 struct TextEdit {
   /// The range of the text document to be manipulated. To insert
   /// text into a document create a range where start === end.
   Range range;
 
   /// The string to be inserted. For delete operations use an
   /// empty string.
   std::string newText;
 };
 inline bool operator==(const TextEdit &L, const TextEdit &R) {
   return std::tie(L.newText, L.range) == std::tie(R.newText, R.range);
 }
 bool fromJSON(const llvm::json::Value &, TextEdit &, llvm::json::Path);
 llvm::json::Value toJSON(const TextEdit &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const TextEdit &);
 
 struct TextDocumentItem {
   /// The text document's URI.
   URIForFile uri;
 
   /// The text document's language identifier.
   std::string languageId;

clang-tools-extra/clangd/Protocol.cpp

 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Range &R) {
   return OS << R.start << '-' << R.end;
 }
 
 llvm::json::Value toJSON(const Location &P) {
   return llvm::json::Object{
       {"uri", P.uri},
       {"range", P.range},
+      {"container", P.container},
   };
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Location &L) {
-  return OS << L.range << '@' << L.uri;
+  return OS << L.range << '@' << L.uri << "(inside " << L.container << ")";
 }
 
 bool fromJSON(const llvm::json::Value &Params, TextDocumentItem &R,
               llvm::json::Path P) {
   llvm::json::ObjectMapper O(Params, P);
   return O && O.map("uri", R.uri) && O.map("languageId", R.languageId) &&
          O.map("version", R.version) && O.map("text", R.text);
 }

clang-tools-extra/clangd/XRefs.cpp

 llvm::Optional<LocatedSymbol> locateFileReferent(const Position &Pos,
                                                  ParsedAST &AST,
                                                  llvm::StringRef MainFilePath) {
   for (auto &Inc : AST.getIncludeStructure().MainFileIncludes) {
     if (!Inc.Resolved.empty() && Inc.HashLine == Pos.line) {
       LocatedSymbol File;
       File.Name = std::string(llvm::sys::path::filename(Inc.Resolved));
       File.PreferredDeclaration = {
-          URIForFile::canonicalize(Inc.Resolved, MainFilePath), Range{}};
+          URIForFile::canonicalize(Inc.Resolved, MainFilePath), Range{}, ""};
       File.Definition = File.PreferredDeclaration;
       // We're not going to find any further symbols on #include lines.
       return File;
     }
   }
   return llvm::None;
 }
 

 
 /// Collects references to symbols within the main file.
 class ReferenceFinder : public index::IndexDataConsumer {
 public:
   struct Reference {
     syntax::Token SpelledTok;
     index::SymbolRoleSet Role;
     SymbolID Target;
+    const DeclContext *Context;
 
     Range range(const SourceManager &SM) const {
       return halfOpenToRange(SM, SpelledTok.range(SM).toCharRange(SM));
     }
   };
 
   ReferenceFinder(const ParsedAST &AST,
                   const llvm::ArrayRef<const NamedDecl *> Targets,
if (PerToken) {
         Locs.clear(); // First token doesn't match, assume our guess was wrong.
     }
     if (Locs.empty())
       Locs.push_back(Loc);
 
     for (SourceLocation L : Locs) {
       L = SM.getFileLoc(L);
       if (const auto *Tok = TB.spelledTokenAt(L))
-        References.push_back({*Tok, Roles, DeclID->getSecond()});
+        References.push_back(
+            {*Tok, Roles, DeclID->getSecond(), ASTNode.ContainerDC});
     }
     return true;
   }
 
 private:
   bool PerToken; // If true, report 3 references for split ObjC selector names.
   std::vector<Reference> References;
   const ParsedAST &AST;
void getOverriddenMethods(const CXXMethodDecl *CMD,
   for (const CXXMethodDecl *Base : CMD->overridden_methods()) {
     if (auto ID = getSymbolID(Base))
       OverriddenMethods.insert(ID);
     getOverriddenMethods(Base, OverriddenMethods);
   }
 }
 } // namespace
 
 ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
                                 const SymbolIndex *Index) {
   ReferencesResult Results;
   const SourceManager &SM = AST.getSourceManager();
   auto MainFilePath = AST.tuPath();
   auto URIMainFile = URIForFile::canonicalize(MainFilePath, MainFilePath);
   auto CurLoc = sourceLocationInMainFile(SM, Pos);
   if (!CurLoc) {
     llvm::consumeError(CurLoc.takeError());
     return {};
   }
 
MainFileRefs.erase(std::unique(MainFileRefs.begin(), MainFileRefs.end(),
                                      return L.SpelledTok.location() ==
                                             R.SpelledTok.location();
                                    }),
                        MainFileRefs.end());
     for (const auto &Ref : MainFileRefs) {
       ReferencesResult::Reference Result;
       Result.Loc.range = Ref.range(SM);
       Result.Loc.uri = URIMainFile;
+      if (const auto *ND = llvm::dyn_cast<NamedDecl>(Ref.Context))
+        Result.Loc.container = ND->getQualifiedNameAsString();
       if (Ref.Role & static_cast<unsigned>(index::SymbolRole::Declaration))
         Result.Attributes |= ReferencesResult::Declaration;
       // clang-index doesn't report definitions as declarations, but they are.
       if (Ref.Role & static_cast<unsigned>(index::SymbolRole::Definition))
         Result.Attributes |=
             ReferencesResult::Definition | ReferencesResult::Declaration;
       Results.References.push_back(std::move(Result));
     }
     // Add decl/def of overridding methods.
     if (Index && !OverriddenBy.Subjects.empty()) {
       Index->relations(OverriddenBy, [&](const SymbolID &Subject,
                                          const Symbol &Object) {
         if (Limit && Results.References.size() >= Limit) {
           Results.HasMore = true;
           return;
         }
         const auto LSPLocDecl =
             toLSPLocation(Object.CanonicalDeclaration, MainFilePath);
         const auto LSPLocDef = toLSPLocation(Object.Definition, MainFilePath);
         if (LSPLocDecl && LSPLocDecl != LSPLocDef) {
if (Index && !OverriddenBy.Subjects.empty()) {
           Results.References.push_back(std::move(Result));
         }
       });
     }
   }
   // Now query the index for references from other files.
   auto QueryIndex = [&](llvm::DenseSet<SymbolID> IDs, bool AllowAttributes,
                         bool AllowMainFileSymbols) {
     if (IDs.empty() || !Index || Results.HasMore)
       return;
     RefsRequest Req;
     Req.IDs = std::move(IDs);
     if (Limit) {
       if (Limit < Results.References.size()) {
         // We've already filled our quota, still check the index to correctly
         // return the `HasMore` info.
         Req.Limit = 0;
       } else {
         // Query index only for the remaining size.
         Req.Limit = Limit - Results.References.size();
       }
     }
+    LookupRequest ContainerLookup;
+    llvm::DenseMap<SymbolID, size_t> RefIndexForContainer;
     Results.HasMore |= Index->refs(Req, [&](const Ref &R) {
       auto LSPLoc = toLSPLocation(R.Location, MainFilePath);
       // Avoid indexed results for the main file - the AST is authoritative.
       if (!LSPLoc ||
           (!AllowMainFileSymbols && LSPLoc->uri.file() == MainFilePath))
         return;
       ReferencesResult::Reference Result;
       Result.Loc = std::move(*LSPLoc);
       if (AllowAttributes) {
         if ((R.Kind & RefKind::Declaration) == RefKind::Declaration)
           Result.Attributes |= ReferencesResult::Declaration;
         // FIXME: our index should definitely store def | decl separately!
         if ((R.Kind & RefKind::Definition) == RefKind::Definition)
           Result.Attributes |=
               ReferencesResult::Declaration | ReferencesResult::Definition;
       }
+      SymbolID Container = R.Container;
+      ContainerLookup.IDs.insert(Container);
       Results.References.push_back(std::move(Result));
+      RefIndexForContainer[Container] = Results.References.size() - 1;
+    });
+
+    Index->lookup(ContainerLookup, [&](const Symbol &Container) {
+      auto Ref = RefIndexForContainer.find(Container.ID);
+      assert(Ref != RefIndexForContainer.end());
+      Results.References[Ref->getSecond()].Loc.container =
+          Container.Scope.str() + Container.Name.str() +
+          Container.Signature.str();
     });
   };
   QueryIndex(std::move(IDsToQuery), /*AllowAttributes=*/true,
              /*AllowMainFileSymbols=*/false);
   // For a virtual method: Occurrences of BaseMethod should be treated as refs
   // and not as decl/def. Allow symbols from main file since AST does not report
   // these.
   QueryIndex(std::move(OverriddenMethods), /*AllowAttributes=*/false,