sammccall updated this revision to Diff 131564.
sammccall added a comment.

Rebase on the Span changes that landed as r323511


Repository:
  rCTE Clang Tools Extra

https://reviews.llvm.org/D42517

Files:
  clangd/ClangdLSPServer.cpp
  clangd/ClangdLSPServer.h
  clangd/ClangdServer.cpp
  clangd/ClangdServer.h
  clangd/ClangdUnit.cpp
  clangd/ClangdUnit.h
  clangd/CodeComplete.cpp
  clangd/CodeComplete.h
  clangd/Context.cpp
  clangd/Context.h
  clangd/GlobalCompilationDatabase.cpp
  clangd/JSONRPCDispatcher.cpp
  clangd/JSONRPCDispatcher.h
  clangd/Logger.cpp
  clangd/Logger.h
  clangd/ProtocolHandlers.cpp
  clangd/ProtocolHandlers.h
  clangd/Trace.cpp
  clangd/Trace.h
  clangd/XRefs.cpp
  clangd/XRefs.h
  clangd/index/FileIndex.cpp
  clangd/index/FileIndex.h
  clangd/index/Index.h
  clangd/index/MemIndex.cpp
  clangd/index/MemIndex.h
  clangd/index/Merge.cpp
  unittests/clangd/FileIndexTests.cpp
  unittests/clangd/IndexTests.cpp
  unittests/clangd/TraceTests.cpp

Index: unittests/clangd/TraceTests.cpp
===================================================================
--- unittests/clangd/TraceTests.cpp
+++ unittests/clangd/TraceTests.cpp
@@ -78,8 +78,8 @@
     auto JSONTracer = trace::createJSONTracer(OS);
     trace::Session Session(*JSONTracer);
     {
-      trace::Span Tracer(Context::empty(), "A");
-      trace::log(Tracer.Ctx, "B");
+      trace::Span Tracer("A");
+      trace::log("B");
     }
   }
 
