[clangd] Extend dexp to support remote index

Summary:
* Merge clangd-remote-client into dexp
* Implement `clangd::remote::IndexClient` that is derived from `SymbolIndex`
* Upgrade remote mode-related CMake infrastructure

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: mgorny, ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D78521

Author: kirillbobyrev

clang-tools-extra/clangd/CMakeLists.txt

@@ -140,7 +140,6 @@ if (${CMAKE_SYSTEM_NAME} STREQUAL "Linux")
 endif()
 add_subdirectory(tool)
 add_subdirectory(indexer)
-add_subdirectory(index/dex/dexp)
 
 if (LLVM_INCLUDE_BENCHMARKS)
   add_subdirectory(benchmarks)
@@ -160,5 +159,6 @@ set(GRPC_INSTALL_PATH "" CACHE PATH "Path to gRPC library manual installation.")
 
 if (CLANGD_ENABLE_REMOTE)
   include(FindGRPC)
-  add_subdirectory(index/remote)
 endif()
+add_subdirectory(index/remote)
+add_subdirectory(index/dex/dexp)

clang-tools-extra/clangd/Features.inc.in

@@ -1 +1,2 @@
 #define CLANGD_BUILD_XPC @CLANGD_BUILD_XPC@
+#define CLANGD_ENABLE_REMOTE @CLANGD_ENABLE_REMOTE@

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

@@ -77,6 +77,13 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const IndexFileOut &O);
 std::string toYAML(const Symbol &);
 std::string toYAML(const std::pair<SymbolID, ArrayRef<Ref>> &);
 std::string toYAML(const Relation &);
+std::string toYAML(const Ref &);
+
+// Deserialize a single symbol from YAML.
+llvm::Expected<clangd::Symbol> symbolFromYAML(StringRef YAML,
+                                              llvm::UniqueStringSaver *Strings);
+llvm::Expected<clangd::Ref> refFromYAML(StringRef YAML,
+                                        llvm::UniqueStringSaver *Strings);
 
 // Build an in-memory static index from an index file.
 // The size should be relatively small, so data can be managed in memory.

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

@@ -517,5 +517,40 @@ std::string toYAML(const Relation &R) {
   return Buf;
 }
 
