[clangd] Make functions of ClangdServer callback-based

[ilya-biryukov]

Summary

As a consequence, all LSP operations are now handled asynchronously,
i.e. they never block the main processing thread. However, if
-run-synchronously flag is specified, clangd still runs everything on
the main thread.

[ioeric]

clangd/ClangdLSPServer.cpp

194

nit: since we are not spelling out the return type here, it might be clearer if we do:

replyError(...);
return;

return replyError(...) makes me wonder what the return type is.

clangd/ClangdServer.cpp

209

nit: I'd probably use a different name than Callback for this parameter for clarity.

441–442

Consider spelling out the captured values, just in case new variables which are not expected to be captured are added in the future.

unittests/clangd/SyncAPI.cpp

107

I'd expect this to be checked by callers. Would return std::move(Result); work?

[ilya-biryukov]

• Capture only needed vars in lambda, not everything.
• Rebase onto head.

clangd/ClangdLSPServer.cpp

194

This trick was used in the code before (there are usages below), so I figured it's ok to do this.
If it's confusing I'll change it everywhere.

clangd/ClangdServer.cpp

209

I actually think we should keep it this way. It's a bit tricky to grasp, but it has two important advantages:

• Makes referencing Callback from outer scope impossible
• saves us from coming up with names for that superficial variable, i.e. should it be called InnerCallback, Callback2, C or something else?

Is it that too confusing or do you feel we can keep it?

441–442

Thanks for spotting this!

unittests/clangd/SyncAPI.cpp

107

The reason why we need it is because capture(Result) writes return value of the callback to the Result variable. But we have to first check the default-constructed value that was there before calling Server.signatureHelp

[ioeric]

clangd/ClangdLSPServer.cpp

194

I guess it's not confusing after figuring out what replyError returns. There is obviously tradeoff between LOC and readability. Either is not too bad. So up to you :)

clangd/ClangdServer.cpp

209

Makes referencing Callback from outer scope impossible.

For lambdas, I think we should rely on captures for this instead of the parameter names.

saves us from coming up with names for that superficial variable, i.e. should it be called InnerCallback, Callback2, C or something else?

I thought this is a common practice with llvm coding style :P I would vote for C :)

Is it that too confusing or do you feel we can keep it?

It is a bit confusing when I first looked at it, but nothing is *too* confusing as long as the code is correct ;) I just think it would make the code easier to read.

unittests/clangd/SyncAPI.cpp

107

I think we should probably fix the behavior of capture since this changes the behavior of llvm::Expected in user code - the check is there to make sure users always check the status. But here we always do this for users.

[sammccall]

LG, suggest a tweak to capture() though.

I wonder whether we want to introduce using Callback<T...> = UniqueFunction<void(T...)> for readability at some point...

clangd/ClangdServer.cpp

209

I agree with you both :-)
Conceptually, this *is* a capture - BindWithForward simulates move-capture that C++11 doesn't have native syntax for. So the same name seems appropriate to me - you want shadowing for the same reasons you want captures to shadow.

unittests/clangd/SyncAPI.cpp

107

+1

I suggested capture(T&) was the right API because I thought it'd be used directly by test code (but we wrap it here) and because I thought T would be default-constructible without problems (but it's not).

I'd suggest changing to capture(Optional<T>&) instead, so the caller can just create an empty optional. That makes the callsites here slightly less obscure.

[ilya-biryukov]

• Use llvm::Optional<> in capture() helper to avoid confusion with llvm::Expected.

[ilya-biryukov]

In D43227#1006584, @sammccall wrote:

I wonder whether we want to introduce using Callback<T...> = UniqueFunction<void(T...)> for readability at some point...

I agree that's a good idea, will come up with a CL doing that.

clangd/ClangdServer.cpp

209

I'll leave as is in this review, we have other callsites, doing the same thing, that we don't touch here. If we decide to change it, we should change all of them in one go.
I'm not opposed to the change, but like the current style a bit more.

unittests/clangd/SyncAPI.cpp

107

The users still have to check the returned value. We only check the default-constructed value. After capture() runs Result is reassigned and now has to be checked again (and this should be done by the users).
Replaced with capture(Optional<T>&) to make it absolutely clear we're doing the right thing there.

[ioeric]

lg

clangd/ClangdServer.cpp

209

Feel free to land without changing this; this is a nit anyway ;)

[ilya-biryukov]

• Rebase onto head

[CarlosAlbertoEnciso]

Hi,

It seems that your submission broke the Windows/Linux builds for 'clangd'.

The generated error message is:

C:\Program Files (x86)\Microsoft Visual Studio 14.0\VC\include\future(263): error C2512: 'llvm::Expected<std::vector<clang::tooling::Replacement,std::allocator<_Kty>>>': no appropriate default constructor available [yyy\build\tools\clang\tools\extra\unittests\clangd\ClangdTests.vcxproj]

