Revert "Address comments."

This reverts commit b43c4d1c731e07172a382567f3146b3c461c5b69.

llvm-svn: 346942
diff --git a/clang-tools-extra/clangd/index/Background.cpp b/clang-tools-extra/clangd/index/Background.cpp
index 94708d6..ef83abf 100644
--- a/clang-tools-extra/clangd/index/Background.cpp
+++ b/clang-tools-extra/clangd/index/Background.cpp
@@ -48,49 +48,54 @@
   return digest(Content);
 }
 
-std::string getShardPathFromFilePath(llvm::StringRef ShardRoot,
-                                     llvm::StringRef FilePath) {
-  llvm::SmallString<128> ShardRootSS(ShardRoot);
-  sys::path::append(ShardRootSS, sys::path::filename(FilePath) +
-                                     toHex(digest(FilePath)) + ".idx");
-  return ShardRoot.str();
+llvm::SmallString<128>
+getShardPathFromFilePath(llvm::SmallString<128> ShardRoot,
+                         llvm::StringRef FilePath) {
+  sys::path::append(ShardRoot, sys::path::filename(FilePath) +
+                                   toHex(digest(FilePath)) + ".idx");
+  return ShardRoot;
 }
 
 // Uses disk as a storage for index shards. Creates a directory called
-// ".clangd-index/" under the path provided during construction.
+// ".clangd-index/" under the path provided during initialize.
+// Note: Requires initialize to be called before storing or retrieval.
 // This class is thread-safe.
 class DiskBackedIndexStorage : public BackgroundIndexStorage {
-  std::string DiskShardRoot;
+  llvm::SmallString<128> DiskShardRoot;
 
 public:
-  llvm::Expected<IndexFileIn> loadShard(llvm::StringRef ShardIdentifier) const {
-    const std::string ShardPath =
-        getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
+  llvm::Expected<IndexFileIn>
+  retrieveShard(llvm::StringRef ShardIdentifier) const override {
+    auto ShardPath = getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
     auto Buffer = MemoryBuffer::getFile(ShardPath);
     if (!Buffer) {
-      elog("Couldn't load {0}: {1}", ShardPath, Buffer.getError().message());
+      elog("Couldn't retrieve {0}: {1}", ShardPath,
+           Buffer.getError().message());
       return llvm::make_error<llvm::StringError>(Buffer.getError());
     }
+    // FIXME: Change readIndexFile to also look at Hash of the source that
+    // generated index and skip if there is a mismatch.
     return readIndexFile(Buffer->get()->getBuffer());
   }
 
-  llvm::Error storeShard(llvm::StringRef ShardIdentifier,
-                         IndexFileOut Shard) const override {
+  bool storeShard(llvm::StringRef ShardIdentifier,
+                  IndexFileOut Shard) const override {
     auto ShardPath = getShardPathFromFilePath(DiskShardRoot, ShardIdentifier);
     std::error_code EC;
     llvm::raw_fd_ostream OS(ShardPath, EC);
-    if (EC)
-      return errorCodeToError(EC);
+    if (EC) {
+      elog("Failed to open {0} for writing: {1}", ShardPath, EC.message());
+      return false;
+    }
     OS << Shard;
-    return llvm::Error::success();
+    return true;
   }
 
-  // Sets DiskShardRoot to (Directory + ".clangd-index/") which is the base
-  // directory for all shard files.
-  DiskBackedIndexStorage(llvm::StringRef Directory) {
-    llvm::SmallString<128> CDBDirectory(Directory);
-    sys::path::append(CDBDirectory, ".clangd-index/");
-    DiskShardRoot = CDBDirectory.str();
+  // Initializes DiskShardRoot to (Directory + ".clangd-index/") which is the
+  // base directory for all shard files. After the initialization succeeds all
+  // subsequent calls are no-op.
+  DiskBackedIndexStorage(llvm::StringRef Directory) : DiskShardRoot(Directory) {
+    sys::path::append(DiskShardRoot, ".clangd-index/");
     if (!llvm::sys::fs::exists(DiskShardRoot)) {
       std::error_code OK;
       std::error_code EC = llvm::sys::fs::create_directory(DiskShardRoot);
@@ -101,7 +106,7 @@
   }
 };
 
-std::string getAbsoluteFilePath(const tooling::CompileCommand &Cmd) {
+SmallString<128> getAbsolutePath(const tooling::CompileCommand &Cmd) {
   SmallString<128> AbsolutePath;
   if (sys::path::is_absolute(Cmd.Filename)) {
     AbsolutePath = Cmd.Filename;
@@ -109,7 +114,7 @@
     AbsolutePath = Cmd.Directory;
     sys::path::append(AbsolutePath, Cmd.Filename);
   }
-  return AbsolutePath.str();
+  return AbsolutePath;
 }
 
 } // namespace
@@ -118,12 +123,10 @@
                                  StringRef ResourceDir,
                                  const FileSystemProvider &FSProvider,
                                  ArrayRef<std::string> URISchemes,
-                                 IndexStorageFactory IndexStorageCreator,
                                  size_t ThreadPoolSize)
     : SwapIndex(make_unique<MemIndex>()), ResourceDir(ResourceDir),
       FSProvider(FSProvider), BackgroundContext(std::move(BackgroundContext)),
-      URISchemes(URISchemes),
-      IndexStorageCreator(std::move(IndexStorageCreator)) {
+      URISchemes(URISchemes) {
   assert(ThreadPoolSize > 0 && "Thread pool size can't be zero.");
   while (ThreadPoolSize--) {
     ThreadPool.emplace_back([this] { run(); });
@@ -182,24 +185,57 @@
 
 void BackgroundIndex::enqueue(StringRef Directory,
                               tooling::CompileCommand Cmd) {
-  BackgroundIndexStorage *IndexStorage = getIndexStorage(Directory);
-  if (!IndexStorage)
+  auto IndexStorage = BackgroundIndexStorage::getForDirectory(Directory);
+  if (IndexStorage)
+    loadShard(IndexStorage.get(), Cmd);
+  else
     elog("No index storage for: {0}", Directory);
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
-    enqueueLocked(std::move(Cmd), IndexStorage);
+    enqueueLocked(std::move(Cmd), std::move(IndexStorage));
   }
   QueueCV.notify_all();
 }
 
+void BackgroundIndex::loadShard(BackgroundIndexStorage *IndexStorage,
+                                const tooling::CompileCommand &Cmd) {
+  assert(IndexStorage && "No index storage to load from?");
+  auto AbsolutePath = getAbsolutePath(Cmd);
+  auto Shard = IndexStorage->retrieveShard(AbsolutePath);
+  if (Shard) {
+    // FIXME: Updated hashes once we have them in serialized format.
+    // IndexedFileDigests[AbsolutePath] = Hash;
+    IndexedSymbols.update(AbsolutePath,
+                          make_unique<SymbolSlab>(std::move(*Shard->Symbols)),
+                          make_unique<RefSlab>(std::move(*Shard->Refs)));
+
+    vlog("Loaded {0} from storage", AbsolutePath);
+  }
+}
+
+void BackgroundIndex::loadShards(
+    BackgroundIndexStorage *IndexStorage,
+    const std::vector<tooling::CompileCommand> &Cmds) {
+  assert(IndexStorage && "No index storage to load from?");
+  for (const auto &Cmd : Cmds)
+    loadShard(IndexStorage, Cmd);
+  // FIXME: Maybe we should get rid of this one once we start building index
+  // periodically? Especially if we also offload this task onto the queue.
+  vlog("Rebuilding automatic index");
+  reset(IndexedSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge,
+                                  URISchemes));
+}
+
 void BackgroundIndex::enqueueAll(StringRef Directory,
                                  const tooling::CompilationDatabase &CDB) {
   trace::Span Tracer("BackgroundIndexEnqueueCDB");
   // FIXME: this function may be slow. Perhaps enqueue a task to re-read the CDB
   // from disk and enqueue the commands asynchronously?
   auto Cmds = CDB.getAllCompileCommands();
-  BackgroundIndexStorage *IndexStorage = getIndexStorage(Directory);
-  if (!IndexStorage)
+  auto IndexStorage = BackgroundIndexStorage::getForDirectory(Directory);
+  if (IndexStorage)
+    loadShards(IndexStorage.get(), Cmds);
+  else
     elog("No index storage for: {0}", Directory);
   SPAN_ATTACH(Tracer, "commands", int64_t(Cmds.size()));
   std::mt19937 Generator(std::random_device{}());
@@ -213,16 +249,18 @@
   QueueCV.notify_all();
 }
 
-void BackgroundIndex::enqueueLocked(tooling::CompileCommand Cmd,
-                                    BackgroundIndexStorage *IndexStorage) {
+void BackgroundIndex::enqueueLocked(
+    tooling::CompileCommand Cmd,
+    std::shared_ptr<BackgroundIndexStorage> IndexStorage) {
   Queue.push_back(Bind(
-      [this, IndexStorage](tooling::CompileCommand Cmd) {
+      [this](tooling::CompileCommand Cmd,
+             std::shared_ptr<BackgroundIndexStorage> IndexStorage) {
         std::string Filename = Cmd.Filename;
         Cmd.CommandLine.push_back("-resource-dir=" + ResourceDir);
-        if (auto Error = index(std::move(Cmd), IndexStorage))
+        if (auto Error = index(std::move(Cmd), IndexStorage.get()))
           log("Indexing {0} failed: {1}", Filename, std::move(Error));
       },
-      std::move(Cmd)));
+      std::move(Cmd), std::move(IndexStorage)));
 }
 
 // Resolves URI to file paths with cache.
@@ -318,8 +356,7 @@
       IndexFileOut Shard;
       Shard.Symbols = SS.get();
       Shard.Refs = RS.get();
-      if (auto Error = IndexStorage->storeShard(Path, Shard))
-        elog("Failed to store shard for {0}: {1}", Path, std::move(Error));
+      IndexStorage->storeShard(Path, Shard);
     }
 
     std::lock_guard<std::mutex> Lock(DigestsMu);
@@ -367,7 +404,7 @@
                              BackgroundIndexStorage *IndexStorage) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
-  const std::string AbsolutePath = getAbsoluteFilePath(Cmd);
+  SmallString<128> AbsolutePath = getAbsolutePath(Cmd);
 
   auto FS = FSProvider.getFileSystem();
   auto Buf = FS->getBufferForFile(AbsolutePath);
@@ -449,18 +486,8 @@
   return Error::success();
 }
 
-BackgroundIndexStorage *
-BackgroundIndex::getIndexStorage(llvm::StringRef CDBDirectory) {
-  if (!IndexStorageCreator)
-    return nullptr;
-  auto IndexStorageIt = IndexStorageMap.find(CDBDirectory);
-  if (IndexStorageIt == IndexStorageMap.end())
-    IndexStorageIt =
-        IndexStorageMap
-            .insert({CDBDirectory, IndexStorageCreator(CDBDirectory)})
-            .first;
-  return IndexStorageIt->second.get();
-}
+std::function<std::shared_ptr<BackgroundIndexStorage>(llvm::StringRef)>
+    BackgroundIndexStorage::Factory = nullptr;
 
 } // namespace clangd
 } // namespace clang