Index: unittests/clangd/IndexTests.cpp
===================================================================
--- unittests/clangd/IndexTests.cpp
+++ unittests/clangd/IndexTests.cpp
@@ -92,8 +92,7 @@
 std::vector<std::string> match(const SymbolIndex &I,
                                const FuzzyFindRequest &Req) {
   std::vector<std::string> Matches;
-  auto Ctx = Context::empty();
-  I.fuzzyFind(Ctx, Req, [&](const Symbol &Sym) {
+  I.fuzzyFind(Req, [&](const Symbol &Sym) {
     Matches.push_back(
         (Sym.Scope + (Sym.Scope.empty() ? "" : "::") + Sym.Name).str());
   });
Index: unittests/clangd/FileIndexTests.cpp
===================================================================
--- unittests/clangd/FileIndexTests.cpp
+++ unittests/clangd/FileIndexTests.cpp
@@ -76,35 +76,32 @@
 std::vector<std::string> match(const SymbolIndex &I,
                                const FuzzyFindRequest &Req) {
   std::vector<std::string> Matches;
-  auto Ctx = Context::empty();
-  I.fuzzyFind(Ctx, Req, [&](const Symbol &Sym) {
+  I.fuzzyFind(Req, [&](const Symbol &Sym) {
     Matches.push_back((Sym.Scope + Sym.Name).str());
   });
   return Matches;
 }
 
 /// Create an ParsedAST for \p Code. Returns None if \p Code is empty.
 llvm::Optional<ParsedAST> build(std::string Path, llvm::StringRef Code) {
-  Context Ctx = Context::empty();
   if (Code.empty())
     return llvm::None;
   const char *Args[] = {"clang", "-xc++", Path.c_str()};
 
   auto CI = createInvocationFromCommandLine(Args);
 
   auto Buf = llvm::MemoryBuffer::getMemBuffer(Code);
-  auto AST = ParsedAST::Build(Ctx, std::move(CI), nullptr, std::move(Buf),
+  auto AST = ParsedAST::Build(std::move(CI), nullptr, std::move(Buf),
                               std::make_shared<PCHContainerOperations>(),
                               vfs::getRealFileSystem());
   assert(AST.hasValue());
   return std::move(*AST);
 }
 
 TEST(FileIndexTest, IndexAST) {
   FileIndex M;
-  auto Ctx = Context::empty();
   M.update(
-      Ctx, "f1",
+      "f1",
       build("f1", "namespace ns { void f() {} class X {}; }").getPointer());
 
   FuzzyFindRequest Req;
@@ -115,9 +112,8 @@
 
 TEST(FileIndexTest, NoLocal) {
   FileIndex M;
-  auto Ctx = Context::empty();
   M.update(
-      Ctx, "f1",
+      "f1",
       build("f1", "namespace ns { void f() { int local = 0; } class X {}; }")
           .getPointer());
 
@@ -128,12 +124,11 @@
 
 TEST(FileIndexTest, IndexMultiASTAndDeduplicate) {
   FileIndex M;
-  auto Ctx = Context::empty();
   M.update(
-      Ctx, "f1",
+      "f1",
       build("f1", "namespace ns { void f() {} class X {}; }").getPointer());
   M.update(
-      Ctx, "f2",
+      "f2",
       build("f2", "namespace ns { void ff() {} class X {}; }").getPointer());
 
   FuzzyFindRequest Req;
@@ -144,31 +139,28 @@
 
 TEST(FileIndexTest, RemoveAST) {
   FileIndex M;
-  auto Ctx = Context::empty();
   M.update(
-      Ctx, "f1",
+      "f1",
       build("f1", "namespace ns { void f() {} class X {}; }").getPointer());
 
   FuzzyFindRequest Req;
   Req.Query = "";
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(M, Req), UnorderedElementsAre("ns::f", "ns::X"));
 
-  M.update(Ctx, "f1", nullptr);
+  M.update("f1", nullptr);
   EXPECT_THAT(match(M, Req), UnorderedElementsAre());
 }
 
 TEST(FileIndexTest, RemoveNonExisting) {
   FileIndex M;
-  auto Ctx = Context::empty();
-  M.update(Ctx, "no", nullptr);
+  M.update("no", nullptr);
   EXPECT_THAT(match(M, FuzzyFindRequest()), UnorderedElementsAre());
 }
 
 TEST(FileIndexTest, IgnoreClassMembers) {
   FileIndex M;
-  auto Ctx = Context::empty();
-  M.update(Ctx, "f1",
+  M.update("f1",
            build("f1", "class X { static int m1; int m2; static void f(); };")
                .getPointer());
 
Index: clangd/index/Merge.cpp
===================================================================
--- clangd/index/Merge.cpp
+++ clangd/index/Merge.cpp
@@ -24,7 +24,7 @@
    //          - find the generating file from each Symbol which is Static-only
    //          - ask Dynamic if it has that file (needs new SymbolIndex method)
    //          - if so, drop the Symbol.
-   bool fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+   bool fuzzyFind(const FuzzyFindRequest &Req,
                   function_ref<void(const Symbol &)> Callback) const override {
      // We can't step through both sources in parallel. So:
      //  1) query all dynamic symbols, slurping results into a slab
@@ -34,13 +34,12 @@
      //  3) now yield all the dynamic symbols we haven't processed.
      bool More = false; // We'll be incomplete if either source was.
      SymbolSlab::Builder DynB;
-     More |=
-         Dynamic->fuzzyFind(Ctx, Req, [&](const Symbol &S) { DynB.insert(S); });
+     More |= Dynamic->fuzzyFind(Req, [&](const Symbol &S) { DynB.insert(S); });
      SymbolSlab Dyn = std::move(DynB).build();
 
      DenseSet<SymbolID> SeenDynamicSymbols;
      Symbol::Details Scratch;
-     More |= Static->fuzzyFind(Ctx, Req, [&](const Symbol &S) {
+     More |= Static->fuzzyFind(Req, [&](const Symbol &S) {
        auto DynS = Dyn.find(S.ID);
        if (DynS == Dyn.end())
          return Callback(S);
Index: clangd/index/MemIndex.h
===================================================================
--- clangd/index/MemIndex.h
+++ clangd/index/MemIndex.h
@@ -28,7 +28,7 @@
   static std::unique_ptr<SymbolIndex> build(SymbolSlab Slab);
 
   bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 private:
Index: clangd/index/MemIndex.cpp
===================================================================
--- clangd/index/MemIndex.cpp
+++ clangd/index/MemIndex.cpp
@@ -29,7 +29,7 @@
 }
 
 bool MemIndex::fuzzyFind(
-    const Context &Ctx, const FuzzyFindRequest &Req,
+    const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   assert(!StringRef(Req.Query).contains("::") &&
          "There must be no :: in query.");
Index: clangd/index/Index.h
===================================================================
--- clangd/index/Index.h
+++ clangd/index/Index.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_INDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_INDEX_H
 
-#include "../Context.h"
 #include "clang/Index/IndexSymbol.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
@@ -244,7 +243,7 @@
   /// Returns true if the result list is complete, false if it was truncated due
   /// to MaxCandidateCount
   virtual bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const = 0;
 
   // FIXME: add interfaces for more index use cases:
Index: clangd/index/FileIndex.h
===================================================================
--- clangd/index/FileIndex.h
+++ clangd/index/FileIndex.h
@@ -17,7 +17,6 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "../ClangdUnit.h"
-#include "../Context.h"
 #include "Index.h"
 #include "MemIndex.h"
 
@@ -58,10 +57,10 @@
 public:
   /// \brief Update symbols in \p Path with symbols in \p AST. If \p AST is
   /// nullptr, this removes all symbols in the file
-  void update(const Context &Ctx, PathRef Path, ParsedAST *AST);
+  void update(PathRef Path, ParsedAST *AST);
 
   bool
-  fuzzyFind(const Context &Ctx, const FuzzyFindRequest &Req,
+  fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
 private:
Index: clangd/index/FileIndex.cpp
===================================================================
--- clangd/index/FileIndex.cpp
+++ clangd/index/FileIndex.cpp
@@ -69,7 +69,7 @@
   return {std::move(Snap), Pointers};
 }
 
-void FileIndex::update(const Context &Ctx, PathRef Path, ParsedAST *AST) {
+void FileIndex::update(PathRef Path, ParsedAST *AST) {
   if (!AST) {
     FSymbols.update(Path, nullptr);
   } else {
@@ -82,9 +82,9 @@
 }
 
 bool FileIndex::fuzzyFind(
-    const Context &Ctx, const FuzzyFindRequest &Req,
+    const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
-  return Index.fuzzyFind(Ctx, Req, Callback);
+  return Index.fuzzyFind(Req, Callback);
 }
 
 } // namespace clangd
Index: clangd/XRefs.h
===================================================================
--- clangd/XRefs.h
+++ clangd/XRefs.h
@@ -14,20 +14,18 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_XREFS_H
 
 #include "ClangdUnit.h"
-#include "Context.h"
 #include "Protocol.h"
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
 /// Get definition of symbol at a specified \p Pos.
-std::vector<Location> findDefinitions(const Context &Ctx, ParsedAST &AST,
-                                      Position Pos);
+std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos);
 
 /// Returns highlights for all usages of a symbol at \p Pos.
-std::vector<DocumentHighlight>
-findDocumentHighlights(const Context &Ctx, ParsedAST &AST, Position Pos);
+std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
+                                                      Position Pos);
 
 } // namespace clangd
 } // namespace clang
Index: clangd/XRefs.cpp
===================================================================
--- clangd/XRefs.cpp
+++ clangd/XRefs.cpp
@@ -146,8 +146,7 @@
 
 } // namespace
 
-std::vector<Location> findDefinitions(const Context &Ctx, ParsedAST &AST,
-                                      Position Pos) {
+std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
   const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
   if (!FE)
@@ -258,8 +257,8 @@
 
 } // namespace
 
-std::vector<DocumentHighlight>
-findDocumentHighlights(const Context &Ctx, ParsedAST &AST, Position Pos) {
+std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
+                                                      Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
   const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
   if (!FE)
Index: clangd/Trace.h
===================================================================
--- clangd/Trace.h
+++ clangd/Trace.h
@@ -43,8 +43,7 @@
                             json::obj *Args) = 0;
 
   /// Called for instant events.
-  virtual void instant(const Context &Ctx, llvm::StringRef Name,
-                       json::obj &&Args) = 0;
+  virtual void instant(llvm::StringRef Name, json::obj &&Args) = 0;
 };
 
 /// Sets up a global EventTracer that consumes events produced by Span and
@@ -65,10 +64,10 @@
                                               bool Pretty = false);
 
 /// Records a single instant event, associated with the current thread.
-void log(const Context &Ctx, const llvm::Twine &Name);
+void log(const llvm::Twine &Name);
 
 /// Records an event whose duration is the lifetime of the Span object.
-/// This lifetime is extended when the span's context is reused.
+/// This lifetime is extended if the current Context escapes the Span lifetime.
 ///
 /// This is the main public interface for producing tracing events.
 ///
@@ -78,18 +77,15 @@
 /// SomeJSONExpr is evaluated and copied only if actually needed.
 class Span {
 public:
-  Span(const Context &Ctx, llvm::StringRef Name);
+  Span(llvm::StringRef Name);
 
   /// Mutable metadata, if this span is interested.
   /// Prefer to use SPAN_ATTACH rather than accessing this directly.
   json::obj *const Args;
-  /// Propagating this context will keep the span alive.
-  const Context Ctx;
-};
 
-/// Returns mutable span metadata if this span is interested.
-/// Prefer to use SPAN_ATTACH rather than accessing this directly.
-json::obj *spanArgs(const Context &Ctx);
+private:
+  llvm::Optional<WithContext> RestoreOriginalContext;
+};
 
 /// Attach a key-value pair to a Span event.
 /// This is not threadsafe when used with the same Span.
Index: clangd/Trace.cpp
===================================================================
--- clangd/Trace.cpp
+++ clangd/Trace.cpp
@@ -7,6 +7,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Context.h"
+#include "Function.h"
 #include "Trace.h"
 #include "Context.h"
 #include "Function.h"
@@ -54,8 +56,7 @@
     }));
   }
 
-  void instant(const Context &Ctx, llvm::StringRef Name,
-               json::obj &&Args) override {
+  void instant(llvm::StringRef Name, json::obj &&Args) override {
     jsonEvent("i", json::obj{{"name", Name}, {"args", std::move(Args)}});
   }
 
@@ -120,20 +121,23 @@
   return llvm::make_unique<JSONTracer>(OS, Pretty);
 }
 
-void log(const Context &Ctx, const Twine &Message) {
+void log(const Twine &Message) {
   if (!T)
     return;
-  T->instant(Ctx, "Log", json::obj{{"Message", Message.str()}});
+  T->instant("Log", json::obj{{"Message", Message.str()}});
 }
 
-// Span keeps a non-owning pointer to the args, which is how users access them.
-// The args are owned by the context though. They stick around until the
-// beginSpan() context is destroyed, when the tracing engine will consume them.
-Span::Span(const Context &Ctx, llvm::StringRef Name)
-    : Args(T ? new json::obj() : nullptr),
-      Ctx(T ? T->beginSpan(Ctx.derive(std::unique_ptr<json::obj>(Args)), Name,
-                           Args)
-            : Ctx.clone()) {}
+Span::Span(llvm::StringRef Name)
+    // Span keeps a non-owning pointer to the args, for users to access.
+    : Args(T ? new json::obj() /* SpanCtx will take ownership */ : nullptr) {
+  if (T) {
+    // The args are owned by the context though. They must outlive SpanCtx,
+    // so that the tracing engine can consume them when the event ends.
+    auto ArgsCtx = Context::current().derive(std::unique_ptr<json::obj>(Args));
+    auto SpanCtx = T->beginSpan(std::move(ArgsCtx), Name, Args);
+    RestoreOriginalContext.emplace(std::move(SpanCtx));
+  }
+}
 
 } // namespace trace
 } // namespace clangd
Index: clangd/ProtocolHandlers.h
===================================================================
--- clangd/ProtocolHandlers.h
+++ clangd/ProtocolHandlers.h
@@ -29,34 +29,28 @@
 // The interface implemented by ClangLSPServer to handle incoming requests.
 class ProtocolCallbacks {
 public:
-  // FIXME(ibiryukov): remove this typedef, inline its usages.
-  using Ctx = Context;
   virtual ~ProtocolCallbacks() = default;
 
-  virtual void onInitialize(Ctx C, InitializeParams &Params) = 0;
-  virtual void onShutdown(Ctx C, ShutdownParams &Params) = 0;
-  virtual void onExit(Ctx C, ExitParams &Params) = 0;
-  virtual void onDocumentDidOpen(Ctx C, DidOpenTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidChange(Ctx C,
-                                   DidChangeTextDocumentParams &Params) = 0;
-  virtual void onDocumentDidClose(Ctx C,
-                                  DidCloseTextDocumentParams &Params) = 0;
-  virtual void onDocumentFormatting(Ctx C,
-                                    DocumentFormattingParams &Params) = 0;
+  virtual void onInitialize(InitializeParams &Params) = 0;
+  virtual void onShutdown(ShutdownParams &Params) = 0;
+  virtual void onExit(ExitParams &Params) = 0;
+  virtual void onDocumentDidOpen(DidOpenTextDocumentParams &Params) = 0;
+  virtual void onDocumentDidChange(DidChangeTextDocumentParams &Params) = 0;
+  virtual void onDocumentDidClose(DidCloseTextDocumentParams &Params) = 0;
+  virtual void onDocumentFormatting(DocumentFormattingParams &Params) = 0;
   virtual void
-  onDocumentOnTypeFormatting(Ctx C, DocumentOnTypeFormattingParams &Params) = 0;
+  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) = 0;
   virtual void
-  onDocumentRangeFormatting(Ctx C, DocumentRangeFormattingParams &Params) = 0;
-  virtual void onCodeAction(Ctx C, CodeActionParams &Params) = 0;
-  virtual void onCompletion(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onSignatureHelp(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onGoToDefinition(Ctx C, TextDocumentPositionParams &Params) = 0;
-  virtual void onSwitchSourceHeader(Ctx C, TextDocumentIdentifier &Params) = 0;
-  virtual void onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) = 0;
-  virtual void onCommand(Ctx C, ExecuteCommandParams &Params) = 0;
-  virtual void onRename(Ctx C, RenameParams &Parames) = 0;
-  virtual void onDocumentHighlight(Ctx C,
-                                   TextDocumentPositionParams &Params) = 0;
+  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) = 0;
+  virtual void onCodeAction(CodeActionParams &Params) = 0;
+  virtual void onCompletion(TextDocumentPositionParams &Params) = 0;
+  virtual void onSignatureHelp(TextDocumentPositionParams &Params) = 0;
+  virtual void onGoToDefinition(TextDocumentPositionParams &Params) = 0;
+  virtual void onSwitchSourceHeader(TextDocumentIdentifier &Params) = 0;
+  virtual void onFileEvent(DidChangeWatchedFilesParams &Params) = 0;
+  virtual void onCommand(ExecuteCommandParams &Params) = 0;
+  virtual void onRename(RenameParams &Parames) = 0;
+  virtual void onDocumentHighlight(TextDocumentPositionParams &Params) = 0;
 };
 
 void registerCallbackHandlers(JSONRPCDispatcher &Dispatcher, JSONOutput &Out,
Index: clangd/ProtocolHandlers.cpp
===================================================================
--- clangd/ProtocolHandlers.cpp
+++ clangd/ProtocolHandlers.cpp
@@ -24,17 +24,16 @@
 // FooParams should have a fromJSON function.
 struct HandlerRegisterer {
   template <typename Param>
-  void operator()(StringRef Method,
-                  void (ProtocolCallbacks::*Handler)(Context, Param)) {
+  void operator()(StringRef Method, void (ProtocolCallbacks::*Handler)(Param)) {
     // Capture pointers by value, as the lambda will outlive this object.
     auto *Callbacks = this->Callbacks;
     Dispatcher.registerHandler(
-        Method, [=](Context C, const json::Expr &RawParams) {
+        Method, [=](const json::Expr &RawParams) {
           typename std::remove_reference<Param>::type P;
           if (fromJSON(RawParams, P)) {
-            (Callbacks->*Handler)(std::move(C), P);
+            (Callbacks->*Handler)(P);
           } else {
-            log(C, "Failed to decode " + Method + " request.");
+            log("Failed to decode " + Method + " request.");
           }
         });
   }
Index: clangd/Logger.h
===================================================================
--- clangd/Logger.h
+++ clangd/Logger.h
@@ -10,24 +10,23 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_LOGGER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_LOGGER_H
 
-#include "Context.h"
 #include "llvm/ADT/Twine.h"
 
 namespace clang {
 namespace clangd {
 
 /// Main logging function.
 /// Logs messages to a global logger, which can be set up by LoggingSesssion.
 /// If no logger is registered, writes to llvm::errs().
-void log(const Context &Ctx, const llvm::Twine &Message);
+void log(const llvm::Twine &Message);
 
 /// Interface to allow custom logging in clangd.
 class Logger {
 public:
   virtual ~Logger() = default;
 
   /// Implementations of this method must be thread-safe.
-  virtual void log(const Context &Ctx, const llvm::Twine &Message) = 0;
+  virtual void log(const llvm::Twine &Message) = 0;
 };
 
 /// Only one LoggingSession can be active at a time.
Index: clangd/Logger.cpp
===================================================================
--- clangd/Logger.cpp
+++ clangd/Logger.cpp
@@ -25,9 +25,9 @@
 
 LoggingSession::~LoggingSession() { L = nullptr; }
 
-void log(const Context &Ctx, const llvm::Twine &Message) {
+void log(const llvm::Twine &Message) {
   if (L)
-    L->log(Ctx, Message);
+    L->log(Message);
   else {
     static std::mutex Mu;
     std::lock_guard<std::mutex> Guard(Mu);
Index: clangd/JSONRPCDispatcher.h
===================================================================
--- clangd/JSONRPCDispatcher.h
+++ clangd/JSONRPCDispatcher.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_JSONRPCDISPATCHER_H
 
-#include "Context.h"
 #include "JSONExpr.h"
 #include "Logger.h"
 #include "Protocol.h"
@@ -38,7 +37,7 @@
   void writeMessage(const json::Expr &Result);
 
   /// Write a line to the logging stream.
-  void log(const Context &Ctx, const Twine &Message) override;
+  void log(const Twine &Message) override;
 
   /// Mirror \p Message into InputMirror stream. Does nothing if InputMirror is
   /// null.
@@ -56,23 +55,22 @@
   std::mutex StreamMutex;
 };
 
-/// Sends a successful reply. \p Ctx must either be the Context accepted by
-/// JSONRPCDispatcher::Handler or be derived from it.
-void reply(const Context &Ctx, json::Expr &&Result);
-/// Sends an error response to the client, and logs it. \p Ctx must either be
-/// the Context accepted by JSONRPCDispatcher::Handler or be derived from it.
-void replyError(const Context &Ctx, ErrorCode code,
-                const llvm::StringRef &Message);
-/// Sends a request to the client. \p Ctx must either be the Context accepted by
-/// JSONRPCDispatcher::Handler or be derived from it.
-void call(const Context &Ctx, llvm::StringRef Method, json::Expr &&Params);
+/// Sends a successful reply.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void reply(json::Expr &&Result);
+/// Sends an error response to the client, and logs it.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void replyError(ErrorCode code, const llvm::StringRef &Message);
+/// Sends a request to the client.
+/// Current context must derive from JSONRPCDispatcher::Handler.
+void call(llvm::StringRef Method, json::Expr &&Params);
 
 /// Main JSONRPC entry point. This parses the JSONRPC "header" and calls the
 /// registered Handler for the method received.
 class JSONRPCDispatcher {
 public:
   // A handler responds to requests for a particular method name.
-  using Handler = std::function<void(Context, const json::Expr &)>;
+  using Handler = std::function<void(const json::Expr &)>;
 
   /// Create a new JSONRPCDispatcher. UnknownHandler is called when an unknown
   /// method is received.
Index: clangd/JSONRPCDispatcher.cpp
===================================================================
--- clangd/JSONRPCDispatcher.cpp
+++ clangd/JSONRPCDispatcher.cpp
@@ -34,14 +34,15 @@
 public:
   // Return a context that's aware of the enclosing request, identified by Span.
   static Context stash(const trace::Span &Span) {
-    return Span.Ctx.derive(RequestSpan::Key, std::unique_ptr<RequestSpan>(
-                                                 new RequestSpan(Span.Args)));
+    return Context::current().derive(
+        RequestSpan::Key,
+        std::unique_ptr<RequestSpan>(new RequestSpan(Span.Args)));
   }
 
   // If there's an enclosing request and the tracer is interested, calls \p F
   // with a json::obj where request info can be added.
-  template <typename Func> static void attach(const Context &Ctx, Func &&F) {
-    auto *RequestArgs = Ctx.get(RequestSpan::Key);
+  template <typename Func> static void attach(Func &&F) {
+    auto *RequestArgs = Context::current().get(RequestSpan::Key);
     if (!RequestArgs || !*RequestArgs || !(*RequestArgs)->Args)
       return;
     std::lock_guard<std::mutex> Lock((*RequestArgs)->Mu);
@@ -70,8 +71,8 @@
   Outs.flush();
 }
 
-void JSONOutput::log(const Context &Ctx, const Twine &Message) {
-  trace::log(Ctx, Message);
+void JSONOutput::log(const Twine &Message) {
+  trace::log(Message);
   std::lock_guard<std::mutex> Guard(StreamMutex);
   Logs << Message << '\n';
   Logs.flush();
@@ -85,33 +86,32 @@
   InputMirror->flush();
 }
 
-void clangd::reply(const Context &Ctx, json::Expr &&Result) {
-  auto ID = Ctx.get(RequestID);
+void clangd::reply(json::Expr &&Result) {
+  auto ID = Context::current().get(RequestID);
   if (!ID) {
-    log(Ctx, "Attempted to reply to a notification!");
+    log("Attempted to reply to a notification!");
     return;
   }
+  RequestSpan::attach([&](json::obj &Args) { Args["Reply"] = Result; });
 
-  RequestSpan::attach(Ctx, [&](json::obj &Args) { Args["Reply"] = Result; });
-
-  Ctx.getExisting(RequestOut)
+  Context::current().getExisting(RequestOut)
       ->writeMessage(json::obj{
           {"jsonrpc", "2.0"},
           {"id", *ID},
           {"result", std::move(Result)},
       });
 }
 
-void clangd::replyError(const Context &Ctx, ErrorCode code,
-                        const llvm::StringRef &Message) {
-  log(Ctx, "Error " + Twine(static_cast<int>(code)) + ": " + Message);
-  RequestSpan::attach(Ctx, [&](json::obj &Args) {
+void clangd::replyError(ErrorCode code, const llvm::StringRef &Message) {
+  log("Error " + Twine(static_cast<int>(code)) + ": " + Message);
+  RequestSpan::attach([&](json::obj &Args) {
     Args["Error"] =
         json::obj{{"code", static_cast<int>(code)}, {"message", Message.str()}};
   });
 
-  if (auto ID = Ctx.get(RequestID)) {
-    Ctx.getExisting(RequestOut)
+  if (auto ID = Context::current().get(RequestID)) {
+    Context::current()
+        .getExisting(RequestOut)
         ->writeMessage(json::obj{
             {"jsonrpc", "2.0"},
             {"id", *ID},
@@ -121,13 +121,14 @@
   }
 }
 
-void clangd::call(const Context &Ctx, StringRef Method, json::Expr &&Params) {
+void clangd::call(StringRef Method, json::Expr &&Params) {
   // FIXME: Generate/Increment IDs for every request so that we can get proper
   // replies once we need to.
-  RequestSpan::attach(Ctx, [&](json::obj &Args) {
+  RequestSpan::attach([&](json::obj &Args) {
     Args["Call"] = json::obj{{"method", Method.str()}, {"params", Params}};
   });
-  Ctx.getExisting(RequestOut)
+  Context::current()
+      .getExisting(RequestOut)
       ->writeMessage(json::obj{
           {"jsonrpc", "2.0"},
           {"id", 1},
@@ -163,18 +164,20 @@
   auto &Handler = I != Handlers.end() ? I->second : UnknownHandler;
 
   // Create a Context that contains request information.
-  auto Ctx = Context::empty().derive(RequestOut, &Out);
+  WithContextValue WithRequestOut(RequestOut, &Out);
+  llvm::Optional<WithContextValue> WithID;
   if (ID)
-    Ctx = std::move(Ctx).derive(RequestID, *ID);
+    WithID.emplace(RequestID, *ID);
 
   // Create a tracing Span covering the whole request lifetime.
-  trace::Span Tracer(Ctx, *Method);
+  trace::Span Tracer(*Method);
   if (ID)
     SPAN_ATTACH(Tracer, "ID", *ID);
   SPAN_ATTACH(Tracer, "Params", Params);
-
   // Stash a reference to the span args, so later calls can add metadata.
-  Handler(RequestSpan::stash(Tracer), std::move(Params));
+  WithContext WithRequestSpan(RequestSpan::stash(Tracer));
+
+  Handler(std::move(Params));
   return true;
 }
 
@@ -216,10 +219,9 @@
       // The end of headers is signified by an empty line.
       if (LineRef.consume_front("Content-Length: ")) {
         if (ContentLength != 0) {
-          log(Context::empty(),
-              "Warning: Duplicate Content-Length header received. "
+          log("Warning: Duplicate Content-Length header received. "
               "The previous value for this message (" +
-                  llvm::Twine(ContentLength) + ") was ignored.\n");
+              llvm::Twine(ContentLength) + ") was ignored.\n");
         }
 
         llvm::getAsUnsignedInteger(LineRef.trim(), 0, ContentLength);
@@ -238,8 +240,8 @@
     // and we don't want to crash downstream because of it.
     if (ContentLength > 1 << 30) { // 1024M
       In.ignore(ContentLength);
-      log(Context::empty(), "Skipped overly large message of " +
-                                Twine(ContentLength) + " bytes.\n");
+      log("Skipped overly large message of " + Twine(ContentLength) +
+          " bytes.\n");
       continue;
     }
 
@@ -253,35 +255,32 @@
         // If the stream is aborted before we read ContentLength bytes, In
         // will have eofbit and failbit set.
         if (!In) {
-          log(Context::empty(),
-              "Input was aborted. Read only " + llvm::Twine(In.gcount()) +
-                  " bytes of expected " + llvm::Twine(ContentLength) + ".\n");
+          log("Input was aborted. Read only " + llvm::Twine(In.gcount()) +
+              " bytes of expected " + llvm::Twine(ContentLength) + ".\n");
           break;
         }
 
         JSONRef = StringRef(JSON.data(), ContentLength);
       }
 
       if (auto Doc = json::parse(JSONRef)) {
         // Log the formatted message.
-        log(Context::empty(),
-            llvm::formatv(Out.Pretty ? "<-- {0:2}\n" : "<-- {0}\n", *Doc));
+        log(llvm::formatv(Out.Pretty ? "<-- {0:2}\n" : "<-- {0}\n", *Doc));
         // Finally, execute the action for this JSON message.
         if (!Dispatcher.call(*Doc, Out))
-          log(Context::empty(), "JSON dispatch failed!\n");
+          log("JSON dispatch failed!\n");
       } else {
         // Parse error. Log the raw message.
-        log(Context::empty(), "<-- " + JSONRef + "\n");
-        log(Context::empty(), llvm::Twine("JSON parse error: ") +
-                                  llvm::toString(Doc.takeError()) + "\n");
+        log("<-- " + JSONRef + "\n");
+        log(llvm::Twine("JSON parse error: ") +
+            llvm::toString(Doc.takeError()) + "\n");
       }
 
       // If we're done, exit the loop.
       if (IsDone)
         break;
     } else {
-      log(Context::empty(),
-          "Warning: Missing Content-Length header, or message has zero "
+      log("Warning: Missing Content-Length header, or message has zero "
           "length.\n");
     }
   }
Index: clangd/GlobalCompilationDatabase.cpp
===================================================================
--- clangd/GlobalCompilationDatabase.cpp
+++ clangd/GlobalCompilationDatabase.cpp
@@ -38,8 +38,7 @@
       return std::move(Candidates.front());
     }
   } else {
-    log(Context::empty(), // FIXME(ibiryukov): pass a proper Context here.
-        "Failed to find compilation database for " + Twine(File));
+    log("Failed to find compilation database for " + Twine(File));
   }
   return llvm::None;
 }
Index: clangd/Context.h
===================================================================
--- clangd/Context.h
+++ clangd/Context.h
@@ -16,15 +16,28 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CONTEXT_H_
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Compiler.h"
 #include <memory>
 #include <type_traits>
 
 namespace clang {
 namespace clangd {
 
-/// A key for a value of type \p Type, stored inside a context. Keys are
-/// non-movable and non-copyable. See documentation of the Context class for
-/// more details and usage examples.
+/// Values in a Context are indexed by typed keys.
+/// Key<T> serves two purposes:
+///   - it provides a lookup key for the context (each Key is unique),
+///   - it makes lookup type-safe: a Key<T> can only map to a T (or nothing).
+///
+/// Example:
+///    Key<int> RequestID;
+///    Key<int> Version;
+///
+///    Context Ctx = Context::empty().derive(RequestID, 10).derive(Version, 3);
+///    assert(*Ctx.get(RequestID) == 10);
+///    assert(*Ctx.get(Version) == 3);
+///
+/// Keys are typically used across multiple functions, so most of the time you
+/// would want to make them static class members or global variables.
 template <class Type> class Key {
 public:
   static_assert(!std::is_reference<Type>::value,
@@ -45,53 +58,23 @@
 /// Conceptually, a context is a heterogeneous map<Key<T>, T>. Each key has
 /// an associated value type, which allows the map to be typesafe.
 ///
+/// There is an "ambient" context for each thread, Context::current().
+/// Most functions should read from this, and use WithContextValue or
+/// WithContext to extend or replace the context within a block scope.
+/// Only code dealing with threads and extension points should need to use
+/// other Context objects.
+///
 /// You can't add data to an existing context, instead you create a new
 /// immutable context derived from it with extra data added. When you retrieve
 /// data, the context will walk up the parent chain until the key is found.
-///
-/// Contexts should be:
-///  - passed by reference when calling synchronous functions
-///  - passed by value (move) when calling asynchronous functions. The result
-///    callback of async operations will receive the context again.
-///  - cloned only when 'forking' an asynchronous computation that we don't wait
-///    for.
-///
-/// Copy operations for this class are deleted, use an explicit clone() method
-/// when you need a copy of the context instead.
-///
-/// To derive a child context use derive() function, e.g.
-///     Context ChildCtx = ParentCtx.derive(RequestIdKey, 123);
-///
-/// To create a new root context, derive() from empty Context.
-/// e.g.:
-///     Context Ctx = Context::empty().derive(RequestIdKey, 123);
-///
-/// Values in the context are indexed by typed keys (instances of Key<T> class).
-/// Key<T> serves two purposes:
-///   - it provides a lookup key for the context (each instance of a key is
-///   unique),
-///   - it keeps the type information about the value stored in the context map
-///   in the template arguments.
-/// This provides a type-safe interface to store and access values of multiple
-/// types inside a single context.
-/// For example,
-///    Key<int> RequestID;
-///    Key<int> Version;
-///
-///    Context Ctx = Context::empty().derive(RequestID, 10).derive(Version, 3);
-///    assert(*Ctx.get(RequestID) == 10);
-///    assert(*Ctx.get(Version) == 3);
-///
-/// Keys are typically used across multiple functions, so most of the time you
-/// would want to make them static class members or global variables.
-///
-/// FIXME: Rather than manual plumbing, pass Context using thread-local storage
-/// by default, and make thread boundaries deal with propagation explicitly.
 class Context {
 public:
-  /// Returns an empty context that contains no data. Useful for calling
-  /// functions that require a context when no explicit context is available.
+  /// Returns an empty root context that contains no data.
   static Context empty();
+  /// Returns the context for the current thread, creating it if needed.
+  /// Assigning to current() will permanently change the thread-local context.
+  /// See WithContext and WithContextValue for how to do this safely.
+  static Context &current();
 
 private:
   struct Data;
@@ -103,7 +86,8 @@
   /// (arguments of std::future<> must be default-construcitble in MSVC).
   Context() = default;
 
-  /// Move-only.
+  /// Copy operations for this class are deleted, use an explicit clone() method
+  /// when you need a copy of the context instead.
   Context(Context const &) = delete;
   Context &operator=(const Context &) = delete;
 
@@ -131,9 +115,8 @@
   }
 
   /// Derives a child context
-  /// It is safe to move or destroy a parent context after calling derive() from
-  /// it. The child context will continue to have access to the data stored in
-  /// the parent context.
+  /// It is safe to move or destroy a parent context after calling derive().
+  /// The child will keep its parent alive, and its data remains accessible.
   template <class Type>
   Context derive(const Key<Type> &Key,
                  typename std::decay<Type>::type Value) const & {
@@ -198,7 +181,46 @@
   };
 
   std::shared_ptr<const Data> DataPtr;
-}; // namespace clangd
+};
+
+/// WithContext replaces Context::current() with a provided scope.
+/// When the WithContext is destroyed, the original scope is restored.
+/// For extending the current context with new value, prefer WithContextValue.
+class LLVM_NODISCARD WithContext {
+public:
+  WithContext(Context C)
+      : Restore(llvm::make_unique<Context>(std::move(Context::current()))) {
+    Context::current() = std::move(C);
+  }
+  ~WithContext() {
+    if (Restore)
+      Context::current() = std::move(*Restore);
+  }
+  WithContext(const WithContext&) = delete;
+  WithContext &operator=(const WithContext&) = delete;
+  WithContext(WithContext &&) = default; // Allow return-by-value.
+  WithContext &operator=(WithContext&&) = delete;
+
+private:
+  std::unique_ptr<Context> Restore; // unique_ptr for move semantics.
+};
+
+/// WithContextValue extends Context::current() with a single value.
+/// When the WithContextValue is destroyed, the original scope is restored.
+class LLVM_NODISCARD WithContextValue {
+public:
+  template <typename T>
+  WithContextValue(const Key<T> &K, typename std::decay<T>::type V)
+      : Restore(Context::current().derive(K, std::move(V))) {}
+
+  // Anonymous values can be used for the destructor side-effect.
+  template <typename T>
+  WithContextValue(T &&V)
+      : Restore(Context::current().derive(std::forward<T>(V))) {}
+
+private:
+  WithContext Restore;
+};
 
 } // namespace clangd
 } // namespace clang
Index: clangd/Context.cpp
===================================================================
--- clangd/Context.cpp
+++ clangd/Context.cpp
@@ -20,5 +20,10 @@
 
 Context Context::clone() const { return Context(DataPtr); }
 
+Context &Context::current() {
+  static thread_local Context C = empty();
+  return C;
+}
+
 } // namespace clangd
 } // namespace clang
Index: clangd/CodeComplete.h
===================================================================
--- clangd/CodeComplete.h
+++ clangd/CodeComplete.h
@@ -15,7 +15,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CODECOMPLETE_H
 
-#include "Context.h"
 #include "Logger.h"
 #include "Path.h"
 #include "Protocol.h"
@@ -67,16 +66,16 @@
 };
 
 /// Get code completions at a specified \p Pos in \p FileName.
-CompletionList codeComplete(const Context &Ctx, PathRef FileName,
+CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts);
 
 /// Get signature help at a specified \p Pos in \p FileName.
-SignatureHelp signatureHelp(const Context &Ctx, PathRef FileName,
+SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
Index: clangd/CodeComplete.cpp
===================================================================
--- clangd/CodeComplete.cpp
+++ clangd/CodeComplete.cpp
@@ -632,8 +632,7 @@
 
 // Invokes Sema code completion on a file.
 // Callback will be invoked once completion is done, but before cleaning up.
-bool semaCodeComplete(const Context &Ctx,
-                      std::unique_ptr<CodeCompleteConsumer> Consumer,
+bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
                       const clang::CodeCompleteOptions &Options,
                       const SemaCompleteInput &Input,
                       llvm::function_ref<void()> Callback = nullptr) {
@@ -688,13 +687,12 @@
 
   SyntaxOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
-    log(Ctx, "BeginSourceFile() failed when running codeComplete for " +
-                 Input.FileName);
+    log("BeginSourceFile() failed when running codeComplete for " +
+        Input.FileName);
     return false;
   }
   if (!Action.Execute()) {
-    log(Ctx,
-        "Execute() failed when running codeComplete for " + Input.FileName);
+    log("Execute() failed when running codeComplete for " + Input.FileName);
     return false;
   }
 
@@ -796,7 +794,6 @@
 //     This score is combined with the result quality score for the final score.
 //   - TopN determines the results with the best score.
 class CodeCompleteFlow {
-  const Context &Ctx;
   const CodeCompleteOptions &Opts;
   // Sema takes ownership of Recorder. Recorder is valid until Sema cleanup.
   std::unique_ptr<CompletionRecorder> RecorderOwner;
@@ -807,25 +804,24 @@
 
 public:
   // A CodeCompleteFlow object is only useful for calling run() exactly once.
-  CodeCompleteFlow(const Context &Ctx, const CodeCompleteOptions &Opts)
-      : Ctx(Ctx), Opts(Opts), RecorderOwner(new CompletionRecorder(Opts)),
+  CodeCompleteFlow(const CodeCompleteOptions &Opts)
+      : Opts(Opts), RecorderOwner(new CompletionRecorder(Opts)),
         Recorder(*RecorderOwner) {}
 
   CompletionList run(const SemaCompleteInput &SemaCCInput) && {
     // We run Sema code completion first. It builds an AST and calculates:
     //   - completion results based on the AST. These are saved for merging.
     //   - partial identifier and context. We need these for the index query.
     CompletionList Output;
-    semaCodeComplete(Ctx, std::move(RecorderOwner), Opts.getClangCompleteOpts(),
+    semaCodeComplete(std::move(RecorderOwner), Opts.getClangCompleteOpts(),
                      SemaCCInput, [&] {
                        if (Recorder.CCSema)
                          Output = runWithSema();
                        else
-                         log(Ctx, "Code complete: no Sema callback, 0 results");
+                         log("Code complete: no Sema callback, 0 results");
                      });
 
-    log(Ctx,
-        llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
+    log(llvm::formatv("Code complete: {0} results from Sema, {1} from Index, "
                       "{2} matched, {3} returned{4}.",
                       NSema, NIndex, NBoth, Output.items.size(),
                       Output.isIncomplete ? " (incomplete)" : ""));
@@ -868,12 +864,12 @@
     Req.Query = Filter->pattern();
     Req.Scopes =
         getQueryScopes(Recorder.CCContext, Recorder.CCSema->getSourceManager());
-    log(Ctx, llvm::formatv(
-                 "Code complete: fuzzyFind(\"{0}\", Scopes: [{1}]", Req.Query,
-                 llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
+    log(llvm::formatv("Code complete: fuzzyFind(\"{0}\", Scopes: [{1}]",
+                      Req.Query,
+                      llvm::join(Req.Scopes.begin(), Req.Scopes.end(), ",")));
     // Run the query against the index.
     Incomplete |= !Opts.Index->fuzzyFind(
-        Ctx, Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); });
+        Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); });
     return std::move(ResultsBuilder).build();
   }
 
@@ -942,18 +938,18 @@
   }
 };
 
-CompletionList codeComplete(const Context &Ctx, PathRef FileName,
+CompletionList codeComplete(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
                             IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                             std::shared_ptr<PCHContainerOperations> PCHs,
                             CodeCompleteOptions Opts) {
-  return CodeCompleteFlow(Ctx, Opts).run(
+  return CodeCompleteFlow(Opts).run(
       {FileName, Command, Preamble, Contents, Pos, VFS, PCHs});
 }
 
-SignatureHelp signatureHelp(const Context &Ctx, PathRef FileName,
+SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             PrecompiledPreamble const *Preamble,
                             StringRef Contents, Position Pos,
@@ -965,10 +961,10 @@
   Options.IncludeMacros = false;
   Options.IncludeCodePatterns = false;
   Options.IncludeBriefComments = true;
-  semaCodeComplete(
-      Ctx, llvm::make_unique<SignatureHelpCollector>(Options, Result), Options,
-      {FileName, Command, Preamble, Contents, Pos, std::move(VFS),
-       std::move(PCHs)});
+  semaCodeComplete(llvm::make_unique<SignatureHelpCollector>(Options, Result),
+                   Options,
+                   {FileName, Command, Preamble, Contents, Pos, std::move(VFS),
+                    std::move(PCHs)});
   return Result;
 }
 
Index: clangd/ClangdUnit.h
===================================================================
--- clangd/ClangdUnit.h
+++ clangd/ClangdUnit.h
@@ -10,7 +10,6 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNIT_H
 
-#include "Context.h"
 #include "Function.h"
 #include "Path.h"
 #include "Protocol.h"
@@ -71,7 +70,7 @@
   /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
   /// it is reused during parsing.
   static llvm::Optional<ParsedAST>
-  Build(const Context &Ctx, std::unique_ptr<clang::CompilerInvocation> CI,
+  Build(std::unique_ptr<clang::CompilerInvocation> CI,
         std::shared_ptr<const PreambleData> Preamble,
         std::unique_ptr<llvm::MemoryBuffer> Buffer,
         std::shared_ptr<PCHContainerOperations> PCHs,
@@ -148,8 +147,7 @@
   mutable llvm::Optional<ParsedAST> AST;
 };
 
-using ASTParsedCallback =
-    std::function<void(const Context &Ctx, PathRef Path, ParsedAST *)>;
+using ASTParsedCallback = std::function<void(PathRef Path, ParsedAST *)>;
 
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
@@ -186,8 +184,7 @@
   /// Returns a list of diagnostics or a llvm::None, if another rebuild was
   /// requested in parallel (effectively cancelling this rebuild) before
   /// diagnostics were produced.
-  llvm::Optional<std::vector<DiagWithFixIts>> rebuild(const Context &Ctx,
-                                                      ParseInputs &&Inputs);
+  llvm::Optional<std::vector<DiagWithFixIts>> rebuild(ParseInputs &&Inputs);
 
   /// Schedule a rebuild and return a deferred computation that will finish the
   /// rebuild, that can be called on a different thread.
@@ -203,7 +200,7 @@
   /// The future to finish rebuild returns a list of diagnostics built during
   /// reparse, or None, if another deferRebuild was called before this
   /// rebuild was finished.
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
   deferRebuild(ParseInputs &&Inputs);
 
   /// Returns a future to get the most fresh PreambleData for a file. The
Index: clangd/ClangdUnit.cpp
===================================================================
--- clangd/ClangdUnit.cpp
+++ clangd/ClangdUnit.cpp
@@ -229,8 +229,7 @@
 }
 
 llvm::Optional<ParsedAST>
-ParsedAST::Build(const Context &Ctx,
-                 std::unique_ptr<clang::CompilerInvocation> CI,
+ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
                  std::shared_ptr<const PreambleData> Preamble,
                  std::unique_ptr<llvm::MemoryBuffer> Buffer,
                  std::shared_ptr<PCHContainerOperations> PCHs,
@@ -252,12 +251,12 @@
   auto Action = llvm::make_unique<ClangdFrontendAction>();
   const FrontendInputFile &MainInput = Clang->getFrontendOpts().Inputs[0];
   if (!Action->BeginSourceFile(*Clang, MainInput)) {
-    log(Ctx, "BeginSourceFile() failed when building AST for " +
+    log("BeginSourceFile() failed when building AST for " +
                  MainInput.getFile());
     return llvm::None;
   }
   if (!Action->Execute())
-    log(Ctx, "Execute() failed when building AST for " + MainInput.getFile());
+    log("Execute() failed when building AST for " + MainInput.getFile());
 
   // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
   // has a longer lifetime.
@@ -385,8 +384,7 @@
       RebuildCounter(0), RebuildInProgress(false), ASTMemUsage(0),
       PreambleMemUsage(0), PCHs(std::move(PCHs)),
       ASTCallback(std::move(ASTCallback)) {
-  // FIXME(ibiryukov): we should pass a proper Context here.
-  log(Context::empty(), "Created CppFile for " + FileName);
+  log("Created CppFile for " + FileName);
 
   std::lock_guard<std::mutex> Lock(Mutex);
   LatestAvailablePreamble = nullptr;
@@ -440,11 +438,11 @@
 }
 
 llvm::Optional<std::vector<DiagWithFixIts>>
-CppFile::rebuild(const Context &Ctx, ParseInputs &&Inputs) {
-  return deferRebuild(std::move(Inputs))(Ctx);
+CppFile::rebuild(ParseInputs &&Inputs) {
+  return deferRebuild(std::move(Inputs))();
 }
 
-UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
 CppFile::deferRebuild(ParseInputs &&Inputs) {
   std::shared_ptr<const PreambleData> OldPreamble;
   std::shared_ptr<PCHContainerOperations> PCHs;
@@ -481,13 +479,11 @@
   std::shared_ptr<CppFile> That = shared_from_this();
   auto FinishRebuild =
       [OldPreamble, RequestRebuildCounter, PCHs,
-       That](ParseInputs Inputs,
-             const Context &Ctx) mutable /* to allow changing OldPreamble. */
+       That](ParseInputs Inputs) mutable /* to allow changing OldPreamble. */
       -> llvm::Optional<std::vector<DiagWithFixIts>> {
-    log(Context::empty(),
-        "Rebuilding file " + That->FileName + " with command [" +
-            Inputs.CompileCommand.Directory + "] " +
-            llvm::join(Inputs.CompileCommand.CommandLine, " "));
+    log("Rebuilding file " + That->FileName + " with command [" +
+        Inputs.CompileCommand.Directory + "] " +
+        llvm::join(Inputs.CompileCommand.CommandLine, " "));
 
     // Only one execution of this method is possible at a time.
     // RebuildGuard will wait for any ongoing rebuilds to finish and will put us
@@ -530,16 +526,16 @@
       if (OldPreamble &&
           OldPreamble->Preamble.CanReuse(*CI, ContentsBuffer.get(), Bounds,
                                          Inputs.FS.get())) {
-        log(Ctx, "Reusing preamble for file " + Twine(That->FileName));
+        log("Reusing preamble for file " + Twine(That->FileName));
         return OldPreamble;
       }
-      log(Ctx, "Premble for file " + Twine(That->FileName) +
-                   " cannot be reused. Attempting to rebuild it.");
+      log("Preamble for file " + Twine(That->FileName) +
+          " cannot be reused. Attempting to rebuild it.");
       // We won't need the OldPreamble anymore, release it so it can be
       // deleted (if there are no other references to it).
       OldPreamble.reset();
 
-      trace::Span Tracer(Ctx, "Preamble");
+      trace::Span Tracer("Preamble");
       SPAN_ATTACH(Tracer, "File", That->FileName);
       std::vector<DiagWithFixIts> PreambleDiags;
       StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
@@ -564,17 +560,15 @@
       CI->getFrontendOpts().SkipFunctionBodies = false;
 
       if (BuiltPreamble) {
-        log(Tracer.Ctx, "Built preamble of size " +
-                            Twine(BuiltPreamble->getSize()) + " for file " +
-                            Twine(That->FileName));
+        log("Built preamble of size " + Twine(BuiltPreamble->getSize()) +
+            " for file " + Twine(That->FileName));
 
         return std::make_shared<PreambleData>(
             std::move(*BuiltPreamble),
             SerializedDeclsCollector.takeTopLevelDeclIDs(),
             std::move(PreambleDiags));
       } else {
-        log(Tracer.Ctx,
-            "Could not build a preamble for file " + Twine(That->FileName));
+        log("Could not build a preamble for file " + Twine(That->FileName));
         return nullptr;
       }
     };
@@ -604,18 +598,17 @@
     // Compute updated AST.
     llvm::Optional<ParsedAST> NewAST;
     {
-      trace::Span Tracer(Ctx, "Build");
+      trace::Span Tracer("Build");
       SPAN_ATTACH(Tracer, "File", That->FileName);
-      NewAST =
-          ParsedAST::Build(Tracer.Ctx, std::move(CI), std::move(NewPreamble),
-                           std::move(ContentsBuffer), PCHs, Inputs.FS);
+      NewAST = ParsedAST::Build(std::move(CI), std::move(NewPreamble),
+                                std::move(ContentsBuffer), PCHs, Inputs.FS);
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());
       if (That->ASTCallback)
-        That->ASTCallback(Ctx, That->FileName, NewAST.getPointer());
+        That->ASTCallback(That->FileName, NewAST.getPointer());
     } else {
       // Don't report even Preamble diagnostics if we coulnd't build AST.
       Diagnostics.clear();
Index: clangd/ClangdServer.h
===================================================================
--- clangd/ClangdServer.h
+++ clangd/ClangdServer.h
@@ -10,6 +10,7 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDSERVER_H
 
+#include "Context.h"
 #include "ClangdUnit.h"
 #include "ClangdUnitStore.h"
 #include "CodeComplete.h"
@@ -73,7 +74,7 @@
 
   /// Called by ClangdServer when \p Diagnostics for \p File are ready.
   virtual void
-  onDiagnosticsReady(const Context &Ctx, PathRef File,
+  onDiagnosticsReady(PathRef File,
                      Tagged<std::vector<DiagWithFixIts>> Diagnostics) = 0;
 };
 
@@ -126,7 +127,8 @@
 
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.push_front(
+      RequestQueue.emplace_front(
+          Context::current().clone(),
           BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
     }
     RequestCV.notify_one();
@@ -142,7 +144,8 @@
 
     {
       std::lock_guard<std::mutex> Lock(Mutex);
-      RequestQueue.push_back(
+      RequestQueue.emplace_back(
+          Context::current().clone(),
           BindWithForward(std::forward<Func>(F), std::forward<Args>(As)...));
     }
     RequestCV.notify_one();
@@ -159,7 +162,7 @@
   bool Done = false;
   /// A queue of requests. Elements of this vector are async computations (i.e.
   /// results of calling std::async(std::launch::deferred, ...)).
-  std::deque<UniqueFunction<void()>> RequestQueue;
+  std::deque<std::pair<Context, UniqueFunction<void()>>> RequestQueue;
   /// Condition variable to wake up worker threads.
   std::condition_variable RequestCV;
 };
@@ -221,26 +224,20 @@
   /// constructor will receive onDiagnosticsReady callback.
   /// \return A future that will become ready when the rebuild (including
   /// diagnostics) is finished.
-  std::future<Context> addDocument(Context Ctx, PathRef File,
-                                   StringRef Contents);
+  void addDocument(PathRef File, StringRef Contents);
   /// Remove \p File from list of tracked files, schedule a request to free
   /// resources associated with it.
   /// \return A future that will become ready when the file is removed and all
   /// associated resources are freed.
-  std::future<Context> removeDocument(Context Ctx, PathRef File);
+  void removeDocument(PathRef File);
   /// Force \p File to be reparsed using the latest contents.
   /// Will also check if CompileCommand, provided by GlobalCompilationDatabase
   /// for \p File has changed. If it has, will remove currently stored Preamble
   /// and AST and rebuild them from scratch.
-  std::future<Context> forceReparse(Context Ctx, PathRef File);
+  void forceReparse(PathRef File);
 
-  /// DEPRECATED. Please use a callback-based version, this API is deprecated
-  /// and will soon be removed.
-  ///
   /// Run code completion for \p File at \p Pos.
-  ///
-  /// Request is processed asynchronously. You can use the returned future to
-  /// wait for the results of the async request.
+  /// Request is processed asynchronously.
   ///
   /// If \p OverridenContents is not None, they will used only for code
   /// completion, i.e. no diagnostics update will be scheduled and a draft for
@@ -251,18 +248,12 @@
   /// This method should only be called for currently tracked files. However, it
   /// is safe to call removeDocument for \p File after this method returns, even
   /// while returned future is not yet ready.
-  std::future<std::pair<Context, Tagged<CompletionList>>>
-  codeComplete(Context Ctx, PathRef File, Position Pos,
-               const clangd::CodeCompleteOptions &Opts,
-               llvm::Optional<StringRef> OverridenContents = llvm::None,
-               IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
-
   /// A version of `codeComplete` that runs \p Callback on the processing thread
   /// when codeComplete results become available.
   void
-  codeComplete(Context Ctx, PathRef File, Position Pos,
+  codeComplete(PathRef File, Position Pos,
                const clangd::CodeCompleteOptions &Opts,
-               UniqueFunction<void(Context, Tagged<CompletionList>)> Callback,
+               UniqueFunction<void(Tagged<CompletionList>)> Callback,
                llvm::Optional<StringRef> OverridenContents = llvm::None,
                IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
 
@@ -274,21 +265,21 @@
   /// vfs::FileSystem used for signature help. This method should only be called
   /// for currently tracked files.
   llvm::Expected<Tagged<SignatureHelp>>
-  signatureHelp(const Context &Ctx, PathRef File, Position Pos,
+  signatureHelp(PathRef File, Position Pos,
                 llvm::Optional<StringRef> OverridenContents = llvm::None,
                 IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS = nullptr);
 
   /// Get definition of symbol at a specified \p Line and \p Column in \p File.
   llvm::Expected<Tagged<std::vector<Location>>>
-  findDefinitions(const Context &Ctx, PathRef File, Position Pos);
+  findDefinitions(PathRef File, Position Pos);
 
   /// 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(const Context &Ctx, PathRef File, Position Pos);
+  findDocumentHighlights(PathRef File, Position Pos);
 
   /// Run formatting for \p Rng inside \p File with content \p Code.
   llvm::Expected<tooling::Replacements> formatRange(StringRef Code,
@@ -305,8 +296,7 @@
 
   /// Rename all occurrences of the symbol at the \p Pos in \p File to
   /// \p NewName.
-  Expected<std::vector<tooling::Replacement>> rename(const Context &Ctx,
-                                                     PathRef File, Position Pos,
+  Expected<std::vector<tooling::Replacement>> rename(PathRef File, Position Pos,
                                                      llvm::StringRef NewName);
 
   /// Gets current document contents for \p File. Returns None if \p File is not
@@ -338,13 +328,12 @@
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  std::future<Context>
-  scheduleReparseAndDiags(Context Ctx, PathRef File, VersionedDraft Contents,
+  void
+  scheduleReparseAndDiags(PathRef File, VersionedDraft Contents,
                           std::shared_ptr<CppFile> Resources,
                           Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS);
 
-  std::future<Context>
-  scheduleCancelRebuild(Context Ctx, std::shared_ptr<CppFile> Resources);
+  void scheduleCancelRebuild(std::shared_ptr<CppFile> Resources);
 
   CompileArgsCache CompileArgs;
   DiagnosticsConsumer &DiagConsumer;
Index: clangd/ClangdServer.cpp
===================================================================
--- clangd/ClangdServer.cpp
+++ clangd/ClangdServer.cpp
@@ -90,6 +90,7 @@
       llvm::set_thread_name(llvm::formatv("scheduler/{0}", I));
       while (true) {
         UniqueFunction<void()> Request;
+        llvm::Optional<WithContext> Ctx;
 
         // Pick request from the queue
         {
@@ -106,10 +107,10 @@
           // ClangdScheduler have a way to prioritise their requests by putting
           // them to the either side of the queue (using either addToEnd or
           // addToFront).
-          Request = std::move(RequestQueue.front());
+          Request = std::move(RequestQueue.front().second);
+          Ctx.emplace(std::move(RequestQueue.front().first));
           RequestQueue.pop_front();
         } // unlock Mutex
-
         Request();
       }
     }));
@@ -147,10 +148,9 @@
       // parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
-      Units(FileIdx
-                ? [this](const Context &Ctx, PathRef Path,
-                         ParsedAST *AST) { FileIdx->update(Ctx, Path, AST); }
-                : ASTParsedCallback()),
+      Units(FileIdx ? [this](PathRef Path,
+                             ParsedAST *AST) { FileIdx->update(Path, AST); }
+                    : ASTParsedCallback()),
       PCHs(std::make_shared<PCHContainerOperations>()),
       StorePreamblesInMemory(StorePreamblesInMemory),
       WorkScheduler(AsyncThreadsCount) {
@@ -172,27 +172,25 @@
     this->RootPath = NewRootPath;
 }
 
-std::future<Context> ClangdServer::addDocument(Context Ctx, PathRef File,
-                                               StringRef Contents) {
+void ClangdServer::addDocument(PathRef File, StringRef Contents) {
   DocVersion Version = DraftMgr.updateDraft(File, Contents);
 
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   std::shared_ptr<CppFile> Resources =
       Units.getOrCreateFile(File, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File,
-                                 VersionedDraft{Version, Contents.str()},
+  return scheduleReparseAndDiags(File, VersionedDraft{Version, Contents.str()},
                                  std::move(Resources), std::move(TaggedFS));
 }
 
-std::future<Context> ClangdServer::removeDocument(Context Ctx, PathRef File) {
+void ClangdServer::removeDocument(PathRef File) {
   DraftMgr.removeDraft(File);
   CompileArgs.invalidate(File);
 
   std::shared_ptr<CppFile> Resources = Units.removeIfPresent(File);
-  return scheduleCancelRebuild(std::move(Ctx), std::move(Resources));
+  return scheduleCancelRebuild(std::move(Resources));
 }
 
-std::future<Context> ClangdServer::forceReparse(Context Ctx, PathRef File) {
+void ClangdServer::forceReparse(PathRef File) {
   auto FileContents = DraftMgr.getDraft(File);
   assert(FileContents.Draft &&
          "forceReparse() was called for non-added document");
@@ -204,38 +202,16 @@
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
   std::shared_ptr<CppFile> Resources =
       Units.getOrCreateFile(File, StorePreamblesInMemory, PCHs);
-  return scheduleReparseAndDiags(std::move(Ctx), File, FileContents,
-                                 std::move(Resources), std::move(TaggedFS));
-}
-
-std::future<std::pair<Context, Tagged<CompletionList>>>
-ClangdServer::codeComplete(Context Ctx, PathRef File, Position Pos,
-                           const clangd::CodeCompleteOptions &Opts,
-                           llvm::Optional<StringRef> OverridenContents,
-                           IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
-  using ResultType = std::pair<Context, Tagged<CompletionList>>;
-
-  std::promise<ResultType> ResultPromise;
-
-  auto Callback = [](std::promise<ResultType> ResultPromise, Context Ctx,
-                     Tagged<CompletionList> Result) -> void {
-    ResultPromise.set_value({std::move(Ctx), std::move(Result)});
-  };
-
-  std::future<ResultType> ResultFuture = ResultPromise.get_future();
-  codeComplete(std::move(Ctx), File, Pos, Opts,
-               BindWithForward(Callback, std::move(ResultPromise)),
-               OverridenContents, UsedFS);
-  return ResultFuture;
+  return scheduleReparseAndDiags(File, FileContents, std::move(Resources),
+                                 std::move(TaggedFS));
 }
 
 void ClangdServer::codeComplete(
-    Context Ctx, PathRef File, Position Pos,
-    const clangd::CodeCompleteOptions &Opts,
-    UniqueFunction<void(Context, Tagged<CompletionList>)> Callback,
+    PathRef File, Position Pos, const clangd::CodeCompleteOptions &Opts,
+    UniqueFunction<void(Tagged<CompletionList>)> Callback,
     llvm::Optional<StringRef> OverridenContents,
     IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
-  using CallbackType = UniqueFunction<void(Context, Tagged<CompletionList>)>;
+  using CallbackType = UniqueFunction<void(Tagged<CompletionList>)>;
 
   std::string Contents;
   if (OverridenContents) {
@@ -276,28 +252,26 @@
   auto Task =
       // 'mutable' to reassign Preamble variable.
       [FileStr, Preamble, Resources, Contents, Pos, CodeCompleteOpts, TaggedFS,
-       PCHs, CompileCommand](Context Ctx, CallbackType Callback) mutable {
+       PCHs, CompileCommand](CallbackType Callback) mutable {
         if (!Preamble) {
           // Maybe we built some preamble before processing this request.
           Preamble = Resources->getPossiblyStalePreamble();
         }
         // FIXME(ibiryukov): even if Preamble is non-null, we may want to check
         // both the old and the new version in case only one of them matches.
         CompletionList Result = clangd::codeComplete(
-            Ctx, FileStr, CompileCommand,
+            FileStr, CompileCommand,
             Preamble ? &Preamble->Preamble : nullptr, Contents, Pos,
             TaggedFS.Value, PCHs, CodeCompleteOpts);
 
-        Callback(std::move(Ctx),
-                 make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
+        Callback(make_tagged(std::move(Result), std::move(TaggedFS.Tag)));
       };
 
-  WorkScheduler.addToFront(std::move(Task), std::move(Ctx),
-                           std::move(Callback));
+  WorkScheduler.addToFront(std::move(Task), std::move(Callback));
 }
 
 llvm::Expected<Tagged<SignatureHelp>>
-ClangdServer::signatureHelp(const Context &Ctx, PathRef File, Position Pos,
+ClangdServer::signatureHelp(PathRef File, Position Pos,
                             llvm::Optional<StringRef> OverridenContents,
                             IntrusiveRefCntPtr<vfs::FileSystem> *UsedFS) {
   std::string DraftStorage;
@@ -324,7 +298,7 @@
 
   auto Preamble = Resources->getPossiblyStalePreamble();
   auto Result =
-      clangd::signatureHelp(Ctx, File, CompileArgs.getCompileCommand(File),
+      clangd::signatureHelp(File, CompileArgs.getCompileCommand(File),
                             Preamble ? &Preamble->Preamble : nullptr,
                             *OverridenContents, Pos, TaggedFS.Value, PCHs);
   return make_tagged(std::move(Result), TaggedFS.Tag);
@@ -357,8 +331,7 @@
 }
 
 Expected<std::vector<tooling::Replacement>>
-ClangdServer::rename(const Context &Ctx, PathRef File, Position Pos,
-                     llvm::StringRef NewName) {
+ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName) {
   std::shared_ptr<CppFile> Resources = Units.getFile(File);
   RefactoringResultCollector ResultCollector;
   Resources->getAST().get()->runUnderLock([&](ParsedAST *AST) {
@@ -430,7 +403,7 @@
 }
 
 llvm::Expected<Tagged<std::vector<Location>>>
-ClangdServer::findDefinitions(const Context &Ctx, PathRef File, Position Pos) {
+ClangdServer::findDefinitions(PathRef File, Position Pos) {
   auto TaggedFS = FSProvider.getTaggedFileSystem(File);
 
   std::shared_ptr<CppFile> Resources = Units.getFile(File);
@@ -440,10 +413,10 @@
         llvm::errc::invalid_argument);
 
   std::vector<Location> Result;
-  Resources->getAST().get()->runUnderLock([Pos, &Result, &Ctx](ParsedAST *AST) {
+  Resources->getAST().get()->runUnderLock([Pos, &Result](ParsedAST *AST) {
     if (!AST)
       return;
-    Result = clangd::findDefinitions(Ctx, *AST, Pos);
+    Result = clangd::findDefinitions(*AST, Pos);
   });
   return make_tagged(std::move(Result), TaggedFS.Tag);
 }
@@ -522,8 +495,7 @@
 }
 
 llvm::Expected<Tagged<std::vector<DocumentHighlight>>>
-ClangdServer::findDocumentHighlights(const Context &Ctx, PathRef File,
-                                     Position Pos) {
+ClangdServer::findDocumentHighlights(PathRef File, Position Pos) {
   auto FileContents = DraftMgr.getDraft(File);
   if (!FileContents.Draft)
     return llvm::make_error<llvm::StringError>(
@@ -540,51 +512,43 @@
 
   std::vector<DocumentHighlight> Result;
   llvm::Optional<llvm::Error> Err;
-  Resources->getAST().get()->runUnderLock([Pos, &Ctx, &Err,
-                                           &Result](ParsedAST *AST) {
+  Resources->getAST().get()->runUnderLock([Pos, &Err, &Result](ParsedAST *AST) {
     if (!AST) {
       Err = llvm::make_error<llvm::StringError>("Invalid AST",
                                                 llvm::errc::invalid_argument);
       return;
     }
-    Result = clangd::findDocumentHighlights(Ctx, *AST, Pos);
+    Result = clangd::findDocumentHighlights(*AST, Pos);
   });
 
   if (Err)
     return std::move(*Err);
   return make_tagged(Result, TaggedFS.Tag);
 }
 
-std::future<Context> ClangdServer::scheduleReparseAndDiags(
-    Context Ctx, PathRef File, VersionedDraft Contents,
+void ClangdServer::scheduleReparseAndDiags(
+    PathRef File, VersionedDraft Contents,
     std::shared_ptr<CppFile> Resources,
     Tagged<IntrusiveRefCntPtr<vfs::FileSystem>> TaggedFS) {
   assert(Contents.Draft && "Draft must have contents");
   ParseInputs Inputs = {CompileArgs.getCompileCommand(File),
                         std::move(TaggedFS.Value), *std::move(Contents.Draft)};
 
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(const Context &)>
+  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
       DeferredRebuild = Resources->deferRebuild(std::move(Inputs));
-  std::promise<Context> DonePromise;
-  std::future<Context> DoneFuture = DonePromise.get_future();
 
   DocVersion Version = Contents.Version;
   Path FileStr = File;
   VFSTag Tag = TaggedFS.Tag;
   auto ReparseAndPublishDiags =
       [this, FileStr, Version,
-       Tag](UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>(
-                const Context &)>
-                DeferredRebuild,
-            std::promise<Context> DonePromise, Context Ctx) -> void {
-    auto Guard =
-        llvm::make_scope_exit([&]() { DonePromise.set_value(std::move(Ctx)); });
-
+       Tag](UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
+                DeferredRebuild) -> void {
     auto CurrentVersion = DraftMgr.getVersion(FileStr);
     if (CurrentVersion != Version)
       return; // This request is outdated
 
-    auto Diags = DeferredRebuild(Ctx);
+    auto Diags = DeferredRebuild();
     if (!Diags)
       return; // A new reparse was requested before this one completed.
 
@@ -600,37 +564,20 @@
       return;
     LastReportedDiagsVersion = Version;
 
-    DiagConsumer.onDiagnosticsReady(Ctx, FileStr,
+    DiagConsumer.onDiagnosticsReady(FileStr,
                                     make_tagged(std::move(*Diags), Tag));
   };
 
   WorkScheduler.addToFront(std::move(ReparseAndPublishDiags),
-                           std::move(DeferredRebuild), std::move(DonePromise),
-                           std::move(Ctx));
-  return DoneFuture;
+                           std::move(DeferredRebuild));
 }
 
-std::future<Context>
-ClangdServer::scheduleCancelRebuild(Context Ctx,
-                                    std::shared_ptr<CppFile> Resources) {
-  std::promise<Context> DonePromise;
-  std::future<Context> DoneFuture = DonePromise.get_future();
+void ClangdServer::scheduleCancelRebuild(std::shared_ptr<CppFile> Resources) {
   if (!Resources) {
     // No need to schedule any cleanup.
-    DonePromise.set_value(std::move(Ctx));
-    return DoneFuture;
+    return;
   }
-
-  UniqueFunction<void()> DeferredCancel = Resources->deferCancelRebuild();
-  auto CancelReparses = [Resources](std::promise<Context> DonePromise,
-                                    UniqueFunction<void()> DeferredCancel,
-                                    Context Ctx) {
-    DeferredCancel();
-    DonePromise.set_value(std::move(Ctx));
-  };
-  WorkScheduler.addToFront(std::move(CancelReparses), std::move(DonePromise),
-                           std::move(DeferredCancel), std::move(Ctx));
-  return DoneFuture;
+  WorkScheduler.addToFront(Resources->deferCancelRebuild());
 }
 
 void ClangdServer::onFileEvent(const DidChangeWatchedFilesParams &Params) {
Index: clangd/ClangdLSPServer.h
===================================================================
--- clangd/ClangdLSPServer.h
+++ clangd/ClangdLSPServer.h
@@ -50,32 +50,30 @@
 private:
   // Implement DiagnosticsConsumer.
   virtual void
-  onDiagnosticsReady(const Context &Ctx, PathRef File,
+  onDiagnosticsReady(PathRef File,
                      Tagged<std::vector<DiagWithFixIts>> Diagnostics) override;
 
   // Implement ProtocolCallbacks.
-  void onInitialize(Ctx C, InitializeParams &Params) override;
-  void onShutdown(Ctx C, ShutdownParams &Params) override;
-  void onExit(Ctx C, ExitParams &Params) override;
-  void onDocumentDidOpen(Ctx C, DidOpenTextDocumentParams &Params) override;
-  void onDocumentDidChange(Ctx C, DidChangeTextDocumentParams &Params) override;
-  void onDocumentDidClose(Ctx C, DidCloseTextDocumentParams &Params) override;
+  void onInitialize(InitializeParams &Params) override;
+  void onShutdown(ShutdownParams &Params) override;
+  void onExit(ExitParams &Params) override;
+  void onDocumentDidOpen(DidOpenTextDocumentParams &Params) override;
+  void onDocumentDidChange(DidChangeTextDocumentParams &Params) override;
+  void onDocumentDidClose(DidCloseTextDocumentParams &Params) override;
   void
-  onDocumentOnTypeFormatting(Ctx C,
-                             DocumentOnTypeFormattingParams &Params) override;
+  onDocumentOnTypeFormatting(DocumentOnTypeFormattingParams &Params) override;
   void
-  onDocumentRangeFormatting(Ctx C,
-                            DocumentRangeFormattingParams &Params) override;
-  void onDocumentFormatting(Ctx C, DocumentFormattingParams &Params) override;
-  void onCodeAction(Ctx C, CodeActionParams &Params) override;
-  void onCompletion(Ctx C, TextDocumentPositionParams &Params) override;
-  void onSignatureHelp(Ctx C, TextDocumentPositionParams &Params) override;
-  void onGoToDefinition(Ctx C, TextDocumentPositionParams &Params) override;
-  void onSwitchSourceHeader(Ctx C, TextDocumentIdentifier &Params) override;
-  void onDocumentHighlight(Ctx C, TextDocumentPositionParams &Params) override;
-  void onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) override;
-  void onCommand(Ctx C, ExecuteCommandParams &Params) override;
-  void onRename(Ctx C, RenameParams &Parames) override;
+  onDocumentRangeFormatting(DocumentRangeFormattingParams &Params) override;
+  void onDocumentFormatting(DocumentFormattingParams &Params) override;
+  void onCodeAction(CodeActionParams &Params) override;
+  void onCompletion(TextDocumentPositionParams &Params) override;
+  void onSignatureHelp(TextDocumentPositionParams &Params) override;
+  void onGoToDefinition(TextDocumentPositionParams &Params) override;
+  void onSwitchSourceHeader(TextDocumentIdentifier &Params) override;
+  void onDocumentHighlight(TextDocumentPositionParams &Params) override;
+  void onFileEvent(DidChangeWatchedFilesParams &Params) override;
+  void onCommand(ExecuteCommandParams &Params) override;
+  void onRename(RenameParams &Parames) override;
 
   std::vector<TextEdit> getFixIts(StringRef File, const clangd::Diagnostic &D);
 
Index: clangd/ClangdLSPServer.cpp
===================================================================
--- clangd/ClangdLSPServer.cpp
+++ clangd/ClangdLSPServer.cpp
@@ -45,8 +45,8 @@
 
 } // namespace
 
-void ClangdLSPServer::onInitialize(Ctx C, InitializeParams &Params) {
-  reply(C, json::obj{
+void ClangdLSPServer::onInitialize(InitializeParams &Params) {
+  reply(json::obj{
       {{"capabilities",
         json::obj{
             {"textDocumentSync", 1},
@@ -81,38 +81,35 @@
     Server.setRootPath(*Params.rootPath);
 }
 
-void ClangdLSPServer::onShutdown(Ctx C, ShutdownParams &Params) {
+void ClangdLSPServer::onShutdown(ShutdownParams &Params) {
   // Do essentially nothing, just say we're ready to exit.
   ShutdownRequestReceived = true;
-  reply(C, nullptr);
+  reply(nullptr);
 }
 
-void ClangdLSPServer::onExit(Ctx C, ExitParams &Params) { IsDone = true; }
+void ClangdLSPServer::onExit(ExitParams &Params) { IsDone = true; }
 
-void ClangdLSPServer::onDocumentDidOpen(Ctx C,
-                                        DidOpenTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidOpen(DidOpenTextDocumentParams &Params) {
   if (Params.metadata && !Params.metadata->extraFlags.empty())
     CDB.setExtraFlagsForFile(Params.textDocument.uri.file,
                              std::move(Params.metadata->extraFlags));
-  Server.addDocument(std::move(C), Params.textDocument.uri.file,
-                     Params.textDocument.text);
+  Server.addDocument(Params.textDocument.uri.file, Params.textDocument.text);
 }
 
-void ClangdLSPServer::onDocumentDidChange(Ctx C,
-                                          DidChangeTextDocumentParams &Params) {
+void ClangdLSPServer::onDocumentDidChange( DidChangeTextDocumentParams &Params) {
   if (Params.contentChanges.size() != 1)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "can only apply one change at a time");
   // We only support full syncing right now.
-  Server.addDocument(std::move(C), Params.textDocument.uri.file,
+  Server.addDocument(Params.textDocument.uri.file,
                      Params.contentChanges[0].text);
 }
 
-void ClangdLSPServer::onFileEvent(Ctx C, DidChangeWatchedFilesParams &Params) {
+void ClangdLSPServer::onFileEvent(DidChangeWatchedFilesParams &Params) {
   Server.onFileEvent(Params);
 }
 
-void ClangdLSPServer::onCommand(Ctx C, ExecuteCommandParams &Params) {
+void ClangdLSPServer::onCommand(ExecuteCommandParams &Params) {
   if (Params.command == ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND &&
       Params.workspaceEdit) {
     // The flow for "apply-fix" :
@@ -126,100 +123,98 @@
 
     ApplyWorkspaceEditParams ApplyEdit;
     ApplyEdit.edit = *Params.workspaceEdit;
-    reply(C, "Fix applied.");
+    reply("Fix applied.");
     // We don't need the response so id == 1 is OK.
     // Ideally, we would wait for the response and if there is no error, we
     // would reply success/failure to the original RPC.
-    call(C, "workspace/applyEdit", ApplyEdit);
+    call("workspace/applyEdit", ApplyEdit);
   } else {
     // We should not get here because ExecuteCommandParams would not have
     // parsed in the first place and this handler should not be called. But if
     // more commands are added, this will be here has a safe guard.
     replyError(
-        C, ErrorCode::InvalidParams,
+        ErrorCode::InvalidParams,
         llvm::formatv("Unsupported command \"{0}\".", Params.command).str());
   }
 }
 
-void ClangdLSPServer::onRename(Ctx C, RenameParams &Params) {
+void ClangdLSPServer::onRename(RenameParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onRename called for non-added file");
 
-  auto Replacements = Server.rename(C, File, Params.position, Params.newName);
+  auto Replacements = Server.rename(File, Params.position, Params.newName);
   if (!Replacements) {
-    replyError(C, ErrorCode::InternalError,
+    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(C, WE);
+  reply(WE);
 }
 
-void ClangdLSPServer::onDocumentDidClose(Ctx C,
-                                         DidCloseTextDocumentParams &Params) {
-  Server.removeDocument(std::move(C), Params.textDocument.uri.file);
+void ClangdLSPServer::onDocumentDidClose(DidCloseTextDocumentParams &Params) {
+  Server.removeDocument(Params.textDocument.uri.file);
 }
 
 void ClangdLSPServer::onDocumentOnTypeFormatting(
-    Ctx C, DocumentOnTypeFormattingParams &Params) {
+    DocumentOnTypeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentOnTypeFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatOnType(*Code, File, Params.position);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
 void ClangdLSPServer::onDocumentRangeFormatting(
-    Ctx C, DocumentRangeFormattingParams &Params) {
+    DocumentRangeFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentRangeFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatRange(*Code, File, Params.range);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
-void ClangdLSPServer::onDocumentFormatting(Ctx C,
-                                           DocumentFormattingParams &Params) {
+void ClangdLSPServer::onDocumentFormatting(DocumentFormattingParams &Params) {
   auto File = Params.textDocument.uri.file;
   auto Code = Server.getDocument(File);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onDocumentFormatting called for non-added file");
 
   auto ReplacementsOrError = Server.formatFile(*Code, File);
   if (ReplacementsOrError)
-    reply(C, json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
+    reply(json::ary(replacementsToEdits(*Code, ReplacementsOrError.get())));
   else
-    replyError(C, ErrorCode::UnknownErrorCode,
+    replyError(ErrorCode::UnknownErrorCode,
                llvm::toString(ReplacementsOrError.takeError()));
 }
 
-void ClangdLSPServer::onCodeAction(Ctx C, CodeActionParams &Params) {
+void ClangdLSPServer::onCodeAction(CodeActionParams &Params) {
   // We provide a code action for each diagnostic at the requested location
   // which has FixIts available.
   auto Code = Server.getDocument(Params.textDocument.uri.file);
   if (!Code)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       "onCodeAction called for non-added file");
 
   json::ary Commands;
@@ -235,68 +230,54 @@
       });
     }
   }
-  reply(C, std::move(Commands));
+  reply(std::move(Commands));
 }
 
-void ClangdLSPServer::onCompletion(Ctx C, TextDocumentPositionParams &Params) {
-  auto Reply = Server
-                   .codeComplete(std::move(C), Params.textDocument.uri.file,
-                                 Position{Params.position.line,
-                                          Params.position.character},
-                                 CCOpts)
-                   .get(); // FIXME(ibiryukov): This could be made async if we
-                           // had an API that would allow to attach callbacks to
-                           // futures returned by ClangdServer.
-
-  // We have std::move'd from C, now restore it from response of codeComplete.
-  C = std::move(Reply.first);
-  auto List = std::move(Reply.second.Value);
-  reply(C, List);
+void ClangdLSPServer::onCompletion(TextDocumentPositionParams &Params) {
+  Server.codeComplete(Params.textDocument.uri.file,
+                      Position{Params.position.line, Params.position.character},
+                      CCOpts,
+                      [](Tagged<CompletionList> List) { reply(List.Value); });
 }
 
-void ClangdLSPServer::onSignatureHelp(Ctx C,
-                                      TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onSignatureHelp(TextDocumentPositionParams &Params) {
   auto SignatureHelp = Server.signatureHelp(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
   if (!SignatureHelp)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       llvm::toString(SignatureHelp.takeError()));
-  reply(C, SignatureHelp->Value);
+  reply(SignatureHelp->Value);
 }
 
-void ClangdLSPServer::onGoToDefinition(Ctx C,
-                                       TextDocumentPositionParams &Params) {
+void ClangdLSPServer::onGoToDefinition(TextDocumentPositionParams &Params) {
   auto Items = Server.findDefinitions(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
   if (!Items)
-    return replyError(C, ErrorCode::InvalidParams,
+    return replyError(ErrorCode::InvalidParams,
                       llvm::toString(Items.takeError()));
-  reply(C, json::ary(Items->Value));
+  reply(json::ary(Items->Value));
 }
 
-void ClangdLSPServer::onSwitchSourceHeader(Ctx C,
-                                           TextDocumentIdentifier &Params) {
+void ClangdLSPServer::onSwitchSourceHeader(TextDocumentIdentifier &Params) {
   llvm::Optional<Path> Result = Server.switchSourceHeader(Params.uri.file);
   std::string ResultUri;
-  reply(C, Result ? URI::fromFile(*Result).uri : "");
+  reply(Result ? URI::fromFile(*Result).uri : "");
 }
 
-void ClangdLSPServer::onDocumentHighlight(Ctx C,
-                                          TextDocumentPositionParams &Params) {
-
+void ClangdLSPServer::onDocumentHighlight(TextDocumentPositionParams &Params) {
   auto Highlights = Server.findDocumentHighlights(
-      C, Params.textDocument.uri.file,
+      Params.textDocument.uri.file,
       Position{Params.position.line, Params.position.character});
 
   if (!Highlights) {
-    replyError(C, ErrorCode::InternalError,
+    replyError(ErrorCode::InternalError,
                llvm::toString(Highlights.takeError()));
     return;
   }
 
-  reply(C, json::ary(Highlights->Value));
+  reply(json::ary(Highlights->Value));
 }
 
 ClangdLSPServer::ClangdLSPServer(JSONOutput &Out, unsigned AsyncThreadsCount,
@@ -315,8 +296,8 @@
   assert(!IsDone && "Run was called before");
 
   // Set up JSONRPCDispatcher.
-  JSONRPCDispatcher Dispatcher([](Context Ctx, const json::Expr &Params) {
-    replyError(Ctx, ErrorCode::MethodNotFound, "method not found");
+  JSONRPCDispatcher Dispatcher([](const json::Expr &Params) {
+    replyError(ErrorCode::MethodNotFound, "method not found");
   });
   registerCallbackHandlers(Dispatcher, Out, /*Callbacks=*/*this);
 
@@ -346,8 +327,7 @@
 }
 
 void ClangdLSPServer::onDiagnosticsReady(
-    const Context &Ctx, PathRef File,
-    Tagged<std::vector<DiagWithFixIts>> Diagnostics) {
+    PathRef File, Tagged<std::vector<DiagWithFixIts>> Diagnostics) {
   json::ary DiagnosticsJSON;
 
   DiagnosticToReplacementMap LocalFixIts; // Temporary storage
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
http://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to