1 files changed, 127 insertions, 16 deletions
diff --git a/include/clang/CrossTU/CrossTranslationUnit.h b/include/clang/CrossTU/CrossTranslationUnit.h
index d64329cdff3e6..4d2b7109c62a6 100644
--- a/include/clang/CrossTU/CrossTranslationUnit.h
+++ b/include/clang/CrossTU/CrossTranslationUnit.h
@@ -17,6 +17,7 @@
 #include "clang/AST/ASTImporterSharedState.h"
 #include "clang/Basic/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Error.h"
@@ -153,18 +154,34 @@ public:
   ///        was passed to the constructor.
   ///
   /// \return Returns the resulting definition or an error.
-  llvm::Expected<const FunctionDecl *> importDefinition(const FunctionDecl *FD);
-  llvm::Expected<const VarDecl *> importDefinition(const VarDecl *VD);
+  llvm::Expected<const FunctionDecl *> importDefinition(const FunctionDecl *FD,
+                                                        ASTUnit *Unit);
+  llvm::Expected<const VarDecl *> importDefinition(const VarDecl *VD,
+                                                   ASTUnit *Unit);
 
   /// Get a name to identify a named decl.
-  static std::string getLookupName(const NamedDecl *ND);
+  static llvm::Optional<std::string> getLookupName(const NamedDecl *ND);
 
   /// Emit diagnostics for the user for potential configuration errors.
   void emitCrossTUDiagnostics(const IndexError &IE);
 
+  /// Determine the original source location in the original TU for an
+  /// imported source location.
+  /// \p ToLoc Source location in the imported-to AST.
+  /// \return Source location in the imported-from AST and the corresponding
+  /// ASTUnit object (the AST was loaded from a file using an internal ASTUnit
+  /// object that is returned here).
+  /// If any error happens (ToLoc is a non-imported source location) empty is
+  /// returned.
+  llvm::Optional<std::pair<SourceLocation /*FromLoc*/, ASTUnit *>>
+  getImportedFromSourceLocation(const clang::SourceLocation &ToLoc) const;
+
 private:
+  using ImportedFileIDMap =
+      llvm::DenseMap<FileID, std::pair<FileID, ASTUnit *>>;
+
   void lazyInitImporterSharedSt(TranslationUnitDecl *ToTU);
-  ASTImporter &getOrCreateASTImporter(ASTContext &From);
+  ASTImporter &getOrCreateASTImporter(ASTUnit *Unit);
   template <typename T>
   llvm::Expected<const T *> getCrossTUDefinitionImpl(const T *D,
                                                      StringRef CrossTUDir,
@@ -174,20 +191,114 @@ private:
   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;
+  llvm::Expected<const T *> importDefinitionImpl(const T *D, ASTUnit *Unit);
+
+  using ImporterMapTy =
+      llvm::DenseMap<TranslationUnitDecl *, std::unique_ptr<ASTImporter>>;
+
+  ImporterMapTy ASTUnitImporterMap;
+
   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};
+  /// Map of imported FileID's (in "To" context) to FileID in "From" context
+  /// and the ASTUnit for the From context.
+  /// This map is used by getImportedFromSourceLocation to lookup a FileID and
+  /// its Preprocessor when knowing only the FileID in the 'To' context. The
+  /// FileID could be imported by any of multiple 'From' ASTImporter objects.
+  /// we do not want to loop over all ASTImporter's to find the one that
+  /// 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;
+  };
+
+  /// 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 FunctionName 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,
+                                                    bool DisplayCTUProgress);
+    /// Identifies the path of the file which can be used to load the ASTUnit
+    /// for a given function.
+    ///
+    /// \param FunctionName 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,
+                                                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;
+
+    /// 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;
+
 };
 
 } // namespace cross_tu