+std::string toYAML(const Ref &R) {
+  std::string Buf;
+  {
+    llvm::raw_string_ostream OS(Buf);
+    llvm::yaml::Output Yout(OS);
+    Ref Reference = R; // copy: Yout<< requires mutability.
+    Yout << Reference;
+  }
+  return Buf;
+}
+
+llvm::Expected<clangd::Symbol>
+symbolFromYAML(StringRef YAML, llvm::UniqueStringSaver *Strings) {
+  clangd::Symbol Deserialized;
+  llvm::yaml::Input YAMLInput(YAML, Strings);
+  if (YAMLInput.error())
+    return llvm::make_error<llvm::StringError>(
+        llvm::formatv("Unable to deserialize Symbol from YAML: {0}", YAML),
+        llvm::inconvertibleErrorCode());
+  YAMLInput >> Deserialized;
+  return Deserialized;
+}
+
+llvm::Expected<clangd::Ref> refFromYAML(StringRef YAML,
+                                        llvm::UniqueStringSaver *Strings) {
+  clangd::Ref Deserialized;
+  llvm::yaml::Input YAMLInput(YAML, Strings);
+  if (YAMLInput.error())
+    return llvm::make_error<llvm::StringError>(
+        llvm::formatv("Unable to deserialize Symbol from YAML: {0}", YAML),
+        llvm::inconvertibleErrorCode());
+  YAMLInput >> Deserialized;
+  return Deserialized;
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/dex/dexp/CMakeLists.txt

@@ -1,4 +1,5 @@
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../../)
+include_directories(${CMAKE_CURRENT_BINARY_DIR}/../../../)
 
 set(LLVM_LINK_COMPONENTS
   LineEditor
@@ -16,4 +17,5 @@ clang_target_link_libraries(dexp
 target_link_libraries(dexp
   PRIVATE
   clangDaemon
+  clangdRemoteIndex
   )

clang-tools-extra/clangd/index/dex/dexp/Dexp.cpp

@@ -11,9 +11,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Features.inc"
 #include "SourceCode.h"
 #include "index/Serialization.h"
 #include "index/dex/Dex.h"
+#include "index/remote/Client.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -26,8 +28,9 @@ namespace clang {
 namespace clangd {
 namespace {
 
-llvm::cl::opt<std::string> IndexPath(llvm::cl::desc("<INDEX FILE>"),
-                                     llvm::cl::Positional, llvm::cl::Required);
+llvm::cl::opt<std::string> IndexLocation(
+    llvm::cl::desc("<path to index file | remote:server.address>"),
+    llvm::cl::Positional);
 
 llvm::cl::opt<std::string>
     ExecCommand("c", llvm::cl::desc("Command to execute and then exit"));
@@ -38,6 +41,10 @@ queries over given symbol collection obtained via clangd-indexer. The
 tool can be used to evaluate search quality of existing index implementations
 and manually construct non-trivial test cases.
 
+You can connect to remote index by passing remote:address to dexp. Example:
+
+$ dexp remote:0.0.0.0:9000
+
 Type use "help" request to get information about the details.
 )";
 
@@ -150,7 +157,7 @@ class FuzzyFind : public Command {
     }
     Request.AnyScope = Request.Scopes.empty();
     // FIXME(kbobyrev): Print symbol final scores to see the distribution.
-    static const auto OutputFormat = "{0,-4} | {1,-40} | {2,-25}\n";
+    static const auto *OutputFormat = "{0,-4} | {1,-40} | {2,-25}\n";
     llvm::outs() << llvm::formatv(OutputFormat, "Rank", "Symbol ID",
                                   "Symbol Name");
     size_t Rank = 0;
@@ -316,13 +323,14 @@ struct {
     {"find", "Search for symbols with fuzzyFind", std::make_unique<FuzzyFind>},
     {"lookup", "Dump symbol details by ID or qualified name",
      std::make_unique<Lookup>},
-    {"refs", "Find references by ID or qualified name",
-     std::make_unique<Refs>},
+    {"refs", "Find references by ID or qualified name", std::make_unique<Refs>},
     {"export", "Export index", std::make_unique<Export>},
 };
 
 std::unique_ptr<SymbolIndex> openIndex(llvm::StringRef Index) {
-  return loadIndex(Index, /*UseDex=*/true);
+  return Index.startswith("remote:")
+             ? remote::getClient(Index.drop_front(strlen("remote:")))
+             : loadIndex(Index, /*UseDex=*/true);
 }
 
 bool runCommand(std::string Request, const SymbolIndex &Index) {
@@ -365,9 +373,10 @@ int main(int argc, const char *argv[]) {
   llvm::sys::PrintStackTraceOnErrorSignal(argv[0]);
 
   std::unique_ptr<SymbolIndex> Index;
-  reportTime("Dex build", [&]() {
-    Index = openIndex(IndexPath);
-  });
+  reportTime(llvm::StringRef(IndexLocation).startswith("remote:")
+                 ? "Remote index client creation"
+                 : "Dex build",
+             [&]() { Index = openIndex(IndexLocation); });
 
   if (!Index) {
     llvm::outs() << "Failed to open the index.\n";

clang-tools-extra/clangd/index/remote/CMakeLists.txt

@@ -1,7 +1,28 @@
-generate_grpc_protos(RemoteIndexProtos "Index.proto")
+if (CLANGD_ENABLE_REMOTE)
+  generate_grpc_protos(RemoteIndexProtos "Index.proto")
+  include_directories(${CMAKE_CURRENT_BINARY_DIR})
+  include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../)
 
-include_directories("${CMAKE_CURRENT_BINARY_DIR}")
-include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../)
+  # FIXME(kirillbobyrev): target_compile_definitions is not working with
+  # add_clang_library for some reason. Is there any way to make this
+  # target-local?
+  add_definitions(-DGOOGLE_PROTOBUF_NO_RTTI=1)
 
-add_subdirectory(client)
-add_subdirectory(server)
+  add_clang_library(clangdRemoteIndex
+    Client.cpp
+
+    LINK_LIBS
+    RemoteIndexProtos
+    clangdRemoteMarshalling
+
+    protobuf
+    grpc++
+    clangDaemon
+    )
+
+  add_subdirectory(marshalling)
+  add_subdirectory(server)
+else()
+  # Provides a dummy implementation of clangdRemoteIndex.
+  add_subdirectory(unimplemented)
+endif()

clang-tools-extra/clangd/index/remote/Client.cpp

@@ -0,0 +1,105 @@
+//===--- Client.cpp ----------------------------------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <grpcpp/grpcpp.h>
+
+#include "Client.h"
+#include "Index.grpc.pb.h"
+#include "Logger.h"
+#include "Trace.h"
+#include "index/Serialization.h"
+#include "marshalling/Marshalling.h"
+#include "llvm/Support/YAMLTraits.h"
+
+namespace clang {
+namespace clangd {
+namespace remote {
+namespace {
+
+class IndexClient : public clangd::SymbolIndex {
+  template <typename RequestT, typename ReplyT>
+  using StreamingCall = std::unique_ptr<grpc::ClientReader<ReplyT>> (
+      remote::SymbolIndex::Stub::*)(grpc::ClientContext *, const RequestT &);
+
+  // FIXME(kirillbobyrev): Set deadlines for requests.
+  template <typename RequestT, typename ReplyT, typename ClangdRequestT,
+            typename CallbackT>
+  bool streamRPC(ClangdRequestT Request,
+                 StreamingCall<RequestT, ReplyT> RPCCall,
+                 CallbackT Callback) const {
+    bool FinalResult = false;
+    trace::Span Tracer(RequestT::descriptor()->name());
+    const auto RPCRequest = toProtobuf(Request);
+    grpc::ClientContext Context;
+    auto Reader = (Stub.get()->*RPCCall)(&Context, RPCRequest);
+    llvm::BumpPtrAllocator Arena;
+    llvm::UniqueStringSaver Strings(Arena);
+    ReplyT Reply;
+    while (Reader->Read(&Reply)) {
+      if (!Reply.has_stream_result()) {
+        FinalResult = Reply.final_result();
+        continue;
+      }
+      auto Sym = fromProtobuf(Reply.stream_result(), &Strings);
+      if (!Sym)
+        elog("Received invalid {0}: {1}", ReplyT::descriptor()->name(),
+             Reply.stream_result().yaml_serialization());
+      Callback(*Sym);
+    }
+    SPAN_ATTACH(Tracer, "status", Reader->Finish().ok());
+    return FinalResult;
+  }
+
+public:
+  IndexClient(std::shared_ptr<grpc::Channel> Channel)
+      : Stub(remote::SymbolIndex::NewStub(Channel)) {}
+
+  void lookup(const clangd::LookupRequest &Request,
+              llvm::function_ref<void(const clangd::Symbol &)> Callback) const {
+    streamRPC<LookupRequest, LookupReply>(
+        Request, &remote::SymbolIndex::Stub::Lookup, Callback);
+  }
+
+  bool
+  fuzzyFind(const clangd::FuzzyFindRequest &Request,
+            llvm::function_ref<void(const clangd::Symbol &)> Callback) const {
+    return streamRPC<FuzzyFindRequest, FuzzyFindReply>(
+        Request, &remote::SymbolIndex::Stub::FuzzyFind, Callback);
+  }
+
+  bool refs(const clangd::RefsRequest &Request,
+            llvm::function_ref<void(const clangd::Ref &)> Callback) const {
+    return streamRPC<RefsRequest, RefsReply>(
+        Request, &remote::SymbolIndex::Stub::Refs, Callback);
+  }
+
+  // FIXME(kirillbobyrev): Implement this.
+  void
+  relations(const clangd::RelationsRequest &,
+            llvm::function_ref<void(const SymbolID &, const clangd::Symbol &)>)
+      const {}
+
+  // IndexClient does not take any space since the data is stored on the server.
+  size_t estimateMemoryUsage() const { return 0; }
+
+private:
+  std::unique_ptr<remote::SymbolIndex::Stub> Stub;
+};
+
+} // namespace
+
+std::unique_ptr<clangd::SymbolIndex> getClient(llvm::StringRef Address) {
+  const auto Channel =
+      grpc::CreateChannel(Address.str(), grpc::InsecureChannelCredentials());
+  Channel->GetState(true);
+  return std::unique_ptr<clangd::SymbolIndex>(new IndexClient(Channel));
+}
+
+} // namespace remote
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/index/remote/Client.h

@@ -0,0 +1,31 @@
+//===--- Client.h - Connect to a remote index via gRPC -----------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REMOTE_INDEX_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REMOTE_INDEX_H
+
+#include "index/Index.h"
+
+namespace clang {
+namespace clangd {
+namespace remote {
+
+/// Returns an SymbolIndex client that passes requests to remote index located
+/// at \p Address. The client allows synchronous RPC calls.
+///
+/// This method attempts to resolve the address and establish the connection.
+///
+/// \returns nullptr if the address is not resolved during the function call or
+/// if the project was compiled without Remote Index support.
+std::unique_ptr<clangd::SymbolIndex> getClient(llvm::StringRef Address);
+
+} // namespace remote
+} // namespace clangd
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_REMOTE_INDEX_H

clang-tools-extra/clangd/index/remote/Index.proto

@@ -10,10 +10,60 @@ syntax = "proto3";
 
 package clang.clangd.remote;
 
-service Index {
+service SymbolIndex {
   rpc Lookup(LookupRequest) returns (stream LookupReply) {}
+
+  rpc FuzzyFind(FuzzyFindRequest) returns (stream FuzzyFindReply) {}
+
+  rpc Refs(RefsRequest) returns (stream RefsReply) {}
+}
+
+message LookupRequest { repeated string ids = 1; }
+
+// The response is a stream of symbol messages and the terminating message
+// indicating the end of stream.
+message LookupReply {
+  oneof kind {
+    Symbol stream_result = 1;
+    bool final_result = 2;
+  }
+}
+
+message FuzzyFindRequest {
+  string query = 1;
+  repeated string scopes = 2;
+  bool any_scope = 3;
+  uint32 limit = 4;
+  bool resricted_for_code_completion = 5;
+  repeated string proximity_paths = 6;
+  repeated string preferred_types = 7;
+}
+
+// The response is a stream of symbol messages, and one terminating has_more
+// message.
+message FuzzyFindReply {
+  oneof kind {
+    Symbol stream_result = 1;
+    bool final_result = 2; // HasMore
+  }
 }
 
-message LookupRequest { string id = 1; }
+message RefsRequest {
+  repeated string ids = 1;
+  uint32 filter = 2;
+  uint32 limit = 3;
+}
+
+// The response is a stream of reference messages, and one terminating has_more
+// message.
+message RefsReply {
+  oneof kind {
+    Ref stream_result = 1;
+    bool final_result = 2; // HasMore
+  }
+}
 
-message LookupReply { string symbol_yaml = 1; }
+// FIXME(kirillbobyrev): Properly serialize symbols and refs instead of passing
+// YAML.
+message Ref { string yaml_serialization = 1; }
+message Symbol { string yaml_serialization = 1; }