[CrossTU][NFCI] Refactor loadExternalAST function

Summary:
Refactor loadExternalAST method of CrossTranslationUnitContext in order to
reduce maintenance burden and so that features are easier to add in the future.

Reviewers: martong

Subscribers: rnkovacs, dkrupp, Szelethus, cfe-commits

Tags: #clang

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

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@367829 91177308-0d34-0410-b5e6-96231b3b80d8

include/clang/CrossTU/CrossTranslationUnit.h

@@ -192,11 +192,11 @@ class CrossTranslationUnitContext {
   template <typename T>
   llvm::Expected<const T *> importDefinitionImpl(const T *D, ASTUnit *Unit);
 
-  llvm::StringMap<std::unique_ptr<clang::ASTUnit>> FileASTUnitMap;
-  llvm::StringMap<clang::ASTUnit *> NameASTUnitMap;
-  llvm::StringMap<std::string> NameFileMap;
-  llvm::DenseMap<TranslationUnitDecl *, std::unique_ptr<ASTImporter>>
-      ASTUnitImporterMap;
+  using ImporterMapTy =
+      llvm::DenseMap<TranslationUnitDecl *, std::unique_ptr<ASTImporter>>;
+
+  ImporterMapTy ASTUnitImporterMap;
+
   CompilerInstance &CI;
   ASTContext &Context;
   std::shared_ptr<ASTImporterSharedState> ImporterSharedSt;
@@ -209,10 +209,105 @@ class CrossTranslationUnitContext {
   /// imported the FileID.
   ImportedFileIDMap ImportedFileIDs;
 
+  /// Functor for loading ASTUnits from AST-dump files.
+  class ASTFileLoader {
+  public:
+    ASTFileLoader(const CompilerInstance &CI);
+    std::unique_ptr<ASTUnit> operator()(StringRef ASTFilePath);
+
+  private:
+    const CompilerInstance &CI;
+  };
+
+  /// Storage for ASTUnits, cached access, and providing searchability are the
+  /// concerns of ASTUnitStorage class.
+  class ASTUnitStorage {
+  public:
+    ASTUnitStorage(const CompilerInstance &CI);
+    /// Loads an ASTUnit for a function.
+    ///
+    /// \param FuncitionName USR name of the function.
+    /// \param CrossTUDir Path to the directory used to store CTU related files.
+    /// \param IndexName Name of the file inside \p CrossTUDir which maps
+    /// function USR names to file paths. These files contain the corresponding
+    /// AST-dumps.
+    ///
+    /// \return An Expected instance which contains the ASTUnit pointer or the
+    /// error occured during the load.
+    llvm::Expected<ASTUnit *> getASTUnitForFunction(StringRef FunctionName,
+                                                    StringRef CrossTUDir,
+                                                    StringRef IndexName);
+    /// Identifies the path of the file which can be used to load the ASTUnit
+    /// for a given function.
+    ///
+    /// \param FuncitionName USR name of the function.
+    /// \param CrossTUDir Path to the directory used to store CTU related files.
+    /// \param IndexName Name of the file inside \p CrossTUDir which maps
+    /// function USR names to file paths. These files contain the corresponding
+    /// AST-dumps.
+    ///
+    /// \return An Expected instance containing the filepath.
+    llvm::Expected<std::string> getFileForFunction(StringRef FunctionName,
+                                                   StringRef CrossTUDir,
+                                                   StringRef IndexName);
+
+  private:
+    llvm::Error ensureCTUIndexLoaded(StringRef CrossTUDir, StringRef IndexName);
+    llvm::Expected<ASTUnit *> getASTUnitForFile(StringRef FileName);
+
+    template <typename... T> using BaseMapTy = llvm::StringMap<T...>;
+    using OwningMapTy = BaseMapTy<std::unique_ptr<clang::ASTUnit>>;
+    using NonOwningMapTy = BaseMapTy<clang::ASTUnit *>;
+
+    OwningMapTy FileASTUnitMap;
+    NonOwningMapTy NameASTUnitMap;
+
+    using IndexMapTy = BaseMapTy<std::string>;
+    IndexMapTy NameFileMap;
+
+    ASTFileLoader FileAccessor;
+  };
+
+  ASTUnitStorage ASTStorage;
+
   /// \p CTULoadTreshold should serve as an upper limit to the number of TUs
   /// imported in order to reduce the memory footprint of CTU analysis.
   const unsigned CTULoadThreshold;
-  unsigned NumASTLoaded{0u};
+
+  /// The number successfully loaded ASTs. Used to indicate, and  - with the
+  /// appropriate threshold value - limit the  memory usage of the
+  /// CrossTranslationUnitContext.
+  unsigned NumASTLoaded;
+
+  /// RAII counter to signal 'threshold reached' condition, and to increment the
+  /// NumASTLoaded counter upon a successful load.
+  class LoadGuard {
+  public:
+    LoadGuard(unsigned Limit, unsigned &Counter)
+        : Counter(Counter), Enabled(Counter < Limit){};
+    ~LoadGuard() {
+      if (StoreSuccess)
+        ++Counter;
+    }
+    /// Flag the LoadGuard instance as successful, meaning that the load
+    /// operation succeeded, and the memory footprint of the AST storage
+    /// actually increased. In this case, \p Counter should be incremented upon
+    /// destruction.
+    void storedSuccessfully() { StoreSuccess = true; }
+    /// Indicates, whether a new load operation is permitted, it is within the
+    /// threshold.
+    operator bool() const { return Enabled; };
+
+  private:
+    /// The number of ASTs actually imported. LoadGuard does not own the
+    /// counter, just uses on given to it at construction time.
+    unsigned &Counter;
+    /// Indicates whether a load operation can begin, which is equivalent to the
+    /// 'threshold not reached' condition.
+    bool Enabled;
+    /// Shows the state of the current load operation.
+    bool StoreSuccess;
+  };
 };
 
 } // namespace cross_tu

