Implement remote index support for the new containedRefs request

Author: HighCommander4

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

@@ -146,6 +146,13 @@ class IndexClient : public clangd::SymbolIndex {
     return streamRPC(Request, &remote::v1::SymbolIndex::Stub::Refs, Callback);
   }
 
+  bool containedRefs(const clangd::ContainedRefsRequest &Request,
+                     llvm::function_ref<void(const ContainedRefsResult &)>
+                         Callback) const override {
+    return streamRPC(Request, &remote::v1::SymbolIndex::Stub::ContainedRefs,
+                     Callback);
+  }
+
   void
   relations(const clangd::RelationsRequest &Request,
             llvm::function_ref<void(const SymbolID &, const clangd::Symbol &)>

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

@@ -131,3 +131,21 @@ message Relation {
   optional string subject_id = 1;
   optional Symbol object = 2;
 }
+
+message ContainedRefsRequest {
+  required string id = 1;
+  optional uint32 limit = 2;
+}
+
+message ContainedRefsReply {
+  oneof kind {
+    ContainedRef stream_result = 1;
+    FinalResult final_result = 2;
+  }
+}
+
+message ContainedRef {
+  required SymbolLocation location = 1;
+  required uint32 kind = 2;
+  required string symbol = 3;
+}

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

@@ -21,5 +21,7 @@ service SymbolIndex {
 
   rpc Refs(RefsRequest) returns (stream RefsReply) {}
 
+  rpc ContainedRefs(ContainedRefsRequest) returns (stream ContainedRefsReply) {}
+
   rpc Relations(RelationsRequest) returns (stream RelationsReply) {}
 }

clang-tools-extra/clangd/index/remote/marshalling/Marshalling.cpp

@@ -126,6 +126,18 @@ Marshaller::fromProtobuf(const RefsRequest *Message) {
   return Req;
 }
 
+llvm::Expected<clangd::ContainedRefsRequest>
+Marshaller::fromProtobuf(const ContainedRefsRequest *Message) {
+  clangd::ContainedRefsRequest Req;
+  auto ID = SymbolID::fromStr(Message->id());
+  if (!ID)
+    return ID.takeError();
+  Req.ID = *ID;
+  if (Message->has_limit())
+    Req.Limit = Message->limit();
+  return Req;
+}
+
 llvm::Expected<clangd::RelationsRequest>
 Marshaller::fromProtobuf(const RelationsRequest *Message) {
   clangd::RelationsRequest Req;
@@ -192,6 +204,21 @@ llvm::Expected<clangd::Ref> Marshaller::fromProtobuf(const Ref &Message) {
   return Result;
 }
 
+llvm::Expected<clangd::ContainedRefsResult>
+Marshaller::fromProtobuf(const ContainedRef &Message) {
+  clangd::ContainedRefsResult Result;
+  auto Location = fromProtobuf(Message.location());
+  if (!Location)
+    return Location.takeError();
+  Result.Location = *Location;
+  Result.Kind = static_cast<RefKind>(Message.kind());
+  auto Symbol = SymbolID::fromStr(Message.symbol());
+  if (!Symbol)
+    return Symbol.takeError();
+  Result.Symbol = *Symbol;
+  return Result;
+}
+
 llvm::Expected<std::pair<clangd::SymbolID, clangd::Symbol>>
 Marshaller::fromProtobuf(const Relation &Message) {
   auto SubjectID = SymbolID::fromStr(Message.subject_id());
@@ -244,6 +271,15 @@ RefsRequest Marshaller::toProtobuf(const clangd::RefsRequest &From) {
   return RPCRequest;
 }
 
+ContainedRefsRequest
+Marshaller::toProtobuf(const clangd::ContainedRefsRequest &From) {
+  ContainedRefsRequest RPCRequest;
+  RPCRequest.set_id(From.ID.str());
+  if (From.Limit)
+    RPCRequest.set_limit(*From.Limit);
+  return RPCRequest;
+}
+
 RelationsRequest Marshaller::toProtobuf(const clangd::RelationsRequest &From) {
   RelationsRequest RPCRequest;
   for (const auto &ID : From.Subjects)
@@ -299,6 +335,18 @@ llvm::Expected<Ref> Marshaller::toProtobuf(const clangd::Ref &From) {
   return Result;
 }
 
+llvm::Expected<ContainedRef>
+Marshaller::toProtobuf(const clangd::ContainedRefsResult &From) {
+  ContainedRef Result;
+  auto Location = toProtobuf(From.Location);
+  if (!Location)
+    return Location.takeError();
+  *Result.mutable_location() = *Location;
+  Result.set_kind(static_cast<uint32_t>(From.Kind));
+  *Result.mutable_symbol() = From.Symbol.str();
+  return Result;
+}
+
 llvm::Expected<Relation> Marshaller::toProtobuf(const clangd::SymbolID &Subject,
                                                 const clangd::Symbol &Object) {
   Relation Result;

clang-tools-extra/clangd/index/remote/marshalling/Marshalling.h

@@ -40,6 +40,8 @@ class Marshaller {
 
   llvm::Expected<clangd::Symbol> fromProtobuf(const Symbol &Message);
   llvm::Expected<clangd::Ref> fromProtobuf(const Ref &Message);
+  llvm::Expected<clangd::ContainedRefsResult>
+  fromProtobuf(const ContainedRef &Message);
   llvm::Expected<std::pair<clangd::SymbolID, clangd::Symbol>>
   fromProtobuf(const Relation &Message);
 
@@ -48,6 +50,8 @@ class Marshaller {
   llvm::Expected<clangd::FuzzyFindRequest>
   fromProtobuf(const FuzzyFindRequest *Message);
   llvm::Expected<clangd::RefsRequest> fromProtobuf(const RefsRequest *Message);
+  llvm::Expected<clangd::ContainedRefsRequest>
+  fromProtobuf(const ContainedRefsRequest *Message);
   llvm::Expected<clangd::RelationsRequest>
   fromProtobuf(const RelationsRequest *Message);
 
@@ -58,10 +62,13 @@ class Marshaller {
   LookupRequest toProtobuf(const clangd::LookupRequest &From);
   FuzzyFindRequest toProtobuf(const clangd::FuzzyFindRequest &From);
   RefsRequest toProtobuf(const clangd::RefsRequest &From);
+  ContainedRefsRequest toProtobuf(const clangd::ContainedRefsRequest &From);
   RelationsRequest toProtobuf(const clangd::RelationsRequest &From);
 
   llvm::Expected<Symbol> toProtobuf(const clangd::Symbol &From);
   llvm::Expected<Ref> toProtobuf(const clangd::Ref &From);
+  llvm::Expected<ContainedRef>
+  toProtobuf(const clangd::ContainedRefsResult &From);
   llvm::Expected<Relation> toProtobuf(const clangd::SymbolID &Subject,
                                       const clangd::Symbol &Object);
 

clang-tools-extra/clangd/index/remote/server/Server.cpp

@@ -258,6 +258,53 @@ class RemoteIndexServer final : public v1::SymbolIndex::Service {
     return grpc::Status::OK;
   }
 
+  grpc::Status
+  ContainedRefs(grpc::ServerContext *Context,
+                const ContainedRefsRequest *Request,
+                grpc::ServerWriter<ContainedRefsReply> *Reply) override {
+    auto StartTime = stopwatch::now();
+    WithContextValue WithRequestContext(CurrentRequest, Context);
+    logRequest(*Request);
+    trace::Span Tracer("ContainedRefsRequest");
+    auto Req = ProtobufMarshaller->fromProtobuf(Request);
+    if (!Req) {
+      elog("Can not parse ContainedRefsRequest from protobuf: {0}",
+           Req.takeError());
+      return grpc::Status::CANCELLED;
+    }
+    if (!Req->Limit || *Req->Limit > LimitResults) {
+      log("[public] Limiting result size for ContainedRefs request from {0} to "
+          "{1}.",
+          Req->Limit, LimitResults);
+      Req->Limit = LimitResults;
+    }
+    unsigned Sent = 0;
+    unsigned FailedToSend = 0;
+    bool HasMore =
+        Index.containedRefs(*Req, [&](const clangd::ContainedRefsResult &Item) {
+          auto SerializedItem = ProtobufMarshaller->toProtobuf(Item);
+          if (!SerializedItem) {
+            elog("Unable to convert ContainedRefsResult to protobuf: {0}",
+                 SerializedItem.takeError());
+            ++FailedToSend;
+            return;
+          }
+          ContainedRefsReply NextMessage;
+          *NextMessage.mutable_stream_result() = *SerializedItem;
+          logResponse(NextMessage);
+          Reply->Write(NextMessage);
+          ++Sent;
+        });
+    ContainedRefsReply LastMessage;
+    LastMessage.mutable_final_result()->set_has_more(HasMore);
+    logResponse(LastMessage);
+    Reply->Write(LastMessage);
+    SPAN_ATTACH(Tracer, "Sent", Sent);
+    SPAN_ATTACH(Tracer, "Failed to send", FailedToSend);
+    logRequestSummary("v1/ContainedRefs", Sent, StartTime);
+    return grpc::Status::OK;
+  }
+
   grpc::Status Relations(grpc::ServerContext *Context,
                          const RelationsRequest *Request,
                          grpc::ServerWriter<RelationsReply> *Reply) override {