[clangd] Remove threading-related code from ClangdUnit.h

[ioeric]

Nice! The code looks much simpler!

Just some drive-by nits. I don't know the threading work well enough to give useful comments. Will leave the approval to others.

clangd/ClangdUnit.cpp

399

Do we still need this block?

434

nit: no braces

clangd/ClangdUnit.h

161

s/latest/last/?

[ilya-biryukov]

• Removed braces
• s/latest/last/

[ilya-biryukov]

Thanks for the NITs :-)

clangd/ClangdUnit.cpp

399

I added it to avoid referencing local variables declared there.

I would move this to be a separate function rather than removing the braces here, but other may have different opinions :-)
WDYT? Does a separate function sound ok?

[ioeric]

clangd/ClangdUnit.cpp

399

I don't think it's necessary though? But up to you :)

[sammccall]

:-D

clangd/ClangdServer.h

 #include "TUScheduler.h"
 #include "index/FileIndex.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include <functional>
+#include <future>
 #include <string>
 #include <type_traits>
 #include <utility>
 
 namespace clang {
 class PCHContainerOperations;
 
 namespace clangd {

clangd/ClangdUnit.h

 #include "Function.h"
 #include "Path.h"
 #include "Protocol.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Serialization/ASTBitCodes.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
-#include <atomic>
-#include <future>
 #include <memory>
-#include <mutex>
+#include <string>
+#include <vector>
 
 namespace llvm {
 class raw_ostream;
 }
 
 namespace clang {
 class PCHContainerOperations;
 
private:
   std::unique_ptr<FrontendAction> Action;
 
   // Data, stored after parsing.
   std::vector<DiagWithFixIts> Diags;
   std::vector<const Decl *> TopLevelDecls;
   bool PreambleDeclsDeserialized;
 };
 
-// Provides thread-safe access to ParsedAST.
-class ParsedASTWrapper {
-public:
-  ParsedASTWrapper(ParsedASTWrapper &&Wrapper);
-  ParsedASTWrapper(llvm::Optional<ParsedAST> AST);
-
-  /// Runs \p F on wrapped ParsedAST under lock. Ensures it is not accessed
-  /// concurrently.
-  template <class Func> void runUnderLock(Func F) const {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    F(AST ? AST.getPointer() : nullptr);
-  }
-
-private:
-  // This wrapper is used as an argument to std::shared_future (and it returns a
-  // const ref in get()), but we need to have non-const ref in order to
-  // implement some features.
-  mutable std::mutex Mutex;
-  mutable llvm::Optional<ParsedAST> AST;
-};
-
 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.
-/// NOTE: Threading-related bits of CppFile are now deprecated and will be
-/// removed soon.
-class CppFile : public std::enable_shared_from_this<CppFile> {
+class CppFile {
 public:
-  // We only allow to create CppFile as shared_ptr, because a future returned by
-  // deferRebuild will hold references to it.
-  static std::shared_ptr<CppFile>
-  Create(PathRef FileName, bool StorePreamblesInMemory,
-         std::shared_ptr<PCHContainerOperations> PCHs,
-         ASTParsedCallback ASTCallback);
-
-private:
   CppFile(PathRef FileName, bool StorePreamblesInMemory,
           std::shared_ptr<PCHContainerOperations> PCHs,
           ASTParsedCallback ASTCallback);
 
-public:
-  CppFile(CppFile const &) = delete;
-  CppFile(CppFile &&) = delete;
-
-  /// Cancels a scheduled rebuild, if any, and sets AST and Preamble to nulls.
-  /// If a rebuild is in progress, will wait for it to finish.
-  void cancelRebuild();
-
-  /// Similar to deferRebuild, but sets both Preamble and AST to nulls instead
-  /// of doing an actual parsing. Returned function is a deferred computation
-  /// that will wait for any ongoing rebuilds to finish and actually set the AST
-  /// and Preamble to nulls. It can be run on a different thread. This function
-  /// is useful to cancel ongoing rebuilds, if any, before removing CppFile.
-  /// DEPRECATED. This function will be removed soon, please do not use it.
-  UniqueFunction<void()> deferCancelRebuild();
-
-  /// Rebuild AST and Preamble synchronously on the calling thread.
-  /// Returns a list of diagnostics or a llvm::None, if another rebuild was
-  /// requested in parallel (effectively cancelling this rebuild) before
-  /// diagnostics were produced.
+  /// Rebuild the AST and the preamble.
+  /// Returns a list of diagnostics or llvm::None, if an error occured.
   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.
-  /// After calling this method, resources, available via futures returned by
-  /// getPreamble() and getAST(), will be waiting for rebuild to finish. A
-  /// continuation fininshing rebuild, returned by this function, must be
-  /// computed(i.e., operator() must be called on it) in order to make those
-  /// resources ready. If deferRebuild is called again before the rebuild is
-  /// finished (either because returned future had not been called or because it
-  /// had not returned yet), the previous rebuild request is cancelled and the
-  /// resource futures (returned by getPreamble() or getAST()) that were not
-  /// ready will be waiting for the last rebuild to finish instead.
-  /// 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.
-  /// DEPRECATED. This function will be removed soon, please do not use it.
-  UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
-  deferRebuild(ParseInputs &&Inputs);
-
-  /// Returns a future to get the most fresh PreambleData for a file. The
-  /// future will wait until the Preamble is rebuilt.
-  std::shared_future<std::shared_ptr<const PreambleData>> getPreamble() const;
-  /// Return some preamble for a file. It might be stale, but won't wait for
-  /// rebuild to finish.
-  std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
-
-  /// Returns a future to get the most fresh AST for a file. Returned AST is
-  /// wrapped to prevent concurrent accesses.
-  /// We use std::shared_ptr here because MVSC fails to compile non-copyable
-  /// classes as template arguments of promise/future. It is guaranteed to
-  /// always be non-null.
-  std::shared_future<std::shared_ptr<ParsedASTWrapper>> getAST() const;
-
+  /// Returns the last built preamble.
+  const std::shared_ptr<const PreambleData> &getPreamble() const;
+  /// Returns the last built AST.
+  ParsedAST *getAST() const;
   /// Returns an estimated size, in bytes, currently occupied by the AST and the
   /// Preamble.
   std::size_t getUsedBytes() const;
 
 private:
-  /// A helper guard that manages the state of CppFile during rebuild.
-  class RebuildGuard {
-  public:
-    RebuildGuard(CppFile &File, unsigned RequestRebuildCounter);
-    ~RebuildGuard();
-
-    bool wasCancelledBeforeConstruction() const;
-
-  private:
-    CppFile &File;
-    unsigned RequestRebuildCounter;
-    bool WasCancelledBeforeConstruction;
-  };
+  /// Build a new preamble for \p Inputs. If the current preamble can be reused,
+  /// it is returned instead.
+  /// This method is const to ensure we don't incidentally modify any fields.
+  std::shared_ptr<const PreambleData>
+  rebuildPreamble(CompilerInvocation &CI,
+                  IntrusiveRefCntPtr<vfs::FileSystem> FS,
+                  llvm::MemoryBuffer &ContentsBuffer) const;
 
   Path FileName;
   bool StorePreamblesInMemory;
 
-  /// Mutex protects all fields, declared below it, FileName and Command are not
-  /// mutated.
-  mutable std::mutex Mutex;
-  /// A counter to cancel old rebuilds.
-  unsigned RebuildCounter;
-  /// Used to wait when rebuild is finished before starting another one.
-  bool RebuildInProgress;
-  /// Condition variable to indicate changes to RebuildInProgress.
-  std::condition_variable RebuildCond;
-
-  /// Size of the last built AST, in bytes.
-  std::size_t ASTMemUsage;
-  /// Size of the last build Preamble, in bytes.
-  std::size_t PreambleMemUsage;
-
-  /// Promise and future for the latests AST. Fulfilled during rebuild.
-  /// We use std::shared_ptr here because MVSC fails to compile non-copyable
-  /// classes as template arguments of promise/future.
-  std::promise<std::shared_ptr<ParsedASTWrapper>> ASTPromise;
-  std::shared_future<std::shared_ptr<ParsedASTWrapper>> ASTFuture;
-
-  /// Promise and future for the latests Preamble. Fulfilled during rebuild.
-  std::promise<std::shared_ptr<const PreambleData>> PreamblePromise;
-  std::shared_future<std::shared_ptr<const PreambleData>> PreambleFuture;
-  /// Latest preamble that was built. May be stale, but always available without
-  /// waiting for rebuild to finish.
-  std::shared_ptr<const PreambleData> LatestAvailablePreamble;
-  /// Utility class, required by clang.
+  /// The latest parsed AST.
+  llvm::Optional<ParsedAST> AST;
+  /// The latest built Preamble.
+  std::shared_ptr<const PreambleData> Preamble;
+  /// Utility class required by clang
   std::shared_ptr<PCHContainerOperations> PCHs;
   /// This is called after the file is parsed. This can be nullptr if there is
   /// no callback.
   ASTParsedCallback ASTCallback;
 };
 
 /// Get the beginning SourceLocation at a specified \p Pos.
 SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,

clangd/ClangdUnit.cpp

 #include "clang/Serialization/ASTWriter.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
-#include <chrono>
 
 using namespace clang::clangd;
 using namespace clang;
 
 namespace {
 
 template <class T> std::size_t getUsedBytes(const std::vector<T> &Vec) {
   return Vec.capacity() * sizeof(T);
if (LangOpts)
         Output.push_back(std::move(*D));
   }
 
 private:
   std::vector<DiagWithFixIts> &Output;
   llvm::Optional<LangOptions> LangOpts;
 };
 
-template <class T> bool futureIsReady(std::shared_future<T> const &Future) {
-  return Future.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
-}
-
 } // namespace
 
 void clangd::dumpAST(ParsedAST &AST, llvm::raw_ostream &OS) {
   AST.getASTContext().getTranslationUnitDecl()->dump(OS, true);
 }
 
 llvm::Optional<ParsedAST>
 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,
                  IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
-
   std::vector<DiagWithFixIts> ASTDiags;
   StoreDiagsConsumer UnitDiagsConsumer(/*ref*/ ASTDiags);
 
   const PrecompiledPreamble *PreamblePCH =
       Preamble ? &Preamble->Preamble : nullptr;
   auto Clang = prepareCompilerInstance(
       std::move(CI), PreamblePCH, std::move(Buffer), std::move(PCHs),
       std::move(VFS), /*ref*/ UnitDiagsConsumer);
ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Diags(std::move(Diags)),
       TopLevelDecls(std::move(TopLevelDecls)),
       PreambleDeclsDeserialized(false) {
   assert(this->Clang);
   assert(this->Action);
 }
 
-ParsedASTWrapper::ParsedASTWrapper(ParsedASTWrapper &&Wrapper)
-    : AST(std::move(Wrapper.AST)) {}
-
-ParsedASTWrapper::ParsedASTWrapper(llvm::Optional<ParsedAST> AST)
-    : AST(std::move(AST)) {}
-
-std::shared_ptr<CppFile>
-CppFile::Create(PathRef FileName, bool StorePreamblesInMemory,
-                std::shared_ptr<PCHContainerOperations> PCHs,
-                ASTParsedCallback ASTCallback) {
-  return std::shared_ptr<CppFile>(new CppFile(FileName, StorePreamblesInMemory,
-                                              std::move(PCHs),
-                                              std::move(ASTCallback)));
-}
-
 CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
                  std::shared_ptr<PCHContainerOperations> PCHs,
                  ASTParsedCallback ASTCallback)
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
-      RebuildCounter(0), RebuildInProgress(false), ASTMemUsage(0),
-      PreambleMemUsage(0), PCHs(std::move(PCHs)),
-      ASTCallback(std::move(ASTCallback)) {
+      PCHs(std::move(PCHs)), ASTCallback(std::move(ASTCallback)) {
   log("Created CppFile for " + FileName);
-
-  std::lock_guard<std::mutex> Lock(Mutex);
-  LatestAvailablePreamble = nullptr;
-  PreamblePromise.set_value(nullptr);
-  PreambleFuture = PreamblePromise.get_future();
-
-  ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
-  ASTFuture = ASTPromise.get_future();
-}
-
-void CppFile::cancelRebuild() { deferCancelRebuild()(); }
-
-UniqueFunction<void()> CppFile::deferCancelRebuild() {
-  std::unique_lock<std::mutex> Lock(Mutex);
-  // Cancel an ongoing rebuild, if any, and wait for it to finish.
-  unsigned RequestRebuildCounter = ++this->RebuildCounter;
-  // Rebuild asserts that futures aren't ready if rebuild is cancelled.
-  // We want to keep this invariant.
-  if (futureIsReady(PreambleFuture)) {
-    PreamblePromise = std::promise<std::shared_ptr<const PreambleData>>();
-    PreambleFuture = PreamblePromise.get_future();
-  }
-  if (futureIsReady(ASTFuture)) {
-    ASTPromise = std::promise<std::shared_ptr<ParsedASTWrapper>>();
-    ASTFuture = ASTPromise.get_future();
-  }
-
-  Lock.unlock();
-  // Notify about changes to RebuildCounter.
-  RebuildCond.notify_all();
-
-  std::shared_ptr<CppFile> That = shared_from_this();
-  return [That, RequestRebuildCounter]() {
-    std::unique_lock<std::mutex> Lock(That->Mutex);
-    CppFile *This = &*That;
-    This->RebuildCond.wait(Lock, [This, RequestRebuildCounter]() {
-      return !This->RebuildInProgress ||
-             This->RebuildCounter != RequestRebuildCounter;
-    });
-
-    // This computation got cancelled itself, do nothing.
-    if (This->RebuildCounter != RequestRebuildCounter)
-      return;
-
-    // Set empty results for Promises.
-    That->PreambleMemUsage = 0;
-    That->PreamblePromise.set_value(nullptr);
-    That->ASTMemUsage = 0;
-    That->ASTPromise.set_value(std::make_shared<ParsedASTWrapper>(llvm::None));
-  };
 }
 
 llvm::Optional<std::vector<DiagWithFixIts>>
 CppFile::rebuild(ParseInputs &&Inputs) {
-  return deferRebuild(std::move(Inputs))();
-}
-
-UniqueFunction<llvm::Optional<std::vector<DiagWithFixIts>>()>
-CppFile::deferRebuild(ParseInputs &&Inputs) {
-  std::shared_ptr<const PreambleData> OldPreamble;
-  std::shared_ptr<PCHContainerOperations> PCHs;
-  unsigned RequestRebuildCounter;
-  {
-    std::unique_lock<std::mutex> Lock(Mutex);
-    // Increase RebuildCounter to cancel all ongoing FinishRebuild operations.
-    // They will try to exit as early as possible and won't call set_value on
-    // our promises.
-    RequestRebuildCounter = ++this->RebuildCounter;
-    PCHs = this->PCHs;
-
-    // Remember the preamble to be used during rebuild.
-    OldPreamble = this->LatestAvailablePreamble;
-    // Setup std::promises and std::futures for Preamble and AST. Corresponding
-    // futures will wait until the rebuild process is finished.
-    if (futureIsReady(this->PreambleFuture)) {
-      this->PreamblePromise =
-          std::promise<std::shared_ptr<const PreambleData>>();
-      this->PreambleFuture = this->PreamblePromise.get_future();
-    }
-    if (futureIsReady(this->ASTFuture)) {
-      this->ASTPromise = std::promise<std::shared_ptr<ParsedASTWrapper>>();
-      this->ASTFuture = this->ASTPromise.get_future();
-    }
-  } // unlock Mutex.
-  // Notify about changes to RebuildCounter.
-  RebuildCond.notify_all();
-
-  // A helper to function to finish the rebuild. May be run on a different
-  // thread.
-
-  // Don't let this CppFile die before rebuild is finished.
-  std::shared_ptr<CppFile> That = shared_from_this();
-  auto FinishRebuild =
-      [OldPreamble, RequestRebuildCounter, PCHs,
-       That](ParseInputs Inputs) mutable /* to allow changing OldPreamble. */
-      -> llvm::Optional<std::vector<DiagWithFixIts>> {
-    log("Rebuilding file " + That->FileName + " with command [" +
-        Inputs.CompileCommand.Directory + "] " +
-        llvm::join(Inputs.CompileCommand.CommandLine, " "));
+  log("Rebuilding file " + FileName + " with command [" +
 
-    // Only one execution of this method is possible at a time.
-    // RebuildGuard will wait for any ongoing rebuilds to finish and will put us
-    // into a state for doing a rebuild.
-    RebuildGuard Rebuild(*That, RequestRebuildCounter);
-    if (Rebuild.wasCancelledBeforeConstruction())
-      return llvm::None;
-
-    std::vector<const char *> ArgStrs;
-    for (const auto &S : Inputs.CompileCommand.CommandLine)
-      ArgStrs.push_back(S.c_str());
 
-    Inputs.FS->setCurrentWorkingDirectory(Inputs.CompileCommand.Directory);
 
-    std::unique_ptr<CompilerInvocation> CI;
-    {
-      // FIXME(ibiryukov): store diagnostics from CommandLine when we start
-      // reporting them.
-      IgnoreDiagnostics IgnoreDiagnostics;
-      IntrusiveRefCntPtr<DiagnosticsEngine> CommandLineDiagsEngine =
-          CompilerInstance::createDiagnostics(new DiagnosticOptions,
-                                              &IgnoreDiagnostics, false);
-      CI = createInvocationFromCommandLine(ArgStrs, CommandLineDiagsEngine,
-                                           Inputs.FS);
-      // createInvocationFromCommandLine sets DisableFree.
-      CI->getFrontendOpts().DisableFree = false;
-    }
-    assert(CI && "Couldn't create CompilerInvocation");
+  // Prepare CompilerInvocation.
 
-    std::unique_ptr<llvm::MemoryBuffer> ContentsBuffer =
-        llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents, That->FileName);
-
-    // A helper function to rebuild the preamble or reuse the existing one. Does
-    // not mutate any fields of CppFile, only does the actual computation.
-    // Lamdba is marked mutable to call reset() on OldPreamble.
-    auto DoRebuildPreamble =
-        [&]() mutable -> std::shared_ptr<const PreambleData> {
-      auto Bounds =
-          ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
-      if (OldPreamble &&
-          OldPreamble->Preamble.CanReuse(*CI, ContentsBuffer.get(), Bounds,
-                                         Inputs.FS.get())) {
-        log("Reusing preamble for file " + Twine(That->FileName));
-        return OldPreamble;
-      }
-      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("Preamble");
-      SPAN_ATTACH(Tracer, "File", That->FileName);
-      std::vector<DiagWithFixIts> PreambleDiags;
-      StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
-      IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
-          CompilerInstance::createDiagnostics(
-              &CI->getDiagnosticOpts(), &PreambleDiagnosticsConsumer, false);
-
-      // Skip function bodies when building the preamble to speed up building
-      // the preamble and make it smaller.
-      assert(!CI->getFrontendOpts().SkipFunctionBodies);
-      CI->getFrontendOpts().SkipFunctionBodies = true;
-
-      CppFilePreambleCallbacks SerializedDeclsCollector;
-      auto BuiltPreamble = PrecompiledPreamble::Build(
-          *CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine, Inputs.FS,
-          PCHs,
-          /*StoreInMemory=*/That->StorePreamblesInMemory,
-          SerializedDeclsCollector);
-
-      // When building the AST for the main file, we do want the function
-      // bodies.
-      CI->getFrontendOpts().SkipFunctionBodies = false;
-
-      if (BuiltPreamble) {
-        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("Could not build a preamble for file " + Twine(That->FileName));
-        return nullptr;
-      }
-    };
+      llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents, FileName);
 
-    // Compute updated Preamble.
-    std::shared_ptr<const PreambleData> NewPreamble = DoRebuildPreamble();
-    // Publish the new Preamble.
-    {
-      std::lock_guard<std::mutex> Lock(That->Mutex);
-      // We always set LatestAvailablePreamble to the new value, hoping that it
-      // will still be usable in the further requests.
-      That->LatestAvailablePreamble = NewPreamble;
-      if (RequestRebuildCounter != That->RebuildCounter)
-        return llvm::None; // Our rebuild request was cancelled, do nothing.
-      That->PreambleMemUsage =
-          NewPreamble ? NewPreamble->Preamble.getSize() : 0;
-      That->PreamblePromise.set_value(NewPreamble);
-    } // unlock Mutex
-
-    // Prepare the Preamble and supplementary data for rebuilding AST.
-    std::vector<DiagWithFixIts> Diagnostics;
-    if (NewPreamble) {
-      Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
-                         NewPreamble->Diags.end());
-    }
+  std::shared_ptr<const PreambleData> NewPreamble =
+      rebuildPreamble(*CI, Inputs.FS, *ContentsBuffer);
 
-    // Compute updated AST.
-    llvm::Optional<ParsedAST> NewAST;
-    {
-      trace::Span Tracer("Build");
-      SPAN_ATTACH(Tracer, "File", That->FileName);
-      NewAST = ParsedAST::Build(std::move(CI), std::move(NewPreamble),
-                                std::move(ContentsBuffer), PCHs, Inputs.FS);
-    }
+  // Remove current AST to avoid wasting memory.
+  AST = llvm::None;
+    SPAN_ATTACH(Tracer, "File", FileName);
+    NewAST = ParsedAST::Build(std::move(CI), NewPreamble,
 
-    if (NewAST) {
-      Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
-                         NewAST->getDiagnostics().end());
-      if (That->ASTCallback)
-        That->ASTCallback(That->FileName, NewAST.getPointer());
-    } else {
-      // Don't report even Preamble diagnostics if we couldn't build AST.
-      Diagnostics.clear();
-    }
+  std::vector<DiagWithFixIts> Diagnostics;
+    // Collect diagnostics from both the preamble and the AST.
+    if (NewPreamble) {
+      Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
+                         NewPreamble->Diags.end());
+    }
+  if (ASTCallback && NewAST)
+    ASTCallback(FileName, NewAST.getPointer());
 
-    // Publish the new AST.
-    {
-      std::lock_guard<std::mutex> Lock(That->Mutex);
-      if (RequestRebuildCounter != That->RebuildCounter)
-        return Diagnostics; // Our rebuild request was cancelled, don't set
-                            // ASTPromise.
-
-      That->ASTMemUsage = NewAST ? NewAST->getUsedBytes() : 0;
-      That->ASTPromise.set_value(
-          std::make_shared<ParsedASTWrapper>(std::move(NewAST)));
-    } // unlock Mutex
-
-    return Diagnostics;
-  };
-
-  return BindWithForward(FinishRebuild, std::move(Inputs));
+  // Write the results of rebuild into class fields.
+  Preamble = std::move(NewPreamble);
+  AST = std::move(NewAST);
 }
 
