[clangd] [WIP] Find references of local symbols

[malaperle]

Summary

We do not have a global index of references but we can find the references
of local symbols within the AST in the mean time. Also, since we will not
record local symbol references in the index, we will need that logic anyway.

Signed-off-by: Marc-Andre Laperle <[e-mail]>

[malaperle]

Fix silly bug I introduced in last minute clean-up.

[malaperle]

Needs tests.

clangd/ClangdLSPServer.h

private:
   void onSwitchSourceHeader(TextDocumentIdentifier &Params) override;
   void onDocumentHighlight(TextDocumentPositionParams &Params) override;
   void onFileEvent(DidChangeWatchedFilesParams &Params) override;
   void onCommand(ExecuteCommandParams &Params) override;
   void onWorkspaceSymbol(WorkspaceSymbolParams &Params) override;
   void onRename(RenameParams &Parames) override;
   void onHover(TextDocumentPositionParams &Params) override;
   void onChangeConfiguration(DidChangeConfigurationParams &Params) override;
+  void onReferences(ReferenceParams &Params) override;
 
   std::vector<Fix> getFixes(StringRef File, const clangd::Diagnostic &D);
 
   /// Forces a reparse of all currently opened files.  As a result, this method
   /// may be very expensive.  This method is normally called when the
   /// compilation database is changed.
   void reparseOpenedFiles();
 

clangd/ClangdLSPServer.cpp

reply(json::Object{
             {"hoverProvider", true},
             {"renameProvider", true},
             {"documentSymbolProvider", true},
             {"workspaceSymbolProvider", true},
             {"executeCommandProvider",
              json::Object{
                  {"commands", {ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND}},
              }},
+            {"referencesProvider", true},
         }}}});
 }
 
 void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
   // Do essentially nothing, just say we're ready to exit.
   ShutdownRequestReceived = true;
   reply(nullptr);
 }
if (Settings.compilationDatabasePath.hasValue()) {
     NonCachedCDB.setCompileCommandsDir(
         Settings.compilationDatabasePath.getValue());
     CDB.clear();
 
     reparseOpenedFiles();
   }
 }
 
+void ClangdLSPServer::onReferences(ReferenceParams &Params) {
+  Server.findReferences(Params.textDocument.uri.file(), Params.position,
+                        Params.context.includeDeclaration,
+                        [](llvm::Expected<std::vector<Location>> Items) {
+                          if (!Items)
+                            return replyError(
+                                ErrorCode::InvalidParams,
+                                llvm::toString(Items.takeError()));
+                          reply(json::Array(*Items));
+                        });
+}
+
 ClangdLSPServer::ClangdLSPServer(JSONOutput &Out,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  const ClangdServer::Options &Opts)
     : Out(Out), NonCachedCDB(std::move(CompileCommandsDir)), CDB(NonCachedCDB),
       CCOpts(CCOpts), SupportedSymbolKinds(defaultSymbolKinds()),
       Server(CDB, FSProvider, /*DiagConsumer=*/*this, Opts) {}
 

clangd/ClangdServer.h

public:
   /// Provide signature help for \p File at \p Pos.  This method should only be
   /// called for tracked files.
   void signatureHelp(PathRef File, Position Pos, Callback<SignatureHelp> CB);
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
   void findDefinitions(PathRef File, Position Pos,
                        Callback<std::vector<Location>> CB);
 
+  /// Get references of symbol at a specified \p Line and
+  /// \p Column in \p File.
+  void findReferences(PathRef File, Position Pos, bool IncludeDeclaration,
+                      Callback<std::vector<Location>> CB);
+
   /// Helper function that returns a path to the corresponding source file when
   /// given a header file and vice versa.
   llvm::Optional<Path> switchSourceHeader(PathRef Path);
 
   /// Get document highlights for a given position.
   void findDocumentHighlights(PathRef File, Position Pos,
                               Callback<std::vector<DocumentHighlight>> CB);
 

