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

[balazske]

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.

[martong]

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

@balazske Could you please prepare the upcoming patch and put that on the Stack? I think that may ease the work of the reviewers if they could see how you want to use this function.

[martong]

include/clang/AST/ASTImporter.h

234

Unit->getASTContext must be equal to FromContext, right?
Then FromUnit would be a better naming.
The importer does handle two ASTContexts, so it would be great to know which ASTUnit's ASTContext is which. Or vica-versa, we must know precisely to which context does this unit belong. Also a comment would be useful that we must store the unit because there is no way to get that from the context.

[martong]

include/clang/AST/ASTImporter.h

317

What if we provided an additional constructor where we take over the ASTUnits instead of the ASTContexts?
Then we would not need to pass the FileManagers neither.

ASTImporter(ASTUnit &ToUnit, 
            ASTUnit &FromUnit,
            bool MinimalImport,
            std::shared_ptr<ASTImporterSharedState> SharedState = nullptr,

[balazske]

include/clang/AST/ASTImporter.h

317

Is the SharedState==nullptr case only for LLDB? If yes then it is possible to have a "LLDB" constructor when no shared state and no ASTUnit is needed. And another for CTU case when a From and To ASTUnit is specified and a shared state (theoretically minimal can be true in any case but probably only true for LLDB?).

[balazske]

include/clang/AST/ASTImporter.h

317

It is not trivial to make and use this new kind of constructor (with ToUnit), there is no ToUnit in the CrossTU context. Is it OK to have the original constructor, or one with FromUnit (but ToContext and ToFileManager)?

[martong]

include/clang/AST/ASTImporter.h

317

Is the SharedState==nullptr case only for LLDB?

Yes.

Is it OK to have the original constructor, or one with FromUnit (but ToContext and ToFileManager)?

I think it is OK to have a new ctor with FromUnit (but ToContext and ToFileManager). Because ASTUnit is just a utility class for loading an ASTContext from an AST file. So, one of the "to" or the "from" context could be initialized by an ASTUnit.

In the future the API of the ASTImporter could be extended with such a ctor which takes two ASTUnits, also we could add the counterpart ctor where we have ToUnit, FromContext and FromFileManager as params.

[martong]

include/clang/AST/ASTImporter.h

317

Note that if we have a ctor which takes FromUnit as a param then we won't need the assertions which check whether the Unit provides the same context or not.

[balazske]

• Renamed 'Unit', changed constructors.

[martong]

include/clang/AST/ASTImporter.h

286

Why do we still need FromUnit in this ctor? I thought the other ctor would handle that case.

[balazske]

include/clang/AST/ASTImporter.h

286

This is a private and common constructor that can handle both cases (called by the others).

[martong]

LGTM.

include/clang/AST/ASTImporter.h

286

Ok, I've missed that.

[balazske]

Commit was reverted because we were not aware of that clangAST library can not use clangFrontend library. So a different solution is needed, probably pass a SourceLocation import callback to ASTImporter. (A Preprocessor that can be used in place of ASTUnit here can not be get from a SourceManager or ASTContext.) And then no ImportedFileIDs is needed in ASTImporterSharedState, instead in CrossTranslationUnitContext.

[balazske]

New patch:
https://reviews.llvm.org/D65064

include/clang/AST/ASTImporter.h

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Error.h"
 #include <utility>
 
 namespace clang {
 
 class ASTContext;
 class ASTImporterSharedState;
+class ASTUnit;
 class Attr;
 class CXXBaseSpecifier;
 class CXXCtorInitializer;
 class Decl;
 class DeclContext;
 class Expr;
 class FileManager;
 class NamedDecl;
private:
     SavedImportPathsTy SavedImportPaths;
 
     /// The contexts we're importing to and from.
     ASTContext &ToContext, &FromContext;
 
     /// The file managers we're importing to and from.
     FileManager &ToFileManager, &FromFileManager;
 
+    /// The ASTUnit that this importer belongs to, if any.
+    ASTUnit *Unit;
+
     /// Whether to perform a minimal import.
     bool Minimal;
 
     /// Whether the last diagnostic came from the "from" context.
     bool LastDiagFromFrom = false;
 
     /// Mapping from the already-imported types in the "from" context
     /// to the corresponding types in the "to" context.
private:
 
     using FoundDeclsTy = SmallVector<NamedDecl *, 2>;
     FoundDeclsTy findDeclsInToCtx(DeclContext *DC, DeclarationName Name);
 
     void AddToLookupTable(Decl *ToD);
 
   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
     /// decl on its own.
     virtual Expected<Decl *> ImportImpl(Decl *From);
 
     /// Used only in unittests to verify the behaviour of the error handling.
     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.
     ///
     /// \param FromContext The context we'll be importing from.
     ///
     /// \param FromFileManager The file manager we'll be importing into.
     ///
     /// \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.
+    ///
+    /// \param Unit Pointer to an ASTUnit that contains this importer, if any.
+    /// A mapping of imported FileID's to ASTUnit is stored in the shared state,
+    /// if both are set.
     ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
                 ASTContext &FromContext, FileManager &FromFileManager,
                 bool MinimalImport,
-                std::shared_ptr<ASTImporterSharedState> SharedState = nullptr);
+                std::shared_ptr<ASTImporterSharedState> SharedState = nullptr,
+                ASTUnit *Unit = nullptr);
 
     virtual ~ASTImporter();
 
     /// Whether the importer will perform a minimal import, creating
     /// to-be-completed forward declarations when possible.
     bool isMinimalImport() const { return Minimal; }
 
     /// \brief Import the given object, returns the result.

