From 664f952f387cb725faaef17d9b070374b4e46bc3 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Tue, 5 Nov 2024 08:56:57 -0800
Subject: [PATCH] Implement textDocument/switchSourceheader

When the current file is X.cc, there might be multiple X.h.
Use a heuristic to find the best X.h.

Vote for each interesting symbol's definitions (for header) or
declarations (for non-header). Select the file with the most votes.

If `file_id2cnt` is empty, use a simpler heuristic.
---
 src/message_handler.cc                |  1 +
 src/message_handler.hh                |  1 +
 src/messages/textDocument_document.cc | 69 +++++++++++++++++++++++++++
 src/query.hh                          |  2 +-
 4 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/src/message_handler.cc b/src/message_handler.cc
index efc3ba60..4b796bdf 100644
--- a/src/message_handler.cc
+++ b/src/message_handler.cc
@@ -216,6 +216,7 @@ MessageHandler::MessageHandler() {
   bind("textDocument/semanticTokens/full", &MessageHandler::textDocument_semanticTokensFull);
   bind("textDocument/semanticTokens/range", &MessageHandler::textDocument_semanticTokensRange);
   bind("textDocument/signatureHelp", &MessageHandler::textDocument_signatureHelp);
+  bind("textDocument/switchSourceHeader", &MessageHandler::textDocument_switchSourceHeader);
   bind("textDocument/typeDefinition", &MessageHandler::textDocument_typeDefinition);
   bind("workspace/didChangeConfiguration", &MessageHandler::workspace_didChangeConfiguration);
   bind("workspace/didChangeWatchedFiles", &MessageHandler::workspace_didChangeWatchedFiles);
diff --git a/src/message_handler.hh b/src/message_handler.hh
index 07044c68..8f09dc7c 100644
--- a/src/message_handler.hh
+++ b/src/message_handler.hh
@@ -311,6 +311,7 @@ private:
   void textDocument_semanticTokensFull(TextDocumentParam &, ReplyOnce &);
   void textDocument_semanticTokensRange(SemanticTokensRangeParams &, ReplyOnce &);
   void textDocument_signatureHelp(TextDocumentPositionParam &, ReplyOnce &);
+  void textDocument_switchSourceHeader(TextDocumentIdentifier &, ReplyOnce &);
   void textDocument_typeDefinition(TextDocumentPositionParam &, ReplyOnce &);
   void workspace_didChangeConfiguration(EmptyParam &);
   void workspace_didChangeWatchedFiles(DidChangeWatchedFilesParam &);
diff --git a/src/messages/textDocument_document.cc b/src/messages/textDocument_document.cc
index 5d34668c..d43a9c22 100644
--- a/src/messages/textDocument_document.cc
+++ b/src/messages/textDocument_document.cc
@@ -3,10 +3,17 @@
 
 #include "message_handler.hh"
 #include "pipeline.hh"
+#include "project.hh"
 #include "query.hh"
 
+#include <llvm/ADT/STLExtras.h>
+#include <llvm/ADT/StringRef.h>
+#include <llvm/Support/Path.h>
+
 #include <algorithm>
 
+using namespace llvm;
+
 MAKE_HASHABLE(ccls::SymbolIdx, t.usr, t.kind);
 
 namespace ccls {
@@ -226,4 +233,66 @@ void MessageHandler::textDocument_documentSymbol(JsonReader &reader,
     reply(result);
   }
 }
+
+void MessageHandler::textDocument_switchSourceHeader(TextDocumentIdentifier &param, ReplyOnce &reply) {
+  QueryFile *file;
+  WorkingFile *wf;
+  std::tie(file, wf) = findOrFail(param.uri.getPath(), reply);
+  if (!wf)
+    return reply(JsonNull{});
+  int file_id = file->id;
+
+  DocumentUri result;
+  const std::string &path = wf->filename;
+  bool is_hdr = lookupExtension(path).second;
+
+  // Vote for each interesting symbol's definitions (for header) or declarations (for non-header).
+  // Select the file with the most votes.
+  // Ignore Type symbols to skip class forward declarations and namespaces.
+  std::unordered_map<int, int> file_id2cnt;
+  for (auto [sym, refcnt] : file->symbol2refcnt) {
+    if (refcnt <= 0 || !sym.extent.valid() || sym.kind == Kind::Type)
+      continue;
+
+    if (is_hdr) {
+      withEntity(db, sym, [&](const auto &entity) {
+        for (auto &def : entity.def)
+          if (def.spell && def.file_id != file_id)
+            ++file_id2cnt[def.file_id];
+      });
+    } else {
+      for (DeclRef dr : getNonDefDeclarations(db, sym))
+        if (dr.file_id != file_id)
+          ++file_id2cnt[dr.file_id];
+    }
+  }
+  if (file_id2cnt.size()) {
+    auto best = file_id2cnt.begin();
+    for (auto it = file_id2cnt.begin(); it != file_id2cnt.end(); ++it)
+      if (it->second > best->second || (it->second == best->second && it->first < best->first))
+        best = it;
+    if (auto &def = db->files[best->first].def)
+      return reply(DocumentUri::fromPath(def->path));
+  }
+
+  if (is_hdr) {
+    // Check if `path` is in a #include entry.
+    for (QueryFile &file1 : db->files) {
+      auto &def = file1.def;
+      if (!def || lookupExtension(def->path).second)
+        continue;
+      for (IndexInclude &include : def->includes)
+        if (path == include.resolved_path)
+          return reply(DocumentUri::fromPath(def->path));
+    }
+    return reply(JsonNull{});
+  }
+
+  // Otherwise, find the #include with the same stem.
+  StringRef stem = sys::path::stem(path);
+  for (IndexInclude &include : file->def->includes)
+    if (sys::path::stem(include.resolved_path) == stem)
+      return reply(DocumentUri::fromPath(std::string(include.resolved_path)));
+  reply(JsonNull{});
+}
 } // namespace ccls
diff --git a/src/query.hh b/src/query.hh
index cd3a19c1..f4a72197 100644
--- a/src/query.hh
+++ b/src/query.hh
@@ -144,7 +144,7 @@ using Lid2file_id = std::unordered_map<int, int>;
 // The query database is heavily optimized for fast queries. It is stored
 // in-memory.
 struct DB {
-  std::vector<QueryFile> files;
+  llvm::SmallVector<QueryFile, 0> files;
   llvm::StringMap<int> name2file_id;
   llvm::DenseMap<Usr, int, DenseMapInfoForUsr> func_usr, type_usr, var_usr;
   llvm::SmallVector<QueryFunc, 0> funcs;