-std::shared_future<std::shared_ptr<const PreambleData>>
-CppFile::getPreamble() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  return PreambleFuture;
+const std::shared_ptr<const PreambleData> &CppFile::getPreamble() const {
+  return Preamble;
 }
 
-std::shared_ptr<const PreambleData> CppFile::getPossiblyStalePreamble() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  return LatestAvailablePreamble;
-}
-
-std::shared_future<std::shared_ptr<ParsedASTWrapper>> CppFile::getAST() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  return ASTFuture;
+ParsedAST *CppFile::getAST() const {
+  // We could add mutable to AST instead of const_cast here, but that would also
+  // allow writing to AST from const methods.
+  return AST ? const_cast<ParsedAST *>(AST.getPointer()) : nullptr;
 }
 
 std::size_t CppFile::getUsedBytes() const {
-  std::lock_guard<std::mutex> Lock(Mutex);
-  // FIXME: We should not store extra size fields. When we store AST and
-  // Preamble directly, not inside futures, we could compute the sizes from the
-  // stored AST and the preamble in this function directly.
-  return ASTMemUsage + PreambleMemUsage;
+  std::size_t Total = 0;
+  if (AST)
+    Total += AST->getUsedBytes();
+  if (StorePreamblesInMemory && Preamble)
+    Total += Preamble->Preamble.getSize();
+  return Total;
 }
 
