Package CMO: Enable serializing decls imported with @_spiOnly or package import.

Starting in Swift 6.0, `package` access level and `@_spiOnly` attribute have been increasingly used in import statements.
However, existing import filtering prevented serialization of package APIs that included such decls, leading to a
significant drop in overall serialization. This PR removes these restrictive filters, and allows decls from SDK or system
modules to be included in serialization.

rdar://130788606

Author: elsh

include/swift/AST/Module.h

@@ -1046,11 +1046,6 @@ class ModuleDecl
   /// This assumes that \p module was imported.
   bool isImportedImplementationOnly(const ModuleDecl *module) const;
 
-  /// Returns true if decl context or its content can be serialized by
-  /// cross-module-optimization.
-  /// The \p ctxt can e.g. be a NominalType or the context of a function.
-  bool canBeUsedForCrossModuleOptimization(DeclContext *ctxt) const;
-
   /// Finds all top-level decls of this module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

lib/AST/Module.cpp

@@ -3047,33 +3047,6 @@ bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module) const {
   return true;
 }
 
-bool ModuleDecl::
-canBeUsedForCrossModuleOptimization(DeclContext *ctxt) const {
-  ModuleDecl *moduleOfCtxt = ctxt->getParentModule();
-
-  // If the context defined in the same module - or is the same module, it's
-  // fine.
-  if (moduleOfCtxt == this)
-    return true;
-
-  // See if context is imported in a "regular" way, i.e. not with
-  // @_implementationOnly, `package import` or @_spiOnly.
-  ModuleDecl::ImportFilter filter = {
-    ModuleDecl::ImportFilterKind::ImplementationOnly,
-    ModuleDecl::ImportFilterKind::PackageOnly,
-    ModuleDecl::ImportFilterKind::SPIOnly
-  };
-  SmallVector<ImportedModule, 4> results;
-  getImportedModules(results, filter);
-
-  auto &imports = getASTContext().getImportCache();
-  for (auto &desc : results) {
-    if (imports.isImportedBy(moduleOfCtxt, desc.importedModule))
-      return false;
-  }
-  return true;
-}
-
 void SourceFile::lookupImportedSPIGroups(
                         const ModuleDecl *importedModule,
                         llvm::SmallSetVector<Identifier, 4> &spiGroups) const {

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -16,6 +16,7 @@
 
 #define DEBUG_TYPE "cross-module-serialization-setup"
 #include "swift/AST/Module.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/IRGen/TBDGen.h"
 #include "swift/SIL/ApplySite.h"
@@ -103,6 +104,11 @@ class CrossModuleOptimization {
   bool canSerializeType(CanType type);
   bool canSerializeDecl(NominalTypeDecl *decl);
 
+  /// Check whether decls imported with certain access levels or attributes
+  /// can be serialized.
+  /// The \p ctxt can e.g. be a NominalType or the context of a function.
+  bool checkImports(DeclContext *ctxt) const;
+
   bool canUseFromInline(DeclContext *declCtxt);
 
   bool canUseFromInline(SILFunction *func);
@@ -734,7 +740,12 @@ static bool couldBeLinkedStatically(DeclContext *funcCtxt, SILModule &module) {
   // The stdlib module is always linked dynamically.
   if (funcModule == module.getASTContext().getStdlibModule())
     return false;
-    
+
+  // An sdk or system module should be linked dynamically.
+  if (isPackageCMOEnabled(module.getSwiftModule()) &&
+      funcModule->isNonUserModule())
+    return false;
+
   // Conservatively assume the function is in a statically linked module.
   return true;
 }
@@ -744,7 +755,7 @@ bool CrossModuleOptimization::canUseFromInline(DeclContext *declCtxt) {
   if (everything)
     return true;
 
-  if (!M.getSwiftModule()->canBeUsedForCrossModuleOptimization(declCtxt))
+  if (!checkImports(declCtxt))
     return false;
 
   /// If we are emitting a TBD file, the TBD file only contains public symbols
@@ -760,6 +771,52 @@ bool CrossModuleOptimization::canUseFromInline(DeclContext *declCtxt) {
   return true;
 }
 
+bool CrossModuleOptimization::checkImports(DeclContext *ctxt) const {
+  ModuleDecl *moduleOfCtxt = ctxt->getParentModule();
+
+  // If the context defined in the same module - or is the same module, it's
+  // fine.
+  if (moduleOfCtxt == M.getSwiftModule())
+    return true;
+
+  ModuleDecl::ImportFilter filter;
+
+  if (isPackageCMOEnabled(M.getSwiftModule())) {
+    // If Package CMO is enabled, decls imported with `package import`
+    // or `@_spiOnly import` into this module should be allowed to be
+    // serialized. They are used in decls with `package` or higher
+    // access level, with or without @_spi; a client of this module
+    // should be able to access them directly if in the same package.
+    filter = { ModuleDecl::ImportFilterKind::ImplementationOnly };
+  } else {
+    // See if context is imported in a "regular" way, i.e. not with
+    // @_implementationOnly, `package import` or @_spiOnly.
+    filter = {
+      ModuleDecl::ImportFilterKind::ImplementationOnly,
+      ModuleDecl::ImportFilterKind::PackageOnly,
+      ModuleDecl::ImportFilterKind::SPIOnly
+    };
+  }
+  SmallVector<ImportedModule, 4> results;
+  M.getSwiftModule()->getImportedModules(results, filter);
+
+  auto &imports = M.getSwiftModule()->getASTContext().getImportCache();
+  for (auto &desc : results) {
+    if (imports.isImportedBy(moduleOfCtxt, desc.importedModule)) {
+      // E.g. `@_implementationOnly import QuartzCore_Private.CALayerPrivate`
+      // imports `Foundation` as its transitive dependency module; use of a
+      // a `public` decl in `Foundation` such as `IndexSet` in a function
+      // signature should not block serialization in Package CMO given the
+      // function has `package` or higher access level.
+      if (isPackageCMOEnabled(M.getSwiftModule()) &&
+          moduleOfCtxt->isNonUserModule())
+          continue;
+      return false;
+    }
+  }
+  return true;
+}
+
 /// Returns true if the function \p func can be used from a serialized function.
 bool CrossModuleOptimization::canUseFromInline(SILFunction *function) {
   if (everything)

test/SILOptimizer/package-cmo-import-filter.swift

@@ -0,0 +1,60 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend %t/CoreA.swift \
+// RUN: -module-name=CoreA -package-name Pkg \
+// RUN: -parse-as-library -emit-module \
+// RUN: -emit-module-path %t/CoreA.swiftmodule -I%t \
+// RUN: -O -wmo -enable-library-evolution
+
+// RUN: %target-swift-frontend %t/CoreB.swift \
+// RUN: -module-name=CoreB -package-name Pkg \
+// RUN: -parse-as-library -emit-module \
+// RUN: -emit-module-path %t/CoreB.swiftmodule -I%t \
+// RUN: -O -wmo -enable-library-evolution
+
+// RUN: %target-swift-frontend %t/Lib.swift \
+// RUN: -module-name=Lib -package-name Pkg \
+// RUN: -parse-as-library -emit-module \
+// RUN: -experimental-spi-only-imports \
+// RUN: -emit-module-path %t/Lib.swiftmodule -I %t \
+// RUN: -experimental-package-cmo -experimental-allow-non-resilient-access \
+// RUN: -O -wmo -enable-library-evolution -Rmodule-loading 2> %t/Lib-result.txt
+// RUN: %target-sil-opt %t/Lib.swiftmodule -I %t -sil-verify-all -o %t/Lib.sil
+// RUN: %FileCheck %s < %t/Lib.sil
+
+// REQUIRES: swift_in_compiler
+
+
+//--- Lib.swift
+package import CoreA
+@_spiOnly public import CoreB
+
+/// PkgStruct is imported with `package import` and should be serialized.
+// CHECK-DAG: sil package [serialized_for_package] [canonical] @$s3Lib7libFuncyy5CoreA9PkgStructVF : $@convention(thin) (@in_guaranteed PkgStruct) -> () {
+package func libFunc(_ arg: PkgStruct) {
+  print(arg.pkgVar)
+}
+
+/// PubStruct is imported with `@_spiOnly public import` and should be serialized.
+// CHECK-DAG: sil [serialized_for_package] [canonical] @$s3Lib7spiFuncyy5CoreB15PubStructForSPIVF : $@convention(thin) (@in_guaranteed PubStructForSPI) -> () {
+@_spi(InCoreB)
+public func spiFunc(_ arg: PubStructForSPI) {
+  print(arg.pubVarForSPI)
+}
+
+//--- CoreA.swift
+package struct PkgStruct {
+  package var pkgVar: Int
+  package init(_ arg: Int) {
+    self.pkgVar = arg
+  }
+}
+
+//--- CoreB.swift
+public struct PubStructForSPI {
+  public var pubVarForSPI: String
+  public init(_ arg: String) {
+    self.pubVarForSPI = arg
+  }
+}