[clangd] Ensure preamble outlives the AST

[ilya-biryukov]

Summary

In-memory preambles will not be copied anymore, so we need to make
sure they outlive the AST.

[sammccall]

clangd/ClangdUnit.h

51

can you move the definition here to avoid the extra decl?
(I tend to find this more readable and the struct is small, but if you prefer "top-down" that's fine too)

[ilya-biryukov]

• Moved PreambleData declaration up to avoid fwd-decl.

clangd/ClangdUnit.h

 class Logger;
 
 /// A diagnostic with its FixIts.
 struct DiagWithFixIts {
   clangd::Diagnostic Diag;
   llvm::SmallVector<tooling::Replacement, 1> FixIts;
 };
 
+// Stores Preamble and associated data.
+struct PreambleData {
+  PreambleData(PrecompiledPreamble Preamble,
+               std::vector<serialization::DeclID> TopLevelDeclIDs,
+               std::vector<DiagWithFixIts> Diags);
+
+  PrecompiledPreamble Preamble;
+  std::vector<serialization::DeclID> TopLevelDeclIDs;
+  std::vector<DiagWithFixIts> Diags;
+};
+
 /// Stores and provides access to parsed AST.
 class ParsedAST {
 public:
   /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
   /// it is reused during parsing.
   static llvm::Optional<ParsedAST>
   Build(std::unique_ptr<clang::CompilerInvocation> CI,
-        const PrecompiledPreamble *Preamble,
-        ArrayRef<serialization::DeclID> PreambleDeclIDs,
+        std::shared_ptr<const PreambleData> Preamble,
         std::unique_ptr<llvm::MemoryBuffer> Buffer,
         std::shared_ptr<PCHContainerOperations> PCHs,
         IntrusiveRefCntPtr<vfs::FileSystem> VFS, clangd::Logger &Logger);
 
   ParsedAST(ParsedAST &&Other);
   ParsedAST &operator=(ParsedAST &&Other);
 
   ~ParsedAST();
 
   ASTContext &getASTContext();
   const ASTContext &getASTContext() const;
 
   Preprocessor &getPreprocessor();
   const Preprocessor &getPreprocessor() const;
 
   /// This function returns all top-level decls, including those that come
   /// from Preamble. Decls, coming from Preamble, have to be deserialized, so
   /// this call might be expensive.
   ArrayRef<const Decl *> getTopLevelDecls();
 
   const std::vector<DiagWithFixIts> &getDiagnostics() const;
 
 private:
-  ParsedAST(std::unique_ptr<CompilerInstance> Clang,
+  ParsedAST(std::shared_ptr<const PreambleData> Preamble,
+            std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action,
             std::vector<const Decl *> TopLevelDecls,
-            std::vector<serialization::DeclID> PendingTopLevelDecls,
             std::vector<DiagWithFixIts> Diags);
 
 private:
   void ensurePreambleDeclsDeserialized();
 
+  // In-memory preambles must outlive the AST, it is important that this member
+  // goes before Clang and Action.
+  std::shared_ptr<const PreambleData> Preamble;
   // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
   // on it) and CompilerInstance used to run it. That way we don't have to do
   // complex memory management of all Clang structures on our own. (They are
   // stored in CompilerInstance and cleaned up by
   // FrontendAction.EndSourceFile).
   std::unique_ptr<CompilerInstance> Clang;
   std::unique_ptr<FrontendAction> Action;
 
   // Data, stored after parsing.
   std::vector<DiagWithFixIts> Diags;
   std::vector<const Decl *> TopLevelDecls;
-  std::vector<serialization::DeclID> PendingTopLevelDecls;
+  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;
 };
 
-// Stores Preamble and associated data.
-struct PreambleData {
-  PreambleData(PrecompiledPreamble Preamble,
-               std::vector<serialization::DeclID> TopLevelDeclIDs,
-               std::vector<DiagWithFixIts> Diags);
-
-  PrecompiledPreamble Preamble;
-  std::vector<serialization::DeclID> TopLevelDeclIDs;
-  std::vector<DiagWithFixIts> Diags;
-};
-
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
 class CppFile : public std::enable_shared_from_this<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, tooling::CompileCommand Command,

clangd/ClangdUnit.cpp

 }
 
 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,
-                 const PrecompiledPreamble *Preamble,
-                 ArrayRef<serialization::DeclID> PreambleDeclIDs,
+                 std::shared_ptr<const PreambleData> Preamble,
                  std::unique_ptr<llvm::MemoryBuffer> Buffer,
                  std::shared_ptr<PCHContainerOperations> PCHs,
                  IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                  clangd::Logger &Logger) {
 
   std::vector<DiagWithFixIts> ASTDiags;
   StoreDiagsConsumer UnitDiagsConsumer(/*ref*/ ASTDiags);
 
+  const PrecompiledPreamble *PreamblePCH =
+      Preamble ? &Preamble->Preamble : nullptr;
   auto Clang = prepareCompilerInstance(
-      std::move(CI), Preamble, std::move(Buffer), std::move(PCHs),
+      std::move(CI), PreamblePCH, std::move(Buffer), std::move(PCHs),
       std::move(VFS), /*ref*/ UnitDiagsConsumer);
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<CompilerInstance> CICleanup(
       Clang.get());
 
   auto Action = llvm::make_unique<ClangdFrontendAction>();
   const FrontendInputFile &MainInput = Clang->getFrontendOpts().Inputs[0];
   if (!Action->BeginSourceFile(*Clang, MainInput)) {
     Logger.log("BeginSourceFile() failed when building AST for " +
                MainInput.getFile());
     return llvm::None;
   }
   if (!Action->Execute())
     Logger.log("Execute() failed when building AST for " + MainInput.getFile());
 
   // UnitDiagsConsumer is local, we can not store it in CompilerInstance that
   // has a longer lifetime.
   Clang->getDiagnostics().setClient(new EmptyDiagsConsumer);
 
   std::vector<const Decl *> ParsedDecls = Action->takeTopLevelDecls();
-  std::vector<serialization::DeclID> PendingDecls;
-  if (Preamble) {
-    PendingDecls.reserve(PreambleDeclIDs.size());
-    PendingDecls.insert(PendingDecls.begin(), PreambleDeclIDs.begin(),
-                        PreambleDeclIDs.end());
-  }
-
-  return ParsedAST(std::move(Clang), std::move(Action), std::move(ParsedDecls),
-                   std::move(PendingDecls), std::move(ASTDiags));
+  return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
+                   std::move(ParsedDecls), std::move(ASTDiags));
 }
 
 namespace {
 
 SourceLocation getMacroArgExpandedLocation(const SourceManager &Mgr,
                                            const FileEntry *FE,
                                            unsigned Offset) {
   SourceLocation FileLoc = Mgr.translateFileLineCol(FE, 1, 1);
std::vector<Location> clangd::findDefinitions(ParsedAST &AST, Position Pos,
 
   indexTopLevelDecls(AST.getASTContext(), AST.getTopLevelDecls(),
                      DeclLocationsFinder, IndexOpts);
 
   return DeclLocationsFinder->takeLocations();
 }
 
 void ParsedAST::ensurePreambleDeclsDeserialized() {
-  if (PendingTopLevelDecls.empty())
+  if (PreambleDeclsDeserialized || !Preamble)
     return;
 
   std::vector<const Decl *> Resolved;
-  Resolved.reserve(PendingTopLevelDecls.size());
+  Resolved.reserve(Preamble->TopLevelDeclIDs.size());
 
   ExternalASTSource &Source = *getASTContext().getExternalSource();
-  for (serialization::DeclID TopLevelDecl : PendingTopLevelDecls) {
+  for (serialization::DeclID TopLevelDecl : Preamble->TopLevelDeclIDs) {
     // Resolve the declaration ID to an actual declaration, possibly
     // deserializing the declaration in the process.
     if (Decl *D = Source.GetExternalDecl(TopLevelDecl))
       Resolved.push_back(D);
   }
 
-  TopLevelDecls.reserve(TopLevelDecls.size() + PendingTopLevelDecls.size());
+  TopLevelDecls.reserve(TopLevelDecls.size() +
+                        Preamble->TopLevelDeclIDs.size());
   TopLevelDecls.insert(TopLevelDecls.begin(), Resolved.begin(), Resolved.end());
 
-  PendingTopLevelDecls.clear();
+  PreambleDeclsDeserialized = true;
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
 ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
 ParsedAST::~ParsedAST() {
   if (Action) {
ArrayRef<const Decl *> ParsedAST::getTopLevelDecls() {
   ensurePreambleDeclsDeserialized();
   return TopLevelDecls;
 }
 
 const std::vector<DiagWithFixIts> &ParsedAST::getDiagnostics() const {
   return Diags;
 }
 
-ParsedAST::ParsedAST(std::unique_ptr<CompilerInstance> Clang,
+PreambleData::PreambleData(PrecompiledPreamble Preamble,
+                           std::vector<serialization::DeclID> TopLevelDeclIDs,
+                           std::vector<DiagWithFixIts> Diags)
+    : Preamble(std::move(Preamble)),
+      TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)) {}
+
+ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
+                     std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      std::vector<const Decl *> TopLevelDecls,
-                     std::vector<serialization::DeclID> PendingTopLevelDecls,
                      std::vector<DiagWithFixIts> Diags)
-    : Clang(std::move(Clang)), Action(std::move(Action)),
-      Diags(std::move(Diags)), TopLevelDecls(std::move(TopLevelDecls)),
-      PendingTopLevelDecls(std::move(PendingTopLevelDecls)) {
+    : 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)) {}
 
-PreambleData::PreambleData(PrecompiledPreamble Preamble,
-                           std::vector<serialization::DeclID> TopLevelDeclIDs,
-                           std::vector<DiagWithFixIts> Diags)
-    : Preamble(std::move(Preamble)),
-      TopLevelDeclIDs(std::move(TopLevelDeclIDs)), Diags(std::move(Diags)) {}
-
 std::shared_ptr<CppFile>
 CppFile::Create(PathRef FileName, tooling::CompileCommand Command,
                 bool StorePreamblesInMemory,
                 std::shared_ptr<PCHContainerOperations> PCHs,
                 clangd::Logger &Logger) {
   return std::shared_ptr<CppFile>(new CppFile(FileName, std::move(Command),
                                               StorePreamblesInMemory,
                                               std::move(PCHs), Logger));
// Publish the new Preamble.
       // 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->PreamblePromise.set_value(NewPreamble);
     } // unlock Mutex
 
     // Prepare the Preamble and supplementary data for rebuilding AST.
-    const PrecompiledPreamble *PreambleForAST = nullptr;
-    ArrayRef<serialization::DeclID> SerializedPreambleDecls = llvm::None;
     std::vector<DiagWithFixIts> Diagnostics;
     if (NewPreamble) {
-      PreambleForAST = &NewPreamble->Preamble;
-      SerializedPreambleDecls = NewPreamble->TopLevelDeclIDs;
       Diagnostics.insert(Diagnostics.begin(), NewPreamble->Diags.begin(),
                          NewPreamble->Diags.end());
     }
 
     // Compute updated AST.
     llvm::Optional<ParsedAST> NewAST;
     {
       trace::Span Tracer("Build");
       SPAN_ATTACH(Tracer, "File", That->FileName);
-      NewAST = ParsedAST::Build(
-          std::move(CI), PreambleForAST, SerializedPreambleDecls,
+      NewAST =
+          ParsedAST::Build(std::move(CI), std::move(NewPreamble),
                            std::move(ContentsBuffer), PCHs, VFS, That->Logger);
     }
 
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());
     } else {
       // Don't report even Preamble diagnostics if we coulnd't build AST.
       Diagnostics.clear();