Index: cfe/trunk/include/clang/AST/ASTImporter.h =================================================================== --- cfe/trunk/include/clang/AST/ASTImporter.h +++ cfe/trunk/include/clang/AST/ASTImporter.h @@ -87,6 +87,8 @@ using NonEquivalentDeclSet = llvm::DenseSet>; using ImportedCXXBaseSpecifierMap = llvm::DenseMap; + using FileIDImportHandlerType = + std::function; // An ImportPath is the list of the AST nodes which we visit during an // Import call. @@ -210,6 +212,8 @@ }; private: + FileIDImportHandlerType FileIDImportHandler; + std::shared_ptr SharedState = nullptr; /// The path which we go through during the import of a given AST node. @@ -310,6 +314,14 @@ virtual ~ASTImporter(); + /// Set a callback function for FileID import handling. + /// The function is invoked when a FileID is imported from the From context. + /// The imported FileID in the To context and the original FileID in the + /// From context is passed to it. + void setFileIDImportHandler(FileIDImportHandlerType H) { + FileIDImportHandler = H; + } + /// Whether the importer will perform a minimal import, creating /// to-be-completed forward declarations when possible. bool isMinimalImport() const { return Minimal; } Index: cfe/trunk/include/clang/CrossTU/CrossTranslationUnit.h =================================================================== --- cfe/trunk/include/clang/CrossTU/CrossTranslationUnit.h +++ cfe/trunk/include/clang/CrossTU/CrossTranslationUnit.h @@ -153,8 +153,10 @@ /// was passed to the constructor. /// /// \return Returns the resulting definition or an error. - llvm::Expected importDefinition(const FunctionDecl *FD); - llvm::Expected importDefinition(const VarDecl *VD); + llvm::Expected importDefinition(const FunctionDecl *FD, + ASTUnit *Unit); + llvm::Expected importDefinition(const VarDecl *VD, + ASTUnit *Unit); /// Get a name to identify a named decl. static std::string getLookupName(const NamedDecl *ND); @@ -162,9 +164,23 @@ /// 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> + getImportedFromSourceLocation(const clang::SourceLocation &ToLoc) const; + private: + using ImportedFileIDMap = + llvm::DenseMap>; + void lazyInitImporterSharedSt(TranslationUnitDecl *ToTU); - ASTImporter &getOrCreateASTImporter(ASTContext &From); + ASTImporter &getOrCreateASTImporter(ASTUnit *Unit); template llvm::Expected getCrossTUDefinitionImpl(const T *D, StringRef CrossTUDir, @@ -174,7 +190,7 @@ const T *findDefInDeclContext(const DeclContext *DC, StringRef LookupName); template - llvm::Expected importDefinitionImpl(const T *D); + llvm::Expected importDefinitionImpl(const T *D, ASTUnit *Unit); llvm::StringMap> FileASTUnitMap; llvm::StringMap NameASTUnitMap; @@ -184,6 +200,15 @@ CompilerInstance &CI; ASTContext &Context; std::shared_ptr ImporterSharedSt; + /// 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; + /// \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; Index: cfe/trunk/lib/AST/ASTImporter.cpp =================================================================== --- cfe/trunk/lib/AST/ASTImporter.cpp +++ cfe/trunk/lib/AST/ASTImporter.cpp @@ -8421,6 +8421,10 @@ assert(ToID.isValid() && "Unexpected invalid fileID was created."); ImportedFileIDs[FromID] = ToID; + + if (FileIDImportHandler) + FileIDImportHandler(ToID, FromID); + return ToID; } Index: cfe/trunk/lib/CrossTU/CrossTranslationUnit.cpp =================================================================== --- cfe/trunk/lib/CrossTU/CrossTranslationUnit.cpp +++ cfe/trunk/lib/CrossTU/CrossTranslationUnit.cpp @@ -295,7 +295,7 @@ TranslationUnitDecl *TU = Unit->getASTContext().getTranslationUnitDecl(); if (const T *ResultDecl = findDefInDeclContext(TU, LookupName)) - return importDefinition(ResultDecl); + return importDefinition(ResultDecl, Unit); return llvm::make_error(index_error_code::failed_import); } @@ -411,10 +411,13 @@ template llvm::Expected -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 @@ } llvm::Expected -CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD) { - return importDefinitionImpl(FD); +CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD, + ASTUnit *Unit) { + return importDefinitionImpl(FD, Unit); } llvm::Expected -CrossTranslationUnitContext::importDefinition(const VarDecl *VD) { - return importDefinitionImpl(VD); +CrossTranslationUnitContext::importDefinition(const VarDecl *VD, + ASTUnit *Unit) { + return importDefinitionImpl(VD, Unit); } void CrossTranslationUnitContext::lazyInitImporterSharedSt( @@ -457,7 +462,9 @@ } 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; @@ -465,9 +472,32 @@ ASTImporter *NewImporter = new ASTImporter( Context, Context.getSourceManager().getFileManager(), From, From.getSourceManager().getFileManager(), false, ImporterSharedSt); + NewImporter->setFileIDImportHandler([this, Unit](FileID ToID, FileID FromID) { + assert(ImportedFileIDs.find(ToID) == ImportedFileIDs.end() && + "FileID already imported, should not happen."); + ImportedFileIDs[ToID] = std::make_pair(FromID, Unit); + }); ASTUnitImporterMap[From.getTranslationUnitDecl()].reset(NewImporter); return *NewImporter; } +llvm::Optional> +CrossTranslationUnitContext::getImportedFromSourceLocation( + const clang::SourceLocation &ToLoc) const { + const SourceManager &SM = Context.getSourceManager(); + auto DecToLoc = SM.getDecomposedLoc(ToLoc); + + auto I = ImportedFileIDs.find(DecToLoc.first); + if (I == ImportedFileIDs.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 Index: cfe/trunk/unittests/CrossTU/CrossTranslationUnitTest.cpp =================================================================== --- cfe/trunk/unittests/CrossTU/CrossTranslationUnitTest.cpp +++ cfe/trunk/unittests/CrossTU/CrossTranslationUnitTest.cpp @@ -28,13 +28,18 @@ : 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(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(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 @@ if (NewFDorError) { const FunctionDecl *NewFD = *NewFDorError; *Success = NewFD && NewFD->hasBody() && !OrigFDHasBody; + + if (NewFD) { + // Check GetImportedFromSourceLocation. + llvm::Optional> 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()); + } + } + } } }