[clangd] Provide a helper to report estimated memory usage per-file

Reviewers: sammccall, ioeric, hokein

Reviewed By: ioeric

Subscribers: klimek, cfe-commits, jkorous-apple

Differential Revision: https://reviews.llvm.org/D42480

1 files changed, 27 insertions, 1 deletions

diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index 90dd221ec42..ed06607ef3c 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -35,6 +35,10 @@ using namespace clang;
 
 namespace {
 
+template <class T> std::size_t getUsedBytes(const std::vector<T> &Vec) {
+  return Vec.capacity() * sizeof(T);
+}
+
 class DeclTrackingASTConsumer : public ASTConsumer {
 public:
   DeclTrackingASTConsumer(std::vector<const Decl *> &TopLevelDecls)
@@ -332,6 +336,14 @@ const std::vector<DiagWithFixIts> &ParsedAST::getDiagnostics() const {
   return Diags;
 }
 
+std::size_t ParsedAST::getUsedBytes() const {
+  auto &AST = getASTContext();
+  // FIXME(ibiryukov): we do not account for the dynamically allocated part of
+  // SmallVector<FixIt> inside each Diag.
+  return AST.getASTAllocatedMemory() + AST.getSideTableAllocatedMemory() +
+         ::getUsedBytes(TopLevelDecls) + ::getUsedBytes(Diags);
+}
+
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<serialization::DeclID> TopLevelDeclIDs,
                            std::vector<DiagWithFixIts> Diags)
@@ -370,7 +382,8 @@ CppFile::CppFile(PathRef FileName, bool StorePreamblesInMemory,
                  std::shared_ptr<PCHContainerOperations> PCHs,
                  ASTParsedCallback ASTCallback)
     : FileName(FileName), StorePreamblesInMemory(StorePreamblesInMemory),
-      RebuildCounter(0), RebuildInProgress(false), PCHs(std::move(PCHs)),
+      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);
@@ -419,7 +432,9 @@ UniqueFunction<void()> CppFile::deferCancelRebuild() {
       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));
   };
 }
@@ -573,6 +588,8 @@ CppFile::deferRebuild(ParseInputs &&Inputs) {
       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
 
@@ -609,6 +626,7 @@ CppFile::deferRebuild(ParseInputs &&Inputs) {
         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
@@ -635,6 +653,14 @@ std::shared_future<std::shared_ptr<ParsedASTWrapper>> CppFile::getAST() const {
   return ASTFuture;
 }
 
+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;
+}
+
 CppFile::RebuildGuard::RebuildGuard(CppFile &File,
                                     unsigned RequestRebuildCounter)
     : File(File), RequestRebuildCounter(RequestRebuildCounter) {