include/clang/AST/ASTImporterSharedState.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_AST_ASTIMPORTERSHAREDSTATE_H
 #define LLVM_CLANG_AST_ASTIMPORTERSHAREDSTATE_H
 
 #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;
 
   /// Mapping from the already-imported declarations in the "to"
   /// context to the error status of the import of that declaration.
   /// This map contains only the declarations that were not correctly
   /// imported. The same declaration may or may not be included in
   /// ImportedFromDecls. This map is updated continuously during imports and
   /// never cleared (like ImportedFromDecls).
   llvm::DenseMap<Decl *, ImportError> ImportErrors;
 
   // 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;
 
   ASTImporterSharedState(TranslationUnitDecl &ToTU) {
     LookupTable = llvm::make_unique<ASTImporterLookupTable>(ToTU);
   }
 
   ASTImporterLookupTable *getLookupTable() { return LookupTable.get(); }
if (Pos != ImportErrors.end())
       return Pos->second;
     else
       return Optional<ImportError>();
   }
 
   void setImportDeclError(Decl *To, ImportError Error) {
     ImportErrors[To] = Error;
   }
+
+  ImportedFileIDMap &getImportedFileIDs() { return ImportedFileIDs; }
+
+  const ImportedFileIDMap &getImportedFileIDs() const {
+    return ImportedFileIDs;
+  }
 };
 
 } // namespace clang
 #endif // LLVM_CLANG_AST_ASTIMPORTERSHAREDSTATE_H

include/clang/CrossTU/CrossTranslationUnit.h

llvm::Expected<ASTUnit *> loadExternalAST(StringRef LookupName,
                                             StringRef IndexName,
                                             bool DisplayCTUProgress = false);
 
   /// This function merges a definition from a separate AST Unit into
   ///        the current one which was created by the compiler instance that
   ///        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);
 
   /// 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,
                                                      StringRef IndexName,
                                                      bool DisplayCTUProgress);
   template <typename T>
   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;
   llvm::StringMap<std::string> NameFileMap;
   llvm::DenseMap<TranslationUnitDecl *, std::unique_ptr<ASTImporter>>
       ASTUnitImporterMap;
   CompilerInstance &CI;
   ASTContext &Context;

lib/AST/ASTImporter.cpp

for (auto *FromOverriddenMethod : FromMethod->overridden_methods()) {
     else
       consumeError(ImportedOrErr.takeError());
   }
 }
 
 ASTImporter::ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
                          ASTContext &FromContext, FileManager &FromFileManager,
                          bool MinimalImport,
-                         std::shared_ptr<ASTImporterSharedState> SharedState)
+                         std::shared_ptr<ASTImporterSharedState> SharedState,
+                         ASTUnit *Unit)
     : SharedState(SharedState), ToContext(ToContext), FromContext(FromContext),
       ToFileManager(ToFileManager), FromFileManager(FromFileManager),
