[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.

Reviewers: martong, shafik

Reviewed By: martong

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

Tags: #clang

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

llvm-svn: 366884

1 files changed, 33 insertions, 6 deletions

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());
+          }
+        }
+      }
     }
   }