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

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.

Reviewers: martong, gamesh411

Reviewed By: martong

Subscribers: rnkovacs, dkrupp, Szelethus, gamesh411, cfe-commits

Tags: #clang

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

llvm-svn: 368545

clang/include/clang/CrossTU/CrossTranslationUnit.h

@@ -219,8 +219,27 @@ private:
     const CompilerInstance &CI;
   };
 
-  /// Storage for ASTUnits, cached access, and providing searchability are the
+  /// Maintain number of AST loads and check for reaching the load limit.
-  /// concerns of ASTUnitStorage class.
+  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);
@@ -231,12 +250,14 @@ private:
     /// \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.
     ///
@@ -253,7 +274,8 @@ private:
 
   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>>;
@@ -266,48 +288,17 @@ private:
     IndexMapTy NameFileMap;
 
     ASTFileLoader FileAccessor;
+
+    /// Limit the number of loaded ASTs. Used to limit the  memory usage of the
+    /// CrossTranslationUnitContext.
+    /// 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;
   };
 
   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
-  /// CrossTranslationUnitContext.
-  unsigned NumASTLoaded{0u};
-
-  /// 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; }
-
-  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

clang/lib/CrossTU/CrossTranslationUnit.cpp

@@ -188,8 +188,7 @@ template <typename T> static bool hasBodyOrInit(const T *D) {
 }
 
 CrossTranslationUnitContext::CrossTranslationUnitContext(CompilerInstance &CI)
-    : Context(CI.getASTContext()), ASTStorage(CI),
+    : Context(CI.getASTContext()), ASTStorage(CI) {}
-      CTULoadThreshold(CI.getAnalyzerOpts()->CTUImportThreshold) {}
 
 CrossTranslationUnitContext::~CrossTranslationUnitContext() {}
 
@@ -363,21 +362,38 @@ CrossTranslationUnitContext::ASTFileLoader::operator()(StringRef ASTFilePath) {
 
 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 {
@@ -388,7 +404,8 @@ CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFile(StringRef FileNam
 
 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()) {
@@ -408,7 +425,7 @@ CrossTranslationUnitContext::ASTUnitStorage::getASTUnitForFunction(
     // 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;
@@ -462,19 +479,9 @@ llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
   //        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 =
+  llvm::Expected<ASTUnit *> Unit = ASTStorage.getASTUnitForFunction(
-      ASTStorage.getASTUnitForFunction(LookupName, CrossTUDir, IndexName);
+      LookupName, CrossTUDir, IndexName, DisplayCTUProgress);
 
   if (!Unit)
     return Unit.takeError();
@@ -484,19 +491,6 @@ llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
     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;
 }