-      Minimal(MinimalImport) {
+      Unit(Unit), Minimal(MinimalImport) {
 
   // Create a default state without the lookup table: LLDB case.
   if (!SharedState) {
     this->SharedState = std::make_shared<ASTImporterSharedState>();
   }
 
   ImportedDecls[FromContext.getTranslationUnitDecl()] =
       ToContext.getTranslationUnitDecl();
if (ToID.isInvalid() || IsBuiltin) {
       ToID = ToSM.createFileID(std::move(ToBuf),
                                FromSLoc.getFile().getFileCharacteristic());
     }
   }
 
   assert(ToID.isValid() && "Unexpected invalid fileID was created.");
 
   ImportedFileIDs[FromID] = ToID;
+  if (Unit) {
+    assert(SharedState->getImportedFileIDs().find(ToID) ==
+               SharedState->getImportedFileIDs().end() &&
+           "FileID already imported!");
+    SharedState->getImportedFileIDs()[ToID] = std::make_pair(FromID, Unit);
+  }
+
   return ToID;
 }
 
 Expected<CXXCtorInitializer *> ASTImporter::Import(CXXCtorInitializer *From) {
   ExpectedExpr ToExprOrErr = Import(From->getInit());
   if (!ToExprOrErr)
     return ToExprOrErr.takeError();
 

lib/CrossTU/CrossTranslationUnit.cpp

if (LangTo.CPlusPlus11 != LangFrom.CPlusPlus11 ||
       LangTo.CPlusPlus2a != LangFrom.CPlusPlus2a) {
     ++NumLangDialectMismatch;
     return llvm::make_error<IndexError>(
         index_error_code::lang_dialect_mismatch);
   }
 
   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);
 }
 
 llvm::Expected<const FunctionDecl *>
 CrossTranslationUnitContext::getCrossTUDefinition(const FunctionDecl *FD,
                                                   StringRef CrossTUDir,
                                                   StringRef IndexName,
                                                   bool DisplayCTUProgress) {
llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
   if (!Unit)
     return llvm::make_error<IndexError>(
         index_error_code::failed_to_get_external_ast);
   return Unit;
 }
 
 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(),
                     [&](const ImportError &IE) {
                       switch (IE.Error) {
                       case ImportError::NameConflict:
                         ++NumNameConflicts;
                          break;
CrossTranslationUnitContext::importDefinitionImpl(const T *D, ASTUnit *Unit) {
   auto *ToDecl = cast<T>(*ToDeclOrError);
   assert(hasBodyOrInit(ToDecl) && "Imported Decl should have body or init.");
   ++NumGetCTUSuccess;
 
   return ToDecl;
 }
 
 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(
     TranslationUnitDecl *ToTU) {
   if (!ImporterSharedSt)
     ImporterSharedSt = std::make_shared<ASTImporterSharedState>(*ToTU);
 }
 
 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);
+      From.getSourceManager().getFileManager(), false, ImporterSharedSt, Unit);
   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

unittests/CrossTU/CrossTranslationUnitTest.cpp

 namespace {
 
 class CTUASTConsumer : public clang::ASTConsumer {
 public:
   explicit CTUASTConsumer(clang::CompilerInstance &CI, bool *Success)
       : CTU(CI), Success(Success) {}
 
   void HandleTranslationUnit(ASTContext &Ctx) {
-    const TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
+    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 = FindFInTU(TU);
     assert(FD && FD->getName() == "f");
     bool OrigFDHasBody = FD->hasBody();
 
     // Prepare the index file and the AST file.
     int ASTFD;
     llvm::SmallString<256> ASTFileName;
     ASSERT_FALSE(
         llvm::sys::fs::createTemporaryFile("f_ast", "ast", ASTFD, ASTFileName));
void HandleTranslationUnit(ASTContext &Ctx) {
     // Load the definition from the AST file.
     llvm::Expected<const FunctionDecl *> NewFDorError = handleExpected(
         CTU.getCrossTUDefinition(FD, "", IndexFileName, false),
         []() { return nullptr; }, [](IndexError &) {});
 
     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());
+          }
+        }
+      }
     }
   }
 
 private:
   CrossTranslationUnitContext CTU;
   bool *Success;
 };