[CrossTU][NFCI] Refactor loadExternalAST function

[gamesh411]

Summary

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

[martong]

Thank you Endre, this patch is a great initiative.
However, I think we can do better encapsulation then just the reorganization of the functions:
We could encapsulate into a nested class NameASTUnitMap and the functions which operate on this (getCachedASTUnitForName,
loadFromASTFileCached).
We could do the same with NameFileMap, lazyInitCTUIndex, getASTFileNameForLookup.
And we could also encapsulate NumASTLoaded and its related function (checkThresholdReached). In this case perhaps we could use RAII to increase the counter.

clang/include/clang/CrossTU/CrossTranslationUnit.h

167

Could you please provide comments for these new functions? E.g. it would be useful to know what is the IndexName param.

168

LookupName is the USR?

185

Could you please add comments for these maps (caches) about what we use them for?

clang/lib/CrossTU/CrossTranslationUnit.cpp

364

I don't think we should use auto here, because it is not obvious what will be the return type. It is important to see that we deal with an Expected<...>.
Perhaps we should have a type alias for llvm::StringMap<std::string>.

448

This comment might be also placed in the .h file as a oneliner above the function's prototype.

[gamesh411]

Refactor functionality into local classes

[gamesh411]

Too much autoformat fixed

[martong]

I have a feeling that LoadPass and LoagGuard could be (should be) merged together, other than that we are getting close.

clang/include/clang/CrossTU/CrossTranslationUnit.h

197

Perhaps this comment is not necessary, because it is not obvious where does the mentioned "section" ends and this is a bit redundant with the comments for the classes below.

241

We may not need this comment.

245

Could you please add these descriptions above the member declarations?

257

This name CanBegin sounds strange to me. Perhaps Enabled?

263

I have a feeling that LoadPass and LoagGuard could be (should be) merged together.

clang/lib/CrossTU/CrossTranslationUnit.cpp

352

I think you could leave the implementation of LoadPass:: member functions in the header, because they are so small. And that way I think code comprehension could improve.

[gamesh411]

Thanks for pointing out these issues. Most of them are agreed. Merging the RAII counter with the threshold checker class, however, does not seem like a good decision for me. What would be the benefits of merging the two?

clang/include/clang/CrossTU/CrossTranslationUnit.h

245

Good point! Will do that.

257

First I was thinking about making the variable names "read" as an english sentence. However CanBegin does not appear during use, so it is fair that Enable is a better name. Changing it...

263

LoadPass is the RAII object, the other is decider whether the threshold is reached. Note also that LoadGuard owns the loaded-counter variable. IMHO these should be two separate classes, and I think if they were to be implemented in the header these two would be simple, and easy-to-understand classes.

clang/lib/CrossTU/CrossTranslationUnit.cpp

352

Good point!

[gamesh411]

• Merge RAII class
• Update comments

[gamesh411]

Updated the revision.
I find this solution definitely more compact with responsibilities more separated. One more thing that comes to mind is that maybe the whole explicit passing of CTUDir and IndexName arguments is a bit verbose. What are your thoughts about these being injected earlier than the query operations (eg.: in getASTUnitForFunction and getFileForFunction)?

[martong]

clang/lib/CrossTU/CrossTranslationUnit.cpp

473

LoadOperation should not be successful if the condition here is true.
Perhaps LoadOperation.wasSuccessful(); should be done right before the last return stmt in the function.

[gamesh411]

Specify the exact meaning of successful storage

[martong]

LGTM! Thanks! Please do the commit.

[gamesh411]

• Remove unused member Limit
• Rebase to current master

[balazske]

I discovered that there are problems with this change. The number of loaded AST units is not incremented correctly and the "CTU loaded AST file" is displayed for every access to a AST unit (even if it was got from cache). The problem is that the counting and print of AST loads does not handle the case when the load returned a cached result (and no real loading happened). This causes "annoying" messages and the AST unit limit is not handled correctly.

clang/include/clang/CrossTU/CrossTranslationUnit.h

public:
 
   /// Get a name to identify a named decl.
   static std::string getLookupName(const NamedDecl *ND);
 
   /// Emit diagnostics for the user for potential configuration errors.
   void emitCrossTUDiagnostics(const IndexError &IE);
 
 private:
+  bool checkThresholdReached() const;
+  llvm::Error lazyInitCTUIndex(StringRef CrossTUDir, StringRef IndexName);
+  ASTUnit *getCachedASTUnitForName(StringRef LookupName) const;
+  std::unique_ptr<ASTUnit> loadFromASTFile(StringRef ASTFileName) const;
+  ASTUnit *loadFromASTFileCached(StringRef LookupName, StringRef ASTFileName);
+  llvm::Expected<std::string>
+  getASTFileNameForLookup(StringRef LookupName) const;
+
   void lazyInitImporterSharedSt(TranslationUnitDecl *ToTU);
   ASTImporter &getOrCreateASTImporter(ASTContext &From);
   template <typename T>
-  llvm::Expected<const T *> getCrossTUDefinitionImpl(const T *D,
-                                                     StringRef CrossTUDir,
-                                                     StringRef IndexName,
-                                                     bool DisplayCTUProgress);
+  llvm::Expected<const T *>
+  getCrossTUDefinitionImpl(const T *D, StringRef CrossTUDir,
+                           StringRef IndexName, bool DisplayCTUProgress);
   template <typename T>
