[ASTImporter] Add an ImportImpl method to allow customizing Import behavior.

[teemperor]

Summary

We are currently implementing support in LLDB that reconstructs the STL templates from
the target program in the expression evaluator. This reconstruction happens during the
import process from our debug info AST into the expression evaluation AST, which means
we need a way to intercept the ASTImporter import process.

This patch adds an protected ImportImpl method that we can overwrite in LLDB to implement
our special importing logic (which is essentially just looking into a C++ module that is attached to
the target context). Because ImportImpl has to call MapImported/AddToLookup for the decls it
creates, this patch also exposes those via a new unified method and checks that we call it when
importing decls.

[balazske]

I do not know if "Shim" is the correct name for this, maybe "Strategy" is better?

[balazske]

The replace of Import calls may work if the internal state of the ASTImporter is updated correctly from the 'shim' object (especially if not every Decl is handled by the shim). For this to work some of the state of ASTImporter that was intended to be internal must be made public, for example map of imported decls and lookup table. The shim may corrupt the ASTImporter state for example when the lookup table is not updated correctly. (Probably for the LLDB special use cases this is not a problem?)

[teemperor]

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.

[balazske]

It looks better now.
One "problem" is that now there is the strategy and there is the Imported function. It would be possible to unify these into a strategy that contains maybe a "import" and a "post-import" callback. (But this requires big changes in the ASTImporter user code.)

clang/include/clang/AST/ASTImporter.h

154 ↗

(On Diff #191257)

'Shim' should be replaced with 'strategy' (and the nullptr means "no user-specified strategy" or "default strategy").

[teemperor]

I think we could also refactor the strategy into a PreImport call. That way we don't break all the user-code out there and it seems less intrusive to the ASTImporter code. I'll update the patch.

[balazske]

Another observation: The Import function of the strategy now has no way to return an error. An even better version of it would be to include the possibility of import error (with ImportError, or other error type). Or the "PreImport" function could indicate if the Decl is handled by the strategy, and Import is called only if yes and Import can return an Expected<Decl *>. (The PostImport is called for every successfully imported Decl.)

[balazske]

By the way if the Import of the strategy uses recursive import of other things this can cause same problems as it was in ASTImporter before the GetImportedOrCreateDecl was introduced. So this should be avoided or something similar to GetImportedOrCreateDecl must be performed, maybe that function can be made public somehow.

But if this strategy is used only for special cases like in the test case (replace a specific Decl with something other) probably implementing this functionality in another way is a more simple solution.

[martong]

We can't reliable import templates with the ASTImporter, so actually reconstructing the template in the debug info AST and then importing it doesn't work.

Could you please elaborate how the import of templates fails in ASTImporter? Is it because the AST you build from the DWARF is incomplete? Or is there a lookup problem? Is there an assertion in one of the CodeGen functions once the import is finished?
If we could fix that error in the ASTImporter then other clients of it (like CTU) could benefit from the fix too. Perhaps, the best would be to have a test which reproduces the template import related errors, maybe a lit test with clang-import-test?

If the crux of the problem is because of the incomplete AST from DWARF then probably we cannot fix that here in clang::ASTImporter. However, if that is not the case then I think we should try to fix the error here instead of duplicating lookup and import logic in LLDB. (As for the lookup problems, quite recently I just have discovered, that lookup during expression evaluation in LLDB does not work properly, I'll communicate this in a separate email/patch.)

It would also be much slower.

Could you please explain why it would be slower?

[teemperor]

So I replaced the whole Strategy/Shim thing with a simple ImportInternal method that we can overwrite and is also no longer public. That also gets rid of all the boilerplate code from the first patches and also makes the follow-up LLDB patch nicer.

@martong It's not related to lookup or anything like that, it's just that we don't have enough information in our debug info AST to allow people to use meta programming. The custom logic we have in LLDB looks into the std C++ module to fill out these gaps in the AST by hand when they are used in an expression.

The remark about the alternative being slow just means that fixing all the templates we have in our debug info AST seems like a costly approach. I'm also not sure how well it would work, as I never tried mixing a C++ module in an ASTContext that isn't currently being parsed by Clang.

[balazske]

clang/lib/AST/ASTImporter.cpp

7738 ↗

(On Diff #191765)

Comment can be better: Import it using ASTNodeImporter.

clang/unittests/AST/ASTImporterTest.cpp

581 ↗

(On Diff #191765)

Is this using needed?

583 ↗

(On Diff #191765)

Is it possible to make this a static variable instead of function?

588 ↗

(On Diff #191765)

The static_cast should not be needed.

[martong]

@martong It's not related to lookup or anything like that, it's just that we don't have enough information in our debug info AST to allow people to use meta programming. The custom logic we have in LLDB looks into the std C++ module to fill out these gaps in the AST by hand when they are used in an expression.

The remark about the alternative being slow just means that fixing all the templates we have in our debug info AST seems like a costly approach. I'm also not sure how well it would work, as I never tried mixing a C++ module in an ASTContext that isn't currently being parsed by Clang.

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

[teemperor]

• Addressed feedback.

[teemperor]

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

There are no C++ modules involved in our use case beside the generic std module. No user-written code is read from a module and we don't have any PCM file beside the std module we build for the expression evaluator.

We use the ASTReader to directly read the std module contents into the expression evaluation ASTContext, but this doesn't give us the decls for the instantiations the user has made (e.g. std::vector<Foo>). We only have these user instantiations in the 'normal' debug info where we have a rather minimal AST (again, no modules are not involved in building this debug info AST). The InternalImport in the LLDB patch just stitches the module AST and the debug info AST together when we import something that we also have (in better quality from the C++ module) in the target ASTContext.

clang/unittests/AST/ASTImporterTest.cpp

581 ↗

(On Diff #191765)

Yeah, otherwise we have to provide our own constructor that essentially just forwards to the ASTImporter constructor.

588 ↗

(On Diff #191765)

Thanks!

[balazske]

clang/unittests/AST/ASTImporterTest.cpp

348 ↗

(On Diff #191859)

This function is not needed?

482 ↗

(On Diff #191859)

This is unused, can be removed.

583 ↗

(On Diff #191765)

Another small thing: The comment above is now not updated ("Returns").

[martong]

In D59485#1439390, @teemperor wrote:

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

There are no C++ modules involved in our use case beside the generic std module. No user-written code is read from a module and we don't have any PCM file beside the std module we build for the expression evaluator.

We use the ASTReader to directly read the std module contents into the expression evaluation ASTContext, but this doesn't give us the decls for the instantiations the user has made (e.g. std::vector<Foo>). We only have these user instantiations in the 'normal' debug info where we have a rather minimal AST (again, no modules are not involved in building this debug info AST). The InternalImport in the LLDB patch just stitches the module AST and the debug info AST together when we import something that we also have (in better quality from the C++ module) in the target ASTContext.

Thank you for the explanation.
Now I understand you directly create specializations from the std module and intercept the import to avoid importing broken specializations from the debug info, this makes sense.

[martong]

In D59485#1439570, @martong wrote:

In D59485#1439390, @teemperor wrote:

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

There are no C++ modules involved in our use case beside the generic std module. No user-written code is read from a module and we don't have any PCM file beside the std module we build for the expression evaluator.

We use the ASTReader to directly read the std module contents into the expression evaluation ASTContext, but this doesn't give us the decls for the instantiations the user has made (e.g. std::vector<Foo>). We only have these user instantiations in the 'normal' debug info where we have a rather minimal AST (again, no modules are not involved in building this debug info AST). The InternalImport in the LLDB patch just stitches the module AST and the debug info AST together when we import something that we also have (in better quality from the C++ module) in the target ASTContext.

Thank you for the explanation.
Now I understand you directly create specializations from the std module and intercept the import to avoid importing broken specializations from the debug info, this makes sense.

Really, just one last question to see the whole picture: If clang::ASTImporter were capable of handling a specialization/instantiation with missing info then we could use that. So the reason for this interception is that clang::ASTImporter::VisitClassTemplateSpecializationDecl cannot handle the specialization it receives because that or its dependent nodes has too many missing data, right?

[teemperor]

• Removed unused functions.
• Fixed two comments that I missed to update.

[a_sidorin]

Hello Raphael,

Thank you for the explanation. I have +1 to Gabor's question to understand if this functionality can actually be added to the common ASTImporter.

clang/include/clang/AST/ASTImporter.h

149 ↗

(On Diff #191765)

I'd suggest to rename it to ImportImpl.

clang/unittests/AST/ASTImporterTest.cpp

578 ↗

(On Diff #191946)

I think it's better to merge these conditions: if (!ND || ND->getName() != "shouldNotBeImported")

583 ↗

(On Diff #191946)

auto * (pointer sign is needed).

590 ↗

(On Diff #191765)

LookupTable?

[teemperor]

In D59485#1439628, @martong wrote:

In D59485#1439570, @martong wrote:

In D59485#1439390, @teemperor wrote:

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

There are no C++ modules involved in our use case beside the generic std module. No user-written code is read from a module and we don't have any PCM file beside the std module we build for the expression evaluator.

We use the ASTReader to directly read the std module contents into the expression evaluation ASTContext, but this doesn't give us the decls for the instantiations the user has made (e.g. std::vector<Foo>). We only have these user instantiations in the 'normal' debug info where we have a rather minimal AST (again, no modules are not involved in building this debug info AST). The InternalImport in the LLDB patch just stitches the module AST and the debug info AST together when we import something that we also have (in better quality from the C++ module) in the target ASTContext.

Thank you for the explanation.
Now I understand you directly create specializations from the std module and intercept the import to avoid importing broken specializations from the debug info, this makes sense.

Really, just one last question to see the whole picture: If clang::ASTImporter were capable of handling a specialization/instantiation with missing info then we could use that. So the reason for this interception is that clang::ASTImporter::VisitClassTemplateSpecializationDecl cannot handle the specialization it receives because that or its dependent nodes has too many missing data, right?

Feel free to ask! I'm just traveling so my internet access (and my reply rate) is a bit low at the moment.

Not sure if we can easily merge the logic of D59537 into the ASTImporter. It has no logic at all to actually determine if a source AST node has missing data but solely relies on our knowledge that our AST in LLDB isn't very useful for std templates. Also instantiating a template instead of importing an existing instantiation seems like a very rare scenario. I'm not even sure how we would test having a broken AST with Clang (the parser won't produce one, so we would have to hack our own AST builder for broken nodes).

If there is a reasonable way to get this logic into the ASTImporter I have no objection against doing so. The only problem I see is that I can't come up with a reasonable way of merging/testing the logic in the ASTImporter (and what other application would benefit from it).

[martong]

In D59485#1444673, @teemperor wrote:

In D59485#1439628, @martong wrote:

In D59485#1439570, @martong wrote:

In D59485#1439390, @teemperor wrote:

Well, I still don't understand how LLDB synthesis the AST.
So you have a C++ module in a .pcm file. This means you could create an AST from that with ASTReader from it's .clang_ast section, right? In that case the AST should be complete with all type information. If this was the case then I don't see why it is not possible to use clang::ASTImporter to import templates and specializations, since we do exactly that in CTU.

Or do you guys create the AST from the debug info, from the .debug_info section of .pcm module file? And this is why AST is incomplete? I've got this from https://youtu.be/EgkZ8PTNSHQ?list=PL85Cf-ok7HpppFWlp2wX_-A1dkyFVlaEB&t=1565
If this is the case, then comes my naiive quiestion: Why don't you use the ASTReader?

There are no C++ modules involved in our use case beside the generic std module. No user-written code is read from a module and we don't have any PCM file beside the std module we build for the expression evaluator.

We use the ASTReader to directly read the std module contents into the expression evaluation ASTContext, but this doesn't give us the decls for the instantiations the user has made (e.g. std::vector<Foo>). We only have these user instantiations in the 'normal' debug info where we have a rather minimal AST (again, no modules are not involved in building this debug info AST). The InternalImport in the LLDB patch just stitches the module AST and the debug info AST together when we import something that we also have (in better quality from the C++ module) in the target ASTContext.

Thank you for the explanation.
Now I understand you directly create specializations from the std module and intercept the import to avoid importing broken specializations from the debug info, this makes sense.

Really, just one last question to see the whole picture: If clang::ASTImporter were capable of handling a specialization/instantiation with missing info then we could use that. So the reason for this interception is that clang::ASTImporter::VisitClassTemplateSpecializationDecl cannot handle the specialization it receives because that or its dependent nodes has too many missing data, right?

Feel free to ask! I'm just traveling so my internet access (and my reply rate) is a bit low at the moment.

Not sure if we can easily merge the logic of D59537 into the ASTImporter. It has no logic at all to actually determine if a source AST node has missing data but solely relies on our knowledge that our AST in LLDB isn't very useful for std templates. Also instantiating a template instead of importing an existing instantiation seems like a very rare scenario. I'm not even sure how we would test having a broken AST with Clang (the parser won't produce one, so we would have to hack our own AST builder for broken nodes).

If there is a reasonable way to get this logic into the ASTImporter I have no objection against doing so. The only problem I see is that I can't come up with a reasonable way of merging/testing the logic in the ASTImporter (and what other application would benefit from it).

Raphael, thank you for your answer. This explanation makes so much more easier for me to understand the need for this patch.

[teemperor]

• Addressed (most of) Aleksei's comments.

clang/unittests/AST/ASTImporterTest.cpp

590 ↗

(On Diff #191765)

Not sure if I can follow?

[martong]

clang/unittests/AST/ASTImporterTest.cpp

590 ↗

(On Diff #191765)

I think Alexey's comment relates to the parameter in ASTImporterOptionSpecificTestBase::ImporterConstructor. There is a typo in the name of the parameter an 'e' is missing: ASTImporterLookupTable *LookupTabl

[teemperor]

Ah, I get it now. Fixed!

[a_sidorin]

Hello Raphael,
I think we should accept this change. I don't see an easy way to merge the LLDB stuff into ASTImporter; also, this patch provides a good extension point for ASTImporter since it is designed to be a parent class. @martong @shafik Gabor, Shafik, what do you think?

[martong]

In D59485#1446950, @a_sidorin wrote:

Hello Raphael,
I think we should accept this change. I don't see an easy way to merge the LLDB stuff into ASTImporter; also, this patch provides a good extension point for ASTImporter since it is designed to be a parent class. @martong @shafik Gabor, Shafik, what do you think?

Yes, I agree.

[teemperor]

• Added RegisterImportedDecl call as suggested in D59537
• Added an assert that checks that we register any decl that an overwritten ImportInternal would create.
• Fixed the tests that they no longer trigger that new assert.

[martong]

Looks good, thanks!

[martong]

clang/unittests/AST/ASTImporterTest.cpp

625 ↗

(On Diff #196955)

I just realized that maybe we could simplify and make this part more elegant:
If ASTImporterTestBase had a data member with the type ImporterConstructor then we could overwrite that data member in the constructor of CustomImporter.
(The ASTImporter is created lazily when the first getImportedDecl is called, but by that time the creator is already set, so this will work.)
This way we could elide the last argument here (RedirectingImporter::Constructor) and in all subsequent TEST_P tests, it would be enough to set up the creator per test fixture.

[teemperor]

• Refactored test according to Gábor's feedback.

[teemperor]

I'll land this as it seems only the tests are subject to change and I want to also land the dependencies of this patch. Please let me know if you want any other changes changes to the test and thanks for the review!

include/clang/AST/ASTImporter.h

private:
     /// (which we have already complained about).
     NonEquivalentDeclSet NonEquivalentDecls;
 
     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);
+
   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.
     ///
public:
     ///
     /// \param D A declaration in the "to" context.
     virtual void CompleteDecl(Decl* D);
 
     /// Subclasses can override this function to observe all of the \c From ->
     /// \c To declaration mappings as they are imported.
     virtual void Imported(Decl *From, Decl *To) {}
 
+    void RegisterImportedDecl(Decl *FromD, Decl *ToD);
+
     /// Store and assign the imported declaration to its counterpart.
     Decl *MapImported(Decl *From, Decl *To);
 
     /// Called by StructuralEquivalenceContext.  If a RecordDecl is
     /// being compared to another RecordDecl as part of import, completing the
     /// other RecordDecl may trigger importation of the first RecordDecl. This
     /// happens especially for anonymous structs.  If the original of the second
     /// RecordDecl can be found, we can complete it without the need for

lib/AST/ASTImporter.cpp

GetImportedOrCreateSpecialDecl(ToDeclT *&ToD, CreateFunT CreateFun,
       //  ToD = nullptr;
       //  return true; // Already imported but with error.
       //}
       ToD = cast_or_null<ToDeclT>(Importer.GetAlreadyImportedOrNull(FromD));
       if (ToD)
         return true; // Already imported.
       ToD = CreateFun(std::forward<Args>(args)...);
       // Keep track of imported Decls.
-      Importer.MapImported(FromD, ToD);
-      Importer.AddToLookupTable(ToD);
+      Importer.RegisterImportedDecl(FromD, ToD);
       InitializeImportedDecl(FromD, ToD);
       return false; // A new Decl is created.
     }
 
     void InitializeImportedDecl(Decl *FromD, Decl *ToD) {
       ToD->IdentifierNamespace = FromD->IdentifierNamespace;
       if (FromD->hasAttrs())
         for (const Attr *FromAttr : FromD->getAttrs()) {

 }
 
 void ASTImporter::AddToLookupTable(Decl *ToD) {
   if (LookupTable)
     if (auto *ToND = dyn_cast<NamedDecl>(ToD))
       LookupTable->add(ToND);
 }
 
+Expected<Decl *> ASTImporter::ImportImpl(Decl *FromD) {
+  // Import the decl using ASTNodeImporter.
+  ASTNodeImporter Importer(*this);
+  return Importer.Visit(FromD);
+}
+
+void ASTImporter::RegisterImportedDecl(Decl *FromD, Decl *ToD) {
+  MapImported(FromD, ToD);
+  AddToLookupTable(ToD);
+}
+
 Expected<QualType> ASTImporter::Import_New(QualType FromT) {
   if (FromT.isNull())
     return QualType{};
 
   const Type *FromTy = FromT.getTypePtr();
 
   // Check whether we've already imported this type.
   llvm::DenseMap<const Type *, const Type *>::iterator Pos
if (FromDPos == ImportedFromDecls.end())
     return nullptr;
   return FromDPos->second->getTranslationUnitDecl();
 }
 
 Expected<Decl *> ASTImporter::Import_New(Decl *FromD) {
   if (!FromD)
     return nullptr;
 
-  ASTNodeImporter Importer(*this);
 
   // Check whether we've already imported this declaration.
   Decl *ToD = GetAlreadyImportedOrNull(FromD);
   if (ToD) {
     // If FromD has some updated flags after last import, apply it
     updateFlags(FromD, ToD);
     return ToD;
   }
 
   // Import the declaration.
-  ExpectedDecl ToDOrErr = Importer.Visit(FromD);
+  ExpectedDecl ToDOrErr = ImportImpl(FromD);
   if (!ToDOrErr)
     return ToDOrErr;
   ToD = *ToDOrErr;
 
   // FIXME Use getImportDeclErrorIfAny() here (and return with the error) once
   // the error handling is finished in GetImportedOrCreateSpecialDecl().
   if (!ToD) {
     return nullptr;
   }
 
+  // Make sure that ImportImpl registered the imported decl.
+  assert(ImportedDecls.count(FromD) != 0 && "Missing call to MapImported?");
+
   // Once the decl is connected to the existing declarations, i.e. when the
   // redecl chain is properly set then we populate the lookup again.
   // This way the primary context will be able to find all decls.
   AddToLookupTable(ToD);
 
   // Notify subclasses.
   Imported(FromD, ToD);
 

unittests/AST/ASTImporterTest.cpp

 // This class provides generic methods to write tests which can check internal
 // attributes of AST nodes like getPreviousDecl(), isVirtual(), etc. Also,
 // this fixture makes it possible to import from several "From" contexts.
 class ASTImporterTestBase : public CompilerOptionSpecificTest {
 
   const char *const InputFileName = "input.cc";
   const char *const OutputFileName = "output.cc";
 
+public:
+  /// Allocates an ASTImporter (or one of its subclasses).
+  typedef std::function<ASTImporter *(ASTContext &, FileManager &, ASTContext &,
+                                      FileManager &, bool,
+                                      ASTImporterLookupTable *)>
+      ImporterConstructor;
+
+  // The lambda that constructs the ASTImporter we use in this test.
+  ImporterConstructor Creator;
+
+private:
   // Buffer for the To context, must live in the test scope.
   std::string ToCode;
 
   // Represents a "From" translation unit and holds an importer object which we
   // use to import from this translation unit.
   struct TU {
     // Buffer for the context, must live in the test scope.
     std::string Code;
     std::string FileName;
     std::unique_ptr<ASTUnit> Unit;
     TranslationUnitDecl *TUDecl = nullptr;
     std::unique_ptr<ASTImporter> Importer;
-    TU(StringRef Code, StringRef FileName, ArgVector Args)
+    ImporterConstructor Creator;
+    TU(StringRef Code, StringRef FileName, ArgVector Args,
+       ImporterConstructor C = ImporterConstructor())
         : Code(Code), FileName(FileName),
           Unit(tooling::buildASTFromCodeWithArgs(this->Code, Args,
                                                  this->FileName)),
-          TUDecl(Unit->getASTContext().getTranslationUnitDecl()) {
+          TUDecl(Unit->getASTContext().getTranslationUnitDecl()), Creator(C) {
       Unit->enableSourceFileDiagnostics();
+
+      // If the test doesn't need a specific ASTImporter, we just create a
+      // normal ASTImporter with it.
+      if (!Creator)
+        Creator = [](ASTContext &ToContext, FileManager &ToFileManager,
+                     ASTContext &FromContext, FileManager &FromFileManager,
+                     bool MinimalImport, ASTImporterLookupTable *LookupTable) {
+          return new ASTImporter(ToContext, ToFileManager, FromContext,
+                                 FromFileManager, MinimalImport, LookupTable);
+        };
     }
 
     void lazyInitImporter(ASTImporterLookupTable &LookupTable, ASTUnit *ToAST) {
       assert(ToAST);
-      if (!Importer) {
-        Importer.reset(
-            new ASTImporter(ToAST->getASTContext(), ToAST->getFileManager(),
+      if (!Importer)
+        Importer.reset(Creator(ToAST->getASTContext(), ToAST->getFileManager(),
                                Unit->getASTContext(), Unit->getFileManager(),
                                false, &LookupTable));
-      }
       assert(&ToAST->getASTContext() == &Importer->getToContext());
       createVirtualFileIfNeeded(ToAST, FileName, Code);
     }
 
     Decl *import(ASTImporterLookupTable &LookupTable, ASTUnit *ToAST,
                  Decl *FromDecl) {
       lazyInitImporter(LookupTable, ToAST);
       if (auto ImportedOrErr = Importer->Import_New(FromDecl))
public:
   // of the identifier into the To context.
   // 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) {
     ArgVector FromArgs = getArgVectorForLanguage(FromLang),
               ToArgs = getArgVectorForLanguage(ToLang);
 
-    FromTUs.emplace_back(FromSrcCode, InputFileName, FromArgs);
+    FromTUs.emplace_back(FromSrcCode, InputFileName, FromArgs, Creator);
     TU &FromTU = FromTUs.back();
 
     assert(!ToAST);
     lazyInitToAST(ToLang, ToSrcCode, OutputFileName);
 
     ASTContext &FromCtx = FromTU.Unit->getASTContext();
 
     IdentifierInfo *ImportedII = &FromCtx.Idents.get(Identifier);
TEST_P(CanonicalRedeclChain, ShouldBeSameForAllDeclInTheChain) {
   auto RedeclsD0 = getCanonicalForwardRedeclChain(D0);
   auto RedeclsD1 = getCanonicalForwardRedeclChain(D1);
   auto RedeclsD2 = getCanonicalForwardRedeclChain(D2);
 
   EXPECT_THAT(RedeclsD0, ::testing::ContainerEq(RedeclsD1));
   EXPECT_THAT(RedeclsD1, ::testing::ContainerEq(RedeclsD2));
 }
 
+namespace {
+struct RedirectingImporter : public ASTImporter {
+  using ASTImporter::ASTImporter;
+
+protected:
+  llvm::Expected<Decl *> ImportImpl(Decl *FromD) override {
+    auto *ND = dyn_cast<NamedDecl>(FromD);
+    if (!ND || ND->getName() != "shouldNotBeImported")
+      return ASTImporter::ImportImpl(FromD);
+    for (Decl *D : getToContext().getTranslationUnitDecl()->decls()) {
+      if (auto *ND = dyn_cast<NamedDecl>(D))
+        if (ND->getName() == "realDecl") {
+          RegisterImportedDecl(FromD, ND);
+          return ND;
+        }
+    }
+    return ASTImporter::ImportImpl(FromD);
+  }
+};
+
+} // namespace
+
+struct RedirectingImporterTest : ASTImporterOptionSpecificTestBase {
+  RedirectingImporterTest() {
+    Creator = [](ASTContext &ToContext, FileManager &ToFileManager,
+                 ASTContext &FromContext, FileManager &FromFileManager,
+                 bool MinimalImport, ASTImporterLookupTable *LookupTable) {
+      return new RedirectingImporter(ToContext, ToFileManager, FromContext,
+                                     FromFileManager, MinimalImport,
+                                     LookupTable);
+    };
+  }
+};
+
+// Test that an ASTImporter subclass can intercept an import call.
+TEST_P(RedirectingImporterTest, InterceptImport) {
+  Decl *From, *To;
+  std::tie(From, To) =
+      getImportedDecl("class shouldNotBeImported {};", Lang_CXX,
+                      "class realDecl {};", Lang_CXX, "shouldNotBeImported");
+  auto *Imported = cast<CXXRecordDecl>(To);
+  EXPECT_EQ(Imported->getQualifiedNameAsString(), "realDecl");
+
+  // Make sure our importer 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 custom ASTImporter
+// is still intercepting the import.
+TEST_P(RedirectingImporterTest, InterceptIndirectImport) {
+  Decl *From, *To;
+  std::tie(From, To) =
+      getImportedDecl("class shouldNotBeImported {};"
+                      "class F { shouldNotBeImported f; };",
+                      Lang_CXX, "class realDecl {};", Lang_CXX, "F");
+
+  // Make sure our ASTImporter 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);
+}
+
 TEST_P(ImportExpr, ImportStringLiteral) {
   MatchVerifier<Decl> Verifier;
   testImport(
       "void declToImport() { (void)\"foo\"; }",
       Lang_CXX, "", Lang_CXX, Verifier,
       functionDecl(hasDescendant(
           stringLiteral(hasType(asString("const char [4]"))))));
   testImport(
INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportType,
                         DefaultTestValuesForRunOptions, );
 
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportDecl,
                         DefaultTestValuesForRunOptions, );
 
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ASTImporterOptionSpecificTestBase,
                         DefaultTestValuesForRunOptions, );
 
+INSTANTIATE_TEST_CASE_P(ParameterizedTests, RedirectingImporterTest,
+                        DefaultTestValuesForRunOptions, );
+
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportFunctions,
                         DefaultTestValuesForRunOptions, );
 
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportFriendFunctionTemplates,
                         DefaultTestValuesForRunOptions, );
 
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportClasses,
                         DefaultTestValuesForRunOptions, );