[CrossTU] Add a function to retrieve original source location.

Summary:
A new function will be added to get the original SourceLocation
for a SourceLocation that was imported as result of getCrossTUDefinition.
The returned SourceLocation is in the context of the (original)
SourceManager for the original source file. Additionally the
ASTUnit object for that source file is returned. This is needed
to get a SourceManager to operate on with the returned source location.

The new function works if multiple different source files are loaded
with the same CrossTU context.

This patch can be treated as part of a bigger change that is needed to
improve macro expansion handliong at plist generation.

Reviewers: martong, shafik, a_sidorin, xazax.hun

Reviewed By: martong

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

Tags: #clang

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

llvm-svn: 366449

6 files changed, 162 insertions, 25 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;
diff --git a/clang/lib/AST/ASTImporter.cpp b/clang/lib/AST/ASTImporter.cpp
index 9d5dd84161d..8c917f67beb 100644
--- a/clang/lib/AST/ASTImporter.cpp
+++ b/clang/lib/AST/ASTImporter.cpp
@@ -12,9 +12,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTImporter.h"
-#include "clang/AST/ASTImporterSharedState.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTDiagnostic.h"
+#include "clang/AST/ASTImporterSharedState.h"
 #include "clang/AST/ASTStructuralEquivalence.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
@@ -52,13 +52,14 @@
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
+#include "clang/Frontend/ASTUnit.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
-#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -7716,9 +7717,22 @@ ASTImporter::ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
                          ASTContext &FromContext, FileManager &FromFileManager,
                          bool MinimalImport,
                          std::shared_ptr<ASTImporterSharedState> SharedState)
+    : ASTImporter(ToContext, ToFileManager, FromContext, FromFileManager,
+                  nullptr, MinimalImport, SharedState){}
+ASTImporter::ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
+                         ASTUnit &FromUnit, bool MinimalImport,
+                         std::shared_ptr<ASTImporterSharedState> SharedState)
+    : ASTImporter(ToContext, ToFileManager, FromUnit.getASTContext(),
+                  FromUnit.getFileManager(), &FromUnit, MinimalImport,
+                  SharedState){}
+
+ASTImporter::ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
+                         ASTContext &FromContext, FileManager &FromFileManager,
+                         ASTUnit *FromUnit, bool MinimalImport,
+                         std::shared_ptr<ASTImporterSharedState> SharedState)
     : SharedState(SharedState), ToContext(ToContext), FromContext(FromContext),
       ToFileManager(ToFileManager), FromFileManager(FromFileManager),