-CppFile::RebuildGuard::RebuildGuard(CppFile &File,
-                                    unsigned RequestRebuildCounter)
-    : File(File), RequestRebuildCounter(RequestRebuildCounter) {
-  std::unique_lock<std::mutex> Lock(File.Mutex);
-  WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
-  if (WasCancelledBeforeConstruction)
-    return;
+std::shared_ptr<const PreambleData>
+CppFile::rebuildPreamble(CompilerInvocation &CI,
+                         IntrusiveRefCntPtr<vfs::FileSystem> FS,
+                         llvm::MemoryBuffer &ContentsBuffer) const {
+  const auto &OldPreamble = this->Preamble;
+  auto Bounds = ComputePreambleBounds(*CI.getLangOpts(), &ContentsBuffer, 0);
+  if (OldPreamble &&
+      OldPreamble->Preamble.CanReuse(CI, &ContentsBuffer, Bounds, FS.get())) {
+    log("Reusing preamble for file " + Twine(FileName));
+    return OldPreamble;
+  }
+  log("Preamble for file " + Twine(FileName) +
+      " cannot be reused. Attempting to rebuild it.");
 
-  File.RebuildCond.wait(Lock, [&File, RequestRebuildCounter]() {
-    return !File.RebuildInProgress ||
-           File.RebuildCounter != RequestRebuildCounter;
-  });
+  trace::Span Tracer("Preamble");
+  SPAN_ATTACH(Tracer, "File", FileName);
+  std::vector<DiagWithFixIts> PreambleDiags;
+  StoreDiagsConsumer PreambleDiagnosticsConsumer(/*ref*/ PreambleDiags);
+  IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
+      CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
+                                          &PreambleDiagnosticsConsumer, false);
 
-  WasCancelledBeforeConstruction = File.RebuildCounter != RequestRebuildCounter;
-  if (WasCancelledBeforeConstruction)
-    return;
+  // Skip function bodies when building the preamble to speed up building
+  // the preamble and make it smaller.
+  assert(!CI.getFrontendOpts().SkipFunctionBodies);
+  CI.getFrontendOpts().SkipFunctionBodies = true;
 
-  File.RebuildInProgress = true;
-}
+  CppFilePreambleCallbacks SerializedDeclsCollector;
+  auto BuiltPreamble = PrecompiledPreamble::Build(
+      CI, &ContentsBuffer, Bounds, *PreambleDiagsEngine, FS, PCHs,
+      /*StoreInMemory=*/StorePreamblesInMemory, SerializedDeclsCollector);
 