clangd/ClangdServer.cpp

 }
 
 void ClangdServer::workspaceSymbols(
     StringRef Query, int Limit, Callback<std::vector<SymbolInformation>> CB) {
   CB(clangd::getWorkspaceSymbols(Query, Limit, Index,
                                  RootPath ? *RootPath : ""));
 }
 
+void ClangdServer::findReferences(PathRef File, Position Pos,
+                                  bool IncludeDeclaration,
+                                  Callback<std::vector<Location>> CB) {
+  auto FS = FSProvider.getFileSystem();
+  auto Action = [this, Pos, FS,
+                 IncludeDeclaration](Callback<std::vector<Location>> CB,
+                                     llvm::Expected<InputsAndAST> InpAST) {
+    if (!InpAST)
+      return CB(InpAST.takeError());
+    CB(clangd::findReferences(InpAST->AST, Pos, IncludeDeclaration));
+  };
+
+  WorkScheduler.runWithAST("Definitions", File, Bind(Action, std::move(CB)));
+}
+
 void ClangdServer::documentSymbols(
     StringRef File, Callback<std::vector<SymbolInformation>> CB) {
   auto Action = [](Callback<std::vector<SymbolInformation>> CB,
                    llvm::Expected<InputsAndAST> InpAST) {
     if (!InpAST)
       return CB(InpAST.takeError());
     CB(clangd::getDocumentSymbols(InpAST->AST));
   };

clangd/Protocol.h

struct Hover {
   MarkupContent contents;
 
   /// An optional range is a range inside a text document
   /// that is used to visualize a hover, e.g. by changing the background color.
   llvm::Optional<Range> range;
 };
 llvm::json::Value toJSON(const Hover &H);
 
+struct ReferenceContext {
+  /// Include the declaration of the current symbol.
+  bool includeDeclaration;
+};
+bool fromJSON(const llvm::json::Value &Params, ReferenceContext &R);
+
+struct ReferenceParams : public TextDocumentPositionParams {
+  ReferenceContext context;
+};
+bool fromJSON(const llvm::json::Value &Params, ReferenceParams &R);
+
 /// The kind of a completion entry.
 enum class CompletionItemKind {
   Missing = 0,
   Text = 1,
   Method = 2,
   Function = 3,
   Constructor = 4,
   Field = 5,

clangd/Protocol.cpp

if (CI.insertTextFormat != InsertTextFormat::Missing)
     Result["insertTextFormat"] = static_cast<int>(CI.insertTextFormat);
   if (CI.textEdit)
     Result["textEdit"] = *CI.textEdit;
   if (!CI.additionalTextEdits.empty())
     Result["additionalTextEdits"] = json::Array(CI.additionalTextEdits);
   return std::move(Result);
 }
 
+bool fromJSON(const json::Value &Params, ReferenceContext &R) {
+  json::ObjectMapper O(Params);
+  return O && O.map("includeDeclaration", R.includeDeclaration);
+}
+
+bool fromJSON(const json::Value &Params, ReferenceParams &R) {
+  json::ObjectMapper O(Params);
+  return O && O.map("context", R.context) &&
+         O.map("textDocument", R.textDocument) && O.map("position", R.position);
+}
+
 llvm::raw_ostream &operator<<(llvm::raw_ostream &O, const CompletionItem &I) {
   O << I.label << " - " << toJSON(I);
   return O;
 }
 
 bool operator<(const CompletionItem &L, const CompletionItem &R) {
   return (L.sortText.empty() ? L.label : L.sortText) <
          (R.sortText.empty() ? R.label : R.sortText);

clangd/ProtocolHandlers.h

public:
   virtual void onSwitchSourceHeader(TextDocumentIdentifier &Params) = 0;
   virtual void onFileEvent(DidChangeWatchedFilesParams &Params) = 0;
   virtual void onCommand(ExecuteCommandParams &Params) = 0;
   virtual void onWorkspaceSymbol(WorkspaceSymbolParams &Params) = 0;
   virtual void onRename(RenameParams &Parames) = 0;
   virtual void onDocumentHighlight(TextDocumentPositionParams &Params) = 0;
   virtual void onHover(TextDocumentPositionParams &Params) = 0;
   virtual void onChangeConfiguration(DidChangeConfigurationParams &Params) = 0;
+  virtual void onReferences(ReferenceParams &Params) = 0;
 };
 
 void registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
                               ProtocolCallbacks &Callbacks);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ProtocolHandlers.cpp

