//===--- ClangdUnitStore.h - A ClangdUnits container -------------*-C++-*-===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===---------------------------------------------------------------------===//

#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H
#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_CLANGDUNITSTORE_H

#include <mutex>

#include "ClangdUnit.h"
#include "GlobalCompilationDatabase.h"
#include "Path.h"
#include "clang/Tooling/CompilationDatabase.h"

namespace clang {
namespace clangd {

/// Thread-safe collection of ASTs built for specific files. Provides
/// synchronized access to ASTs.
class ClangdUnitStore {
public:
  /// Run specified \p Action on the ClangdUnit for \p File.
  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
  /// created from the \p FileContents. If the file is already present in the
  /// store, ClangdUnit::reparse will be called with the new contents before
  /// running \p Action.
  template <class Func>
  void runOnUnit(PathRef File, StringRef FileContents, StringRef ResourceDir,
                 GlobalCompilationDatabase &CDB,
                 std::shared_ptr<PCHContainerOperations> PCHs,
                 IntrusiveRefCntPtr<vfs::FileSystem> VFS, Func Action) {
    runOnUnitImpl(File, FileContents, ResourceDir, CDB, PCHs,
                  /*ReparseBeforeAction=*/true, VFS,
                  std::forward<Func>(Action));
  }

  /// Run specified \p Action on the ClangdUnit for \p File.
  /// If the file is not present in ClangdUnitStore, a new ClangdUnit will be
  /// created from the \p FileContents. If the file is already present in the
  /// store, the \p Action will be run directly on it.
  template <class Func>
  void runOnUnitWithoutReparse(PathRef File, StringRef FileContents,
                               StringRef ResourceDir,
                               GlobalCompilationDatabase &CDB,
                               std::shared_ptr<PCHContainerOperations> PCHs,
                               IntrusiveRefCntPtr<vfs::FileSystem> VFS,
                               Func Action) {
    runOnUnitImpl(File, FileContents, ResourceDir, CDB, PCHs,
                  /*ReparseBeforeAction=*/false, VFS,
                  std::forward<Func>(Action));
  }

  /// Run the specified \p Action on the ClangdUnit for \p File.
  /// Unit for \p File should exist in the store.
  template <class Func> void runOnExistingUnit(PathRef File, Func Action) {
    std::lock_guard<std::mutex> Lock(Mutex);

    auto It = OpenedFiles.find(File);
    assert(It != OpenedFiles.end() && "File is not in OpenedFiles");

    Action(It->second);
  }

  /// Remove ClangdUnit for \p File, if any
  void removeUnitIfPresent(PathRef File);

private:
  /// Run specified \p Action on the ClangdUnit for \p File.
  template <class Func>
  void runOnUnitImpl(PathRef File, StringRef FileContents,
                     StringRef ResourceDir, GlobalCompilationDatabase &CDB,
                     std::shared_ptr<PCHContainerOperations> PCHs,
                     bool ReparseBeforeAction,
                     IntrusiveRefCntPtr<vfs::FileSystem> VFS, Func Action) {
    std::lock_guard<std::mutex> Lock(Mutex);

    auto Commands = getCompileCommands(CDB, File);
    assert(!Commands.empty() &&
           "getCompileCommands should add default command");
    VFS->setCurrentWorkingDirectory(Commands.front().Directory);

    auto It = OpenedFiles.find(File);
    if (It == OpenedFiles.end()) {
      It = OpenedFiles
               .insert(std::make_pair(File, ClangdUnit(File, FileContents,
                                                       ResourceDir, PCHs,
                                                       Commands, VFS)))
               .first;
    } else if (ReparseBeforeAction) {
      It->second.reparse(FileContents, VFS);
    }
    return Action(It->second);
  }

  std::vector<tooling::CompileCommand>
  getCompileCommands(GlobalCompilationDatabase &CDB, PathRef File);

  std::mutex Mutex;
  llvm::StringMap<ClangdUnit> OpenedFiles;
};
} // namespace clangd
} // namespace clang

#endif