lib/CrossTU/CrossTranslationUnit.cpp

@@ -188,7 +188,7 @@ template <typename T> static bool hasBodyOrInit(const T *D) {
 }
 
 CrossTranslationUnitContext::CrossTranslationUnitContext(CompilerInstance &CI)
-    : CI(CI), Context(CI.getASTContext()),
+    : CI(CI), Context(CI.getASTContext()), ASTStorage(CI),
       CTULoadThreshold(CI.getAnalyzerOpts()->CTUImportThreshold) {}
 
 CrossTranslationUnitContext::~CrossTranslationUnitContext() {}
@@ -237,8 +237,8 @@ llvm::Expected<const T *> CrossTranslationUnitContext::getCrossTUDefinitionImpl(
   if (LookupName.empty())
     return llvm::make_error<IndexError>(
         index_error_code::failed_to_generate_usr);
-  llvm::Expected<ASTUnit *> ASTUnitOrError = loadExternalAST(
-      LookupName, CrossTUDir, IndexName, DisplayCTUProgress);
+  llvm::Expected<ASTUnit *> ASTUnitOrError =
+      loadExternalAST(LookupName, CrossTUDir, IndexName, DisplayCTUProgress);
   if (!ASTUnitOrError)
     return ASTUnitOrError.takeError();
   ASTUnit *Unit = *ASTUnitOrError;
@@ -340,6 +340,118 @@ void CrossTranslationUnitContext::emitCrossTUDiagnostics(const IndexError &IE) {
   }
 }
 
