3 files changed, 66 insertions, 5 deletions

diff --git a/clang/include/clang/AST/ASTImporter.h b/clang/include/clang/AST/ASTImporter.h
index 4a55c120a45..38e9b03aea5 100644
--- a/clang/include/clang/AST/ASTImporter.h
+++ b/clang/include/clang/AST/ASTImporter.h
@@ -34,6 +34,7 @@ namespace clang {
 
 class ASTContext;
 class ASTImporterSharedState;
+class ASTUnit;
 class Attr;
 class CXXBaseSpecifier;
 class CXXCtorInitializer;
@@ -229,6 +230,11 @@ class TypeSourceInfo;
     /// The file managers we're importing to and from.
     FileManager &ToFileManager, &FromFileManager;
 
+    /// The ASTUnit for the From context, if any.
+    /// This is used to create a mapping of imported ('To') FileID's to the
+    /// original ('From') FileID and ASTUnit.
+    ASTUnit *FromUnit;
+
     /// Whether to perform a minimal import.
     bool Minimal;
 
@@ -277,6 +283,11 @@ class TypeSourceInfo;
 
     void AddToLookupTable(Decl *ToD);
 
+    ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
+                ASTContext &FromContext, FileManager &FromFileManager,
+                ASTUnit *FromUnit, bool MinimalImport,
+                std::shared_ptr<ASTImporterSharedState> SharedState);
+
   protected:
     /// Can be overwritten by subclasses to implement their own import logic.
     /// The overwritten method should call this method if it didn't import the
@@ -287,7 +298,6 @@ class TypeSourceInfo;
     virtual bool returnWithErrorInTest() { return false; };
 
   public:
-
     /// \param ToContext The context we'll be importing into.
     ///
     /// \param ToFileManager The file manager we'll be importing into.
@@ -307,6 +317,23 @@ class TypeSourceInfo;
                 ASTContext &FromContext, FileManager &FromFileManager,
                 bool MinimalImport,
                 std::shared_ptr<ASTImporterSharedState> SharedState = nullptr);
+    /// \param ToContext The context we'll be importing into.
+    ///
+    /// \param ToFileManager The file manager we'll be importing into.
+    ///
+    /// \param FromUnit Pointer to an ASTUnit that contains the context and
+    /// file manager to import from.
+    ///
+    /// \param MinimalImport If true, the importer will attempt to import
+    /// as little as it can, e.g., by importing declarations as forward
+    /// declarations that can be completed at a later point.
+    ///
+    /// \param SharedState The importer specific lookup table which may be
+    /// shared amongst several ASTImporter objects.
+    /// If not set then the original C/C++ lookup is used.
+    ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
+                ASTUnit &FromUnit, bool MinimalImport,
+                std::shared_ptr<ASTImporterSharedState> SharedState = nullptr);
 
     virtual ~ASTImporter();
 
diff --git a/clang/include/clang/AST/ASTImporterSharedState.h b/clang/include/clang/AST/ASTImporterSharedState.h
index 3635a62deef..73b3282497c 100644
--- a/clang/include/clang/AST/ASTImporterSharedState.h
+++ b/clang/include/clang/AST/ASTImporterSharedState.h
@@ -17,18 +17,24 @@
 
 #include "clang/AST/ASTImporterLookupTable.h"
 #include "clang/AST/Decl.h"
+#include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/DenseMap.h"
 // FIXME We need this because of ImportError.
 #include "clang/AST/ASTImporter.h"
 
 namespace clang {
 
+class ASTUnit;
 class TranslationUnitDecl;
 
 /// Importer specific state, which may be shared amongst several ASTImporter
 /// objects.
 class ASTImporterSharedState {
+public:
+  using ImportedFileIDMap =
+      llvm::DenseMap<FileID, std::pair<FileID, ASTUnit *>>;
 
+private:
   /// Pointer to the import specific lookup table.
   std::unique_ptr<ASTImporterLookupTable> LookupTable;
 
@@ -43,6 +49,16 @@ class ASTImporterSharedState {
   // FIXME put ImportedFromDecls here!
   // And from that point we can better encapsulate the lookup table.
 
+  /// Map of imported FileID's (in "To" context) to FileID in "From" context
+  /// and the ASTUnit that contains the preprocessor and source manager for the
+  /// "From" FileID. This map is used to lookup a FileID and its SourceManager
+  /// when knowing only the FileID in the 'To' context. The FileID could be
+  /// imported by any of multiple ASTImporter objects. The map is used because
+  /// we do not want to loop over all ASTImporter's to find the one that
+  /// imported the FileID. (The ASTUnit is usable to get the SourceManager and
+  /// additional data.)
+  ImportedFileIDMap ImportedFileIDs;
+
 public:
   ASTImporterSharedState() = default;
 
@@ -75,6 +91,12 @@ public:
   void setImportDeclError(Decl *To, ImportError Error) {
     ImportErrors[To] = Error;
   }
+
+  ImportedFileIDMap &getImportedFileIDs() { return ImportedFileIDs; }
+
+  const ImportedFileIDMap &getImportedFileIDs() const {
+    return ImportedFileIDs;
+  }
 };
 
 } // namespace clang
diff --git a/clang/include/clang/CrossTU/CrossTranslationUnit.h b/clang/include/clang/CrossTU/CrossTranslationUnit.h
index d64329cdff3..7b689bdf330 100644
--- a/clang/include/clang/CrossTU/CrossTranslationUnit.h
+++ b/clang/include/clang/CrossTU/CrossTranslationUnit.h
@@ -153,8 +153,10 @@ 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);
@@ -162,9 +164,19 @@ public:
   /// 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.
+  /// 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:
   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,7 +186,7 @@ private:
   const T *findDefInDeclContext(const DeclContext *DC,
                                 StringRef LookupName);
   template <typename T>
-  llvm::Expected<const T *> importDefinitionImpl(const T *D);
+  llvm::Expected<const T *> importDefinitionImpl(const T *D, ASTUnit *Unit);
 
   llvm::StringMap<std::unique_ptr<clang::ASTUnit>> FileASTUnitMap;
   llvm::StringMap<clang::ASTUnit *> NameASTUnitMap;