Merge pull request #40860 from apple/QuietMisdreavus/display-decl-recur

[AST] check modules before recursing for display decls

rdar://87601741

Author: QuietMisdreavus

include/swift/AST/FileUnit.h

@@ -16,6 +16,7 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/RawComment.h"
 #include "swift/Basic/BasicSourceInfo.h"
+#include "swift/Basic/Debug.h"
 
 #include "llvm/ADT/PointerIntPair.h"
 
@@ -188,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.
@@ -308,6 +311,9 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
     return getParentModule()->getRealName().str();
   }
 
+  SWIFT_DEBUG_DUMPER(dumpDisplayDecls());
+  SWIFT_DEBUG_DUMPER(dumpTopLevelDecls());
+
   /// Traverse the decls within this file.
   ///
   /// \returns true if traversal was aborted, false if it completed

include/swift/AST/Module.h

@@ -26,6 +26,7 @@
 #include "swift/AST/Type.h"
 #include "swift/Basic/BasicSourceInfo.h"
 #include "swift/Basic/Compiler.h"
+#include "swift/Basic/Debug.h"
 #include "swift/Basic/OptionSet.h"
 #include "swift/Basic/STLExtras.h"
 #include "swift/Basic/SourceLoc.h"
@@ -759,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.
   ///
@@ -856,6 +866,9 @@ class ModuleDecl
   /// transferred from module files to the dSYMs, remove this.
   bool isExternallyConsumed() const;
 
+  SWIFT_DEBUG_DUMPER(dumpDisplayDecls());
+  SWIFT_DEBUG_DUMPER(dumpTopLevelDecls());
+
   SourceRange getSourceRange() const { return SourceRange(); }
 
   static bool classof(const DeclContext *DC) {

include/swift/AST/SourceFile.h

@@ -288,6 +288,10 @@ class SourceFile final : public FileUnit {
 
   ~SourceFile();
 
+  bool hasImports() const {
+    return Imports.hasValue();
+  }
+
   /// Retrieve an immutable view of the source file's imports.
   ArrayRef<AttributedImport<ImportedModule>> getImports() const {
     return *Imports;

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,6 +780,30 @@ 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) &&
+              !Imports.contains(import.module.importedModule) &&
+              import.module.importedModule->shouldCollectDisplayDecls()) {
+            Imports.insert(import.module.importedModule);
+          }
+        }
+      }
+    }
+  }
+}
+
 void ModuleDecl::getLocalTypeDecls(SmallVectorImpl<TypeDecl*> &Results) const {
   FORWARD(getLocalTypeDecls, (Results));
 }
@@ -788,6 +812,24 @@ void ModuleDecl::getTopLevelDecls(SmallVectorImpl<Decl*> &Results) const {
   FORWARD(getTopLevelDecls, (Results));
 }
 
+void ModuleDecl::dumpDisplayDecls() const {
+  SmallVector<Decl *, 32> Decls;
+  getDisplayDecls(Decls);
+  for (auto *D : Decls) {
+    D->dump(llvm::errs());
+    llvm::errs() << "\n";
+  }
+}
+
+void ModuleDecl::dumpTopLevelDecls() const {
+  SmallVector<Decl *, 32> Decls;
+  getTopLevelDecls(Decls);
+  for (auto *D : Decls) {
+    D->dump(llvm::errs());
+    llvm::errs() << "\n";
+  }
+}
+
 void ModuleDecl::getExportedPrespecializations(
     SmallVectorImpl<Decl *> &Results) const {
   FORWARD(getExportedPrespecializations, (Results));
@@ -908,8 +950,23 @@ SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
 }
 
 void ModuleDecl::getDisplayDecls(SmallVectorImpl<Decl*> &Results) const {
+  if (isParsedModule(this)) {
+    SmallPtrSet<ModuleDecl *, 4> Modules;
+    collectParsedExportedImports(this, Modules);
+    for (const ModuleDecl *import : Modules) {
+      import->getDisplayDecls(Results);
+    }
+  }
   // FIXME: Should this do extra access control filtering?
   FORWARD(getDisplayDecls, (Results));
+
+#ifndef NDEBUG
+  llvm::DenseSet<Decl *> visited;
+  for (auto *D : Results) {
+    auto inserted = visited.insert(D).second;
+    assert(inserted && "there should be no duplicate decls");
+  }
+#endif
 }
 
 ProtocolConformanceRef