C:\xxx\llvm\tools\clang\tools\extra\clangd\ClangdServer.h(231): note: see declaration of 'llvm::Expected<std::vector<clang::tooling::Replacement,std::allocator<_Kty>>>'
C:\Program Files (x86)\Microsoft Visual Studio 14.0\VC\include\future(255): note: while compiling class template member function 'std::_Associated_state<_Ty>::_Associated_state(std::_Deleter_base<_Ty> *)'
C:\Program Files (x86)\Microsoft Visual Studio 14.0\VC\include\future(1491): note: see reference to function template instantiation 'std::_Associated_state<_Ty>::_Associated_state(std::_Deleter_base<_Ty> *)' being compiled
C:\Program Files (x86)\Microsoft Visual Studio 14.0\VC\include\future(989): note: see reference to class template instantiation 'std::_Associated_state<_Ty>' being compiled
C:\Program Files (x86)\Microsoft Visual Studio 14.0\VC\include\future(987): note: while compiling class template member function 'bool std::_State_manager<_Ty>::valid(void) throw() const'

Thanks

[CarlosAlbertoEnciso]

Build now is OK. Thanks @ilya-biryukov

clangd/ClangdLSPServer.cpp

if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
     // more commands are added, this will be here has a safe guard.
     replyError(
         ErrorCode::InvalidParams,
         llvm::formatv("Unsupported command \"{0}\".", Params.command).str());
   }
 }
 
 void ClangdLSPServer::onRename(RenameParams &Params) {
-  auto File = Params.textDocument.uri.file;
-  auto Code = Server.getDocument(File);
+  Path File = Params.textDocument.uri.file;
+  llvm::Optional<std::string> Code = Server.getDocument(File);
   if (!Code)
     return replyError(ErrorCode::InvalidParams,
                       "onRename called for non-added file");
 
-  auto Replacements = Server.rename(File, Params.position, Params.newName);
-  if (!Replacements) {
-    replyError(ErrorCode::InternalError,
+  Server.rename(
+      File, Params.position, Params.newName,
+      [File, Code,
+       Params](llvm::Expected<std::vector<tooling::Replacement>> Replacements) {
+        if (!Replacements)
+          return replyError(ErrorCode::InternalError,
                             llvm::toString(Replacements.takeError()));
-    return;
-  }
 
         std::vector<TextEdit> Edits = replacementsToEdits(*Code, *Replacements);
         WorkspaceEdit WE;
         WE.changes = {{Params.textDocument.uri.uri(), Edits}};
         reply(WE);
+      });
 }
 
 void ClangdLSPServer::onDocumentDidClose(DidCloseTextDocumentParams &Params) {
   Server.removeDocument(Params.textDocument.uri.file);
 }
 
 void ClangdLSPServer::onDocumentOnTypeFormatting(
     DocumentOnTypeFormattingParams &Params) {

 }
 
 void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
   Server.codeComplete(Params.textDocument.uri.file, Params.position, CCOpts,
                       [](Tagged<CompletionList> List) { reply(List.Value); });
 }
 
 void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
-  auto SignatureHelp =
-      Server.signatureHelp(Params.textDocument.uri.file, Params.position);
+  Server.signatureHelp(Params.textDocument.uri.file, Params.position,
+                       [](llvm::Expected<Tagged<SignatureHelp>> SignatureHelp) {
                          if (!SignatureHelp)
-    return replyError(ErrorCode::InvalidParams,
+                           return replyError(
+                               ErrorCode::InvalidParams,
                                llvm::toString(SignatureHelp.takeError()));
                          reply(SignatureHelp->Value);
+                       });
 }
 
 void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