void clangd::registerCallbackHandlers(JSONRPCDispatcher &Dispatcher,
   Register("textDocument/documentSymbol", &ProtocolCallbacks::onDocumentSymbol);
   Register("workspace/didChangeWatchedFiles", &ProtocolCallbacks::onFileEvent);
   Register("workspace/executeCommand", &ProtocolCallbacks::onCommand);
   Register("textDocument/documentHighlight",
            &ProtocolCallbacks::onDocumentHighlight);
   Register("workspace/didChangeConfiguration",
            &ProtocolCallbacks::onChangeConfiguration);
   Register("workspace/symbol", &ProtocolCallbacks::onWorkspaceSymbol);
+  Register("textDocument/references", &ProtocolCallbacks::onReferences);
 }

clangd/XRefs.h

 /// Get definition of symbol at a specified \p Pos.
 std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
                                       const SymbolIndex *Index = nullptr);
 
 /// Returns highlights for all usages of a symbol at \p Pos.
 std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
                                                       Position Pos);
 
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     bool IncludeDeclaration);
+
 /// Get the hover information when hovering at \p Pos.
 llvm::Optional<Hover> getHover(ParsedAST &AST, Position Pos);
 
 } // namespace clangd
 } // namespace clang
 #endif

clangd/XRefs.cpp

 llvm::Optional<SymbolID> getSymbolID(const Decl *D) {
   llvm::SmallString<128> USR;
   if (index::generateUSRForDecl(D, USR)) {
     return None;
   }
   return SymbolID(USR);
 }
 
+/// Finds declarations locations that a given source Decl refers to, in the main
+/// file.
+class ReferenceLocationsFinder : public index::IndexDataConsumer {
+  std::vector<Location> ReferenceLocations;
+  ParsedAST &AST;
+  const Decl *ReferencedDecl;
+  index::SymbolRoleSet InterestingRoleSet;
+
+public:
+  ReferenceLocationsFinder(ParsedAST &AST, const Decl *D,
+                           bool IncludeDeclaration)
+      : AST(AST), ReferencedDecl(D),
+        InterestingRoleSet(
+            static_cast<index::SymbolRoleSet>(index::SymbolRole::Reference)) {
+    if (IncludeDeclaration)
+      InterestingRoleSet |=
+          static_cast<index::SymbolRoleSet>(index::SymbolRole::Declaration) |
+          static_cast<index::SymbolRoleSet>(index::SymbolRole::Definition);
+  }
+
+  std::vector<Location> takeLocations() {
+    // Don't keep the same location multiple times.
+    // This can happen when nodes in the AST are visited twice.
+    std::sort(ReferenceLocations.begin(), ReferenceLocations.end());
+    auto last =
+        std::unique(ReferenceLocations.begin(), ReferenceLocations.end());
+    ReferenceLocations.erase(last, ReferenceLocations.end());
+    return std::move(ReferenceLocations);
+  }
+
+  bool
+  handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
+                      ArrayRef<index::SymbolRelation> Relations,
+                      SourceLocation Loc,
+                      index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
+    const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+    if (D != ReferencedDecl || !SourceMgr.isWrittenInMainFile(Loc)) {
+      return true;
+    }
+
+    // The end loc is adjusted in makeLocation with getLocForEndOfToken.
+    SourceRange Range(Loc, Loc);
+
+    if (Roles & InterestingRoleSet) {
+      auto L = makeLocation(AST, Range);
+      if (L)
+        ReferenceLocations.push_back(*L);
+    }
+    return true;
+  }
+};
+
 } // namespace
 
 std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
                                       const SymbolIndex *Index) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
 
   std::vector<Location> Result;
   // Handle goto definition for #include.
