[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

1507

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

1507

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

1941–1943

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

[sammccall]

clang-tools-extra/clangd/Protocol.h

231

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`.

236

"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)

244

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

clang-tools-extra/clangd/XRefs.cpp

1311

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

1314

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?

1461

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

1465

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

244

Correct, they are indeed not needed as of now.

clang-tools-extra/clangd/XRefs.cpp

1311

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

1314

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.

1461

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

1426

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.

1484

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.

1501

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

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

2188

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

2238

Did you mean to test the payload here?

2310

nit: the payloads aren't actually used here

2357

nit: we're not actually using the context

2376

likewise

[tom-anders]

Rebase, fix review comments

clang-tools-extra/clangd/XRefs.cpp

1484

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

1501

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

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

2238

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/ClangdLSPServer.h

private:
   void onGoToDeclaration(const TextDocumentPositionParams &,
                          Callback<std::vector<Location>>);
   void onGoToDefinition(const TextDocumentPositionParams &,
                         Callback<std::vector<Location>>);
   void onGoToType(const TextDocumentPositionParams &,
                   Callback<std::vector<Location>>);
   void onGoToImplementation(const TextDocumentPositionParams &,
                             Callback<std::vector<Location>>);
-  void onReference(const ReferenceParams &, Callback<std::vector<Location>>);
+  void onReference(const ReferenceParams &, Callback<std::vector<ReferenceLocation>>);
   void onSwitchSourceHeader(const TextDocumentIdentifier &,
                             Callback<llvm::Optional<URIForFile>>);
   void onDocumentHighlight(const TextDocumentPositionParams &,
                            Callback<std::vector<DocumentHighlight>>);
   void onFileEvent(const DidChangeWatchedFilesParams &);
   void onWorkspaceSymbol(const WorkspaceSymbolParams &,
                          Callback<std::vector<SymbolInformation>>);
   void onPrepareRename(const TextDocumentPositionParams &,
private:
   /// The supported completion item kinds of the client.
   CompletionItemKindBitset SupportedCompletionItemKinds;
   /// Whether the client supports CodeAction response objects.
   bool SupportsCodeAction = false;
   /// From capabilities of textDocument/documentSymbol.
   bool SupportsHierarchicalDocumentSymbol = false;
   /// Whether the client supports showing file status.
   bool SupportFileStatus = false;
+  /// Whether the client supports attaching a container string to references.
+  bool SupportsReferenceContainer = false;
   /// Which kind of markup should we use in textDocument/hover responses.
   MarkupKind HoverContentFormat = MarkupKind::PlainText;
   /// Whether the client supports offsets for parameter info labels.
   bool SupportsOffsetsInSignatureHelp = false;
   std::mutex BackgroundIndexProgressMutex;
   enum class BackgroundIndexProgress {
     // Client doesn't support reporting progress. No transitions possible.
     Unsupported,

clang-tools-extra/clangd/ClangdLSPServer.cpp

DiagOpts.EmitRelatedLocations =
       Params.capabilities.DiagnosticRelatedInformation;
   if (Params.capabilities.WorkspaceSymbolKinds)
     SupportedSymbolKinds |= *Params.capabilities.WorkspaceSymbolKinds;
   if (Params.capabilities.CompletionItemKinds)
     SupportedCompletionItemKinds |= *Params.capabilities.CompletionItemKinds;
   SupportsCodeAction = Params.capabilities.CodeActionStructure;
   SupportsHierarchicalDocumentSymbol =
       Params.capabilities.HierarchicalDocumentSymbol;
+  SupportsReferenceContainer = Params.capabilities.ReferenceContainer;
   SupportFileStatus = Params.initializationOptions.FileStatus;
   HoverContentFormat = Params.capabilities.HoverContentFormat;
   Opts.LineFoldingOnly = Params.capabilities.LineFoldingOnly;
   SupportsOffsetsInSignatureHelp = Params.capabilities.OffsetsInSignatureHelp;
   if (Params.capabilities.WorkDoneProgress)
     BackgroundIndexProgressState = BackgroundIndexProgress::Empty;
   BackgroundIndexSkipCreate = Params.capabilities.ImplicitProgressCreation;
   Opts.ImplicitCancellation = !Params.capabilities.CancelsStaleRequests;

 }
 
 // FIXME: This function needs to be properly tested.
 void ClangdLSPServer::onChangeConfiguration(
     const DidChangeConfigurationParams &Params) {
   applyConfiguration(Params.settings);
 }
 
-void ClangdLSPServer::onReference(const ReferenceParams &Params,
-                                  Callback<std::vector<Location>> Reply) {
-  Server->findReferences(
-      Params.textDocument.uri.file(), Params.position, Opts.ReferencesLimit,
+void ClangdLSPServer::onReference(
+    const ReferenceParams &Params,
+    Callback<std::vector<ReferenceLocation>> Reply) {
+  Server->findReferences(Params.textDocument.uri.file(), Params.position,
+                         Opts.ReferencesLimit, SupportsReferenceContainer,
                          [Reply = std::move(Reply),
                           IncludeDecl(Params.context.includeDeclaration)](
                              llvm::Expected<ReferencesResult> Refs) mutable {
                            if (!Refs)
                              return Reply(Refs.takeError());
                            // Filter out declarations if the client asked.
-        std::vector<Location> Result;
+                           std::vector<ReferenceLocation> Result;
                            Result.reserve(Refs->References.size());
                            for (auto &Ref : Refs->References) {
-          bool IsDecl = Ref.Attributes & ReferencesResult::Declaration;
+                             bool IsDecl =
+                                 Ref.Attributes & ReferencesResult::Declaration;
                              if (IncludeDecl || !IsDecl)
                                Result.push_back(std::move(Ref.Loc));
                            }
                            return Reply(std::move(Result));
                          });
 }
 
 void ClangdLSPServer::onGoToType(const TextDocumentPositionParams &Params,
                                  Callback<std::vector<Location>> Reply) {
   Server->findType(
       Params.textDocument.uri.file(), Params.position,
       [Reply = std::move(Reply)](
           llvm::Expected<std::vector<LocatedSymbol>> Types) mutable {

clang-tools-extra/clangd/ClangdServer.h

void findImplementations(PathRef File, Position Pos,
                            Callback<std::vector<LocatedSymbol>> CB);
 
   /// Retrieve symbols for types referenced at \p Pos.
   void findType(PathRef File, Position Pos,
                 Callback<std::vector<LocatedSymbol>> CB);
 
   /// Retrieve locations for symbol references.
   void findReferences(PathRef File, Position Pos, uint32_t Limit,
-                      Callback<ReferencesResult> CB);
+                      bool AddContainer, Callback<ReferencesResult> CB);
 
   /// Run formatting for the \p File with content \p Code.
   /// If \p Rng is non-null, formats only that region.
   void formatFile(PathRef File, llvm::Optional<Range> Rng,
                   Callback<tooling::Replacements> CB);
 
   /// Run formatting after \p TriggerText was typed at \p Pos in \p File with
   /// content \p Code.

clang-tools-extra/clangd/ClangdServer.cpp

if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::findImplementations(InpAST->AST, Pos, Index));
   };
 
   WorkScheduler->runWithAST("Implementations", File, std::move(Action));
 }
 
 void ClangdServer::findReferences(PathRef File, Position Pos, uint32_t Limit,
+                                  bool AddContainer,
                                   Callback<ReferencesResult> CB) {
-  auto Action = [Pos, Limit, CB = std::move(CB),
+  auto Action = [Pos, Limit, AddContainer, CB = std::move(CB),
                  this](llvm::Expected<InputsAndAST> InpAST) mutable {
     if (!InpAST)
       return CB(InpAST.takeError());
-    CB(clangd::findReferences(InpAST->AST, Pos, Limit, Index));
+    CB(clangd::findReferences(InpAST->AST, Pos, Limit, Index, AddContainer));
   };
 
   WorkScheduler->runWithAST("References", File, std::move(Action));
 }
 
 void ClangdServer::symbolInfo(PathRef File, Position Pos,
                               Callback<std::vector<SymbolDetails>> CB) {
   auto Action =

clang-tools-extra/clangd/Protocol.h

 
 struct Location {
   /// The text document's URI.
   URIForFile uri;
   Range range;
 
   friend bool operator==(const Location &LHS, const Location &RHS) {
     return LHS.uri == RHS.uri && LHS.range == RHS.range;
   }
 
   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);
   }
 };
 llvm::json::Value toJSON(const Location &);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Location &);
 
+/// Extends Locations returned by textDocument/references with extra info.
+/// This is a clangd extension: LSP uses `Location`.
+struct ReferenceLocation : Location {
+  /// clangd extension: contains the name of the function or class in which the
+  /// reference occurs
+  llvm::Optional<std::string> containerName;
+};
+llvm::json::Value toJSON(const ReferenceLocation &);
+llvm::raw_ostream &operator<<(llvm::raw_ostream &, const ReferenceLocation &);
+
 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);
 }
struct ClientCapabilities {
   /// textDocument.completion.completionItem.snippetSupport
   bool CompletionSnippets = false;
 
   /// Client supports completions with additionalTextEdit near the cursor.
   /// This is a clangd extension. (LSP says this is for unrelated text only).
   /// textDocument.completion.editsNearCursor
   bool CompletionFixes = false;
 
+  /// Client supports displaying a container string for results of
+  /// textDocument/reference (clangd extension)
+  bool ReferenceContainer = false;
+
   /// Client supports hierarchical document symbols.
   /// textDocument.documentSymbol.hierarchicalDocumentSymbolSupport
   bool HierarchicalDocumentSymbol = false;
 
   /// Client supports signature help.
   /// textDocument.signatureHelp
   bool HasSignatureHelp = false;
 

clang-tools-extra/clangd/Protocol.cpp

return llvm::json::Object{
       {"range", P.range},
   };
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Location &L) {
   return OS << L.range << '@' << L.uri;
 }
 
+llvm::json::Value toJSON(const ReferenceLocation &P) {
+  llvm::json::Object Result{
+      {"uri", P.uri},
+      {"range", P.range},
+  };
+  if (P.containerName)
+    Result.insert({"containerName", P.containerName});
+  return Result;
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
+                              const ReferenceLocation &L) {
+  return OS << L.range << '@' << L.uri << " (container: " << L.containerName
+            << ")";
+}
+
 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);
 }
 
 bool fromJSON(const llvm::json::Value &Params, TextEdit &R,
if (auto *TextDocument = O->getObject("textDocument")) {
     if (auto *Diagnostics = TextDocument->getObject("publishDiagnostics")) {
       if (auto CategorySupport = Diagnostics->getBoolean("categorySupport"))
         R.DiagnosticCategory = *CategorySupport;
       if (auto CodeActions = Diagnostics->getBoolean("codeActionsInline"))
         R.DiagnosticFixes = *CodeActions;
       if (auto RelatedInfo = Diagnostics->getBoolean("relatedInformation"))
         R.DiagnosticRelatedInformation = *RelatedInfo;
     }
+    if (auto *References = TextDocument->getObject("references"))
+      if (auto ContainerSupport = References->getBoolean("container"))
+        R.ReferenceContainer = *ContainerSupport;
     if (auto *Completion = TextDocument->getObject("completion")) {
       if (auto *Item = Completion->getObject("completionItem")) {
         if (auto SnippetSupport = Item->getBoolean("snippetSupport"))
           R.CompletionSnippets = *SnippetSupport;
         if (const auto *DocumentationFormat =
                 Item->getArray("documentationFormat")) {
           for (const auto &Format : *DocumentationFormat) {
             if (fromJSON(Format, R.CompletionDocumentationFormat, P))

clang-tools-extra/clangd/XRefs.h

struct ReferencesResult {
   // Bitmask describing whether the occurrence is a declaration, definition etc.
   enum ReferenceAttributes : unsigned {
     Declaration = 1 << 0,
     Definition = 1 << 1,
     // The occurrence is an override of the target base method.
     Override = 1 << 2,
   };
   struct Reference {
-    Location Loc;
+    ReferenceLocation Loc;
     unsigned Attributes = 0;
   };
   std::vector<Reference> References;
   bool HasMore = false;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                               const ReferencesResult::Reference &);
 
 /// Returns implementations at a specified \p Pos:
 ///   - overrides for a virtual method;
 ///   - 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);
 
 /// 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);
+                                const SymbolIndex *Index = nullptr,
+                                bool AddContext = false);
 
 /// Get info about symbols at \p Pos.
 std::vector<SymbolDetails> getSymbolInfo(ParsedAST &AST, Position Pos);
 
 /// Find the record types referenced at \p Pos.
 std::vector<const CXXRecordDecl *> findRecordTypeAt(ParsedAST &AST,
                                                     Position Pos);
 

clang-tools-extra/clangd/XRefs.cpp

 #include "Protocol.h"
 #include "Quality.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "URI.h"
 #include "index/Index.h"
 #include "index/Merge.h"
 #include "index/Relation.h"
+#include "index/SymbolCollector.h"
 #include "index/SymbolID.h"
 #include "index/SymbolLocation.h"
 #include "support/Logger.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Attrs.inc"
 #include "clang/AST/Decl.h"

 
 /// 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 Decl *Container;
 
     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) {
       // Sanity check: we expect the *first* token to match the reported loc.
       // Otherwise, maybe it was e.g. some other kind of reference to a Decl.
       if (!Locs.empty() && Locs.front() != Loc)
         Locs.clear(); // First token doesn't match, assume our guess was wrong.
     }
     if (Locs.empty())
       Locs.push_back(Loc);
 
+    SymbolCollector::Options CollectorOpts;
+    CollectorOpts.CollectMainFileSymbols = true;
     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(),
+             SymbolCollector::getRefContainer(ASTNode.Parent, CollectorOpts)});
     }
     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,
   if (!CMD)
     return;
   for (const CXXMethodDecl *Base : CMD->overridden_methods()) {
     if (auto ID = getSymbolID(Base))
       OverriddenMethods.insert(ID);
     getOverriddenMethods(Base, OverriddenMethods);
   }
 }
+
+llvm::Optional<std::string>
+stringifyContainerForMainFileRef(const Decl *Container) {
+  // FIXME We might also want to display the signature here
+  // When doing so, remember to also add the Signature to index results!
+  if (auto *ND = llvm::dyn_cast_if_present<NamedDecl>(Container))
+    return printQualifiedName(*ND);
+  return {};
+}
 } // namespace
 
 ReferencesResult findReferences(ParsedAST &AST, Position Pos, uint32_t Limit,
-                                const SymbolIndex *Index) {
+                                const SymbolIndex *Index, bool AddContext) {
   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 (AddContext)
+        Result.Loc.containerName =
+            stringifyContainerForMainFileRef(Ref.Container);
       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()) {
+      LookupRequest ContainerLookup;
+      llvm::DenseMap<SymbolID, size_t> RefIndexForContainer;
       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) {
           ReferencesResult::Reference Result;
-          Result.Loc = std::move(*LSPLocDecl);
+          Result.Loc = {std::move(*LSPLocDecl), llvm::None};
           Result.Attributes =
               ReferencesResult::Declaration | ReferencesResult::Override;
+          RefIndexForContainer.insert({Object.ID, Results.References.size()});
+          ContainerLookup.IDs.insert(Object.ID);
           Results.References.push_back(std::move(Result));
         }
         if (LSPLocDef) {
           ReferencesResult::Reference Result;
-          Result.Loc = std::move(*LSPLocDef);
+          Result.Loc = {std::move(*LSPLocDef), llvm::None};
           Result.Attributes = ReferencesResult::Declaration |
                               ReferencesResult::Definition |
                               ReferencesResult::Override;
+          RefIndexForContainer.insert({Object.ID, Results.References.size()});
+          ContainerLookup.IDs.insert(Object.ID);
           Results.References.push_back(std::move(Result));
         }
       });
+
+      if (!ContainerLookup.IDs.empty() && AddContext)
+        Index->lookup(ContainerLookup, [&](const Symbol &Container) {
+          auto Ref = RefIndexForContainer.find(Container.ID);
+          assert(Ref != RefIndexForContainer.end());
+          Results.References[Ref->getSecond()].Loc.containerName =
+              Container.Scope.str() + Container.Name.str();
+        });
     }
   }
   // 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);
+      Result.Loc = {std::move(*LSPLoc), llvm::None};
       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;
+    });
+
+    if (!ContainerLookup.IDs.empty() && AddContext)
+      Index->lookup(ContainerLookup, [&](const Symbol &Container) {
+        auto Ref = RefIndexForContainer.find(Container.ID);
+        assert(Ref != RefIndexForContainer.end());
+        Results.References[Ref->getSecond()].Loc.containerName =
+            Container.Scope.str() + Container.Name.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,
              /*AllowMainFileSymbols=*/true);

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

public:
 
   SymbolCollector(Options Opts);
   ~SymbolCollector();
 
   /// Returns true is \p ND should be collected.
   static bool shouldCollectSymbol(const NamedDecl &ND, const ASTContext &ASTCtx,
                                   const Options &Opts, bool IsMainFileSymbol);
 
+  // Given a ref contained in enclosing decl `Enclosing`, return
+  // the decl that should be used as that ref's Ref::Container. This is
+  // usually `Enclosing` itself, but in cases where `Enclosing` is not
+  // indexed, we walk further up because Ref::Container should always be
+  // an indexed symbol.
+  // Note: we don't use DeclContext as the container as in some cases
+  // it's useful to use a Decl which is not a DeclContext. For example,
+  // for a ref occurring in the initializer of a namespace-scope variable,
+  // it's useful to use that variable as the container, as otherwise the
+  // next enclosing DeclContext would be a NamespaceDecl or TranslationUnitDecl,
+  // which are both not indexed and less granular than we'd like for use cases
+  // like call hierarchy.
+  static const Decl *getRefContainer(const Decl *Enclosing,
+                                     const SymbolCollector::Options &Opts);
+
   void initialize(ASTContext &Ctx) override;
 
   void setPreprocessor(std::shared_ptr<Preprocessor> PP) override {
     this->PP = PP.get();
   }
   void setPreprocessor(Preprocessor &PP) { this->PP = &PP; }
 
   bool

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

 llvm::Optional<RelationKind> indexableRelation(const index::SymbolRelation &R) {
   if (R.Roles & static_cast<unsigned>(index::SymbolRole::RelationBaseOf))
     return RelationKind::BaseOf;
   if (R.Roles & static_cast<unsigned>(index::SymbolRole::RelationOverrideOf))
     return RelationKind::OverriddenBy;
   return None;
 }
 
-// Given a ref contained in enclosing decl `Enclosing`, return
-// the decl that should be used as that ref's Ref::Container. This is
-// usually `Enclosing` itself, but in cases where `Enclosing` is not
-// indexed, we walk further up because Ref::Container should always be
-// an indexed symbol.
-// Note: we don't use DeclContext as the container as in some cases
-// it's useful to use a Decl which is not a DeclContext. For example,
-// for a ref occurring in the initializer of a namespace-scope variable,
-// it's useful to use that variable as the container, as otherwise the
-// next enclosing DeclContext would be a NamespaceDecl or TranslationUnitDecl,
-// which are both not indexed and less granular than we'd like for use cases
-// like call hierarchy.
-const Decl *getRefContainer(const Decl *Enclosing,
-                            const SymbolCollector::Options &Opts) {
-  while (Enclosing) {
-    const auto *ND = dyn_cast<NamedDecl>(Enclosing);
-    if (ND && SymbolCollector::shouldCollectSymbol(*ND, ND->getASTContext(),
-                                                   Opts, true)) {
-      break;
-    }
-    Enclosing = dyn_cast_or_null<Decl>(Enclosing->getDeclContext());
-  }
-  return Enclosing;
-}
-
 // Check if there is an exact spelling of \p ND at \p Loc.
 bool isSpelled(SourceLocation Loc, const NamedDecl &ND) {
   auto Name = ND.getDeclName();
   const auto NameKind = Name.getNameKind();
   if (NameKind != DeclarationName::Identifier &&
       NameKind != DeclarationName::CXXConstructorName)
     return false;
   const auto &AST = ND.getASTContext();
if (isPrivateProtoDecl(ND))
     return false;
   if (!Opts.CollectReserved &&
       (hasReservedName(ND) || hasReservedScope(*ND.getDeclContext())))
     return false;
 
   return true;
 }
 
+const Decl *
+SymbolCollector::getRefContainer(const Decl *Enclosing,
+                                 const SymbolCollector::Options &Opts) {
+  while (Enclosing) {
+    const auto *ND = dyn_cast<NamedDecl>(Enclosing);
+    if (ND && shouldCollectSymbol(*ND, ND->getASTContext(), Opts, true)) {
+      break;
+    }
+    Enclosing = dyn_cast_or_null<Decl>(Enclosing->getDeclContext());
+  }
+  return Enclosing;
+}
+
 // Always return true to continue indexing.
 bool SymbolCollector::handleDeclOccurrence(
     const Decl *D, index::SymbolRoleSet Roles,
     llvm::ArrayRef<index::SymbolRelation> Relations, SourceLocation Loc,
     index::IndexDataConsumer::ASTNodeInfo ASTNode) {
   assert(ASTCtx && PP && HeaderFileURIs);
   assert(CompletionAllocator && CompletionTUInfo);
   assert(ASTNode.OrigD);

clang-tools-extra/clangd/test/references-container.test

+# RUN: clangd -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"clangd","capabilities":{"textDocument":{"references":{"container": true}}},"trace":"off"}}
+---
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///main.cpp",
+  "languageId":"cpp",
+  "version":1,
+  "text":"void foo();void bar() { foo(); }"
+}}}
+---
+{"jsonrpc":"2.0","id":1,"method":"textDocument/references","params":{
+  "textDocument":{"uri":"test:///main.cpp"},
+  "position":{"line":0,"character":6},
+  "context":{"includeDeclaration": false}
+}}
+#      CHECK:  "id": 1
+# CHECK-NEXT:  "jsonrpc": "2.0",
+# CHECK-NEXT:  "result": [
+# CHECK-NEXT:    {
+# CHECK-NEXT:      "containerName": "bar",
+# CHECK-NEXT:      "range": {
+# CHECK-NEXT:        "end": {
+# CHECK-NEXT:          "character": 27,
+# CHECK-NEXT:          "line": 0
+# CHECK-NEXT:        },
+# CHECK-NEXT:        "start": {
+# CHECK-NEXT:          "character": 24,
+# CHECK-NEXT:          "line": 0
+# CHECK-NEXT:        }
+# CHECK-NEXT:      },
+# CHECK-NEXT:      "uri": "{{.*}}/main.cpp"
+# CHECK-NEXT:    }
+# CHECK-NEXT:  ]
+---
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+---
+{"jsonrpc":"2.0","method":"exit"}

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

if (arg.Definition->range != *Def) {
     return false;
   }
   return true;
 }
 
 MATCHER_P(sym, Name, "") { return arg.Name == Name; }
 
 MATCHER_P(rangeIs, R, "") { return arg.Loc.range == R; }
+MATCHER_P(containerIs, C, "") {
+  return arg.Loc.containerName.value_or("") == C;
+}
 MATCHER_P(attrsAre, A, "") { return arg.Attributes == A; }
 MATCHER_P(hasID, ID, "") { return arg.ID == ID; }
 
 TEST(LocateSymbol, WithIndex) {
   Annotations SymbolHeader(R"cpp(
         class $forward[[Forward]];
         class $foo[[Foo]] {};
 

 
 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 : T.ranges())
-    ExpectedLocations.push_back(AllOf(rangeIs(R), attrsAre(0u)));
+  for (const auto &[R, Context] : T.rangesWithPayload())
+    ExpectedLocations.push_back(
+        AllOf(rangeIs(R), containerIs(Context), attrsAre(0u)));
   // $def is actually shorthand for both definition and declaration.
   // If we have cases that are definition-only, we should change this.
-  for (const auto &R : T.ranges("def"))
-    ExpectedLocations.push_back(
-        AllOf(rangeIs(R), attrsAre(ReferencesResult::Definition |
+  for (const auto &[R, Context] : T.rangesWithPayload("def"))
+    ExpectedLocations.push_back(AllOf(rangeIs(R), containerIs(Context),
+                                      attrsAre(ReferencesResult::Definition |
                                                ReferencesResult::Declaration)));
-  for (const auto &R : T.ranges("decl"))
-    ExpectedLocations.push_back(
-        AllOf(rangeIs(R), attrsAre(ReferencesResult::Declaration)));
-  for (const auto &R : T.ranges("overridedecl"))
+  for (const auto &[R, Context] : T.rangesWithPayload("decl"))
+    ExpectedLocations.push_back(AllOf(rangeIs(R), containerIs(Context),
+                                      attrsAre(ReferencesResult::Declaration)));
+  for (const auto &[R, Context] : T.rangesWithPayload("overridedecl"))
     ExpectedLocations.push_back(AllOf(
-        rangeIs(R),
+        rangeIs(R), containerIs(Context),
         attrsAre(ReferencesResult::Declaration | ReferencesResult::Override)));
-  for (const auto &R : T.ranges("overridedef"))
-    ExpectedLocations.push_back(
-        AllOf(rangeIs(R), attrsAre(ReferencesResult::Declaration |
+  for (const auto &[R, Context] : T.rangesWithPayload("overridedef"))
+    ExpectedLocations.push_back(AllOf(rangeIs(R), containerIs(Context),
+                                      attrsAre(ReferencesResult::Declaration |
                                                ReferencesResult::Definition |
                                                ReferencesResult::Override)));
   for (const auto &P : T.points()) {
-    EXPECT_THAT(findReferences(AST, P, 0, UseIndex ? TU.index().get() : nullptr)
+    EXPECT_THAT(findReferences(AST, P, 0, UseIndex ? TU.index().get() : nullptr,
+                               /*AddContext*/ true)
                     .References,
                 UnorderedElementsAreArray(ExpectedLocations))
         << "Failed for Refs at " << P << "\n"
         << Test;
   }
 }
 
 TEST(FindReferences, WithinAST) {
   const char *Tests[] = {
       R"cpp(// Local variable
         int main() {
-          int $def[[foo]];
-          [[^foo]] = 2;
-          int test1 = [[foo]];
+          int $def(main)[[foo]];
+          $(main)[[^foo]] = 2;
+          int test1 = $(main)[[foo]];
         }
       )cpp",
 
       R"cpp(// Struct
         namespace ns1 {
-        struct $def[[Foo]] {};
+        struct $def(ns1)[[Foo]] {};
         } // namespace ns1
         int main() {
-          ns1::[[Fo^o]]* Params;
+          ns1::$(main)[[Fo^o]]* Params;
         }
       )cpp",
 
       R"cpp(// Forward declaration
         class $decl[[Foo]];
         class $def[[Foo]] {};
         int main() {
-          [[Fo^o]] foo;
+          $(main)[[Fo^o]] foo;
         }
       )cpp",
 
       R"cpp(// Function
         int $def[[foo]](int) {}
         int main() {
-          auto *X = &[[^foo]];
-          [[foo]](42);
+          auto *X = &$(main)[[^foo]];
+          $(main)[[foo]](42);
         }
       )cpp",
 
       R"cpp(// Field
         struct Foo {
-          int $def[[foo]];
-          Foo() : [[foo]](0) {}
+          int $def(Foo)[[foo]];
+          Foo() : $(Foo::Foo)[[foo]](0) {}
         };
         int main() {
           Foo f;
-          f.[[f^oo]] = 1;
+          f.$(main)[[f^oo]] = 1;
         }
       )cpp",
 
       R"cpp(// Method call
-        struct Foo { int $decl[[foo]](); };
-        int Foo::$def[[foo]]() {}
+        struct Foo { int $decl(Foo)[[foo]](); };
+        int Foo::$def(Foo)[[foo]]() {}
         int main() {
           Foo f;
-          f.[[^foo]]();
+          f.$(main)[[^foo]]();
         }
       )cpp",
 
       R"cpp(// Constructor
         struct Foo {
-          $decl[[F^oo]](int);
+          $decl(Foo)[[F^oo]](int);
         };
         void foo() {
-          Foo f = [[Foo]](42);
+          Foo f = $(foo)[[Foo]](42);
         }
       )cpp",
 
       R"cpp(// Typedef
         typedef int $def[[Foo]];
         int main() {
-          [[^Foo]] bar;
+          $(main)[[^Foo]] bar;
         }
       )cpp",
 
       R"cpp(// Namespace
         namespace $decl[[ns]] { // FIXME: def?
         struct Foo {};
         } // namespace ns
-        int main() { [[^ns]]::Foo foo; }
+        int main() { $(main)[[^ns]]::Foo foo; }
       )cpp",
 
       R"cpp(// Macros
         #define TYPE(X) X
         #define FOO Foo
         #define CAT(X, Y) X##Y
         class $def[[Fo^o]] {};
         void test() {
-          TYPE([[Foo]]) foo;
-          [[FOO]] foo2;
-          TYPE(TYPE([[Foo]])) foo3;
-          [[CAT]](Fo, o) foo4;
+          TYPE($(test)[[Foo]]) foo;
+          $(test)[[FOO]] foo2;
+          TYPE(TYPE($(test)[[Foo]])) foo3;
+          $(test)[[CAT]](Fo, o) foo4;
         }
       )cpp",
 
       R"cpp(// Macros
         #define $def[[MA^CRO]](X) (X+1)
         void test() {
-          int x = [[MACRO]]([[MACRO]](1));
+          int x = $[[MACRO]]($[[MACRO]](1));
         }
       )cpp",
 
       R"cpp(// Macro outside preamble
         int breakPreamble;
         #define $def[[MA^CRO]](X) (X+1)
         void test() {
-          int x = [[MACRO]]([[MACRO]](1));
+          int x = $[[MACRO]]($[[MACRO]](1));
         }
       )cpp",
 
       R"cpp(
         int $def[[v^ar]] = 0;
-        void foo(int s = [[var]]);
+        void foo(int s = $(foo)[[var]]);
       )cpp",
 
       R"cpp(
        template <typename T>
        class $def[[Fo^o]] {};
-       void func([[Foo]]<int>);
+       void func($(func)[[Foo]]<int>);
       )cpp",
 
       R"cpp(
        template <typename T>
        class $def[[Foo]] {};
-       void func([[Fo^o]]<int>);
+       void func($(func)[[Fo^o]]<int>);
       )cpp",
       R"cpp(// Not touching any identifiers.
         struct Foo {
-          $def[[~]]Foo() {};
+          $def(Foo)[[~]]Foo() {};
         };
         void foo() {
           Foo f;
-          f.[[^~]]Foo();
+          f.$(foo)[[^~]]Foo();
         }
       )cpp",
       R"cpp(// Lambda capture initializer
         void foo() {
-          int $def[[w^aldo]] = 42;
-          auto lambda = [x = [[waldo]]](){};
+          int $def(foo)[[w^aldo]] = 42;
+          auto lambda = [x = $(foo)[[waldo]]](){};
         }
       )cpp",
       R"cpp(// Renaming alias
         template <typename> class Vector {};
         using $def[[^X]] = Vector<int>;
-        [[X]] x1;
+        $(x1)[[X]] x1;
         Vector<int> x2;
         Vector<double> y;
       )cpp",
       R"cpp(// Dependent code
         template <typename T> void $decl[[foo]](T t);
-        template <typename T> void bar(T t) { [[foo]](t); } // foo in bar is uninstantiated.
-        void baz(int x) { [[f^oo]](x); }
+        template <typename T> void bar(T t) { $(bar)[[foo]](t); } // foo in bar is uninstantiated. 
+        void baz(int x) { $(baz)[[f^oo]](x); }
       )cpp",
       R"cpp(
         namespace ns {
         struct S{};
-        void $decl[[foo]](S s);
+        void $decl(ns)[[foo]](S s);
         } // namespace ns
         template <typename T> void foo(T t);
         // FIXME: Maybe report this foo as a ref to ns::foo (because of ADL)
         // when bar<ns::S> is instantiated?
         template <typename T> void bar(T t) { foo(t); }
         void baz(int x) {
           ns::S s;
           bar<ns::S>(s);
-          [[f^oo]](s);
+          $(baz)[[f^oo]](s);
         }
       )cpp",
       R"cpp(// unresolved member expression
         struct Foo {
-          template <typename T> void $decl[[b^ar]](T t); 
+          template <typename T> void $decl(Foo)[[b^ar]](T t);
         };
         template <typename T> void test(Foo F, T t) {
-          F.[[bar]](t);
+          F.$(test)[[bar]](t);
         }
       )cpp",
 
       // Enum base
       R"cpp(
         typedef int $def[[MyTypeD^ef]];
-        enum MyEnum : [[MyTy^peDef]] { };
+        enum MyEnum : $(MyEnum)[[MyTy^peDef]] { };
       )cpp",
       R"cpp(
         typedef int $def[[MyType^Def]];
-        enum MyEnum : [[MyTypeD^ef]];
+        enum MyEnum : $(MyEnum)[[MyTypeD^ef]];
       )cpp",
       R"cpp(
         using $def[[MyTypeD^ef]] = int;
-        enum MyEnum : [[MyTy^peDef]] { };
+        enum MyEnum : $(MyEnum)[[MyTy^peDef]] { };
       )cpp",
   };
   for (const char *Test : Tests)
     checkFindRefs(Test);
 }
 
 TEST(FindReferences, ConceptsWithinAST) {
   constexpr llvm::StringLiteral Code = R"cpp(
     template <class T>
     concept $def[[IsSmal^l]] = sizeof(T) <= 8;
 
     template <class T>
     concept IsSmallPtr = requires(T x) {
-      { *x } -> [[IsSmal^l]];
+      { *x } -> $(IsSmallPtr)[[IsSmal^l]];
     };
 
-    [[IsSmall]] auto i = 'c';
-    template<[[IsSmal^l]] U> void foo();
-    template<class U> void bar() requires [[IsSmal^l]]<U>;
-    template<class U> requires [[IsSmal^l]]<U> void baz();
+    $(i)[[IsSmall]] auto i = 'c';
+    template<$(foo)[[IsSmal^l]] U> void foo();
+    template<class U> void bar() requires $(bar)[[IsSmal^l]]<U>;
+    template<class U> requires $(baz)[[IsSmal^l]]<U> void baz();
     static_assert([[IsSma^ll]]<char>);
   )cpp";
   checkFindRefs(Code);
 }
 
 TEST(FindReferences, ConceptReq) {
   constexpr llvm::StringLiteral Code = R"cpp(
     template <class T>
     concept $def[[IsSmal^l]] = sizeof(T) <= 8;
 
     template <class T>
     concept IsSmallPtr = requires(T x) {
-      { *x } -> [[IsSmal^l]];
+      { *x } -> $(IsSmallPtr)[[IsSmal^l]];
     };
   )cpp";
   checkFindRefs(Code);
 }
 
 TEST(FindReferences, RequiresExprParameters) {
   constexpr llvm::StringLiteral Code = R"cpp(
     template <class T>
     concept IsSmall = sizeof(T) <= 8;
 
     template <class T>
     concept IsSmallPtr = requires(T $def[[^x]]) {
-      { *[[^x]] } -> IsSmall;
+      { *$(IsSmallPtr)[[^x]] } -> IsSmall;
     };
   )cpp";
   checkFindRefs(Code);
 }
 
 TEST(FindReferences, IncludeOverrides) {
   llvm::StringRef Test =
       R"cpp(
         class Base {
         public:
-          virtu^al void $decl[[f^unc]]() ^= ^0;
+          virtu^al void $decl(Base)[[f^unc]]() ^= ^0;
         };
         class Derived : public Base {
         public:
-          void $overridedecl[[func]]() override;
+          void $overridedecl(Derived::func)[[func]]() override;
         };
         void Derived::$overridedef[[func]]() {}
         class Derived2 : public Base {
-          void $overridedef[[func]]() override {}
+          void $overridedef(Derived2::func)[[func]]() override {}
         };
-        void test(Derived* D) {
+        void test(Derived* D, Base* B) {
           D->func();  // No references to the overrides.
         })cpp";
   checkFindRefs(Test, /*UseIndex=*/true);
 }
 
 TEST(FindReferences, RefsToBaseMethod) {
   llvm::StringRef Test =
       R"cpp(
         class BaseBase {
         public:
-          virtual void [[func]]();
+          virtual void $(BaseBase)[[func]]();
         };
         class Base : public BaseBase {
         public:
-          void [[func]]() override;
+          void $(Base)[[func]]() override;
         };
         class Derived : public Base {
         public:
-          void $decl[[fu^nc]]() over^ride;
+          void $decl(Derived)[[fu^nc]]() over^ride;
         };
         void test(BaseBase* BB, Base* B, Derived* D) {
           // refs to overridden methods in complete type hierarchy are reported.
-          BB->[[func]]();
-          B->[[func]]();
-          D->[[fu^nc]]();
+          // FIXME BB->func() gets no context here, this seems to be a bug (or feature?)
+          // in the index for multiple levels of inheritance
+          BB->$[[func]]();
+          B->$(test)[[func]]();
+          D->$(test)[[fu^nc]]();
         })cpp";
   checkFindRefs(Test, /*UseIndex=*/true);
 }
 
 TEST(FindReferences, MainFileReferencesOnly) {
   llvm::StringRef Test =
       R"cpp(
         void test() {
-          int [[fo^o]] = 1;
+          int $(test)[[fo^o]] = 1;
           // refs not from main file should not be included.
           #include "foo.inc"
         })cpp";
 
   Annotations Code(Test);
   auto TU = TestTU::withCode(Code.code());
   TU.AdditionalFiles["foo.inc"] = R"cpp(
       foo = 3;
     )cpp";
   auto AST = TU.build();
 
   std::vector<Matcher<ReferencesResult::Reference>> ExpectedLocations;
   for (const auto &R : Code.ranges())
     ExpectedLocations.push_back(rangeIs(R));
   EXPECT_THAT(findReferences(AST, Code.point(), 0).References,
               ElementsAreArray(ExpectedLocations))
       << Test;
 }
 
 TEST(FindReferences, ExplicitSymbols) {
   const char *Tests[] = {
       R"cpp(
-      struct Foo { Foo* $decl[[self]]() const; };
+      struct Foo { Foo* $decl(Foo)[[self]]() const; };
       void f() {
         Foo foo;
-        if (Foo* T = foo.[[^self]]()) {} // Foo member call expr.
+        if (Foo* T = foo.$(f)[[^self]]()) {} // Foo member call expr.
       }
       )cpp",
 
       R"cpp(
       struct Foo { Foo(int); };
       Foo f() {
-        int $def[[b]];
-        return [[^b]]; // Foo constructor expr.
+        int $def(f)[[b]];
+        return $(f)[[^b]]; // Foo constructor expr.
       }
       )cpp",
 
       R"cpp(
       struct Foo {};
       void g(Foo);
       Foo $decl[[f]]();
       void call() {
-        g([[^f]]());  // Foo constructor expr.
+        g($(call)[[^f]]());  // Foo constructor expr.
       }
       )cpp",
 
       R"cpp(
       void $decl[[foo]](int);
       void $decl[[foo]](double);
 
       namespace ns {
-      using ::$decl[[fo^o]];
+      using ::$decl(ns)[[fo^o]];
       }
       )cpp",
 
       R"cpp(
       struct X {
         operator bool();
       };
 
       int test() {
-        X $def[[a]];
-        [[a]].operator bool();
-        if ([[a^]]) {} // ignore implicit conversion-operator AST node
+        X $def(test)[[a]];
+        $(test)[[a]].operator bool();
+        if ($(test)[[a^]]) {} // ignore implicit conversion-operator AST node
       }
     )cpp",
   };
   for (const char *Test : Tests)
     checkFindRefs(Test);
 }
 
 TEST(FindReferences, NeedsIndexForSymbols) {
   const char *Header = "int foo();";
   Annotations Main("int main() { [[f^oo]](); }");
   TestTU TU;
   TU.Code = std::string(Main.code());
   TU.HeaderCode = Header;
   auto AST = TU.build();
 
   // References in main file are returned without index.
   EXPECT_THAT(
       findReferences(AST, Main.point(), 0, /*Index=*/nullptr).References,
       ElementsAre(rangeIs(Main.range())));
   Annotations IndexedMain(R"cpp(
-    int [[foo]]() { return 42; }
+    int $decl[[foo]]() { return 42; }
+    void bar() { $bar(bar)[[foo]](); }
+    struct S { void bar() { $S(S::bar)[[foo]](); } };
+    namespace N { void bar() { $N(N::bar)[[foo]](); } }
   )cpp");
 
   // References from indexed files are included.
   TestTU IndexedTU;
   IndexedTU.Code = std::string(IndexedMain.code());
   IndexedTU.Filename = "Indexed.cpp";
   IndexedTU.HeaderCode = Header;
   EXPECT_THAT(
-      findReferences(AST, Main.point(), 0, IndexedTU.index().get()).References,
-      ElementsAre(rangeIs(Main.range()),
-                  AllOf(rangeIs(IndexedMain.range()),
-                        attrsAre(ReferencesResult::Declaration |
-                                 ReferencesResult::Definition))));
+      findReferences(AST, Main.point(), 0, IndexedTU.index().get(),
+                     /*AddContext*/ true)
+          .References,
+      ElementsAre(
+          rangeIs(Main.range()),
+          AllOf(rangeIs(IndexedMain.range("decl")),
+                         ReferencesResult::Definition)),
+          AllOf(rangeIs(IndexedMain.range("bar")), containerIs("bar")),
+          AllOf(rangeIs(IndexedMain.range("S")), containerIs("S::bar")),
+          AllOf(rangeIs(IndexedMain.range("N")), containerIs("N::bar"))));
   auto LimitRefs =
       findReferences(AST, Main.point(), /*Limit*/ 1, IndexedTU.index().get());
   EXPECT_EQ(1u, LimitRefs.References.size());
   EXPECT_TRUE(LimitRefs.HasMore);
 
   // Avoid indexed results for the main file. Use AST for the mainfile.
   TU.Code = ("\n\n" + Main.code()).str();
   EXPECT_THAT(findReferences(AST, Main.point(), 0, TU.index().get()).References,
               ElementsAre(rangeIs(Main.range())));
 }
 
 TEST(FindReferences, NeedsIndexForMacro) {
   const char *Header = "#define MACRO(X) (X+1)";
   Annotations Main(R"cpp(
     int main() {
       int a = [[MA^CRO]](1);
     }
   )cpp");
   TestTU TU;
   TU.Code = std::string(Main.code());
   TU.HeaderCode = Header;
   auto AST = TU.build();
 
   // References in main file are returned without index.
-  EXPECT_THAT(
-      findReferences(AST, Main.point(), 0, /*Index=*/nullptr).References,
+  EXPECT_THAT(findReferences(AST, Main.point(), 0, /*Index=*/nullptr,
+                             /*AddContext*/ true)
+                  .References,
               ElementsAre(rangeIs(Main.range())));
 
   Annotations IndexedMain(R"cpp(
     int indexed_main() {
       int a = [[MACRO]](1);
     }
   )cpp");
 
   // References from indexed files are included.
   TestTU IndexedTU;
   IndexedTU.Code = std::string(IndexedMain.code());
   IndexedTU.Filename = "Indexed.cpp";
   IndexedTU.HeaderCode = Header;
   EXPECT_THAT(
       findReferences(AST, Main.point(), 0, IndexedTU.index().get()).References,
       ElementsAre(rangeIs(Main.range()), rangeIs(IndexedMain.range())));
-  auto LimitRefs =
-      findReferences(AST, Main.point(), /*Limit*/ 1, IndexedTU.index().get());
+  auto LimitRefs = findReferences(AST, Main.point(), /*Limit*/ 1,
+                                  IndexedTU.index().get(), /*AddContext*/ true);
   EXPECT_EQ(1u, LimitRefs.References.size());
   EXPECT_TRUE(LimitRefs.HasMore);
 }
 
 TEST(FindReferences, NoQueryForLocalSymbols) {
   struct RecordingIndex : public MemIndex {
     mutable Optional<llvm::DenseSet<SymbolID>> RefIDs;
     bool refs(const RefsRequest &Req,