Add isPackageCMOEnabled check for MethodInst and KeyPathInst.

Update lambda in table serialization.
Add more tests.

Author: elsh

lib/SIL/IR/Linker.cpp

@@ -105,7 +105,6 @@ void SILLinkerVisitor::maybeAddFunctionToWorklist(
           F->hasValidLinkageForFragileRef(callerSerializedKind) ||
          hasSharedVisibility(linkage) || F->isExternForwardDeclaration()) &&
          "called function has wrong linkage for serialized function");
-                                         
   if (!F->isExternalDeclaration()) {
     // The function is already in the module, so no need to de-serialized it.
     // But check if we need to set the IsSerialized flag.

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -189,6 +189,10 @@ static bool isPackageOrPublic(AccessLevel accessLevel, SILOptions options) {
   return accessLevel == AccessLevel::Public;
 }
 
+static bool isPackageCMOEnabled(ModuleDecl *mod) {
+  return mod->isResilient() && mod->serializePackageEnabled();
+}
+
 /// Checks wither this function is [serialized_for_package] due to Package CMO
 /// or [serialized] with non-package CMO. The [serialized_for_package] attribute
 /// is used to indicate that a function is serialized because of Package CMO, which
@@ -201,20 +205,21 @@ static bool isPackageOrPublic(AccessLevel accessLevel, SILOptions options) {
 /// is the correct behavior.
 static bool isSerializedWithRightKind(const SILModule &mod,
                                     SILFunction *f) {
-  return mod.getSwiftModule()->serializePackageEnabled() &&
-         mod.getSwiftModule()->isResilient() ?
-         f->isSerializedForPackage() : f->isSerialized();
+  // If Package CMO is enabled in resilient mode, return
+  // true if the function is [serialized] due to @inlinable
+  // (or similar) or [serialized_for_pkg] due to this
+  // optimization.
+  return isPackageCMOEnabled(mod.getSwiftModule()) ? f->isAnySerialized()
+                                                   : f->isSerialized();
 }
 static bool isSerializedWithRightKind(const SILModule &mod,
                                       SILGlobalVariable *g) {
-  return mod.getSwiftModule()->serializePackageEnabled() &&
-         mod.getSwiftModule()->isResilient() ?
-         g->isSerializedForPackage() : g->isSerialized();
+  return isPackageCMOEnabled(mod.getSwiftModule()) ? g->isAnySerialized()
+                                                   : g->isSerialized();
 }
 static SerializedKind_t getRightSerializedKind(const SILModule &mod) {
-  return mod.getSwiftModule()->serializePackageEnabled() &&
-         mod.getSwiftModule()->isResilient() ?
-         IsSerializedForPackage : IsSerialized;
+  return isPackageCMOEnabled(mod.getSwiftModule()) ? IsSerializedForPackage
+                                                   : IsSerialized;
 }
 
 static bool isSerializeCandidate(SILFunction *F, SILOptions options) {
@@ -286,24 +291,28 @@ void CrossModuleOptimization::serializeTablesInModule() {
       // This should not be necessary but is added to ensure all applicable
       // symbols are serialized. Whether serialized or not is cached so
       // this check shouldn't be expensive.
-      auto unserializedClassMethodRange = llvm::make_filter_range(vt->getEntries(), [&](auto &entry) {
-        return entry.getImplementation()->getSerializedKind() != getRightSerializedKind(M);
-      });
+      auto unserializedClassMethodRange = llvm::make_filter_range(
+          vt->getEntries(), [&](const SILVTableEntry &entry) {
+            return entry.getImplementation()->getSerializedKind() !=
+                   getRightSerializedKind(M);
+          });
       std::vector<SILFunction *> classMethodsToSerialize;
       llvm::transform(unserializedClassMethodRange,
                       std::back_inserter(classMethodsToSerialize),
-                      [&](auto entry) {
-        return entry.getImplementation();
-      });
+                      [&](const SILVTableEntry &entry) {
+                        return entry.getImplementation();
+                      });
       trySerializeFunctions(classMethodsToSerialize);
 
-      bool containsInternal = llvm::any_of(classMethodsToSerialize, [&](auto &method) {
-        // If the entry is internal, vtable should not be serialized.
-        // However, if the entry is not serialized but has the right
-        // visibility, it can still be referenced, thus the vtable
-        // should serialized.
-        return !method->hasValidLinkageForFragileRef(getRightSerializedKind(M));
-      });
+      bool containsInternal =
+          llvm::any_of(vt->getEntries(), [&](const SILVTableEntry &entry) {
+            // If the entry is internal, vtable should not be serialized.
+            // However, if the entry is not serialized but has the right
+            // visibility, it can still be referenced, thus the vtable
+            // should serialized.
+            return !entry.getImplementation()->hasValidLinkageForFragileRef(
+                getRightSerializedKind(M));
+          });
       if (!containsInternal)
         vt->setSerializedKind(getRightSerializedKind(M));
     }
@@ -318,25 +327,31 @@ void CrossModuleOptimization::serializeTablesInModule() {
       // This should not be necessary but is added to ensure all applicable
       // symbols are serialized. Whether serialized or not is cached so
       // this check shouldn't be expensive.
-      auto unserializedWTMethodRange = llvm::make_filter_range(wt.getEntries(), [&](auto &entry) {
-        return entry.getKind() == SILWitnessTable::Method &&
-        entry.getMethodWitness().Witness->getSerializedKind() != getRightSerializedKind(M);
-      });
+      auto unserializedWTMethodRange = llvm::make_filter_range(
+          wt.getEntries(), [&](const SILWitnessTable::Entry &entry) {
+            return entry.getKind() == SILWitnessTable::Method &&
+                   entry.getMethodWitness().Witness->getSerializedKind() !=
+                       getRightSerializedKind(M);
+          });
       std::vector<SILFunction *> wtMethodsToSerialize;
       llvm::transform(unserializedWTMethodRange,
                       std::back_inserter(wtMethodsToSerialize),
-                      [&](auto entry) {
-        return entry.getMethodWitness().Witness;
-      });
+                      [&](const SILWitnessTable::Entry &entry) {
+                        return entry.getMethodWitness().Witness;
+                      });
       trySerializeFunctions(wtMethodsToSerialize);
 
-      bool containsInternal = llvm::any_of(wtMethodsToSerialize, [&](auto &method) {
-        // If the entry is internal, wtable should not be serialized.
-        // However, if the entry is not serialized but has the right
-        // visibility, it can still be referenced, thus the vtable
-        // should serialized.
-        return !method->hasValidLinkageForFragileRef(getRightSerializedKind(M));
-      });
+      bool containsInternal = llvm::any_of(
+          wt.getEntries(), [&](const SILWitnessTable::Entry &entry) {
+            // If the entry is internal, wtable should not be serialized.
+            // However, if the entry is not serialized but has the right
+            // visibility, it can still be referenced, thus the vtable
+            // should serialized.
+            return entry.getKind() == SILWitnessTable::Method &&
+                   !entry.getMethodWitness()
+                        .Witness->hasValidLinkageForFragileRef(
+                            getRightSerializedKind(M));
+          });
       if (!containsInternal)
         wt.setSerializedKind(getRightSerializedKind(M));
     }
@@ -485,17 +500,29 @@ bool CrossModuleOptimization::canSerializeInstruction(
         [&](SILDeclRef method) {
           if (method.isForeign)
             canUse = false;
+          else if (isPackageCMOEnabled(method.getModuleContext())) {
+            // If the referenced keypath is internal, do not
+            // serialize.
+            auto methodScope = method.getDecl()->getFormalAccessScope(
+                nullptr,
+                /*treatUsableFromInlineAsPublic*/ true);
+            canUse = methodScope.isPublicOrPackage();
+          }
         });
     return canUse;
   }
   if (auto *MI = dyn_cast<MethodInst>(inst)) {
-    // If a class_method or witness_method is internal, it can't
-    // be serialized.
+    // If a class_method or witness_method is internal,
+    // it can't be serialized.
     auto member = MI->getMember();
-    auto methodAccessScope = member.getDecl()->getFormalAccessScope(nullptr,
-                                               /*treatUsableFromInlineAsPublic*/ true);
-    return methodAccessScope.isPublicOrPackage() &&
-          !member.isForeign;
+    auto canUse = !member.isForeign;
+    if (canUse && isPackageCMOEnabled(member.getModuleContext())) {
+      auto methodScope = member.getDecl()->getFormalAccessScope(
+          nullptr,
+          /*treatUsableFromInlineAsPublic*/ true);
+      canUse = methodScope.isPublicOrPackage();
+    }
+    return canUse;
   }
   if (auto *REAI = dyn_cast<RefElementAddrInst>(inst)) {
     // In conservative mode, we don't support class field accesses of non-public
@@ -737,7 +764,9 @@ void CrossModuleOptimization::serializeInstruction(SILInstruction *inst,
     if (canSerializeGlobal(global)) {
       serializeGlobal(global);
     }
-    if (!hasPublicOrPackageVisibility(global->getLinkage(), M.getSwiftModule()->serializePackageEnabled())) {
+    if (!hasPublicOrPackageVisibility(
+            global->getLinkage(),
+            M.getSwiftModule()->serializePackageEnabled())) {
       global->setLinkage(SILLinkage::Public);
     }
     return;