Index: clang/include/clang/AST/ASTImporter.h =================================================================== --- clang/include/clang/AST/ASTImporter.h +++ clang/include/clang/AST/ASTImporter.h @@ -78,6 +78,20 @@ // the different entries in a given redecl chain. llvm::SmallVector getCanonicalForwardRedeclChain(Decl* D); + class ASTImporter; + /// Allows intercepting the import process of the ASTImporter. + struct ImportStrategy { + virtual ~ImportStrategy() = default; + /// Called from the ASTImporter before the given Decl is imported. + /// + /// If this method returns a Decl, the ASTImporter will abort the current + /// import step and treat the returned decl as if it was just imported. + /// If this method returns nothing, the ASTImporter continues to import the + /// given Decl on its own. + virtual llvm::Optional Import(ASTImporter &Importer, + Decl *FromD) = 0; + }; + /// Imports selected nodes from one AST context into another context, /// merging AST nodes where appropriate. class ASTImporter { @@ -137,6 +151,10 @@ /// (which we have already complained about). NonEquivalentDeclSet NonEquivalentDecls; + /// The Shim that should rewrite the import calls of this ASTImporter, or + /// a nullptr of this ASTImporter has no shim. + ImportStrategy *Strategy = nullptr; + using FoundDeclsTy = SmallVector; FoundDeclsTy findDeclsInToCtx(DeclContext *DC, DeclarationName Name); @@ -170,6 +188,10 @@ /// to-be-completed forward declarations when possible. bool isMinimalImport() const { return Minimal; } + ImportStrategy *getStrategy() { return Strategy; } + + void setStrategy(ImportStrategy *S) { Strategy = S; } + /// \brief Import the given object, returns the result. /// /// \param To Import the object into this variable. Index: clang/lib/AST/ASTImporter.cpp =================================================================== --- clang/lib/AST/ASTImporter.cpp +++ clang/lib/AST/ASTImporter.cpp @@ -7837,6 +7837,21 @@ return ToD; } + // First check if our ImportStrategy wants to rewrite this import call. + if (Strategy) { + llvm::Optional NewD = Strategy->Import(*this, FromD); + // The strategy has found a decl for us and we pretend we successfully + // imported it. + if (NewD) { + // Update LookupTable and notify subclasses. + AddToLookupTable(*NewD); + Imported(FromD, *NewD); + // Update map of already imported decls. + MapImported(FromD, *NewD); + return *NewD; + } + } + // Import the declaration. ExpectedDecl ToDOrErr = Importer.Visit(FromD); if (!ToDOrErr) Index: clang/unittests/AST/ASTImporterTest.cpp =================================================================== --- clang/unittests/AST/ASTImporterTest.cpp +++ clang/unittests/AST/ASTImporterTest.cpp @@ -316,11 +316,14 @@ std::unique_ptr Unit; TranslationUnitDecl *TUDecl = nullptr; std::unique_ptr Importer; - TU(StringRef Code, StringRef FileName, ArgVector Args) + ImportStrategy *Strategy; + TU(StringRef Code, StringRef FileName, ArgVector Args, + ImportStrategy *Strategy = nullptr) : Code(Code), FileName(FileName), Unit(tooling::buildASTFromCodeWithArgs(this->Code, Args, this->FileName)), - TUDecl(Unit->getASTContext().getTranslationUnitDecl()) { + TUDecl(Unit->getASTContext().getTranslationUnitDecl()), + Strategy(Strategy) { Unit->enableSourceFileDiagnostics(); } @@ -331,6 +334,7 @@ new ASTImporter(ToAST->getASTContext(), ToAST->getFileManager(), Unit->getASTContext(), Unit->getFileManager(), false, &LookupTable)); + Importer->setStrategy(Strategy); } assert(&ToAST->getASTContext() == &Importer->getToContext()); createVirtualFileIfNeeded(ToAST, FileName, Code); @@ -401,11 +405,12 @@ // Must not be called more than once within the same test. std::tuple getImportedDecl(StringRef FromSrcCode, Language FromLang, StringRef ToSrcCode, - Language ToLang, StringRef Identifier = DeclToImportID) { + Language ToLang, StringRef Identifier = DeclToImportID, + ImportStrategy *Strategy = nullptr) { ArgVector FromArgs = getArgVectorForLanguage(FromLang), ToArgs = getArgVectorForLanguage(ToLang); - FromTUs.emplace_back(FromSrcCode, InputFileName, FromArgs); + FromTUs.emplace_back(FromSrcCode, InputFileName, FromArgs, Strategy); TU &FromTU = FromTUs.back(); assert(!ToAST); @@ -455,6 +460,12 @@ return ToAST->getASTContext().getTranslationUnitDecl(); } + ASTImporter &getImporter(Decl *From, Language ToLang) { + lazyInitToAST(ToLang, "", OutputFileName); + TU *FromTU = findFromTU(From); + return *FromTU->Importer; + } + // Import the given Decl into the ToCtx. // May be called several times in a given test. // The different instances of the param From may have different ASTContext. @@ -544,6 +555,81 @@ EXPECT_THAT(RedeclsD1, ::testing::ContainerEq(RedeclsD2)); } +namespace { +class RedirectStrategy : public ImportStrategy { + llvm::Optional Import(ASTImporter &Importer, Decl *FromD) override { + auto *ND = dyn_cast(FromD); + if (!ND) + return {}; + if (ND->getName() != "shouldNotBeImported") + return {}; + for (Decl *D : Importer.getToContext().getTranslationUnitDecl()->decls()) + if (auto ND = dyn_cast(D)) + if (ND->getName() == "realDecl") + return ND; + return {}; + } +}; +} // namespace + +struct ImportStrategyTest : ASTImporterOptionSpecificTestBase {}; + +// Test that the ImportStrategy can intercept an import call. +TEST_P(ImportStrategyTest, InterceptImport) { + RedirectStrategy Strategy; + Decl *From, *To; + std::tie(From, To) = getImportedDecl("class shouldNotBeImported {};", + Lang_CXX, "class realDecl {};", Lang_CXX, + "shouldNotBeImported", &Strategy); + auto *Imported = cast(To); + EXPECT_EQ(Imported->getQualifiedNameAsString(), "realDecl"); + + // Make sure the Strategy prevented the importing of the decl. + auto *ToTU = Imported->getTranslationUnitDecl(); + auto Pattern = functionDecl(hasName("shouldNotBeImported")); + unsigned count = + DeclCounterWithPredicate().match(ToTU, Pattern); + EXPECT_EQ(0U, count); +} + +// Test that when we indirectly import a declaration the Strategy still works. +// Also tests that the ImportStrategy allows the import process to continue when +// we try to import a declaration that we ignore in the Strategy. +TEST_P(ImportStrategyTest, InterceptIndirectImport) { + RedirectStrategy Strategy; + Decl *From, *To; + std::tie(From, To) = + getImportedDecl("class shouldNotBeImported {};" + "class F { shouldNotBeImported f; };", + Lang_CXX, "class realDecl {};", Lang_CXX, "F", &Strategy); + + // Make sure the ImportStrategy prevented the importing of the decl. + auto *ToTU = To->getTranslationUnitDecl(); + auto Pattern = functionDecl(hasName("shouldNotBeImported")); + unsigned count = + DeclCounterWithPredicate().match(ToTU, Pattern); + EXPECT_EQ(0U, count); +} + +namespace { +class NoOpStrategy : public ImportStrategy { + llvm::Optional Import(ASTImporter &Importer, Decl *FromD) override { + return {}; + } +}; +} // namespace + +// Test a ImportStrategy that just does nothing. +TEST_P(ImportStrategyTest, NoOpStrategy) { + NoOpStrategy Strategy; + Decl *From, *To; + std::tie(From, To) = + getImportedDecl("class declToImport {};", Lang_CXX, "class realDecl {};", + Lang_CXX, DeclToImportID, &Strategy); + auto *Imported = cast(To); + EXPECT_EQ(Imported->getQualifiedNameAsString(), "declToImport"); +} + TEST_P(ImportExpr, ImportStringLiteral) { MatchVerifier Verifier; testImport( @@ -5512,6 +5598,9 @@ INSTANTIATE_TEST_CASE_P(ParameterizedTests, ASTImporterOptionSpecificTestBase, DefaultTestValuesForRunOptions, ); +INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportStrategyTest, + DefaultTestValuesForRunOptions, ); + INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportFunctions, DefaultTestValuesForRunOptions, );