8 files changed, 48 insertions, 0 deletions
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 3d8e17bfb2f..48d05959dae 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -118,5 +118,9 @@ void FileIndex::findOccurrences(
   log("findOccurrences is not implemented.");
 }
 
+size_t FileIndex::estimateMemoryUsage() const {
+  return Index.estimateMemoryUsage();
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index 6f24f9d2c23..58fad2f62e2 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -81,6 +81,9 @@ public:
   void findOccurrences(const OccurrencesRequest &Req,
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
+
+  size_t estimateMemoryUsage() const override;
+
 private:
   FileSymbols FSymbols;
   MemIndex Index;
diff --git a/clang-tools-extra/clangd/index/Index.h b/clang-tools-extra/clangd/index/Index.h
index eac4c5e7d85..95b081acec7 100644
--- a/clang-tools-extra/clangd/index/Index.h
+++ b/clang-tools-extra/clangd/index/Index.h
@@ -385,6 +385,12 @@ public:
   virtual void findOccurrences(
       const OccurrencesRequest &Req,
       llvm::function_ref<void(const SymbolOccurrence &)> Callback) const = 0;
+
+  /// Returns estimated size of index (in bytes).
+  // FIXME(kbobyrev): Currently, this only returns the size of index itself
+  // excluding the size of actual symbol slab index refers to. We should include
+  // both.
+  virtual size_t estimateMemoryUsage() const = 0;
 };
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/index/MemIndex.cpp b/clang-tools-extra/clangd/index/MemIndex.cpp
index db11244d8ac..19a64ad9525 100644
--- a/clang-tools-extra/clangd/index/MemIndex.cpp
+++ b/clang-tools-extra/clangd/index/MemIndex.cpp
@@ -26,6 +26,9 @@ void MemIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
     Index = std::move(TempIndex);
     Symbols = std::move(Syms); // Relase old symbols.
   }
+
+  vlog("Built MemIndex with estimated memory usage {0} bytes.",
+       estimateMemoryUsage());
 }
 
 std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Slab) {
@@ -98,5 +101,10 @@ getSymbolsFromSlab(SymbolSlab Slab) {
                                                       &Snap->Pointers);
 }
 
+size_t MemIndex::estimateMemoryUsage() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+  return Index.getMemorySize();
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/MemIndex.h b/clang-tools-extra/clangd/index/MemIndex.h
index 8b12f8c7921..feca87af5e4 100644
--- a/clang-tools-extra/clangd/index/MemIndex.h
+++ b/clang-tools-extra/clangd/index/MemIndex.h
@@ -39,7 +39,10 @@ public:
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 
+  size_t estimateMemoryUsage() const override;
+
 private:
+
   std::shared_ptr<std::vector<const Symbol *>> Symbols;
   // Index is a set of symbols that are deduplicated by symbol IDs.
   // FIXME: build smarter index structure.
diff --git a/clang-tools-extra/clangd/index/Merge.cpp b/clang-tools-extra/clangd/index/Merge.cpp
index 4834bfba69f..84928ca1a5e 100644
--- a/clang-tools-extra/clangd/index/Merge.cpp
+++ b/clang-tools-extra/clangd/index/Merge.cpp
@@ -84,6 +84,10 @@ class MergedIndex : public SymbolIndex {
     log("findOccurrences is not implemented.");
   }
 
+  size_t estimateMemoryUsage() const override {
+    return Dynamic->estimateMemoryUsage() + Static->estimateMemoryUsage();
+  }
+
 private:
   const SymbolIndex *Dynamic, *Static;
 };
diff --git a/clang-tools-extra/clangd/index/dex/DexIndex.cpp b/clang-tools-extra/clangd/index/dex/DexIndex.cpp
index 534f51e0c3d..9280cc8fdfc 100644
--- a/clang-tools-extra/clangd/index/dex/DexIndex.cpp
+++ b/clang-tools-extra/clangd/index/dex/DexIndex.cpp
@@ -67,6 +67,9 @@ void DexIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
     InvertedIndex = std::move(TempInvertedIndex);
     SymbolQuality = std::move(TempSymbolQuality);
   }
+
+  vlog("Built DexIndex with estimated memory usage {0} bytes.",
+       estimateMemoryUsage());
 }
 
 std::unique_ptr<SymbolIndex> DexIndex::build(SymbolSlab Slab) {
@@ -171,6 +174,20 @@ void DexIndex::findOccurrences(
   log("findOccurrences is not implemented.");
 }
 
+size_t DexIndex::estimateMemoryUsage() const {
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  size_t Bytes =
+      LookupTable.size() * sizeof(std::pair<SymbolID, const Symbol *>);
+  Bytes += SymbolQuality.size() * sizeof(std::pair<const Symbol *, float>);
+  Bytes += InvertedIndex.size() * sizeof(Token);
+
+  for (const auto &P : InvertedIndex) {
+    Bytes += P.second.size() * sizeof(DocID);
+  }
+  return Bytes;
+}
+
 } // namespace dex
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/index/dex/DexIndex.h b/clang-tools-extra/clangd/index/dex/DexIndex.h
index 4485150b132..d9b2cd14389 100644
--- a/clang-tools-extra/clangd/index/dex/DexIndex.h
+++ b/clang-tools-extra/clangd/index/dex/DexIndex.h
@@ -57,7 +57,10 @@ public:
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 
+  size_t estimateMemoryUsage() const override;
+
 private:
+
   mutable std::mutex Mutex;
 
   std::shared_ptr<std::vector<const Symbol *>> Symbols /*GUARDED_BY(Mutex)*/;