Merge pull request #72606 from apple/elsh/pkg-serialize-tables

Author: elsh

include/swift/SIL/SILLinkage.h

@@ -394,12 +394,16 @@ inline SILLinkage effectiveLinkageForClassMember(SILLinkage linkage,
 // protocol requirement, even if the extended type is not public;
 // then SILGen gives the member private linkage, ignoring the more
 // visible access level it was given in the AST.
-inline bool
-fixmeWitnessHasLinkageThatNeedsToBePublic(SILDeclRef witness) {
+//
+// Despite the FIXME above, this is still used to determine the linkage
+// for witness thunks. In case package serialization is enabled, we need
+// to take the package linkage into account so we can set a proper final
+// linkage to the thunks in the witness table with a package linkage.
+inline bool fixmeWitnessHasLinkageThatNeedsToBePublic(SILDeclRef witness,
+                                                      bool isPackageVisible) {
   auto witnessLinkage = witness.getLinkage(ForDefinition);
-  return !hasPublicVisibility(witnessLinkage)
-         && (!hasSharedVisibility(witnessLinkage)
-             || !witness.isSerialized());
+  return !hasPublicOrPackageVisibility(witnessLinkage, isPackageVisible) &&
+         (!hasSharedVisibility(witnessLinkage) || !witness.isSerialized());
 }
 
 } // end swift namespace

lib/SIL/IR/SILSymbolVisitor.cpp

@@ -289,7 +289,9 @@ class SILSymbolVisitorImpl : public ASTVisitor<SILSymbolVisitorImpl> {
             return;
 
           if (!isa<SelfProtocolConformance>(rootConformance) &&
-              !fixmeWitnessHasLinkageThatNeedsToBePublic(witnessRef)) {
+              !fixmeWitnessHasLinkageThatNeedsToBePublic(
+                  witnessRef,
+                  witnessRef.getASTContext().SILOpts.EnableSerializePackage)) {
             return;
           }
         }

lib/SIL/IR/SILWitnessTable.cpp

@@ -167,15 +167,23 @@ void SILWitnessTable::convertToDefinition(
 
 bool SILWitnessTable::conformanceIsSerialized(
     const RootProtocolConformance *conformance) {
+  // Allow serializing conformance with package or public access level
+  // if package serialization is enabled.
+  auto optInPackage = conformance->getDeclContext()
+                          ->getASTContext()
+                          .SILOpts.EnableSerializePackage;
+  auto accessLevelToCheck =
+      optInPackage ? AccessLevel::Package : AccessLevel::Public;
+
   auto normalConformance = dyn_cast<NormalProtocolConformance>(conformance);
-  if (normalConformance && normalConformance->isResilient())
+  if (normalConformance && normalConformance->isResilient() && !optInPackage)
     return false;
 
-  if (conformance->getProtocol()->getEffectiveAccess() < AccessLevel::Public)
+  if (conformance->getProtocol()->getEffectiveAccess() < accessLevelToCheck)
     return false;
 
   auto *nominal = conformance->getDeclContext()->getSelfNominalTypeDecl();
-  return nominal->getEffectiveAccess() >= AccessLevel::Public;
+  return nominal->getEffectiveAccess() >= accessLevelToCheck;
 }
 
 bool SILWitnessTable::enumerateWitnessTableConditionalConformances(

lib/SILGen/SILGenType.cpp

@@ -630,7 +630,12 @@ class SILGenConformance : public SILGenWitnessTable<SILGenConformance> {
     auto witnessLinkage = witnessRef.getLinkage(ForDefinition);
     auto witnessSerialized = Serialized;
     if (witnessSerialized &&
-        fixmeWitnessHasLinkageThatNeedsToBePublic(witnessRef)) {
+        // If package optimization is enabled, this is false;
+        // witness thunk should get a `shared` linkage in the
+        // else block below.
+        fixmeWitnessHasLinkageThatNeedsToBePublic(
+            witnessRef,
+            witnessRef.getASTContext().SILOpts.EnableSerializePackage)) {
       witnessLinkage = SILLinkage::Public;
       witnessSerialized = IsNotSerialized;
     } else {

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -21,6 +21,8 @@
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h"
+#include "swift/SILOptimizer/Analysis/FunctionOrder.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -70,7 +72,8 @@ class CrossModuleOptimization {
   CrossModuleOptimization(SILModule &M, bool conservative, bool everything)
     : M(M), conservative(conservative), everything(everything) { }
 
-  void serializeFunctionsInModule();
+  void serializeFunctionsInModule(ArrayRef<SILFunction *> functions);
+  void serializeTablesInModule();
 
 private:
   bool canSerializeFunction(SILFunction *function,
@@ -81,7 +84,7 @@ class CrossModuleOptimization {
 
   bool canSerializeGlobal(SILGlobalVariable *global);
 
-  bool canSerializeType(SILType type);
+  bool canSerializeType(SILType type, TypeExpansionContext typeExpCtx);
 
   bool canUseFromInline(DeclContext *declCtxt);
 
@@ -161,29 +164,94 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
   }
 };
 
-static bool isVisible(SILLinkage linkage, SILOptions options) {
+static bool isPackageOrPublic(SILLinkage linkage, SILOptions options) {
   if (options.EnableSerializePackage)
     return linkage == SILLinkage::Public || linkage == SILLinkage::Package;
   return linkage == SILLinkage::Public;
 }
-static bool isVisible(AccessLevel accessLevel, SILOptions options) {
+
+static bool isPackageOrPublic(AccessLevel accessLevel, SILOptions options) {
   if (options.EnableSerializePackage)
     return accessLevel == AccessLevel::Package || accessLevel == AccessLevel::Public;
   return accessLevel == AccessLevel::Public;
 }
 
-/// Select functions in the module which should be serialized.
-void CrossModuleOptimization::serializeFunctionsInModule() {
+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 (options.EnableSerializePackage)
+    return linkage == SILLinkage::Public || linkage == SILLinkage::Package ||
+           (linkage == SILLinkage::Shared && F->isDefinition());
+  return linkage == SILLinkage::Public;
+}
+
+static bool isReferenceSerializeCandidate(SILFunction *F, SILOptions options) {
+  if (options.EnableSerializePackage) {
+    if (F->isSerialized())
+      return true;
+    return hasPublicOrPackageVisibility(F->getLinkage(),
+                                        /*includePackage*/ true);
+  }
+  return hasPublicVisibility(F->getLinkage());
+}
+
+static bool isReferenceSerializeCandidate(SILGlobalVariable *G,
+                                          SILOptions options) {
+  if (options.EnableSerializePackage) {
+    if (G->isSerialized())
+      return true;
+    return hasPublicOrPackageVisibility(G->getLinkage(),
+                                        /*includePackage*/ true);
+  }
+  return hasPublicVisibility(G->getLinkage());
+}
 
+/// Select functions in the module which should be serialized.
+void CrossModuleOptimization::serializeFunctionsInModule(
+    ArrayRef<SILFunction *> functions) {
   FunctionFlags canSerializeFlags;
 
-  // Start with public functions.
-  for (SILFunction &F : M) {
-    if (isVisible(F.getLinkage(), M.getOptions()) ||
-        everything) {
-      if (canSerializeFunction(&F, canSerializeFlags, /*maxDepth*/ 64)) {
-        serializeFunction(&F, canSerializeFlags);
+  // The passed functions are already ordered bottom-up so the most
+  // nested referenced function is checked first.
+  for (SILFunction *F : functions) {
+    if (isSerializeCandidate(F, M.getOptions()) || everything) {
+      if (canSerializeFunction(F, canSerializeFlags, /*maxDepth*/ 64)) {
+        serializeFunction(F, canSerializeFlags);
+      }
+    }
+  }
+}
+
+void CrossModuleOptimization::serializeTablesInModule() {
+  if (!M.getOptions().EnableSerializePackage)
+    return;
+
+  for (const auto &vt : M.getVTables()) {
+    if (!vt->isSerialized() &&
+        vt->getClass()->getEffectiveAccess() >= AccessLevel::Package) {
+      vt->setSerialized(IsSerialized);
+    }
+  }
+
+  for (auto &wt : M.getWitnessTables()) {
+    if (!wt.isSerialized() && hasPublicOrPackageVisibility(
+                                  wt.getLinkage(), /*includePackage*/ true)) {
+      for (auto &entry : wt.getEntries()) {
+        // Witness thunks are not serialized, so serialize them here.
+        if (entry.getKind() == SILWitnessTable::Method &&
+            !entry.getMethodWitness().Witness->isSerialized() &&
+            isSerializeCandidate(entry.getMethodWitness().Witness,
+                                 M.getOptions())) {
+          entry.getMethodWitness().Witness->setSerialized(IsSerialized);
+        }
       }
+      // Then serialize the witness table itself.
+      wt.setSerialized(IsSerialized);
     }
   }
 }
@@ -215,8 +283,10 @@ bool CrossModuleOptimization::canSerializeFunction(
       return false;
   }
 
-  if (function->isSerialized())
+  if (function->isSerialized()) {
+    canSerializeFlags[function] = true;
     return true;
+  }
 
   if (!function->isDefinition() || function->isAvailableExternally())
     return false;
@@ -258,16 +328,17 @@ bool CrossModuleOptimization::canSerializeFunction(
 /// Returns true if \p inst can be serialized.
 ///
 /// If \p inst is a function_ref, recursively visits the referenced function.
-bool CrossModuleOptimization::canSerializeInstruction(SILInstruction *inst,
-                      FunctionFlags &canSerializeFlags, int maxDepth) {
-
+bool CrossModuleOptimization::canSerializeInstruction(
+    SILInstruction *inst, FunctionFlags &canSerializeFlags, int maxDepth) {
   // First check if any result or operand types prevent serialization.
+  auto typeExpCtx = inst->getFunction()->getTypeExpansionContext();
+
   for (SILValue result : inst->getResults()) {
-    if (!canSerializeType(result->getType()))
+    if (!canSerializeType(result->getType(), typeExpCtx))
       return false;
   }
   for (Operand &op : inst->getAllOperands()) {
-    if (!canSerializeType(op.get()->getType()))
+    if (!canSerializeType(op.get()->getType(), typeExpCtx))
       return false;
   }
 
@@ -280,9 +351,9 @@ bool CrossModuleOptimization::canSerializeInstruction(SILInstruction *inst,
     // public functions, because that can increase code size. E.g. if the
     // function is completely inlined afterwards.
     // Also, when emitting TBD files, we cannot introduce a new public symbol.
-    if ((conservative || M.getOptions().emitTBD) &&
-        !hasPublicOrPackageVisibility(callee->getLinkage(), M.getOptions().EnableSerializePackage)) {
-      return false;
+    if (conservative || M.getOptions().emitTBD) {
+      if (!isReferenceSerializeCandidate(callee, M.getOptions()))
+        return false;
     }
 
     // In some project configurations imported C functions are not necessarily
@@ -301,12 +372,13 @@ bool CrossModuleOptimization::canSerializeInstruction(SILInstruction *inst,
     // inline.
     if (!canUseFromInline(callee))
       return false;
+
     return true;
   }
   if (auto *GAI = dyn_cast<GlobalAddrInst>(inst)) {
     SILGlobalVariable *global = GAI->getReferencedGlobal();
     if ((conservative || M.getOptions().emitTBD) &&
-        !hasPublicOrPackageVisibility(global->getLinkage(), M.getOptions().EnableSerializePackage)) {
+        !isReferenceSerializeCandidate(global, M.getOptions())) {
       return false;
     }
 
@@ -354,7 +426,7 @@ bool CrossModuleOptimization::canSerializeGlobal(SILGlobalVariable *global) {
       // function is completely inlined afterwards.
       // Also, when emitting TBD files, we cannot introduce a new public symbol.
       if ((conservative || M.getOptions().emitTBD) &&
-          !hasPublicOrPackageVisibility(referencedFunc->getLinkage(), M.getOptions().EnableSerializePackage)) {
+          !isReferenceSerializeCandidate(referencedFunc, M.getOptions())) {
         return false;
       }
 
@@ -365,16 +437,28 @@ bool CrossModuleOptimization::canSerializeGlobal(SILGlobalVariable *global) {
   return true;
 }
 
-bool CrossModuleOptimization::canSerializeType(SILType type) {
+bool CrossModuleOptimization::canSerializeType(SILType type,
+                                               TypeExpansionContext typeExpCtx) {
   auto iter = typesChecked.find(type);
   if (iter != typesChecked.end())
     return iter->getSecond();
 
+  if (M.getSwiftModule()->isResilient()) {
+    auto minResilientCtx = TypeExpansionContext(ResilienceExpansion::Minimal,
+                                                typeExpCtx.getContext(),
+                                                typeExpCtx.isWholeModuleContext());
+    auto loadableInMinResilientCtx = M.Types.getTypeLowering(type, minResilientCtx).isLoadable();
+    if (!loadableInMinResilientCtx) {
+      typesChecked[type] = false;
+      return false;
+    }
+  }
+
   bool success = !type.getASTType().findIf(
     [this](Type rawSubType) {
       CanType subType = rawSubType->getCanonicalType();
       if (NominalTypeDecl *subNT = subType->getNominalOrBoundGenericNominal()) {
-      
+
         if (conservative && subNT->getEffectiveAccess() < AccessLevel::Package) {
           return true;
         }
@@ -484,13 +568,17 @@ bool CrossModuleOptimization::shouldSerialize(SILFunction *function) {
       return true;
   }
 
-  // Also serialize "small" non-generic functions.
-  int size = 0;
-  for (SILBasicBlock &block : *function) {
-    for (SILInstruction &inst : block) {
-      size += (int)instructionInlineCost(inst);
-      if (size >= CMOFunctionSizeLimit)
-        return false;
+  // If package-cmo is enabled, we don't want to limit inlining
+  // or should at least increase the cap.
+  if (!M.getOptions().EnableSerializePackage) {
+    // Also serialize "small" non-generic functions.
+    int size = 0;
+    for (SILBasicBlock &block : *function) {
+      for (SILInstruction &inst : block) {
+        size += (int)instructionInlineCost(inst);
+        if (size >= CMOFunctionSizeLimit)
+          return false;
+      }
     }
   }
 
@@ -503,7 +591,7 @@ void CrossModuleOptimization::serializeFunction(SILFunction *function,
                                        const FunctionFlags &canSerializeFlags) {
   if (function->isSerialized())
     return;
-  
+
   if (!canSerializeFlags.lookup(function))
     return;
 
@@ -552,9 +640,11 @@ void CrossModuleOptimization::serializeInstruction(SILInstruction *inst,
       }
     }
     serializeFunction(callee, canSerializeFlags);
-    assert(callee->isSerialized() || isVisible(callee->getLinkage(), M.getOptions()));
+    assert(callee->isSerialized() ||
+           isPackageOrPublic(callee->getLinkage(), M.getOptions()));
     return;
   }
+
   if (auto *GAI = dyn_cast<GlobalAddrInst>(inst)) {
     SILGlobalVariable *global = GAI->getReferencedGlobal();
     if (canSerializeGlobal(global)) {
@@ -616,7 +706,7 @@ void CrossModuleOptimization::makeDeclUsableFromInline(ValueDecl *decl) {
   if (M.getSwiftModule() != decl->getDeclContext()->getParentModule())
     return;
 
-  if (!isVisible(decl->getFormalAccess(), M.getOptions()) &&
+  if (!isPackageOrPublic(decl->getFormalAccess(), M.getOptions()) &&
       !decl->isUsableFromInline()) {
     // Mark the nominal type as "usableFromInline".
     // TODO: find a way to do this without modifying the AST. The AST should be
@@ -699,7 +789,8 @@ class CrossModuleOptimizationPass: public SILModuleTransform {
   void run() override {
 
     auto &M = *getModule();
-    if (M.getSwiftModule()->isResilient())
+    if (M.getSwiftModule()->isResilient() &&
+        !M.getOptions().EnableSerializePackage)
       return;
     if (!M.isWholeModule())
       return;
@@ -726,7 +817,17 @@ class CrossModuleOptimizationPass: public SILModuleTransform {
     }
 
     CrossModuleOptimization CMO(M, conservative, everything);
-    CMO.serializeFunctionsInModule();
+
+    // Reorder SIL funtions in the module bottom up so we can serialize
+    // the most nested referenced functions first and avoid unnecessary
+    // recursive checks.
+    BasicCalleeAnalysis *BCA = PM->getAnalysis<BasicCalleeAnalysis>();
+    BottomUpFunctionOrder BottomUpOrder(M, BCA);
+    auto BottomUpFunctions = BottomUpOrder.getFunctions();
+    CMO.serializeFunctionsInModule(BottomUpFunctions);
+
+    // Serialize SIL v-tables and witness-tables if package-cmo is enabled.
+    CMO.serializeTablesInModule();
   }
 };