-  const T *findDefInDeclContext(const DeclContext *DC,
-                                StringRef LookupName);
+  const T *findDefInDeclContext(const DeclContext *DC, StringRef LookupName);
   template <typename T>
   llvm::Expected<const T *> importDefinitionImpl(const T *D);
 
   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;
   CompilerInstance &CI;
   ASTContext &Context;
   std::shared_ptr<ASTImporterSharedState> ImporterSharedSt;
   /// \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};
 };
 
 } // namespace cross_tu
 } // namespace clang
 
 #endif // LLVM_CLANG_CROSSTU_CROSSTRANSLATIONUNIT_H

clang/lib/CrossTU/CrossTranslationUnit.cpp

llvm::Expected<const T *> CrossTranslationUnitContext::getCrossTUDefinitionImpl(
   assert(D && "D is missing, bad call to this function!");
   assert(!hasBodyOrInit(D) &&
          "D has a body or init in current translation unit!");
   ++NumGetCTUCalled;
   const std::string LookupName = getLookupName(D);
   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;
   assert(&Unit->getFileManager() ==
          &Unit->getASTContext().getSourceManager().getFileManager());
 
   const llvm::Triple &TripleTo = Context.getTargetInfo().getTriple();
   const llvm::Triple &TripleFrom =
case index_error_code::triple_mismatch:
     Context.getDiagnostics().Report(diag::warn_ctu_incompat_triple)
         << IE.getFileName() << IE.getTripleToName() << IE.getTripleFromName();
     break;
   default:
     break;
   }
 }
 
-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.
-
+bool CrossTranslationUnitContext::checkThresholdReached() const {
   if (NumASTLoaded >= CTULoadThreshold) {
     ++NumASTLoadThresholdReached;
-    return llvm::make_error<IndexError>(
-        index_error_code::load_threshold_reached);
-  }
+    return true;
+  }
+  return false;
 
-  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();
-    }
+llvm::Error CrossTranslationUnitContext::lazyInitCTUIndex(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.
+
+  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();
+  };
 
-    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));
+ASTUnit *CrossTranslationUnitContext::getCachedASTUnitForName(
+    StringRef LookupName) const {
+  auto CacheEntry = NameASTUnitMap.find(LookupName);
+  if (CacheEntry != NameASTUnitMap.end())
+    return CacheEntry->second;
+  else
+    return nullptr;
+
+std::unique_ptr<ASTUnit>
+CrossTranslationUnitContext::loadFromASTFile(StringRef ASTFileName) const {
+  // Load AST from ast-dump.
 
-      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 ASTUnit::LoadFromASTFile(
+      ASTUnit::LoadEverything, Diags, CI.getFileSystemOpts());
+}
+
+llvm::Expected<std::string>
+CrossTranslationUnitContext::getASTFileNameForLookup(
+    StringRef LookupName) const {
+  auto NameFileMapEntry = NameFileMap.find(LookupName);
+  if (NameFileMapEntry != NameFileMap.end()) {
+    return NameFileMapEntry->second;
+  } else {
+    ++NumNotInOtherTU;
+    return llvm::make_error<IndexError>(index_error_code::missing_definition);
+  }
+}
+
+ASTUnit *
+CrossTranslationUnitContext::loadFromASTFileCached(StringRef LookupName,
+                                                   StringRef ASTFileName) {
+  ASTUnit *Unit = nullptr;
+
+  auto ASTCacheEntry = FileASTUnitMap.find(ASTFileName);
+  if (ASTCacheEntry == FileASTUnitMap.end()) {
+    // Load the ASTUnit from the pre-dumped AST file specified by ASTFileName.
+    std::unique_ptr<ASTUnit> LoadedUnit = loadFromASTFile(ASTFileName);
+
+  return Unit;
+
+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.
+
+  // If import threshold is reached, don't import anything.
+  if (checkThresholdReached())
     return llvm::make_error<IndexError>(
-        index_error_code::failed_to_get_external_ast);
+        index_error_code::load_threshold_reached);
+
+  // First try to access the cache to get the ASTUnit for the function name
+  // specified by LookupName.
+  ASTUnit *Unit = getCachedASTUnitForName(LookupName);
+  if (Unit)
+    return Unit;
+
+  // Lazily initialize the mapping from function names to AST files.
+  if (llvm::Error InitFailed = lazyInitCTUIndex(CrossTUDir, IndexName))
+    return std::move(InitFailed);
+
+  llvm::Expected<std::string> ASTFileName = getASTFileNameForLookup(LookupName);
+  if (!ASTFileName)
+    return ASTFileName.takeError();
+
+  // Try to load from ASTFile but use cache for both file and function names.
+  Unit = loadFromASTFileCached(LookupName, *ASTFileName);
+
+  if (Unit) {
+    ++NumASTLoaded;
+    if (DisplayCTUProgress) {
+      llvm::errs() << "CTU loaded AST file: " << *ASTFileName << "\n";
+    }
     return Unit;
+  } else {
+    return llvm::make_error<IndexError>(
+        index_error_code::failed_to_get_external_ast);
+  }
 }
 
 template <typename T>
 llvm::Expected<const T *>
 CrossTranslationUnitContext::importDefinitionImpl(const T *D) {
   assert(hasBodyOrInit(D) && "Decls to be imported should have body or init.");
 
   ASTImporter &Importer = getOrCreateASTImporter(D->getASTContext());
   auto ToDeclOrError = Importer.Import(D);
   if (!ToDeclOrError) {
     handleAllErrors(ToDeclOrError.takeError(),
                     [&](const ImportError &IE) {
                       switch (IE.Error) {