[CrossTU] Fix problem with CrossTU AST load limit and progress messages.

[balazske]

Summary

Number of loaded ASTs is to be incremented only if the AST was really loaded
but not if it was returned from cache. At the same place the message about
a loaded AST is displayed.

[martong]

LGTM!
Thank you.

[balazske]

• [CrossTU] Fix problem with CrossTU AST load limit and progress messages.
• Reformat.

clang/include/clang/CrossTU/CrossTranslationUnit.h

private:
   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.
+  /// Maintain number of AST loads and check for reaching the load limit.
+  class ASTLoadGuard {
+  public:
+    ASTLoadGuard(unsigned Limit) : Limit(Limit) {}
+
+    /// Indicates, whether a new load operation is permitted, it is within the
+    /// threshold.
+    operator bool() const { return Count < Limit; }
+
+    /// Tell that a new AST was loaded successfully.
+    void indicateLoadSuccess() { ++Count; }
+
+  private:
+    /// The number of ASTs actually imported.
+    unsigned Count{0u};
+    /// The limit (threshold) value for number of loaded ASTs.
+    const unsigned Limit;
+  };
+
+  /// Storage and load of 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.
+    /// \param DisplayCTUProgress Display a message about loading new ASTs.
     ///
     /// \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);
+                                                    StringRef IndexName,
+                                                    bool DisplayCTUProgress);
     /// 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);
+    llvm::Expected<ASTUnit *> getASTUnitForFile(StringRef FileName,
+                                                bool DisplayCTUProgress);
 
     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;
 
-  /// The number successfully loaded ASTs. Used to indicate, and  - with the
-  /// appropriate threshold value - limit the  memory usage of the
+    /// Limit the number of loaded ASTs. Used to limit the  memory usage of the
     /// CrossTranslationUnitContext.
-  unsigned NumASTLoaded{0u};
+    /// The ASTUnitStorage has the knowledge about if the AST to load is
+    /// actually loaded or returned from cache. This information is needed to
+    /// maintain the counter.
+    ASTLoadGuard LoadGuard;
+  };
 
-  /// 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), StoreSuccess(false) {}
-    ~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; }
+  ASTUnitStorage ASTStorage;
 
-  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
 } // namespace clang
 
 #endif // LLVM_CLANG_CROSSTU_CROSSTRANSLATIONUNIT_H

clang/lib/CrossTU/CrossTranslationUnit.cpp

static bool hasBodyOrInit(const VarDecl *D, const VarDecl *&DefD) {
   return D->getAnyInitializer(DefD);
 }
 template <typename T> static bool hasBodyOrInit(const T *D) {
   const T *Unused;
   return hasBodyOrInit(D, Unused);
 }
 
 CrossTranslationUnitContext::CrossTranslationUnitContext(CompilerInstance &CI)
-    : Context(CI.getASTContext()), ASTStorage(CI),
-      CTULoadThreshold(CI.getAnalyzerOpts()->CTUImportThreshold) {}
+    : Context(CI.getASTContext()), ASTStorage(CI) {}
 
 CrossTranslationUnitContext::~CrossTranslationUnitContext() {}
 
 llvm::Optional<std::string>
 CrossTranslationUnitContext::getLookupName(const NamedDecl *ND) {
   SmallString<128> DeclUSR;
   bool Ret = index::generateUSRForDecl(ND, DeclUSR);
   if (Ret)
CrossTranslationUnitContext::ASTFileLoader::operator()(StringRef ASTFilePath) {
 
   return ASTUnit::LoadFromASTFile(
       ASTFilePath, CI.getPCHContainerOperations()->getRawReader(),
       ASTUnit::LoadEverything, Diags, CI.getFileSystemOpts());
 }
 
 CrossTranslationUnitContext::ASTUnitStorage::ASTUnitStorage(
     const CompilerInstance &CI)
-    : FileAccessor(CI) {}
+    : FileAccessor(CI), LoadGuard(const_cast<CompilerInstance &>(CI)
+                                      .getAnalyzerOpts()
+                                      ->CTUImportThreshold) {}
 
 llvm::Expected<ASTUnit *>
-CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFile(StringRef FileName) {
+CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFile(
+    StringRef FileName, bool DisplayCTUProgress) {
   // Try the cache first.
   auto ASTCacheEntry = FileASTUnitMap.find(FileName);
   if (ASTCacheEntry == FileASTUnitMap.end()) {
+
+    // Do not load if the limit is reached.
+    if (!LoadGuard) {
+      ++NumASTLoadThresholdReached;
+      return llvm::make_error<IndexError>(
+          index_error_code::load_threshold_reached);
+    }
+
     // 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();
+    ASTUnit *Unit = LoadedUnit.get();
 
     // Update the cache.
     FileASTUnitMap[FileName] = std::move(LoadedUnit);
+
+    LoadGuard.indicateLoadSuccess();
+
+    if (DisplayCTUProgress)
+      llvm::errs() << "CTU loaded AST file: " << FileName << "\n";
+
     return Unit;
 
   } else {
     // Found in the cache.
     return ASTCacheEntry->second.get();
   }
 }
 
 llvm::Expected<ASTUnit *>
 CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFunction(
-    StringRef FunctionName, StringRef CrossTUDir, StringRef IndexName) {
+    StringRef FunctionName, StringRef CrossTUDir, StringRef IndexName,
+    bool DisplayCTUProgress) {
   // 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])) {
+            getASTUnitForFile(NameFileMap[FunctionName], DisplayCTUProgress)) {
 
       // Update the cache.
       NameASTUnitMap[FunctionName] = *FoundForFile;
       return *FoundForFile;
 
     } else {
       return FoundForFile.takeError();
     }

 llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
     StringRef LookupName, StringRef CrossTUDir, StringRef IndexName,
     bool DisplayCTUProgress) {
   // FIXME: The current implementation only supports loading decls with
   //        a lookup name from a single translation unit. If multiple
   //        translation units contains decls with the same lookup name an
   //        error will be returned.
 
-  // 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);
-  }
-
   // Try to get the value from the heavily cached storage.
-  llvm::Expected<ASTUnit *> Unit =
-      ASTStorage.getASTUnitForFunction(LookupName, CrossTUDir, IndexName);
+  llvm::Expected<ASTUnit *> Unit = ASTStorage.getASTUnitForFunction(
+      LookupName, CrossTUDir, IndexName, DisplayCTUProgress);
 
   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;
 }
 
 template <typename T>
 llvm::Expected<const T *>
 CrossTranslationUnitContext::importDefinitionImpl(const T *D, ASTUnit *Unit) {
   assert(hasBodyOrInit(D) && "Decls to be imported should have body or init.");