+CrossTranslationUnitContext::ASTFileLoader::ASTFileLoader(
+    const CompilerInstance &CI)
+    : CI(CI) {}
+
+std::unique_ptr<ASTUnit>
+CrossTranslationUnitContext::ASTFileLoader::operator()(StringRef ASTFilePath) {
+  // Load AST from ast-dump.
+  IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts = new DiagnosticOptions();
+  TextDiagnosticPrinter *DiagClient =
+      new TextDiagnosticPrinter(llvm::errs(), &*DiagOpts);
+  IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs());
+  IntrusiveRefCntPtr<DiagnosticsEngine> Diags(
+      new DiagnosticsEngine(DiagID, &*DiagOpts, DiagClient));
+
+  return ASTUnit::LoadFromASTFile(
+      ASTFilePath, CI.getPCHContainerOperations()->getRawReader(),
+      ASTUnit::LoadEverything, Diags, CI.getFileSystemOpts());
+}
+
+CrossTranslationUnitContext::ASTUnitStorage::ASTUnitStorage(
+    const CompilerInstance &CI)
+    : FileAccessor(CI) {}
+
+llvm::Expected<ASTUnit *>
+CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFile(StringRef FileName) {
+  // Try the cache first.
+  auto ASTCacheEntry = FileASTUnitMap.find(FileName);
+  if (ASTCacheEntry == FileASTUnitMap.end()) {
+    // Load the ASTUnit from the pre-dumped AST file specified by ASTFileName.
+    std::unique_ptr<ASTUnit> LoadedUnit = FileAccessor(FileName);
+
+    // Need the raw pointer and the unique_ptr as well.
+    ASTUnit* Unit = LoadedUnit.get();
+
+    // Update the cache.
+    FileASTUnitMap[FileName] = std::move(LoadedUnit);
+    return Unit;
+
+  } else {
+    // Found in the cache.
+    return ASTCacheEntry->second.get();
+  }
+}
+
+llvm::Expected<ASTUnit *>
+CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFunction(
+    StringRef FunctionName, StringRef CrossTUDir, StringRef IndexName) {
+  // Try the cache first.
+  auto ASTCacheEntry = NameASTUnitMap.find(FunctionName);
+  if (ASTCacheEntry == NameASTUnitMap.end()) {
+    // Load the ASTUnit from the pre-dumped AST file specified by ASTFileName.
+
+    // Ensure that the Index is loaded, as we need to search in it.
+    if (llvm::Error IndexLoadError =
+            ensureCTUIndexLoaded(CrossTUDir, IndexName))
+      return std::move(IndexLoadError);
+
+    // Check if there is and entry in the index for the function.
+    if (!NameFileMap.count(FunctionName)) {
+      ++NumNotInOtherTU;
+      return llvm::make_error<IndexError>(index_error_code::missing_definition);
+    }
+
+    // Search in the index for the filename where the definition of FuncitonName
+    // resides.
+    if (llvm::Expected<ASTUnit *> FoundForFile =
+            getASTUnitForFile(NameFileMap[FunctionName])) {
+
+      // Update the cache.
+      NameASTUnitMap[FunctionName] = *FoundForFile;
+      return *FoundForFile;
+
+    } else {
+      return FoundForFile.takeError();
+    }
+  } else {
+    // Found in the cache.
+    return ASTCacheEntry->second;
+  }
+}
+
+llvm::Expected<std::string>
+CrossTranslationUnitContext::ASTUnitStorage::getFileForFunction(
+    StringRef FunctionName, StringRef CrossTUDir, StringRef IndexName) {
+  if (llvm::Error IndexLoadError = ensureCTUIndexLoaded(CrossTUDir, IndexName))
+    return std::move(IndexLoadError);
+  return NameFileMap[FunctionName];
+}
+
+llvm::Error CrossTranslationUnitContext::ASTUnitStorage::ensureCTUIndexLoaded(
+    StringRef CrossTUDir, StringRef IndexName) {
+  // Dont initialize if the map is filled.
+  if (!NameFileMap.empty())
+    return llvm::Error::success();
+
+  // Get the absolute path to the index file.
+  SmallString<256> IndexFile = CrossTUDir;
+  if (llvm::sys::path::is_absolute(IndexName))
+    IndexFile = IndexName;
+  else
+    llvm::sys::path::append(IndexFile, IndexName);
+
+  if (auto IndexMapping = parseCrossTUIndex(IndexFile, CrossTUDir)) {
+    // Initialize member map.
+    NameFileMap = *IndexMapping;
+    return llvm::Error::success();
+  } else {
+    // Error while parsing CrossTU index file.
+    return IndexMapping.takeError();
+  };
+}
+
 llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
     StringRef LookupName, StringRef CrossTUDir, StringRef IndexName,
     bool DisplayCTUProgress) {
@@ -348,64 +460,41 @@ llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
   //        translation units contains decls with the same lookup name an
   //        error will be returned.
 
-  if (NumASTLoaded >= CTULoadThreshold) {
+  // RAII incrementing counter is used to count successful loads.
+  LoadGuard LoadOperation(CTULoadThreshold, NumASTLoaded);
+
+  // If import threshold is reached, don't import anything.
+  if (!LoadOperation) {
     ++NumASTLoadThresholdReached;
     return llvm::make_error<IndexError>(
         index_error_code::load_threshold_reached);
   }
 
-  ASTUnit *Unit = nullptr;
-  auto NameUnitCacheEntry = NameASTUnitMap.find(LookupName);
-  if (NameUnitCacheEntry == NameASTUnitMap.end()) {
-    if (NameFileMap.empty()) {
-      SmallString<256> IndexFile = CrossTUDir;
-      if (llvm::sys::path::is_absolute(IndexName))
-        IndexFile = IndexName;
-      else
-        llvm::sys::path::append(IndexFile, IndexName);
-      llvm::Expected<llvm::StringMap<std::string>> IndexOrErr =
-          parseCrossTUIndex(IndexFile, CrossTUDir);
-      if (IndexOrErr)
-        NameFileMap = *IndexOrErr;
-      else
-        return IndexOrErr.takeError();
-    }
+  // Try to get the value from the heavily cached storage.
+  llvm::Expected<ASTUnit *> Unit =
+      ASTStorage.getASTUnitForFunction(LookupName, CrossTUDir, IndexName);
 
-    auto It = NameFileMap.find(LookupName);
-    if (It == NameFileMap.end()) {
-      ++NumNotInOtherTU;
-      return llvm::make_error<IndexError>(index_error_code::missing_definition);
-    }
-    StringRef ASTFileName = It->second;
-    auto ASTCacheEntry = FileASTUnitMap.find(ASTFileName);
-    if (ASTCacheEntry == FileASTUnitMap.end()) {
-      IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts = new DiagnosticOptions();
-      TextDiagnosticPrinter *DiagClient =
-          new TextDiagnosticPrinter(llvm::errs(), &*DiagOpts);
-      IntrusiveRefCntPtr<DiagnosticIDs> DiagID(new DiagnosticIDs());
-      IntrusiveRefCntPtr<DiagnosticsEngine> Diags(
-          new DiagnosticsEngine(DiagID, &*DiagOpts, DiagClient));
-
-      std::unique_ptr<ASTUnit> LoadedUnit(ASTUnit::LoadFromASTFile(
-          ASTFileName, CI.getPCHContainerOperations()->getRawReader(),
-          ASTUnit::LoadEverything, Diags, CI.getFileSystemOpts()));
-      Unit = LoadedUnit.get();
-      FileASTUnitMap[ASTFileName] = std::move(LoadedUnit);
-      ++NumASTLoaded;
-      if (DisplayCTUProgress) {
-        llvm::errs() << "CTU loaded AST file: "
-                     << ASTFileName << "\n";
-      }
-    } else {
-      Unit = ASTCacheEntry->second.get();
-    }
-    NameASTUnitMap[LookupName] = Unit;
-  } else {
-    Unit = NameUnitCacheEntry->second;
-  }
   if (!Unit)
+    return Unit.takeError();
+
+  // Check whether the backing pointer of the Expected is a nullptr.
+  if (!*Unit)
     return llvm::make_error<IndexError>(
         index_error_code::failed_to_get_external_ast);
+
+  // The backing pointer is not null, loading was successful. If anything goes
+  // wrong from this point on, the AST is already stored, so the load part is
+  // finished.
+  LoadOperation.storedSuccessfully();
+
+  if (DisplayCTUProgress) {
+    if (llvm::Expected<std::string> FileName =
+            ASTStorage.getFileForFunction(LookupName, CrossTUDir, IndexName))
+      llvm::errs() << "CTU loaded AST file: " << *FileName << "\n";
+    else
+      return FileName.takeError();
+  }
+
   return Unit;
 }