-  auto Items =
-      Server.findDefinitions(Params.textDocument.uri.file, Params.position);
+  Server.findDefinitions(
+      Params.textDocument.uri.file, Params.position,
+      [](llvm::Expected<Tagged<std::vector<Location>>> Items) {
         if (!Items)
           return replyError(ErrorCode::InvalidParams,
                             llvm::toString(Items.takeError()));
         reply(json::ary(Items->Value));
+      });
 }
 
 void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
   llvm::Optional<Path> Result = Server.switchSourceHeader(Params.uri.file);
   reply(Result ? URI::createFile(*Result).toString() : "");
 }
 
 void ClangdLSPServer::onDocumentHighlight(TextDocumentPositionParams &Params) {
-  auto Highlights = Server.findDocumentHighlights(Params.textDocument.uri.file,
-                                                  Params.position);
-
-  if (!Highlights) {
-    replyError(ErrorCode::InternalError,
+  Server.findDocumentHighlights(
+      Params.textDocument.uri.file, Params.position,
+      [](llvm::Expected<Tagged<std::vector<DocumentHighlight>>> Highlights) {
+        if (!Highlights)
+          return replyError(ErrorCode::InternalError,
                             llvm::toString(Highlights.takeError()));
-    return;
-  }
-
         reply(json::ary(Highlights->Value));
+      });
 }
 
 ClangdLSPServer::ClangdLSPServer(JSONOutput &Out, unsigned AsyncThreadsCount,
                                  bool StorePreamblesInMemory,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<StringRef> ResourceDir,
                                  llvm::Optional<Path> CompileCommandsDir,
                                  bool BuildDynamicSymbolIndex,

clangd/ClangdServer.h

public:
 
   /// Provide signature help for \p File at \p Pos. If \p OverridenContents is
   /// not None, they will used only for signature help, i.e. no diagnostics
   /// update will be scheduled and a draft for \p File will not be updated. If
   /// \p OverridenContents is None, contents of the current draft for \p File
   /// will be used. If \p UsedFS is non-null, it will be overwritten by
   /// vfs::FileSystem used for signature help. This method should only be called
   /// for currently tracked files.
-  llvm::Expected<Tagged<SignatureHelp>>
-  signatureHelp(PathRef File, Position Pos,
-                llvm::Optional<StringRef> OverridenContents = llvm::None,
-                IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
+  void signatureHelp(
+      PathRef File, Position Pos,
+      UniqueFunction<void(llvm::Expected<Tagged<SignatureHelp>>)> Callback,
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
-  llvm::Expected<Tagged<std::vector<Location>>> findDefinitions(PathRef File,
-                                                                Position Pos);
+  void findDefinitions(
+      PathRef File, Position Pos,
+      UniqueFunction<void(llvm::Expected<Tagged<std::vector<Location>>>)>
+          Callback);
 
   /// 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.
-  llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
-  findDocumentHighlights(PathRef File, Position Pos);
+  void findDocumentHighlights(
+      PathRef File, Position Pos,
+      UniqueFunction<
+          void(llvm::Expected<Tagged<std::vector<DocumentHighlight>>>)>
+          Callback);
 
   /// Run formatting for \p Rng inside \p File with content \p Code.
   llvm::Expected<tooling::Replacements> formatRange(StringRef Code,
                                                     PathRef File, Range Rng);
 
   /// Run formatting for the whole \p File with content \p Code.
   llvm::Expected<tooling::Replacements> formatFile(StringRef Code,
                                                    PathRef File);
 
   /// Run formatting after a character was typed at \p Pos in \p File with
   /// content \p Code.
   llvm::Expected<tooling::Replacements>
   formatOnType(StringRef Code, PathRef File, Position Pos);
 
   /// Rename all occurrences of the symbol at the \p Pos in \p File to
   /// \p NewName.
-  Expected<std::vector<tooling::Replacement>> rename(PathRef File, Position Pos,
-                                                     llvm::StringRef NewName);
+  void rename(PathRef File, Position Pos, llvm::StringRef NewName,
+              UniqueFunction<void(Expected<std::vector<tooling::Replacement>>)>
+                  Callback);
 
   /// Gets current document contents for \p File. Returns None if \p File is not
   /// currently tracked.
   /// FIXME(ibiryukov): This function is here to allow offset-to-Position
   /// conversions in outside code, maybe there's a way to get rid of it.
   llvm::Optional<std::string> getDocument(PathRef File);
 
   /// Only for testing purposes.
   /// Waits until all requests to worker thread are finished and dumps AST for
   /// \p File. \p File must be in the list of added documents.
-  std::string dumpAST(PathRef File);
+  void dumpAST(PathRef File, UniqueFunction<void(std::string)> Callback);
   /// Called when an event occurs for a watched file in the workspace.
   void onFileEvent(const DidChangeWatchedFilesParams &Params);
 
   /// Returns estimated memory usage for each of the currently open files.
   /// The order of results is unspecified.
   /// Overall memory usage of clangd may be significantly more than reported
   /// here, as this metric does not account (at least) for:
   ///   - memory occupied by static and dynamic index,

clangd/ClangdServer.cpp

 #include "llvm/Support/raw_ostream.h"
 #include <future>
 
 using namespace clang;
 using namespace clang::clangd;
 
 namespace {
 
-// Issues an async read of AST and waits for results.
-template <class Ret, class Func>
-Ret blockingRunWithAST(TUScheduler &S, PathRef File, Func &&F) {
-  // Using shared_ptr to workaround MSVC bug. It requires future<> arguments to
-  // have default and copy ctor.
-  auto SharedPtrFunc = [&](llvm::Expected<InputsAndAST> Arg) {
-    return std::make_shared<Ret>(F(std::move(Arg)));
-  };
-  std::packaged_task<std::shared_ptr<Ret>(llvm::Expected<InputsAndAST>)> Task(
-      SharedPtrFunc);
-  auto Future = Task.get_future();
-  S.runWithAST(File, std::move(Task));
-  return std::move(*Future.get());
-}
-
-// Issues an async read of preamble and waits for results.
-template <class Ret, class Func>
-Ret blockingRunWithPreamble(TUScheduler &S, PathRef File, Func &&F) {
-  // Using shared_ptr to workaround MSVC bug. It requires future<> arguments to
-  // have default and copy ctor.
-  auto SharedPtrFunc = [&](llvm::Expected<InputsAndPreamble> Arg) {
-    return std::make_shared<Ret>(F(std::move(Arg)));
-  };
-  std::packaged_task<std::shared_ptr<Ret>(llvm::Expected<InputsAndPreamble>)>
-      Task(SharedPtrFunc);
-  auto Future = Task.get_future();
-  S.runWithPreamble(File, std::move(Task));
-  return std::move(*Future.get());
-}
-
 void ignoreError(llvm::Error Err) {
   handleAllErrors(std::move(Err), [](const llvm::ErrorInfoBase &) {});
 }
 
 std::string getStandardResourceDir() {
   static int Dummy; // Just an address in this process.
   return CompilerInvocation::GetResourcesPath("clangd", (void *)&Dummy);
 }
auto Task = [PCHs, Pos, TaggedFS, CodeCompleteOpts](
     Callback(make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
   };
 
   WorkScheduler.runWithPreamble(File, BindWithForward(Task, std::move(Contents),
                                                       File.str(),
                                                       std::move(Callback)));
 }
 
-llvm::Expected<Tagged<SignatureHelp>>
-ClangdServer::signatureHelp(PathRef File, Position Pos,
-                            llvm::Optional<StringRef> OverridenContents,
-                            IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
+void ClangdServer::signatureHelp(
+    PathRef File, Position Pos,
+    UniqueFunction<void(llvm::Expected<Tagged<SignatureHelp>>)> Callback,
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   if (UsedFS)
     *UsedFS = TaggedFS.Value;
 
   std::string Contents;
   if (OverridenContents) {
     Contents = OverridenContents->str();
   } else {
     VersionedDraft Latest = DraftMgr.getDraft(File);
     if (!Latest.Draft)
-      return llvm::make_error<llvm::StringError>(
+      return Callback(llvm::make_error<llvm::StringError>(
           "signatureHelp is called for non-added document",
-          llvm::errc::invalid_argument);
+          llvm::errc::invalid_argument));
     Contents = std::move(*Latest.Draft);
   }
 
-  auto Action = [=](llvm::Expected<InputsAndPreamble> IP)
-      -> Expected<Tagged<SignatureHelp>> {
+  auto PCHs = this->PCHs;
+  auto Action = [Contents, Pos, TaggedFS,
+                 PCHs](Path File, decltype(Callback) Callback,
+                       llvm::Expected<InputsAndPreamble> IP) {
     if (!IP)
-      return IP.takeError();
+      return Callback(IP.takeError());
+
     auto PreambleData = IP->Preamble;
     auto &Command = IP->Inputs.CompileCommand;
-
-    return make_tagged(
+    Callback(make_tagged(
         clangd::signatureHelp(File, Command,
                               PreambleData ? &PreambleData->Preamble : nullptr,
                               Contents, Pos, TaggedFS.Value, PCHs),
-        TaggedFS.Tag);
+        TaggedFS.Tag));
   };
-  return blockingRunWithPreamble<Expected<Tagged<SignatureHelp>>>(WorkScheduler,
-                                                                  File, Action);
+
+  WorkScheduler.runWithPreamble(
+      File, BindWithForward(Action, File.str(), std::move(Callback)));
 }
 
 llvm::Expected<tooling::Replacements>
 ClangdServer::formatRange(StringRef Code, PathRef File, Range Rng) {
   size_t Begin = positionToOffset(Code, Rng.start);
   size_t Len = positionToOffset(Code, Rng.end) - Begin;
   return formatCode(Code, File, {tooling::Range(Begin, Len)});
 }
ClangdServer::formatOnType(StringRef Code, PathRef File, Position Pos) {
   size_t PreviousLBracePos = StringRef(Code).find_last_of('{', CursorPos);
   if (PreviousLBracePos == StringRef::npos)
     PreviousLBracePos = CursorPos;
   size_t Len = CursorPos - PreviousLBracePos;
 
   return formatCode(Code, File, {tooling::Range(PreviousLBracePos, Len)});
 }
 
-Expected<std::vector<tooling::Replacement>>
-ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName) {
-  using RetType = Expected<std::vector<tooling::Replacement>>;
-  auto Action = [=](Expected<InputsAndAST> InpAST) -> RetType {
+void ClangdServer::rename(
+    PathRef File, Position Pos, llvm::StringRef NewName,
+    UniqueFunction<void(Expected<std::vector<tooling::Replacement>>)>
+        Callback) {
+  auto Action = [Pos](Path File, std::string NewName,
+                      decltype(Callback) Callback,
+                      Expected<InputsAndAST> InpAST) {
     if (!InpAST)
-      return InpAST.takeError();
+      return Callback(InpAST.takeError());
     auto &AST = InpAST->AST;
 
     RefactoringResultCollector ResultCollector;
     const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
     const FileEntry *FE =
         SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
     if (!FE)
-      return llvm::make_error<llvm::StringError>(
-          "rename called for non-added document", llvm::errc::invalid_argument);
+      return Callback(llvm::make_error<llvm::StringError>(
+          "rename called for non-added document",
+          llvm::errc::invalid_argument));
     SourceLocation SourceLocationBeg =
         clangd::getBeginningOfIdentifier(AST, Pos, FE);
     tooling::RefactoringRuleContext Context(
         AST.getASTContext().getSourceManager());
     Context.setASTContext(AST.getASTContext());
     auto Rename = clang::tooling::RenameOccurrences::initiate(
-        Context, SourceRange(SourceLocationBeg), NewName.str());
+        Context, SourceRange(SourceLocationBeg), NewName);
     if (!Rename)
-      return Rename.takeError();
+      return Callback(Rename.takeError());
 
     Rename->invoke(ResultCollector, Context);
 
     assert(ResultCollector.Result.hasValue());
     if (!ResultCollector.Result.getValue())
-      return ResultCollector.Result->takeError();
+      return Callback(ResultCollector.Result->takeError());
 
     std::vector<tooling::Replacement> Replacements;
     for (const tooling::AtomicChange &Change : ResultCollector.Result->get()) {
       tooling::Replacements ChangeReps = Change.getReplacements();
       for (const auto &Rep : ChangeReps) {
         // FIXME: Right now we only support renaming the main file, so we
         // drop replacements not for the main file. In the future, we might
         // consider to support:
         //   * rename in any included header
         //   * rename only in the "main" header
         //   * provide an error if there are symbols we won't rename (e.g.
         //     std::vector)
         //   * rename globally in project
         //   * rename in open files
         if (Rep.getFilePath() == File)
           Replacements.push_back(Rep);
       }
     }
-    return Replacements;
+    return Callback(Replacements);
   };
-  return blockingRunWithAST<RetType>(WorkScheduler, File, std::move(Action));
+
+  WorkScheduler.runWithAST(
+      File,
+      BindWithForward(Action, File.str(), NewName.str(), std::move(Callback)));
 }
 
 llvm::Optional<std::string> ClangdServer::getDocument(PathRef File) {
   auto Latest = DraftMgr.getDraft(File);
   if (!Latest.Draft)
     return llvm::None;
   return std::move(*Latest.Draft);
 }
 
-std::string ClangdServer::dumpAST(PathRef File) {
-  auto Action = [](llvm::Expected<InputsAndAST> InpAST) -> std::string {
+void ClangdServer::dumpAST(PathRef File,
+                           UniqueFunction<void(std::string)> Callback) {
+  auto Action = [](decltype(Callback) Callback,
+                   llvm::Expected<InputsAndAST> InpAST) {
     if (!InpAST) {
       ignoreError(InpAST.takeError());
-      return "<no-ast>";
+      return Callback("<no-ast>");
     }
-
     std::string Result;
 
     llvm::raw_string_ostream ResultOS(Result);
     clangd::dumpAST(InpAST->AST, ResultOS);
     ResultOS.flush();
 
-    return Result;
+    Callback(Result);
   };
-  return blockingRunWithAST<std::string>(WorkScheduler, File,
-                                         std::move(Action));
+
+  WorkScheduler.runWithAST(File, BindWithForward(Action, std::move(Callback)));
 }
 
-llvm::Expected<Tagged<std::vector<Location>>>
-ClangdServer::findDefinitions(PathRef File, Position Pos) {
+void ClangdServer::findDefinitions(
+    PathRef File, Position Pos,
+    UniqueFunction<void(llvm::Expected<Tagged<std::vector<Location>>>)>
+        Callback) {
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
-
-  using RetType = llvm::Expected<Tagged<std::vector<Location>>>;
-  auto Action = [=](llvm::Expected<InputsAndAST> InpAST) -> RetType {
+  auto Action = [Pos, TaggedFS](decltype(Callback) Callback,
+                                llvm::Expected<InputsAndAST> InpAST) {
     if (!InpAST)
-      return InpAST.takeError();
+      return Callback(InpAST.takeError());
     auto Result = clangd::findDefinitions(InpAST->AST, Pos);
-    return make_tagged(std::move(Result), TaggedFS.Tag);
+    Callback(make_tagged(std::move(Result), TaggedFS.Tag));
   };
-  return blockingRunWithAST<RetType>(WorkScheduler, File, Action);
+
+  WorkScheduler.runWithAST(File, BindWithForward(Action, std::move(Callback)));
 }
 
 llvm::Optional<Path> ClangdServer::switchSourceHeader(PathRef Path) {
 
   StringRef SourceExtensions[] = {".cpp", ".c", ".cc", ".cxx",
                                   ".c++", ".m", ".mm"};
   StringRef HeaderExtensions[] = {".h", ".hh", ".hpp", ".hxx", ".inc"};
 
auto StyleOrError =
       format::getStyle("file", File, "LLVM", Code, TaggedFS.Value.get());
   if (!StyleOrError) {
     return StyleOrError.takeError();
   } else {
     return format::reformat(StyleOrError.get(), Code, Ranges, File);
   }
 }
 
-llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
-ClangdServer::findDocumentHighlights(PathRef File, Position Pos) {
+void ClangdServer::findDocumentHighlights(
+    PathRef File, Position Pos,
+    UniqueFunction<void(llvm::Expected<Tagged<std::vector<DocumentHighlight>>>)>
+        Callback) {
   auto FileContents = DraftMgr.getDraft(File);
   if (!FileContents.Draft)
-    return llvm::make_error<llvm::StringError>(
+    return Callback(llvm::make_error<llvm::StringError>(
         "findDocumentHighlights called on non-added file",
-        llvm::errc::invalid_argument);
+        llvm::errc::invalid_argument));
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
 
-  using RetType = llvm::Expected<Tagged<std::vector<DocumentHighlight>>>;
-  auto Action = [=](llvm::Expected<InputsAndAST> InpAST) -> RetType {
+  auto Action = [TaggedFS, Pos](decltype(Callback) Callback,
+                                llvm::Expected<InputsAndAST> InpAST) {
     if (!InpAST)
-      return InpAST.takeError();
+      return Callback(InpAST.takeError());
     auto Result = clangd::findDocumentHighlights(InpAST->AST, Pos);
-    return make_tagged(std::move(Result), TaggedFS.Tag);
+    Callback(make_tagged(std::move(Result), TaggedFS.Tag));
   };
-  return blockingRunWithAST<RetType>(WorkScheduler, File, Action);
+
+  WorkScheduler.runWithAST(File, BindWithForward(Action, std::move(Callback)));
 }
 
 void ClangdServer::scheduleReparseAndDiags(
     PathRef File, VersionedDraft Contents,
     Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS) {
   tooling::CompileCommand Command = CompileArgs.getCompileCommand(File);
 
   using OptDiags = llvm::Optional<std::vector<DiagWithFixIts>>;

unittests/clangd/ClangdTests.cpp

while (RE.match(Pending, &Matches)) {
     Pending = Pending.drop_front(MatchPos + Matches[0].size());
   }
   Result += Pending;
 
   return Result;
 }
 
 std::string dumpASTWithoutMemoryLocs(ClangdServer &Server, PathRef File) {
-  auto DumpWithMemLocs = Server.dumpAST(File);
+  auto DumpWithMemLocs = runDumpAST(Server, File);
   return replacePtrsInDump(DumpWithMemLocs);
 }
 
 } // namespace
 
 class ClangdVFSTest : public ::testing::Test {
 protected:
   std::string parseSourceAndDumpAST(
TEST_F(ClangdVFSTest, InvalidCompileCommand) {
   // clang cannot create CompilerInvocation if we pass two files in the
   // CompileCommand. We pass the file in ExtraFlags once and CDB adds another
   // one in getCompileCommand().
   CDB.ExtraClangFlags.push_back(FooCpp.str());
 
   // Clang can't parse command args in that case, but we shouldn't crash.
   Server.addDocument(FooCpp, "int main() {}");
 
-  EXPECT_EQ(Server.dumpAST(FooCpp), "<no-ast>");
-  EXPECT_ERROR(Server.findDefinitions(FooCpp, Position()));
-  EXPECT_ERROR(Server.findDocumentHighlights(FooCpp, Position()));
-  EXPECT_ERROR(Server.rename(FooCpp, Position(), "new_name"));
+  EXPECT_EQ(runDumpAST(Server, FooCpp), "<no-ast>");
+  EXPECT_ERROR(runFindDefinitions(Server, FooCpp, Position()));
+  EXPECT_ERROR(runFindDocumentHighlights(Server, FooCpp, Position()));
+  EXPECT_ERROR(runRename(Server, FooCpp, Position(), "new_name"));
   // FIXME: codeComplete and signatureHelp should also return errors when they
   // can't parse the file.
   EXPECT_THAT(
       runCodeComplete(Server, FooCpp, Position(), clangd::CodeCompleteOptions())
           .Value.items,
       IsEmpty());
-  auto SigHelp = Server.signatureHelp(FooCpp, Position());
+  auto SigHelp = runSignatureHelp(Server, FooCpp, Position());
   ASSERT_TRUE(bool(SigHelp)) << "signatureHelp returned an error";
   EXPECT_THAT(SigHelp->Value.signatures, IsEmpty());
 }
 
 class ClangdThreadingTest : public ClangdVFSTest {};
 
 TEST_F(ClangdThreadingTest, StressTest) {
   // Without 'static' clang gives an error for a usage inside TestDiagConsumer.
auto FindDefinitionsRequest = [&]() {
       // Make sure we don't violate the ClangdServer's contract.
       if (ReqStats[FileIndex].FileIsRemoved)
         AddDocument(FileIndex);
 
       Position Pos;
       Pos.line = LineDist(RandGen);
       Pos.character = ColumnDist(RandGen);
 
-      ASSERT_TRUE(!!Server.findDefinitions(FilePaths[FileIndex], Pos));
+      ASSERT_TRUE(!!runFindDefinitions(Server, FilePaths[FileIndex], Pos));
     };
 
     std::vector<std::function<void()>> AsyncRequests = {
         AddDocumentRequest, ForceReparseRequest, RemoveDocumentRequest};
     std::vector<std::function<void()>> BlockingRequests = {
         CodeCompletionRequest, FindDefinitionsRequest};
 
     // Bash requests to ClangdServer in a loop.
TEST_F(ClangdVFSTest, CheckSourceHeaderSwitch) {
   EXPECT_FALSE(PathResult.hasValue());
 }
 
 TEST_F(ClangdThreadingTest, NoConcurrentDiagnostics) {
   class NoConcurrentAccessDiagConsumer : public DiagnosticsConsumer {
   public:
     std::atomic<int> Count = {0};
 
-     NoConcurrentAccessDiagConsumer(std::promise<void> StartSecondReparse)
-         : StartSecondReparse(std::move(StartSecondReparse)) {}
 
     void onDiagnosticsReady(PathRef,
                             Tagged<std::vector<DiagWithFixIts>>) override {
       ++Count;
       std::unique_lock<std::mutex> Lock(Mutex, std::try_to_lock_t());
       ASSERT_TRUE(Lock.owns_lock())
           << "Detected concurrent onDiagnosticsReady calls for the same file.";
 

unittests/clangd/CodeCompleteTests.cpp

SignatureHelp signatures(StringRef Text) {
   MockCompilationDatabase CDB;
   IgnoreDiagnostics DiagConsumer;
   ClangdServer Server(CDB, DiagConsumer, FS, getDefaultAsyncThreadsCount(),
                       /*StorePreamblesInMemory=*/true);
   auto File = getVirtualTestFilePath("foo.cpp");
   Annotations Test(Text);
   Server.addDocument(File, Test.code());
   EXPECT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for preamble";
-  auto R = Server.signatureHelp(File, Test.point());
+  auto R = runSignatureHelp(Server, File, Test.point());
   assert(R);
   return R.get().Value;
 }
 
 MATCHER_P(ParamsAre, P, "") {
   if (P.size() != arg.parameters.size())
     return false;
   for (unsigned I = 0; I < P.size(); ++I)

unittests/clangd/SyncAPI.h

 
 namespace clang {
 namespace clangd {
 
 Tagged<CompletionList>
 runCodeComplete(ClangdServer &Server, PathRef File, Position Pos,
                 clangd::CodeCompleteOptions Opts,
                 llvm::Optional<StringRef> OverridenContents = llvm::None);
+
+llvm::Expected<Tagged<SignatureHelp>>
+runSignatureHelp(ClangdServer &Server, PathRef File, Position Pos,
+                 llvm::Optional<StringRef> OverridenContents = llvm::None);
+
+llvm::Expected<Tagged<std::vector<Location>>>
+runFindDefinitions(ClangdServer &Server, PathRef File, Position Pos);
+
+llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
+runFindDocumentHighlights(ClangdServer &Server, PathRef File, Position Pos);
+
+llvm::Expected<std::vector<tooling::Replacement>>
+runRename(ClangdServer &Server, PathRef File, Position Pos, StringRef NewName);
+
+std::string runDumpAST(ClangdServer &Server, PathRef File);
+
 } // namespace clangd
 } // namespace clang
 
 #endif

unittests/clangd/SyncAPI.cpp

 namespace clangd {
 
 namespace {
 /// A helper that waits for async callbacks to fire and exposes their result in
 /// the output variable. Intended to be used in the following way:
 ///    T Result;
 ///    someAsyncFunc(Param1, Param2, /*Callback=*/capture(Result));
 template <typename T> struct CaptureProxy {
-  CaptureProxy(T &Target) : Target(&Target) {}
+  CaptureProxy(llvm::Optional<T> &Target) : Target(&Target) {
+    assert(!Target.hasValue());
+  }
 
   CaptureProxy(const CaptureProxy &) = delete;
   CaptureProxy &operator=(const CaptureProxy &) = delete;
   // We need move ctor to return a value from the 'capture' helper.
   CaptureProxy(CaptureProxy &&Other) : Target(Other.Target) {
     Other.Target = nullptr;
   }
   CaptureProxy &operator=(CaptureProxy &&) = delete;
 
   operator UniqueFunction<void(T)>() && {
     assert(!Future.valid() && "conversion to callback called multiple times");
     Future = Promise.get_future();
     return BindWithForward([](std::promise<T> Promise,
                               T Value) { Promise.set_value(std::move(Value)); },
                            std::move(Promise));
   }
 
   ~CaptureProxy() {
     if (!Target)
       return;
     assert(Future.valid() && "conversion to callback was not called");
-    *Target = Future.get();
+    assert(!Target->hasValue());
+    Target->emplace(Future.get());
   }
 
 private:
-  T *Target;
+  llvm::Optional<T> *Target;
   std::promise<T> Promise;
   std::future<T> Future;
 };
 
-template <typename T> CaptureProxy<T> capture(T &Target) {
+template <typename T> CaptureProxy<T> capture(llvm::Optional<T> &Target) {
   return CaptureProxy<T>(Target);
 }
 } // namespace
 
 Tagged<CompletionList>
 runCodeComplete(ClangdServer &Server, PathRef File, Position Pos,
                 clangd::CodeCompleteOptions Opts,
                 llvm::Optional<StringRef> OverridenContents) {
-  Tagged<CompletionList> Result;
+  llvm::Optional<Tagged<CompletionList>> Result;
   Server.codeComplete(File, Pos, Opts, capture(Result), OverridenContents);
-  return Result;
+  return std::move(*Result);
+}
+
+llvm::Expected<Tagged<SignatureHelp>>
+runSignatureHelp(ClangdServer &Server, PathRef File, Position Pos,
+                 llvm::Optional<StringRef> OverridenContents) {
+  llvm::Optional<llvm::Expected<Tagged<SignatureHelp>>> Result;
+  Server.signatureHelp(File, Pos, capture(Result), OverridenContents);
+  return std::move(*Result);
+}
+
+llvm::Expected<Tagged<std::vector<Location>>>
+runFindDefinitions(ClangdServer &Server, PathRef File, Position Pos) {
+  llvm::Optional<llvm::Expected<Tagged<std::vector<Location>>>> Result;
+  Server.findDefinitions(File, Pos, capture(Result));
+  return std::move(*Result);
+}
+
+llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
+runFindDocumentHighlights(ClangdServer &Server, PathRef File, Position Pos) {
+  llvm::Optional<llvm::Expected<Tagged<std::vector<DocumentHighlight>>>> Result;
+  Server.findDocumentHighlights(File, Pos, capture(Result));
+  return std::move(*Result);
+}
+
+llvm::Expected<std::vector<tooling::Replacement>>
+runRename(ClangdServer &Server, PathRef File, Position Pos, StringRef NewName) {
+  llvm::Optional<llvm::Expected<std::vector<tooling::Replacement>>> Result;
+  Server.rename(File, Pos, NewName, capture(Result));
+  return std::move(*Result);
+}
+
+std::string runDumpAST(ClangdServer &Server, PathRef File) {
+  llvm::Optional<std::string> Result;
+  Server.dumpAST(File, capture(Result));
+  return std::move(*Result);
 }
 
 } // namespace clangd
 } // namespace clang

unittests/clangd/XRefsTests.cpp

 //===-- XRefsTests.cpp  ---------------------------*- C++ -*--------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 #include "Annotations.h"
 #include "ClangdUnit.h"
 #include "Compiler.h"
 #include "Matchers.h"
+#include "SyncAPI.h"
 #include "TestFS.h"
 #include "XRefs.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Frontend/Utils.h"
 #include "llvm/Support/Path.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
)cpp");
   MockFSProvider FS;
   ClangdServer Server(CDB, DiagConsumer, FS, /*AsyncThreadsCount=*/0,
                       /*StorePreambleInMemory=*/true);
 
   auto FooCpp = getVirtualTestFilePath("foo.cpp");
   FS.Files[FooCpp] = "";
 
   Server.addDocument(FooCpp, SourceAnnotations.code());
-  auto Locations = Server.findDefinitions(FooCpp, SourceAnnotations.point());
+  auto Locations =
+      runFindDefinitions(Server, FooCpp, SourceAnnotations.point());
   EXPECT_TRUE(bool(Locations)) << "findDefinitions returned an error";
 
   EXPECT_THAT(Locations->Value,
               ElementsAre(Location{URIForFile{FooCpp.str()},
                                    SourceAnnotations.range()}));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang