[AST] check modules before recursing for display decls

Author: QuietMisdreavus

include/swift/AST/FileUnit.h

@@ -189,6 +189,8 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
   virtual Identifier
   getDiscriminatorForPrivateValue(const ValueDecl *D) const = 0;
 
+  virtual bool shouldCollectDisplayDecls() const { return true; }
+
   /// Finds all top-level decls in this file.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/AST/Module.h

@@ -760,6 +760,15 @@ class ModuleDecl
   /// The order of the results is not guaranteed to be meaningful.
   void getPrecedenceGroups(SmallVectorImpl<PrecedenceGroupDecl*> &Results) const;
 
+  /// Determines whether this module should be recursed into when calling
+  /// \c getDisplayDecls.
+  ///
+  /// Some modules should not call \c getDisplayDecls, due to assertions
+  /// in their implementation. These are usually implicit imports that would be
+  /// recursed into for parsed modules. This function provides a guard against
+  /// recusing into modules that should not have decls collected.
+  bool shouldCollectDisplayDecls() const;
+
   /// Finds all top-level decls that should be displayed to a client of this
   /// module.
   ///

include/swift/ClangImporter/ClangModule.h

@@ -87,6 +87,8 @@ class ClangModuleUnit final : public LoadedFile {
          ObjCSelector selector,
          SmallVectorImpl<AbstractFunctionDecl *> &results) const override;
 
+  virtual bool shouldCollectDisplayDecls() const override;
+
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
   virtual void getDisplayDecls(SmallVectorImpl<Decl*> &results) const override;

lib/AST/Module.cpp

@@ -780,15 +780,23 @@ void SourceFile::lookupObjCMethods(
   results.append(known->second.begin(), known->second.end());
 }
 
+bool ModuleDecl::shouldCollectDisplayDecls() const {
+  for (const FileUnit *file : Files) {
+    if (!file->shouldCollectDisplayDecls())
+      return false;
+  }
+  return true;
+}
+
 static void collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports) {
   for (const FileUnit *file : M->getFiles()) {
     if (const SourceFile *source = dyn_cast<SourceFile>(file)) {
       if (source->hasImports()) {
         for (auto import : source->getImports()) {
-          if (import.options.contains(ImportFlags::Exported)) {
-            if (!Imports.contains(import.module.importedModule)) {
-              Imports.insert(import.module.importedModule);
-            }
+          if (import.options.contains(ImportFlags::Exported) &&
+              !Imports.contains(import.module.importedModule) &&
+              import.module.importedModule->shouldCollectDisplayDecls()) {
+            Imports.insert(import.module.importedModule);
           }
         }
       }

lib/ClangImporter/ClangImporter.cpp

@@ -2998,6 +2998,8 @@ class VectorDeclPtrConsumer : public swift::VisibleDeclConsumer {
 };
 } // unnamed namespace
 
+bool ClangModuleUnit::shouldCollectDisplayDecls() const { return isTopLevel(); }
+
 void ClangModuleUnit::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
   VectorDeclPtrConsumer consumer(results);
   FilteringDeclaredDeclConsumer filterConsumer(consumer, this);