-bool CppFile::RebuildGuard::wasCancelledBeforeConstruction() const {
-  return WasCancelledBeforeConstruction;
-}
+  // When building the AST for the main file, we do want the function
+  // bodies.
+  CI.getFrontendOpts().SkipFunctionBodies = false;
 
-CppFile::RebuildGuard::~RebuildGuard() {
-  if (WasCancelledBeforeConstruction)
-    return;
+  if (BuiltPreamble) {
+    log("Built preamble of size " + Twine(BuiltPreamble->getSize()) +
+        " for file " + Twine(FileName));
 
-  std::unique_lock<std::mutex> Lock(File.Mutex);
-  assert(File.RebuildInProgress);
-  File.RebuildInProgress = false;
-
-  if (File.RebuildCounter == RequestRebuildCounter) {
-    // Our rebuild request was successful.
-    assert(futureIsReady(File.ASTFuture));
-    assert(futureIsReady(File.PreambleFuture));
+    return std::make_shared<PreambleData>(
+        std::move(*BuiltPreamble),
+        SerializedDeclsCollector.takeTopLevelDeclIDs(),
+        std::move(PreambleDiags));
   } else {
-    // Our rebuild request was cancelled, because further reparse was requested.
-    assert(!futureIsReady(File.ASTFuture));
-    assert(!futureIsReady(File.PreambleFuture));
+    log("Could not build a preamble for file " + Twine(FileName));
+    return nullptr;
   }
-
-  Lock.unlock();
-  File.RebuildCond.notify_all();
 }
 
 SourceLocation clangd::getBeginningOfIdentifier(ParsedAST &Unit,
                                                 const Position &Pos,
                                                 const FileEntry *FE) {
   // The language server protocol uses zero-based line and column numbers.
   // Clang uses one-based numbers.
 

clangd/TUScheduler.cpp

 //   immediately.
 
 #include "TUScheduler.h"
 #include "Logger.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "llvm/Support/Errc.h"
 #include <memory>
 #include <queue>
+#include <thread>
 
 namespace clang {
 namespace clangd {
 namespace {
 class ASTWorkerHandle;
 
 /// Owns one instance of the AST, schedules updates and reads of it.
 /// Also responsible for building and providing access to the preamble.
 /// Each ASTWorker processes the async requests sent to it on a separate
 /// dedicated thread.
 /// The ASTWorker that manages the AST is shared by both the processing thread
 /// and the TUScheduler. The TUScheduler should discard an ASTWorker when
 /// remove() is called, but its thread may be busy and we don't want to block.
 /// So the workers are accessed via an ASTWorkerHandle. Destroying the handle
 /// signals the worker to exit its run loop and gives up shared ownership of the
 /// worker.
 class ASTWorker {
   friend class ASTWorkerHandle;
-  ASTWorker(Semaphore &Barrier, std::shared_ptr<CppFile> AST, bool RunSync);
+  ASTWorker(Semaphore &Barrier, CppFile AST, bool RunSync);
 
 public:
   /// Create a new ASTWorker and return a handle to it.
   /// The processing thread is spawned using \p Tasks. However, when \p Tasks
   /// is null, all requests will be processed on the calling thread
   /// synchronously instead. \p Barrier is acquired when processing each
   /// request, it is be used to limit the number of actively running threads.
   static ASTWorkerHandle Create(AsyncTaskRunner *Tasks, Semaphore &Barrier,
-                                std::shared_ptr<CppFile> AST);
+                                CppFile AST);
   ~ASTWorker();
 
   void update(ParseInputs Inputs,
               UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
                   OnUpdated);
   void runWithAST(UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action);
 
   std::shared_ptr<const PreambleData> getPossiblyStalePreamble() const;
private:
   void startTask(UniqueFunction<void()> Task, bool isUpdate,
                  llvm::Optional<CancellationFlag> CF);
 
   using RequestWithCtx = std::pair<UniqueFunction<void()>, Context>;
 
   const bool RunSync;
   Semaphore &Barrier;
   // AST and FileInputs are only accessed on the processing thread from run().
-  const std::shared_ptr<CppFile> AST;
+  CppFile AST;
   // Inputs, corresponding to the current state of AST.
   ParseInputs FileInputs;
   // Guards members used by both TUScheduler and the worker thread.
   mutable std::mutex Mutex;
+  std::shared_ptr<const PreambleData> LastBuiltPreamble; /* GUARDED_BY(Mutex) */
+  // Result of getUsedBytes() after the last rebuild or read of AST.
+  std::size_t LastASTSize; /* GUARDED_BY(Mutex) */
   // Set to true to signal run() to finish processing.
   bool Done;                           /* GUARDED_BY(Mutex) */
   std::queue<RequestWithCtx> Requests; /* GUARDED_BY(Mutex) */
   // Only set when last request is an update. This allows us to cancel an update
   // that was never read, if a subsequent update comes in.
   llvm::Optional<CancellationFlag> LastUpdateCF; /* GUARDED_BY(Mutex) */
   std::condition_variable RequestsCV;
 };
public:
   /// are possible.
   std::shared_ptr<const ASTWorker> lock() { return Worker; }
 
 private:
   std::shared_ptr<ASTWorker> Worker;
 };
 
 ASTWorkerHandle ASTWorker::Create(AsyncTaskRunner *Tasks, Semaphore &Barrier,
-                                  std::shared_ptr<CppFile> AST) {
+                                  CppFile AST) {
   std::shared_ptr<ASTWorker> Worker(
       new ASTWorker(Barrier, std::move(AST), /*RunSync=*/!Tasks));
   if (Tasks)
     Tasks->runAsync([Worker]() { Worker->run(); });
 
   return ASTWorkerHandle(std::move(Worker));
 }
 
-ASTWorker::ASTWorker(Semaphore &Barrier, std::shared_ptr<CppFile> AST,
-                     bool RunSync)
+ASTWorker::ASTWorker(Semaphore &Barrier, CppFile AST, bool RunSync)
     : RunSync(RunSync), Barrier(Barrier), AST(std::move(AST)), Done(false) {
   if (RunSync)
     return;
 }
 
 ASTWorker::~ASTWorker() {
 #ifndef NDEBUG
   std::lock_guard<std::mutex> Lock(Mutex);
   assert(Done && "handle was not destroyed");
   assert(Requests.empty() && "unprocessed requests when destroying ASTWorker");
 #endif
 }
 
 void ASTWorker::update(
     ParseInputs Inputs,
     UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
         OnUpdated) {
   auto Task = [=](CancellationFlag CF, decltype(OnUpdated) OnUpdated) mutable {
     if (CF.isCancelled()) {
       OnUpdated(llvm::None);
       return;
     }
     FileInputs = Inputs;
-    auto Diags = AST->rebuild(std::move(Inputs));
+    auto Diags = AST.rebuild(std::move(Inputs));
+
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      if (AST.getPreamble())
+        LastBuiltPreamble = AST.getPreamble();
+      LastASTSize = AST.getUsedBytes();
+    }
     // We want to report the diagnostics even if this update was cancelled.
     // It seems more useful than making the clients wait indefinitely if they
     // spam us with updates.
     OnUpdated(std::move(Diags));
   };
 
   CancellationFlag UpdateCF;
   startTask(BindWithForward(Task, UpdateCF, std::move(OnUpdated)),
             /*isUpdate=*/true, UpdateCF);
 }
 
 void ASTWorker::runWithAST(
     UniqueFunction<void(llvm::Expected<InputsAndAST>)> Action) {
   auto Task = [=](decltype(Action) Action) {
-    auto ASTWrapper = this->AST->getAST().get();
-    // FIXME: no need to lock here, cleanup the CppFile interface to get rid of
-    // them.
-    ASTWrapper->runUnderLock([&](ParsedAST *AST) {
-      if (!AST) {
-        Action(llvm::make_error<llvm::StringError>(
-            "invalid AST", llvm::errc::invalid_argument));
-        return;
-      }
-      Action(InputsAndAST{FileInputs, *AST});
-    });
+    ParsedAST *ActualAST = AST.getAST();
+    if (!ActualAST) {
+      Action(llvm::make_error<llvm::StringError>("invalid AST",
+                                                 llvm::errc::invalid_argument));
+    Action(InputsAndAST{FileInputs, *ActualAST});
+
+    // Size of the AST might have changed after reads too, e.g. if some decls
+    // were deserialized from preamble.
+    std::lock_guard<std::mutex> Lock(Mutex);
+    LastASTSize = ActualAST->getUsedBytes();
   };
 
   startTask(BindWithForward(Task, std::move(Action)), /*isUpdate=*/false,
             llvm::None);
 }
 
 std::shared_ptr<const PreambleData>
 ASTWorker::getPossiblyStalePreamble() const {
-  return AST->getPossiblyStalePreamble();
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return LastBuiltPreamble;
 }
 
 std::size_t ASTWorker::getUsedBytes() const {
-  // FIXME(ibiryukov): we'll need to take locks here after we remove
-  // thread-safety from CppFile. For now, CppFile is thread-safe and we can
-  // safely call methods on it without acquiring a lock.
-  return AST->getUsedBytes();
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return LastASTSize;
 }
 
 void ASTWorker::stop() {
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     assert(!Done && "stop() called twice");
     Done = true;
   }
void TUScheduler::update(
     PathRef File, ParseInputs Inputs,
     UniqueFunction<void(llvm::Optional<std::vector<DiagWithFixIts>>)>
         OnUpdated) {
   std::unique_ptr<FileData> &FD = Files[File];
   if (!FD) {
     // Create a new worker to process the AST-related tasks.
     ASTWorkerHandle Worker = ASTWorker::Create(
         Tasks ? Tasks.getPointer() : nullptr, Barrier,
-        CppFile::Create(File, StorePreamblesInMemory, PCHOps, ASTCallback));
+        CppFile(File, StorePreamblesInMemory, PCHOps, ASTCallback));
     FD = std::unique_ptr<FileData>(new FileData{Inputs, std::move(Worker)});
   } else {
     FD->Inputs = Inputs;
   }
   FD->Worker->update(std::move(Inputs), std::move(OnUpdated));
 }
 
 void TUScheduler::remove(PathRef File) {