for (auto It : ResultCandidates) {
     if (Candidate.Decl &&
         Candidate.Decl != Candidate.Def) // Decl and Def might be the same
       Result.push_back(*Candidate.Decl);
   }
 
   return Result;
 }
 
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     bool IncludeDeclaration) {
+  SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+  SourceLocation SourceLocationBeg =
+      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
+
+  DeclarationAndMacrosFinder DeclMacrosFinder(llvm::errs(), SourceLocationBeg,
+                                              AST.getASTContext(),
+                                              AST.getPreprocessor());
+  index::IndexingOptions IndexOpts;
+  IndexOpts.SystemSymbolFilter =
+      index::IndexingOptions::SystemSymbolFilterKind::All;
+  IndexOpts.IndexFunctionLocals = true;
+  indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
+                     DeclMacrosFinder, IndexOpts);
+  std::vector<const Decl *> Decls = DeclMacrosFinder.takeDecls();
+  if (Decls.empty())
+    return {};
+
+  // Make CXXConstructorDecl lower priority. For example:
+  // MyClass Obj;
+  // We likely want to find the references to Obj not MyClass()
+  std::sort(Decls.begin(), Decls.end(), [](const Decl *&D1, const Decl *&D2) {
+    return !dyn_cast<CXXConstructorDecl>(D1);
+  });
+
+  const Decl *D = Decls[0];
+
+  if (!index::isFunctionLocalSymbol(D)) {
+    return {};
+  }
+
+  ReferenceLocationsFinder ReferencesFinder(AST, D, IncludeDeclaration);
+  indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
+                     ReferencesFinder, IndexOpts);
+
+  return ReferencesFinder.takeLocations();
+}
+
 namespace {
 
 /// Finds document highlights that a given list of declarations refers to.
 class DocumentHighlightsFinder : public index::IndexDataConsumer {
   std::vector<const Decl *> &Decls;
   std::vector<DocumentHighlight> DocumentHighlights;
   const ASTContext &AST;
 

test/clangd/initialize-params-invalid.test

 # CHECK-NEXT:      "documentRangeFormattingProvider": true,
 # CHECK-NEXT:      "documentSymbolProvider": true,
 # CHECK-NEXT:      "executeCommandProvider": {
 # CHECK-NEXT:        "commands": [
 # CHECK-NEXT:          "clangd.applyFix"
 # CHECK-NEXT:        ]
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "hoverProvider": true,
+# CHECK-NEXT:      "referencesProvider": true,
 # CHECK-NEXT:      "renameProvider": true,
 # CHECK-NEXT:      "signatureHelpProvider": {
 # CHECK-NEXT:        "triggerCharacters": [
 # CHECK-NEXT:          "(",
 # CHECK-NEXT:          ","
 # CHECK-NEXT:        ]
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "textDocumentSync": 2,
 # CHECK-NEXT:      "workspaceSymbolProvider": true
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }
 ---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

test/clangd/initialize-params.test

 # CHECK-NEXT:      "documentRangeFormattingProvider": true,
 # CHECK-NEXT:      "documentSymbolProvider": true,
 # CHECK-NEXT:      "executeCommandProvider": {
 # CHECK-NEXT:        "commands": [
 # CHECK-NEXT:          "clangd.applyFix"
 # CHECK-NEXT:        ]
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "hoverProvider": true,
+# CHECK-NEXT:      "referencesProvider": true,
 # CHECK-NEXT:      "renameProvider": true,
 # CHECK-NEXT:      "signatureHelpProvider": {
 # CHECK-NEXT:        "triggerCharacters": [
 # CHECK-NEXT:          "(",
 # CHECK-NEXT:          ","
 # CHECK-NEXT:        ]
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "textDocumentSync": 2,