Serialization: Error on leaked cross-references to @_implementationOnly dependencies

include/swift/AST/DiagnosticsSema.def

@@ -887,6 +887,10 @@ REMARK(serialization_skipped_invalid_type_unknown_name,none,
 ERROR(serialization_failed,none,
       "serialization of module %0 failed due to the errors above",
       (const ModuleDecl *))
+ERROR(serialization_xref_to_hidden_dependency,none,
+      "invalid reference to implementation-only imported module %0"
+      "%select{| for %1}1",
+      (const ModuleDecl *, const Decl *))
 
 WARNING(can_import_invalid_swiftmodule,none,
         "canImport() evaluated to false due to invalid swiftmodule: %0", (StringRef))

include/swift/AST/Module.h

@@ -1064,11 +1064,15 @@ class ModuleDecl
   void
   getImportedModulesForLookup(SmallVectorImpl<ImportedModule> &imports) const;
 
-  /// Has \p module been imported via an '@_implementationOnly' import
-  /// instead of another kind of import?
+  /// Has \p module been imported via an '@_implementationOnly' import and
+  /// not by anything more visible?
   ///
-  /// This assumes that \p module was imported.
-  bool isImportedImplementationOnly(const ModuleDecl *module) const;
+  /// If \p assumeImported, assume that \p module was imported and avoid the
+  /// work to confirm it is imported at all. Transitive modules not reexported
+  /// are not considered imported here and may lead to false positive without
+  /// this setting.
+  bool isImportedImplementationOnly(const ModuleDecl *module,
+      bool assumeImported = true) const;
 
   /// Finds all top-level decls of this module.
   ///

include/swift/Basic/Features.def

@@ -557,6 +557,10 @@ EXPERIMENTAL_FEATURE(AnyAppleOSAvailability, true)
 /// Check @_implementationOnly imports in non-library-evolution mode.
 EXPERIMENTAL_FEATURE(CheckImplementationOnly, true)
 
+/// Extends CheckImplementationOnly with an extra check at serialization to
+/// report references to implementation-only imported modules.
+EXPERIMENTAL_FEATURE(CheckImplementationOnlyStrict, false)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE

include/swift/Serialization/SerializationOptions.h

@@ -163,6 +163,7 @@ class SerializationOptions {
   bool HermeticSealAtLink = false;
   bool EmbeddedSwiftModule = false;
   bool SkipNonExportableDecls = false;
+  bool SkipImplementationOnlyDecls = false;
   bool ExplicitModuleBuild = false;
   bool EnableSerializationRemarks = false;
   bool IsInterfaceSDKRelative = false;

lib/AST/FeatureSet.cpp

@@ -145,6 +145,7 @@ UNINTERESTING_FEATURE(GroupActorErrors)
 UNINTERESTING_FEATURE(SameElementRequirements)
 UNINTERESTING_FEATURE(SendingArgsAndResults)
 UNINTERESTING_FEATURE(CheckImplementationOnly)
+UNINTERESTING_FEATURE(CheckImplementationOnlyStrict)
 
 static bool findUnderscoredLifetimeAttr(Decl *decl) {
   auto hasUnderscoredLifetimeAttr = [](Decl *decl) {

lib/AST/Module.cpp

@@ -3078,7 +3078,8 @@ void ModuleDecl::setPackageName(Identifier name) {
   Package = PackageUnit::create(name, *this, getASTContext());
 }
 
-bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module) const {
+bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module,
+    bool assumeImported) const {
   if (module == this) return false;
 
   auto &imports = getASTContext().getImportCache();
@@ -3099,7 +3100,17 @@ bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module) const {
       return false;
   }
 
-  return true;
+  if (assumeImported)
+    return true;
+
+  results.clear();
+  getImportedModules(results,
+      {ModuleDecl::ImportFilterKind::ImplementationOnly});
+  for (auto &desc : results)
+    if (imports.isImportedBy(module, desc.importedModule))
+      return true;
+
+  return false;
 }
 
 void SourceFile::lookupImportedSPIGroups(

lib/Frontend/CompilerInvocation.cpp

@@ -1897,6 +1897,10 @@ static bool ParseLangArgs(LangOptions &Opts, ArgList &Args,
     Opts.enableFeature(Feature::NoExplicitNonIsolated);
   }
 
+  if (Opts.hasFeature(Feature::CheckImplementationOnlyStrict) &&
+      !::getenv("SWIFT_DISABLE_IMPLICIT_CHECK_IMPLEMENTATION_ONLY"))
+    Opts.enableFeature(Feature::CheckImplementationOnly);
+
 #if !defined(NDEBUG) && SWIFT_ENABLE_EXPERIMENTAL_PARSER_VALIDATION
   /// Enable round trip parsing via the new swift parser unless it is disabled
   /// explicitly. The new Swift parser can have mismatches with C++ parser -

lib/Frontend/Frontend.cpp

@@ -284,6 +284,9 @@ SerializationOptions CompilerInvocation::computeSerializationOptions(
   serializationOpts.SkipNonExportableDecls =
       getLangOptions().SkipNonExportableDecls;
 
+  serializationOpts.SkipImplementationOnlyDecls =
+      getLangOptions().hasFeature(Feature::CheckImplementationOnly);
+
   serializationOpts.ExplicitModuleBuild = FrontendOpts.DisableImplicitModules;
 
   serializationOpts.EnableSerializationRemarks =

lib/Sema/TypeCheckAttr.cpp

@@ -5051,7 +5051,8 @@ AttributeChecker::visitImplementationOnlyAttr(ImplementationOnlyAttr *attr) {
 
   // @_implementationOnly on types only applies to non-public types.
   if (isa<NominalTypeDecl>(D)) {
-    if (!Ctx.LangOpts.hasFeature(Feature::CheckImplementationOnly)) {
+    if (!Ctx.LangOpts.hasFeature(Feature::CheckImplementationOnly) &&
+        !Ctx.LangOpts.hasFeature(Feature::CheckImplementationOnlyStrict)) {
       diagnoseAndRemoveAttr(attr, diag::implementation_only_on_types_feature);
       return;
     }

lib/Serialization/Serialization.cpp

@@ -756,6 +756,16 @@ IdentifierID Serializer::addContainingModuleRef(const DeclContext *DC,
   if (M->isClangHeaderImportModule())
     return OBJC_HEADER_MODULE_ID;
 
+  // Reject references to hidden dependencies.
+  if (getASTContext().LangOpts.hasFeature(
+          Feature::CheckImplementationOnlyStrict) &&
+      !allowCompilerErrors() &&
+      this->M->isImportedImplementationOnly(M, /*assumeImported=*/false)) {
+    getASTContext().Diags.diagnose(SourceLoc(),
+        diag::serialization_xref_to_hidden_dependency,
+        M, crossReferencedDecl);
+  }
+
   auto exportedModuleName = file->getExportedModuleName();
   assert(!exportedModuleName.empty());
   auto moduleID = M->getASTContext().getIdentifier(exportedModuleName);
@@ -2452,6 +2462,8 @@ void Serializer::writeCrossReference(const Decl *D) {
 
   unsigned abbrCode;
 
+  llvm::SaveAndRestore<const Decl *> SaveDecl(crossReferencedDecl, D);
+
   if (auto op = dyn_cast<OperatorDecl>(D)) {
     writeCrossReference(op->getDeclContext(), 1);
 
@@ -5390,6 +5402,19 @@ bool Serializer::shouldSkipDecl(const Decl *D) const {
 void Serializer::writeASTBlockEntity(const Decl *D) {
   using namespace decls_block;
 
+  if (Options.SkipImplementationOnlyDecls) {
+    // Skip @_implementationOnly types.
+    if (D->getAttrs().hasAttribute<ImplementationOnlyAttr>())
+      return;
+
+    // Skip non-public @export(interface) functions.
+    auto FD = dyn_cast<AbstractFunctionDecl>(D);
+    if (FD && FD->isNeverEmittedIntoClient() &&
+        !FD->getFormalAccessScope(/*useDC*/nullptr,
+          /*treatUsableFromInlineAsPublic*/true).isPublicOrPackage())
+      return;
+  }
+
   PrettyStackTraceDecl trace("serializing", D);
   assert(DeclsToSerialize.hasRef(D));
 
@@ -7302,7 +7327,7 @@ void Serializer::writeToStream(
     BCBlockRAII moduleBlock(S.Out, MODULE_BLOCK_ID, 2);
     S.writeHeader();
     S.writeInputBlock();
-    S.writeSIL(SILMod, options.SerializeAllSIL, options.SerializeDebugInfoSIL);
+    S.writeSIL(SILMod);
     S.writeAST(DC);
 
     if (S.hadError)

lib/Serialization/Serialization.h

@@ -118,6 +118,9 @@ class Serializer : public SerializerBase {
 
   bool hadImplementationOnlyImport = false;
 
+  /// Current decl being serialized.
+  const Decl* crossReferencedDecl = nullptr;
+
   /// Helper for serializing entities in the AST block object graph.
   ///
   /// Keeps track of assigning IDs to newly-seen entities, and collecting
@@ -433,8 +436,7 @@ class Serializer : public SerializerBase {
                     const SpecificASTBlockRecordKeeper &entities);
 
   /// Serializes all transparent SIL functions in the SILModule.
-  void writeSIL(const SILModule *M, bool serializeAllSIL,
-                bool serializeDebugInfo);
+  void writeSIL(const SILModule *M);
 
   /// Top-level entry point for serializing a module.
   void writeAST(ModuleOrSourceFile DC);

lib/Serialization/SerializeSIL.cpp

@@ -315,8 +315,7 @@ namespace {
                               << " for layout " << Layout::Code << "\n");
     }
 
-    bool ShouldSerializeAll;
-    bool SerializeDebugInfoSIL;
+    const SerializationOptions &Options;
 
     void addMandatorySILFunction(const SILFunction *F,
                                  bool emitDeclarationsForOnoneSupport);
@@ -399,10 +398,9 @@ namespace {
     IdentifierID addSILFunctionRef(SILFunction *F);
 
   public:
-    SILSerializer(Serializer &S, llvm::BitstreamWriter &Out, bool serializeAll,
-                  bool serializeDebugInfo)
-        : S(S), Out(Out), ShouldSerializeAll(serializeAll),
-          SerializeDebugInfoSIL(serializeDebugInfo) {}
+    SILSerializer(Serializer &S, llvm::BitstreamWriter &Out,
+                  const SerializationOptions &options)
+        : S(S), Out(Out), Options(options) {}
 
     void writeSILModule(const SILModule *SILMod);
   };
@@ -544,7 +542,7 @@ void SILSerializer::writeSILFunction(const SILFunction &F, bool DeclOnly) {
   // Otherwise, the generic specializer fails to remap references to functions
   // in debug scopes to their specialized versions which breaks IRGen.
   // TODO: add an assertion in IRGen when the specializer fails to remap.
-  if (!NoBody || SerializeDebugInfoSIL)
+  if (!NoBody || Options.SerializeDebugInfoSIL)
     if (auto *genericEnv = F.getGenericEnvironment())
       genericSigID = S.addGenericSignatureRef(genericEnv->getGenericSignature());
 
@@ -691,7 +689,7 @@ void SILSerializer::writeSILFunction(const SILFunction &F, bool DeclOnly) {
   DebugScopeMap.clear();
   SourceLocMap.clear();
 
-  if (SerializeDebugInfoSIL)
+  if (Options.SerializeDebugInfoSIL)
     writeDebugScopes(F.getDebugScope(), F.getModule().getSourceManager());
   // Assign a unique ID to each basic block of the SILFunction.
   unsigned BasicID = 0;
@@ -781,13 +779,13 @@ void SILSerializer::writeSILBasicBlock(const SILBasicBlock &BB) {
   const SILDebugScope *Prev = nullptr;
   auto &SM = BB.getParent()->getModule().getSourceManager();
   for (const SILInstruction &SI : BB) {
-    if (SerializeDebugInfoSIL) {
+    if (Options.SerializeDebugInfoSIL) {
       if (SI.getDebugScope() != Prev) {
         Prev = SI.getDebugScope();
         writeDebugScopes(Prev, SM);
       }
     }
-    if (SerializeDebugInfoSIL) {
+    if (Options.SerializeDebugInfoSIL) {
       writeSourceLoc(SI.getLoc(), SM);
     }
 
@@ -1060,7 +1058,7 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
   }
 
   case SILInstructionKind::DebugValueInst: {
-    if (!SerializeDebugInfoSIL)
+    if (!Options.SerializeDebugInfoSIL)
      return;
     auto DVI = cast<DebugValueInst>(&SI);
     unsigned attrs = unsigned(DVI->poisonRefs() & 0x1);
@@ -3287,7 +3285,7 @@ void SILSerializer::writeSILGlobalVar(const SILGlobalVariable &g) {
 void SILSerializer::writeSILVTable(const SILVTable &vt) {
   // Do not emit vtables for non-public classes unless everything has to be
   // serialized.
-  if (!ShouldSerializeAll &&
+  if (!Options.SerializeAllSIL &&
       vt.getClass()->getEffectiveAccess() < swift::AccessLevel::Package)
     return;
 
@@ -3313,7 +3311,7 @@ void SILSerializer::writeSILVTable(const SILVTable &vt) {
     SmallVector<uint64_t, 4> ListOfValues;
     SILFunction *impl = entry.getImplementation();
 
-    if (ShouldSerializeAll ||
+    if (Options.SerializeAllSIL ||
         (vt.isAnySerialized() &&
          impl->hasValidLinkageForFragileRef(vt.getSerializedKind()))) {
       handleSILDeclRef(S, entry.getMethod(), ListOfValues);
@@ -3333,12 +3331,12 @@ void SILSerializer::writeSILVTable(const SILVTable &vt) {
 void SILSerializer::writeSILMoveOnlyDeinit(const SILMoveOnlyDeinit &deinit) {
   // Do not emit deinit for non-public nominal types unless everything has to be
   // serialized.
-  if (!ShouldSerializeAll && deinit.getNominalDecl()->getEffectiveAccess() <
+  if (!Options.SerializeAllSIL && deinit.getNominalDecl()->getEffectiveAccess() <
                                  swift::AccessLevel::Package)
     return;
 
   SILFunction *impl = deinit.getImplementation();
-  if (!ShouldSerializeAll &&
+  if (!Options.SerializeAllSIL &&
       // Package CMO for MoveOnlyDeinit is not supported so
       // pass the IsSerialized argument to keep the behavior
       // consistent with or without the optimization.
@@ -3510,6 +3508,12 @@ void SILSerializer::writeDebugScopes(const SILDebugScope *Scope,
 }
 
 void SILSerializer::writeSILWitnessTable(const SILWitnessTable &wt) {
+  if (Options.SkipImplementationOnlyDecls &&
+      wt.getConformingNominal()->getAttrs().hasAttribute<
+        ImplementationOnlyAttr>()) {
+    return;
+  }
+
   WitnessTableList[wt.getName()] = NextWitnessTableID++;
   WitnessTableOffset.push_back(Out.GetCurrentBitNo());
 
@@ -3734,7 +3738,7 @@ bool SILSerializer::shouldEmitFunctionBody(const SILFunction *F,
   }
 
   // If we are asked to serialize everything, go ahead and do it.
-  if (ShouldSerializeAll)
+  if (Options.SerializeAllSIL)
     return true;
 
   // If F is serialized, we should always emit its body.
@@ -3814,29 +3818,29 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
   // serialize everything.
   // FIXME: Resilience: could write out vtable for fragile classes.
   for (const auto &vt : SILMod->getVTables()) {
-    if ((ShouldSerializeAll || vt->isAnySerialized()) &&
+    if ((Options.SerializeAllSIL || vt->isAnySerialized()) &&
         SILMod->shouldSerializeEntitiesAssociatedWithDeclContext(vt->getClass()))
       writeSILVTable(*vt);
   }
 
   for (const auto &deinit : SILMod->getMoveOnlyDeinits()) {
-    if ((ShouldSerializeAll || deinit->isAnySerialized()) &&
+    if ((Options.SerializeAllSIL || deinit->isAnySerialized()) &&
         SILMod->shouldSerializeEntitiesAssociatedWithDeclContext(
             deinit->getNominalDecl()))
       writeSILMoveOnlyDeinit(*deinit);
   }
 
   // Write out property descriptors.
   for (const SILProperty &prop : SILMod->getPropertyList()) {
-    if ((ShouldSerializeAll || prop.isAnySerialized()) &&
+    if ((Options.SerializeAllSIL || prop.isAnySerialized()) &&
         SILMod->shouldSerializeEntitiesAssociatedWithDeclContext(
                                      prop.getDecl()->getInnermostDeclContext()))
       writeSILProperty(prop);
   }
 
   // Write out fragile WitnessTables.
   for (const SILWitnessTable &wt : SILMod->getWitnessTables()) {
-    if ((ShouldSerializeAll || wt.isAnySerialized()) &&
+    if ((Options.SerializeAllSIL || wt.isAnySerialized()) &&
         SILMod->shouldSerializeEntitiesAssociatedWithDeclContext(
                                          wt.getConformance()->getDeclContext()))
       writeSILWitnessTable(wt);
@@ -3860,7 +3864,7 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
 
   // Add global variables that must be emitted to the list.
   for (const SILGlobalVariable &g : SILMod->getSILGlobals()) {
-    if (g.isAnySerialized() || ShouldSerializeAll)
+    if (g.isAnySerialized() || Options.SerializeAllSIL)
       addReferencedGlobalVariable(&g);
   }
 
@@ -3894,7 +3898,7 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
     // TODO(TF-893): Consider checking
     // `SILMod->shouldSerializeEntitiesAssociatedWithDeclContext` on the JVP/VJP
     // functions.
-    if ((ShouldSerializeAll || diffWitness.isSerialized()))
+    if ((Options.SerializeAllSIL || diffWitness.isSerialized()))
       DifferentiabilityWitnessesToEmit.insert(&diffWitness);
   }
   for (auto *diffWitness : DifferentiabilityWitnessesToEmit)
@@ -3945,11 +3949,10 @@ void SILSerializer::writeSILModule(const SILModule *SILMod) {
   writeIndexTables();
 }
 
-void Serializer::writeSIL(const SILModule *SILMod, bool serializeAllSIL,
-                          bool serializeDebugInfo) {
+void Serializer::writeSIL(const SILModule *SILMod) {
   if (!SILMod)
     return;
 
-  SILSerializer SILSer(*this, Out, serializeAllSIL, serializeDebugInfo);
+  SILSerializer SILSer(*this, Out, Options);
   SILSer.writeSILModule(SILMod);
 }

test/Sema/hidden-memory-layout.swift

@@ -48,8 +48,17 @@
 // RUN:   -verify-additional-prefix embedded-opt-in- \
 // RUN:   -enable-experimental-feature CheckImplementationOnly
 
+/// Same diags with CheckImplementationOnlyStrict
+// RUN: %target-swift-frontend -typecheck -verify -verify-ignore-unrelated %s -I %t \
+// RUN:   -swift-version 5 -target arm64-apple-none-macho \
+// RUN:   -enable-experimental-feature Embedded \
+// RUN:   -verify-additional-prefix opt-in- -DUseImplementationOnly \
+// RUN:   -verify-additional-prefix embedded-opt-in- \
+// RUN:   -enable-experimental-feature CheckImplementationOnlyStrict
+
 // REQUIRES: swift_feature_Embedded
 // REQUIRES: swift_feature_CheckImplementationOnly
+// REQUIRES: swift_feature_CheckImplementationOnlyStrict
 // REQUIRES: embedded_stdlib_cross_compiling
 
 @_implementationOnly import directs