[PackageCMO] Optimize witness thunks.

- Keep witness thunk linkage private for a package protocol member in SILGen.
- Optimize private/hidden functions during Package CMO; if they don't contain
references that have private/hidden symbols, serialize them and set the linkage
to shared. For unserialized witness thunks, set the linkage to package, so the
witness table itself can be serialized.
- Update witness table and vtable serialization.

Resolves rdar://129976582

Author: elsh

lib/SIL/IR/Linker.cpp

@@ -113,7 +113,7 @@ void SILLinkerVisitor::maybeAddFunctionToWorklist(
     if (callerSerializedKind == IsSerialized &&
         hasSharedVisibility(linkage) &&
         !Mod.isSerialized() &&
-        !F->isSerialized()) {
+        !F->isAnySerialized()) {
       F->setSerializedKind(IsSerialized);
 
       // Push the function to the worklist so that all referenced shared functions

lib/SIL/IR/SILWitnessTable.cpp

@@ -108,7 +108,7 @@ SILWitnessTable::SILWitnessTable(
     ArrayRef<ConditionalConformance> conditionalConformances)
     : Mod(M), Name(N), Linkage(Linkage), Conformance(Conformance), Entries(),
       ConditionalConformances(), IsDeclaration(true),
-      SerializedKind(IsNotSerialized) {
+      SerializedKind(SerializedKind) {
   convertToDefinition(entries, conditionalConformances, SerializedKind);
 }
 
@@ -169,21 +169,15 @@ void SILWitnessTable::convertToDefinition(
 
 SerializedKind_t SILWitnessTable::conformanceSerializedKind(
                                                             const RootProtocolConformance *conformance) {
-  // Allow serializing conformance with package or public access level
-  // if package serialization is enabled.
-  auto optInPackage = conformance->getDeclContext()->getParentModule()->serializePackageEnabled();
-  auto accessLevelToCheck =
-  optInPackage ? AccessLevel::Package : AccessLevel::Public;
-
   auto normalConformance = dyn_cast<NormalProtocolConformance>(conformance);
-  if (normalConformance && normalConformance->isResilient() && !optInPackage)
+  if (normalConformance && normalConformance->isResilient())
     return IsNotSerialized;
 
-  if (conformance->getProtocol()->getEffectiveAccess() < accessLevelToCheck)
+  if (conformance->getProtocol()->getEffectiveAccess() < AccessLevel::Public)
     return IsNotSerialized;
 
   auto *nominal = conformance->getDeclContext()->getSelfNominalTypeDecl();
-  if (nominal->getEffectiveAccess() >= accessLevelToCheck)
+  if (nominal->getEffectiveAccess() >= AccessLevel::Public)
     return IsSerialized;
 
   return IsNotSerialized;

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -73,11 +73,19 @@ class CrossModuleOptimization {
   CrossModuleOptimization(SILModule &M, bool conservative, bool everything)
     : M(M), conservative(conservative), everything(everything) { }
 
-  void trySerializeFunctions(ArrayRef<SILFunction *> functions);
   void serializeFunctionsInModule(SILPassManager *manager);
-  void serializeTablesInModule();
+  void serializeWitnessTablesInModule();
+  void serializeVTablesInModule();
 
 private:
+  bool isReferenceSerializeCandidate(SILFunction *F, SILOptions options);
+  bool isReferenceSerializeCandidate(SILGlobalVariable *G, SILOptions options);
+  SerializedKind_t getRightSerializedKind(const SILModule &mod);
+  bool isSerializedWithRightKind(const SILModule &mod, SILFunction *f);
+  bool isSerializedWithRightKind(const SILModule &mod, SILGlobalVariable *g);
+
+  void trySerializeFunctions(ArrayRef<SILFunction *> functions);
+
   bool canSerializeFunction(SILFunction *function,
                     FunctionFlags &canSerializeFlags, int maxDepth);
 
@@ -203,40 +211,50 @@ static bool isPackageCMOEnabled(ModuleDecl *mod) {
 /// function due to `@inlinable`, funtions with [serialized_for_package] from
 /// the imported module are not allowed being inlined into the client function, which
 /// is the correct behavior.
-static bool isSerializedWithRightKind(const SILModule &mod,
+bool CrossModuleOptimization::isSerializedWithRightKind(const SILModule &mod,
                                     SILFunction *f) {
   // 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
+  // (or similar) or [serialized_for_package] due to this
   // optimization.
   return isPackageCMOEnabled(mod.getSwiftModule()) ? f->isAnySerialized()
                                                    : f->isSerialized();
 }
-static bool isSerializedWithRightKind(const SILModule &mod,
+bool CrossModuleOptimization::isSerializedWithRightKind(const SILModule &mod,
                                       SILGlobalVariable *g) {
   return isPackageCMOEnabled(mod.getSwiftModule()) ? g->isAnySerialized()
                                                    : g->isSerialized();
 }
-static SerializedKind_t getRightSerializedKind(const SILModule &mod) {
+SerializedKind_t CrossModuleOptimization::getRightSerializedKind(const SILModule &mod) {
   return isPackageCMOEnabled(mod.getSwiftModule()) ? IsSerializedForPackage
                                                    : IsSerialized;
 }
 
 static bool isSerializeCandidate(SILFunction *F, SILOptions options) {
   auto linkage = F->getLinkage();
-  // We allow serializing a shared definition. For example,
-  // `public func foo() { print("") }` is a function with a
-  // public linkage which only references `print`; the definition
-  // of `print` has a shared linkage and does not reference
-  // non-serializable instructions, so it should be serialized,
-  // thus the public `foo` could be serialized.
+  // If Package CMO is enabled, besides package/public definitions,
+  // we allow serializing private, hidden, or shared definitions
+  // that do not contain private or hidden symbol references (and
+  // their nested references). If private or internal definitions are
+  // serialized, they are set to a shared linkage.
+  //
+  // E.g. `public func foo() { print("") }` is a public function that
+  // references `print`, a shared definition which does not contain
+  // any private or internal symbols, thus is serialized, which in turn
+  // allows `foo` to be serialized.
+  // E.g. a protocol witness method for a package protocol member is
+  // set to a private linkage in SILGen. By allowing such private thunk
+  // to be serialized and set to shared linkage here, functions that
+  // reference the thunk can be serialized as well.
   if (options.EnableSerializePackage)
-    return linkage == SILLinkage::Public || linkage == SILLinkage::Package ||
-           (linkage == SILLinkage::Shared && F->isDefinition());
+    return linkage != SILLinkage::PublicExternal &&
+           linkage != SILLinkage::PackageExternal &&
+           linkage != SILLinkage::HiddenExternal;
   return linkage == SILLinkage::Public;
 }
 
-static bool isReferenceSerializeCandidate(SILFunction *F, SILOptions options) {
+bool CrossModuleOptimization::isReferenceSerializeCandidate(SILFunction *F, 
+                                                            SILOptions options) {
   if (options.EnableSerializePackage) {
     if (isSerializedWithRightKind(F->getModule(), F))
       return true;
@@ -246,8 +264,8 @@ static bool isReferenceSerializeCandidate(SILFunction *F, SILOptions options) {
   return hasPublicVisibility(F->getLinkage());
 }
 
-static bool isReferenceSerializeCandidate(SILGlobalVariable *G,
-                                          SILOptions options) {
+bool CrossModuleOptimization::isReferenceSerializeCandidate(SILGlobalVariable *G,
+                                                            SILOptions options) {
   if (options.EnableSerializePackage) {
     if (isSerializedWithRightKind(G->getModule(), G))
       return true;
@@ -279,85 +297,73 @@ void CrossModuleOptimization::serializeFunctionsInModule(SILPassManager *manager
   trySerializeFunctions(bottomUpFunctions);
 }
 
-void CrossModuleOptimization::serializeTablesInModule() {
-  if (!M.getSwiftModule()->serializePackageEnabled())
+void CrossModuleOptimization::serializeWitnessTablesInModule() {
+  if (!isPackageCMOEnabled(M.getSwiftModule()))
     return;
 
-  for (const auto &vt : M.getVTables()) {
-    if (vt->getSerializedKind() != getRightSerializedKind(M) &&
-        vt->getClass()->getEffectiveAccess() >= AccessLevel::Package) {
-      // This checks if a vtable entry is not serialized and attempts to
-      // serialize (and its references) if they have the right visibility.
-      // 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(), [&](const SILVTableEntry &entry) {
-            return entry.getImplementation()->getSerializedKind() !=
-                   getRightSerializedKind(M);
-          });
-      std::vector<SILFunction *> classMethodsToSerialize;
-      llvm::transform(unserializedClassMethodRange,
-                      std::back_inserter(classMethodsToSerialize),
-                      [&](const SILVTableEntry &entry) {
-                        return entry.getImplementation();
-                      });
-      trySerializeFunctions(classMethodsToSerialize);
-
-      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));
-    }
-  }
-
-  // Witness thunks are not serialized, so serialize them here.
   for (auto &wt : M.getWitnessTables()) {
     if (wt.getSerializedKind() != getRightSerializedKind(M) &&
         hasPublicOrPackageVisibility(wt.getLinkage(), /*includePackage*/ true)) {
-      // This checks if a wtable entry is not serialized and attempts to
-      // serialize (and its references) if they have the right visibility.
-      // 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(), [&](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),
-                      [&](const SILWitnessTable::Entry &entry) {
-                        return entry.getMethodWitness().Witness;
-                      });
-      trySerializeFunctions(wtMethodsToSerialize);
+      // In Package CMO, we try serializing witness thunks that
+      // are private if they don't contain hidden or private
+      // references. If they are serialized, they are set to
+      // a shared linkage. If they can't be serialized, we set
+      // the linkage to package so that the witness table itself
+      // can still be serialized, thus giving a chance for entires
+      // that _are_ serialized to be accessed directly.
+      for (const SILWitnessTable::Entry &entry: unserializedWTMethodRange) {
+        if (entry.getMethodWitness().Witness->getLinkage() == SILLinkage::Private)
+          entry.getMethodWitness().Witness->setLinkage(SILLinkage::Package);
+      }
 
       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));
           });
+      // FIXME: This check shouldn't be necessary but added as a caution
+      // to ensure we don't serialize witness table if it contains an
+      // internal entry.
       if (!containsInternal)
         wt.setSerializedKind(getRightSerializedKind(M));
     }
   }
 }
 
+void CrossModuleOptimization::serializeVTablesInModule() {
+  if (!isPackageCMOEnabled(M.getSwiftModule()))
+    return;
+
+  for (const auto &vt : M.getVTables()) {
+    if (vt->getSerializedKind() != getRightSerializedKind(M) &&
+        vt->getClass()->getEffectiveAccess() >= AccessLevel::Package) {
+       bool containsInternal =
+          llvm::any_of(vt->getEntries(), [&](const SILVTableEntry &entry) {
+            return !entry.getImplementation()->hasValidLinkageForFragileRef(
+                getRightSerializedKind(M));
+          });
+
+      // If the entries are either serialized or have package/public
+      // visibility, the vtable can be serialized; the non-serialized
+      // entries can still be referenced as they have the visibility
+      // outside of their defining module, and the serialized entries
+      // can be directly accessed since the vtable is serialized.
+      // However, if it contains a private/internal entry, we don't
+      // serialize the vtable at all.
+      if (!containsInternal)
+        vt->setSerializedKind(getRightSerializedKind(M));
+    }
+  }
+}
+
 /// Recursively walk the call graph and select functions to be serialized.
 ///
 /// The results are stored in \p canSerializeFlags and the result for \p
@@ -538,7 +544,6 @@ bool CrossModuleOptimization::canSerializeGlobal(SILGlobalVariable *global) {
   for (const SILInstruction &initInst : *global) {
     if (auto *FRI = dyn_cast<FunctionRefInst>(&initInst)) {
       SILFunction *referencedFunc = FRI->getReferencedFunction();
-      
       // In conservative mode we don't want to turn non-public functions into
       // public functions, because that can increase code size. E.g. if the
       // function is completely inlined afterwards.
@@ -701,6 +706,16 @@ void CrossModuleOptimization::serializeFunction(SILFunction *function,
   if (!canSerializeFlags.lookup(function))
     return;
 
+  if (isPackageCMOEnabled(M.getSwiftModule())) {
+    // If a private thunk (such as a protocol witness method for
+    // a package protocol member) does not reference any private
+    // or internal symbols, thus is serialized, it's set to shared
+    // linkage, so that functions that reference the thunk can be
+    // serialized as well.
+    if (function->getLinkage() == SILLinkage::Private ||
+        function->getLinkage() == SILLinkage::Hidden)
+      function->setLinkage(SILLinkage::Shared);
+  }
   function->setSerializedKind(getRightSerializedKind(M));
 
   for (SILBasicBlock &block : *function) {
@@ -900,7 +915,6 @@ void CrossModuleOptimization::makeSubstUsableFromInline(
 
 class CrossModuleOptimizationPass: public SILModuleTransform {
   void run() override {
-
     auto &M = *getModule();
     if (M.getSwiftModule()->isResilient() &&
         !M.getSwiftModule()->serializePackageEnabled())
@@ -933,7 +947,8 @@ class CrossModuleOptimizationPass: public SILModuleTransform {
     CMO.serializeFunctionsInModule(PM);
 
     // Serialize SIL v-tables and witness-tables if package-cmo is enabled.
-    CMO.serializeTablesInModule();
+    CMO.serializeVTablesInModule();
+    CMO.serializeWitnessTablesInModule();
   }
 };