teemperor updated this revision to Diff 191257.
teemperor retitled this revision from "[ASTImporter] Allow adding a shim to the 
ASTImporter" to "[ASTImporter] Allow adding a import strategy to the 
ASTImporter".
teemperor edited the summary of this revision.
teemperor added a comment.

Thanks for the quick review! I changed the code to also update the internal 
ASTImporter state like with a normal decl.

Changes:

- Renamed from Shim -> Strategy
- Now using the ImportedDecls map and also updating the internal ASTImporter 
state.


CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D59485/new/

https://reviews.llvm.org/D59485

Files:
  clang/include/clang/AST/ASTImporter.h
  clang/lib/AST/ASTImporter.cpp
  clang/unittests/AST/ASTImporterTest.cpp

Index: clang/unittests/AST/ASTImporterTest.cpp
===================================================================
--- clang/unittests/AST/ASTImporterTest.cpp
+++ clang/unittests/AST/ASTImporterTest.cpp
@@ -316,11 +316,14 @@
     std::unique_ptr<ASTUnit> Unit;
     TranslationUnitDecl *TUDecl = nullptr;
     std::unique_ptr<ASTImporter> 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<Decl *, Decl *>
   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<Decl *> Import(ASTImporter &Importer, Decl *FromD) override {
+    auto *ND = dyn_cast<NamedDecl>(FromD);
+    if (!ND)
+      return {};
+    if (ND->getName() != "shouldNotBeImported")
+      return {};
+    for (Decl *D : Importer.getToContext().getTranslationUnitDecl()->decls())
+      if (auto ND = dyn_cast<NamedDecl>(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<CXXRecordDecl>(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<CXXRecordDecl>().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<CXXRecordDecl>().match(ToTU, Pattern);
+  EXPECT_EQ(0U, count);
+}
+
+namespace {
+class NoOpStrategy : public ImportStrategy {
+  llvm::Optional<Decl *> 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<CXXRecordDecl>(To);
+  EXPECT_EQ(Imported->getQualifiedNameAsString(), "declToImport");
+}
+
 TEST_P(ImportExpr, ImportStringLiteral) {
   MatchVerifier<Decl> 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, );
 
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<Decl *> 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/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<Decl*, 2> 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<Decl *> 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<NamedDecl *, 2>;
     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.
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to