-      Minimal(MinimalImport) {
+      FromUnit(FromUnit), Minimal(MinimalImport) {
 
   // Create a default state without the lookup table: LLDB case.
   if (!SharedState) {
@@ -8421,6 +8435,13 @@ Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
   assert(ToID.isValid() && "Unexpected invalid fileID was created.");
 
   ImportedFileIDs[FromID] = ToID;
+  if (FromUnit) {
+    assert(SharedState->getImportedFileIDs().find(ToID) ==
+               SharedState->getImportedFileIDs().end() &&
+           "FileID already imported!");
+    SharedState->getImportedFileIDs()[ToID] = std::make_pair(FromID, FromUnit);
+  }
+
   return ToID;
 }
 
diff --git a/clang/lib/CrossTU/CrossTranslationUnit.cpp b/clang/lib/CrossTU/CrossTranslationUnit.cpp
index 977fd4b8dd3..1985d056440 100644
--- a/clang/lib/CrossTU/CrossTranslationUnit.cpp
+++ b/clang/lib/CrossTU/CrossTranslationUnit.cpp
@@ -295,7 +295,7 @@ llvm::Expected<const T *> CrossTranslationUnitContext::getCrossTUDefinitionImpl(
 
   TranslationUnitDecl *TU = Unit->getASTContext().getTranslationUnitDecl();
   if (const T *ResultDecl = findDefInDeclContext<T>(TU, LookupName))
-    return importDefinition(ResultDecl);
+    return importDefinition(ResultDecl, Unit);
   return llvm::make_error<IndexError>(index_error_code::failed_import);
 }
 
@@ -411,10 +411,13 @@ llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
 
 template <typename T>
 llvm::Expected<const T *>
-CrossTranslationUnitContext::importDefinitionImpl(const T *D) {
+CrossTranslationUnitContext::importDefinitionImpl(const T *D, ASTUnit *Unit) {
   assert(hasBodyOrInit(D) && "Decls to be imported should have body or init.");
 
-  ASTImporter &Importer = getOrCreateASTImporter(D->getASTContext());
+  assert(&D->getASTContext() == &Unit->getASTContext() &&
+         "ASTContext of Decl and the unit should match.");
+  ASTImporter &Importer = getOrCreateASTImporter(Unit);
+
   auto ToDeclOrError = Importer.Import(D);
   if (!ToDeclOrError) {
     handleAllErrors(ToDeclOrError.takeError(),
@@ -441,13 +444,15 @@ CrossTranslationUnitContext::importDefinitionImpl(const T *D) {
 }
 
 llvm::Expected<const FunctionDecl *>
-CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD) {
-  return importDefinitionImpl(FD);
+CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD,
+                                              ASTUnit *Unit) {
+  return importDefinitionImpl(FD, Unit);
 }
 
 llvm::Expected<const VarDecl *>
-CrossTranslationUnitContext::importDefinition(const VarDecl *VD) {
-  return importDefinitionImpl(VD);
+CrossTranslationUnitContext::importDefinition(const VarDecl *VD,
+                                              ASTUnit *Unit) {
+  return importDefinitionImpl(VD, Unit);
 }
 
 void CrossTranslationUnitContext::lazyInitImporterSharedSt(
@@ -457,17 +462,40 @@ void CrossTranslationUnitContext::lazyInitImporterSharedSt(
 }
 
 ASTImporter &
-CrossTranslationUnitContext::getOrCreateASTImporter(ASTContext &From) {
+CrossTranslationUnitContext::getOrCreateASTImporter(ASTUnit *Unit) {
+  ASTContext &From = Unit->getASTContext();
+
   auto I = ASTUnitImporterMap.find(From.getTranslationUnitDecl());
   if (I != ASTUnitImporterMap.end())
     return *I->second;
   lazyInitImporterSharedSt(Context.getTranslationUnitDecl());
-  ASTImporter *NewImporter = new ASTImporter(
-      Context, Context.getSourceManager().getFileManager(), From,
-      From.getSourceManager().getFileManager(), false, ImporterSharedSt);
+  ASTImporter *NewImporter =
+      new ASTImporter(Context, Context.getSourceManager().getFileManager(),
+                      *Unit, false, ImporterSharedSt);
   ASTUnitImporterMap[From.getTranslationUnitDecl()].reset(NewImporter);
   return *NewImporter;
 }
 
+llvm::Optional<std::pair<SourceLocation, ASTUnit *>>
+CrossTranslationUnitContext::getImportedFromSourceLocation(
+    const clang::SourceLocation &ToLoc) const {
+  if (!ImporterSharedSt)
+    return {};
+
+  const SourceManager &SM = Context.getSourceManager();
+  auto DecToLoc = SM.getDecomposedLoc(ToLoc);
+
+  auto I = ImporterSharedSt->getImportedFileIDs().find(DecToLoc.first);
+  if (I == ImporterSharedSt->getImportedFileIDs().end())
+    return {};
+
+  FileID FromID = I->second.first;
+  clang::ASTUnit *Unit = I->second.second;
+  SourceLocation FromLoc =
+      Unit->getSourceManager().getComposedLoc(FromID, DecToLoc.second);
+
+  return std::make_pair(FromLoc, Unit);
+}
+
 } // namespace cross_tu
 } // namespace clang
diff --git a/clang/unittests/CrossTU/CrossTranslationUnitTest.cpp b/clang/unittests/CrossTU/CrossTranslationUnitTest.cpp
index b3e7243ce1d..b4f22d5cb03 100644
--- a/clang/unittests/CrossTU/CrossTranslationUnitTest.cpp
+++ b/clang/unittests/CrossTU/CrossTranslationUnitTest.cpp
@@ -28,13 +28,18 @@ public:
       : CTU(CI), Success(Success) {}
 
   void HandleTranslationUnit(ASTContext &Ctx) {
+    auto FindFInTU = [](const TranslationUnitDecl *TU) {
+      const FunctionDecl *FD = nullptr;
+      for (const Decl *D : TU->decls()) {
+        FD = dyn_cast<FunctionDecl>(D);
+        if (FD && FD->getName() == "f")
+          break;
+      }
+      return FD;
+    };
+
     const TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
-    const FunctionDecl *FD = nullptr;
-    for (const Decl *D : TU->decls()) {
-      FD = dyn_cast<FunctionDecl>(D);
-      if (FD && FD->getName() == "f")
-        break;
-    }
+    const FunctionDecl *FD = FindFInTU(TU);
     assert(FD && FD->getName() == "f");
     bool OrigFDHasBody = FD->hasBody();
 
@@ -78,6 +83,28 @@ public:
     if (NewFDorError) {
       const FunctionDecl *NewFD = *NewFDorError;
       *Success = NewFD && NewFD->hasBody() && !OrigFDHasBody;
+
+      if (NewFD) {
+        // Check GetImportedFromSourceLocation.
+        llvm::Optional<std::pair<SourceLocation, ASTUnit *>> SLocResult =
+            CTU.getImportedFromSourceLocation(NewFD->getLocation());
+        EXPECT_TRUE(SLocResult);
+        if (SLocResult) {
+          SourceLocation OrigSLoc = (*SLocResult).first;
+          ASTUnit *OrigUnit = (*SLocResult).second;
+          // OrigUnit is created internally by CTU (is not the
+          // ASTWithDefinition).
+          TranslationUnitDecl *OrigTU =
+              OrigUnit->getASTContext().getTranslationUnitDecl();
+          const FunctionDecl *FDWithDefinition = FindFInTU(OrigTU);
+          EXPECT_TRUE(FDWithDefinition);
+          if (FDWithDefinition) {
+            EXPECT_EQ(FDWithDefinition->getName(), "f");
+            EXPECT_TRUE(FDWithDefinition->isThisDeclarationADefinition());
+            EXPECT_EQ(OrigSLoc, FDWithDefinition->getLocation());
+          }
+        }
+      }
     }
   }