Merge pull request #35472 from eeckstein/cmo-implementationonly-fix

CrossModuleOptimization: fix crash when importing a module as implementationOnly

Author: eeckstein

include/swift/AST/Module.h

@@ -626,6 +626,10 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// This assumes that \p module was imported.
   bool isImportedImplementationOnly(const ModuleDecl *module) const;
 
+  /// Returns true if a function, which is using \p nominal, can be serialized
+  /// by cross-module-optimization.
+  bool canBeUsedForCrossModuleOptimization(NominalTypeDecl *nominal) const;
+
   /// Finds all top-level decls of this module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

lib/AST/Module.cpp

@@ -2055,6 +2055,30 @@ bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module) const {
   return true;
 }
 
+bool ModuleDecl::
+canBeUsedForCrossModuleOptimization(NominalTypeDecl *nominal) const {
+  ModuleDecl *moduleOfNominal = nominal->getParentModule();
+
+  // If the nominal is defined in the same module, it's fine.
+  if (moduleOfNominal == this)
+    return true;
+
+  // See if nominal is imported in a "regular" way, i.e. not with
+  // @_implementationOnly or @_spi.
+  ModuleDecl::ImportFilter filter = {
+    ModuleDecl::ImportFilterKind::Exported,
+    ModuleDecl::ImportFilterKind::Default};
+  SmallVector<ImportedModule, 4> results;
+  getImportedModules(results, filter);
+
+  auto &imports = getASTContext().getImportCache();
+  for (auto &desc : results) {
+    if (imports.isImportedBy(moduleOfNominal, desc.importedModule))
+      return true;
+  }
+  return false;
+}
+
 void SourceFile::lookupImportedSPIGroups(
                         const ModuleDecl *importedModule,
                         llvm::SmallSetVector<Identifier, 4> &spiGroups) const {

lib/SILOptimizer/IPO/CrossModuleSerializationSetup.cpp

@@ -45,6 +45,7 @@ class CrossModuleSerializationSetup {
   llvm::SmallVector<SILFunction *, 16> workList;
   llvm::SmallPtrSet<SILFunction *, 16> functionsHandled;
 
+  llvm::DenseMap<SILType, bool> typesChecked;
   llvm::SmallPtrSet<TypeBase *, 16> typesHandled;
 
   SILModule &M;
@@ -62,6 +63,8 @@ class CrossModuleSerializationSetup {
 
   bool canSerialize(SILInstruction *inst, bool lookIntoThunks);
 
+  bool canSerialize(SILType type);
+
   void setUpForSerialization(SILFunction *F);
 
   void prepareInstructionForSerialization(SILInstruction *inst);
@@ -320,6 +323,12 @@ bool CrossModuleSerializationSetup::canSerialize(SILFunction *F,
 
 bool CrossModuleSerializationSetup::canSerialize(SILInstruction *inst,
                                                  bool lookIntoThunks) {
+
+  for (SILValue result : inst->getResults()) {
+    if (!canSerialize(result->getType()))
+      return false;
+  }
+
   if (auto *FRI = dyn_cast<FunctionRefBaseInst>(inst)) {
     SILFunction *callee = FRI->getReferencedFunctionOrNull();
     return canUseFromInline(callee, lookIntoThunks);
@@ -344,6 +353,31 @@ bool CrossModuleSerializationSetup::canSerialize(SILInstruction *inst,
   return true;
 }
 
+bool CrossModuleSerializationSetup::canSerialize(SILType type) {
+  auto iter = typesChecked.find(type);
+  if (iter != typesChecked.end())
+    return iter->getSecond();
+
+  ModuleDecl *mod = M.getSwiftModule();
+  bool success = !type.getASTType().findIf(
+    [mod](Type rawSubType) {
+      CanType subType = rawSubType->getCanonicalType();
+      if (NominalTypeDecl *subNT = subType->getNominalOrBoundGenericNominal()) {
+      
+        // Exclude types which are defined in an @_implementationOnly imported
+        // module. Such modules are not transitively available.
+        if (!mod->canBeUsedForCrossModuleOptimization(subNT)) {
+          llvm::dbgs() << " === " << mod->getName() << ", " <<
+            subNT->getParentModule()->getName() << ", " << subNT->getName() << '\n';
+          return true;
+        }
+      }
+      return false;
+    });
+  typesChecked[type] = success;
+  return success;
+}
+
 /// Returns true if the function \p func can be used from a serialized function.
 ///
 /// If \p lookIntoThunks is true, serializable shared thunks are also accepted.

test/SILOptimizer/Inputs/cross-module.swift

@@ -1,4 +1,5 @@
 import Submodule
+@_implementationOnly import PrivateSubmodule
 
 private enum PE<T> {
   case A
@@ -267,4 +268,11 @@ public func callUnrelated<T>(_ t: T) -> T {
   return t
 }
 
+public func callImplementationOnly<T>(_ t: T) -> T {
+  let p = PrivateStr(i: 27)
+  print(p.test())
+  return t
+}
+
+
 public let globalLet = 529387

test/SILOptimizer/Inputs/cross-private-submodule.swift

@@ -0,0 +1,14 @@
+
+public struct PrivateStr {
+  var i: Int
+
+  public init(i: Int) {
+    self.i = i
+  }
+
+  @inline(never)
+  public func test() -> Int {
+    return i
+  }
+}
+

test/SILOptimizer/cross-module-optimization.swift

@@ -2,16 +2,17 @@
 
 // RUN: %empty-directory(%t) 
 // RUN: %target-build-swift -O -wmo -parse-as-library -cross-module-optimization -emit-module -emit-module-path=%t/Submodule.swiftmodule -module-name=Submodule %S/Inputs/cross-submodule.swift -c -o %t/submodule.o
+// RUN: %target-build-swift -O -wmo -parse-as-library -cross-module-optimization -emit-module -emit-module-path=%t/PrivateSubmodule.swiftmodule -module-name=PrivateSubmodule %S/Inputs/cross-private-submodule.swift -c -o %t/privatesubmodule.o
 // RUN: %target-build-swift -O -wmo -parse-as-library -cross-module-optimization -emit-module -emit-module-path=%t/Test.swiftmodule -module-name=Test -I%t %S/Inputs/cross-module.swift -c -o %t/test.o
 // RUN: %target-build-swift -O -wmo -module-name=Main -I%t %s -c -o %t/main.o
-// RUN: %target-swiftc_driver %t/main.o %t/test.o %t/submodule.o -o %t/a.out
+// RUN: %target-swiftc_driver %t/main.o %t/test.o %t/submodule.o %t/privatesubmodule.o -o %t/a.out
 // RUN: %target-codesign %t/a.out
 // RUN: %target-run %t/a.out | %FileCheck %s -check-prefix=CHECK-OUTPUT
 
 // Check if it also works if the main module is compiled with -Onone:
 
 // RUN: %target-build-swift -Onone -wmo -module-name=Main -I%t %s -c -o %t/main-onone.o
-// RUN: %target-swiftc_driver %t/main-onone.o %t/test.o %t/submodule.o -o %t/a.out
+// RUN: %target-swiftc_driver %t/main-onone.o %t/test.o %t/submodule.o %t/privatesubmodule.o -o %t/a.out
 // RUN: %target-codesign %t/a.out
 // RUN: %target-run %t/a.out | %FileCheck %s -check-prefix=CHECK-OUTPUT
 
@@ -29,19 +30,19 @@ import Test
 func testNestedTypes() {
   let c = Container()
 
-  // CHECK-OUTPUT: [Test.Container.Base]
+  // CHECK-OUTPUT [Test.Container.Base]
   // CHECK-OUTPUT: 27
   // CHECK-SIL-DAG: sil shared [noinline] @$s4Test9ContainerV9testclassyxxlFSi_Tg5
   print(c.testclass(27))
-  // CHECK-OUTPUT: [Test.Container.Base]
+  // CHECK-OUTPUT [Test.Container.Base]
   // CHECK-OUTPUT: 27
   // CHECK-SIL-DAG: sil shared_external {{.*}} @$s4Test9ContainerV13testclass_genyxxlF
   print(c.testclass_gen(27))
-  // CHECK-OUTPUT: [Test.PE<Swift.Int>.B(27)]
+  // CHECK-OUTPUT [Test.PE<Swift.Int>.B(27)]
   // CHECK-OUTPUT: 27
   // CHECK-SIL-DAG: sil shared [noinline] @$s4Test9ContainerV8testenumyxxlFSi_Tg5
   print(c.testenum(27))
-  // CHECK-OUTPUT: [Test.PE<Swift.Int>.B(27)]
+  // CHECK-OUTPUT [Test.PE<Swift.Int>.B(27)]
   // CHECK-OUTPUT: 27
   // CHECK-SIL-DAG: sil shared_external {{.*}} @$s4Test9ContainerV12testenum_genyxxlF
   print(c.testenum_gen(27))
@@ -141,6 +142,15 @@ func testGlobal() {
   // CHECK-SIL2: } // end sil function '$s4Main10testGlobalyyF'
 }
 
+// CHECK-SIL2-LABEL: sil hidden [noinline] @$s4Main22testImplementationOnlyyyF
+@inline(never)
+func testImplementationOnly() {
+  // CHECK-OUTPUT: 27
+  // CHECK-SIL2: function_ref @$s4Test22callImplementationOnlyyxxlF
+  print(callImplementationOnly(27))
+  // CHECK-SIL2: } // end sil function '$s4Main22testImplementationOnlyyyF'
+}
+
 testNestedTypes()
 testClass()
 testError()
@@ -150,4 +160,5 @@ testClosures()
 testKeypath()
 testMisc()
 testGlobal()
+testImplementationOnly()