@@ -3066,6 +3123,22 @@ void FileUnit::getTopLevelDeclsWhereAttributesMatch(
   Results.erase(newEnd, Results.end());
 }
 
+void FileUnit::dumpDisplayDecls() const {
+  SmallVector<Decl *, 32> Decls;
+  getDisplayDecls(Decls);
+  for (auto *D : Decls) {
+    D->dump(llvm::errs());
+  }
+}
+
+void FileUnit::dumpTopLevelDecls() const {
+  SmallVector<Decl *, 32> Decls;
+  getTopLevelDecls(Decls);
+  for (auto *D : Decls) {
+    D->dump(llvm::errs());
+  }
+}
+
 void swift::simple_display(llvm::raw_ostream &out, const FileUnit *file) {
   if (!file) {
     out << "(null)";

lib/ClangImporter/ClangImporter.cpp

@@ -3124,6 +3124,9 @@ class VectorDeclPtrConsumer : public swift::VisibleDeclConsumer {
 };
 } // unnamed namespace
 
+// FIXME: should submodules still be crawled for the symbol graph? (SR-15753)
+bool ClangModuleUnit::shouldCollectDisplayDecls() const { return isTopLevel(); }
+
 void ClangModuleUnit::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
   VectorDeclPtrConsumer consumer(results);
   FilteringDeclaredDeclConsumer filterConsumer(consumer, this);

test/SymbolGraph/ClangImporter/EmitWhileBuilding.swift

@@ -1,15 +1,19 @@
 // RUN: %empty-directory(%t)
 // RUN: cp -r %S/Inputs/EmitWhileBuilding/EmitWhileBuilding.framework %t
-// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -enable-objc-interop -emit-module-path %t/EmitWhileBuilding.framework/Modules/EmitWhileBuilding.swiftmodule/%target-swiftmodule-name -import-underlying-module -F %t -module-name EmitWhileBuilding -disable-objc-attr-requires-foundation-module %s -emit-symbol-graph -emit-symbol-graph-dir %t
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -enable-objc-interop -emit-module-path %t/EmitWhileBuilding.framework/Modules/EmitWhileBuilding.swiftmodule/%target-swiftmodule-name -import-underlying-module -F %t -module-name EmitWhileBuilding -disable-objc-attr-requires-foundation-module %s %S/Inputs/EmitWhileBuilding/Extra.swift -emit-symbol-graph -emit-symbol-graph-dir %t
 // RUN: %{python} -m json.tool %t/EmitWhileBuilding.symbols.json %t/EmitWhileBuilding.formatted.symbols.json
 // RUN: %FileCheck %s --input-file %t/EmitWhileBuilding.formatted.symbols.json
+// RUN: %FileCheck %s --input-file %t/EmitWhileBuilding.formatted.symbols.json --check-prefix HEADER
 
 // REQUIRES: objc_interop
 
 import Foundation
 
 public enum SwiftEnum {}
 
+// HEADER:       "precise": "c:@testVariable"
+
+// CHECK:        "precise": "s:17EmitWhileBuilding9SwiftEnumO",
 // CHECK:        "declarationFragments": [
 // CHECK-NEXT:       {
 // CHECK-NEXT:           "kind": "keyword",

test/SymbolGraph/ClangImporter/Inputs/EmitWhileBuilding/Extra.swift

@@ -0,0 +1 @@
+public struct SomeStruct {}

test/SymbolGraph/ClangImporter/Inputs/Submodules/Mixed.h

@@ -0,0 +1 @@
+double outerVar = 1.0;

test/SymbolGraph/ClangImporter/Inputs/Submodules/Submodule.h

@@ -0,0 +1 @@
+double innerVar = 2.0;

test/SymbolGraph/ClangImporter/Inputs/Submodules/module.modulemap

@@ -0,0 +1,9 @@
+module Mixed {
+  header "Mixed.h"
+  export *
+
+  explicit module Submodule {
+    header "Submodule.h"
+    export *
+  }
+}

test/SymbolGraph/ClangImporter/Submodules.swift

@@ -0,0 +1,11 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -I %S/Inputs/Submodules -emit-module-path %t/Submodules.swiftmodule -enable-objc-interop -module-name Submodules %s -emit-symbol-graph -emit-symbol-graph-dir %t
+
+// REQUIRES: objc_interop
+
+// Don't crash when a module declared an `@_exported import` for a Clang non-top-level module.
+
+@_exported import Mixed
+@_exported import Mixed.Submodule
+
+public func someFunc() {}