[Package CMO] Use InstructionVisitor to check if inst can be serialized.

Currently, InstructionVisitor is only used during the serialization pass,
that comes after a pass that checks if instructions are serializable. This
causes inconsistencies as some types bypass the first pass but are processed
in the second, leading to assert fails.

This PR uses InstructionVisitor in the first pass to be exhaustive and to
be consistent with the coverage in the second pass. The downside is the visitor
is SILCloner, which introduces overhead of cloning and discarding insts per pass
-- a potential area for future refactoring.

Resolves rdar://130788545.

Author: elsh

lib/SILOptimizer/IPO/CrossModuleOptimization.cpp

@@ -53,7 +53,7 @@ namespace {
 class CrossModuleOptimization {
   friend class InstructionVisitor;
 
-  llvm::DenseMap<SILType, bool> typesChecked;
+  llvm::DenseMap<CanType, bool> canTypesChecked;
   llvm::SmallPtrSet<TypeBase *, 16> typesHandled;
 
   SILModule &M;
@@ -90,14 +90,18 @@ class CrossModuleOptimization {
   void trySerializeFunctions(ArrayRef<SILFunction *> functions);
 
   bool canSerializeFunction(SILFunction *function,
-                    FunctionFlags &canSerializeFlags, int maxDepth);
+                            FunctionFlags &canSerializeFlags,
+                            int maxDepth);
 
-  bool canSerializeInstruction(SILInstruction *inst,
-                    FunctionFlags &canSerializeFlags, int maxDepth);
+  bool canSerializeFieldsByInstructionKind(SILInstruction *inst,
+                               FunctionFlags &canSerializeFlags,
+                               int maxDepth);
 
   bool canSerializeGlobal(SILGlobalVariable *global);
 
   bool canSerializeType(SILType type);
+  bool canSerializeType(CanType type);
+  bool canSerializeDecl(NominalTypeDecl *decl);
 
   bool canUseFromInline(DeclContext *declCtxt);
 
@@ -120,11 +124,10 @@ class CrossModuleOptimization {
   void makeDeclUsableFromInline(ValueDecl *decl);
 
   void makeTypeUsableFromInline(CanType type);
-
-  void makeSubstUsableFromInline(const SubstitutionMap &substs);
 };
 
-/// Visitor for making used types of an instruction inlinable.
+/// Visitor for detecting if an instruction can be serialized and also making used
+/// types of an instruction inlinable if so.
 ///
 /// We use the SILCloner for visiting types, though it sucks that we allocate
 /// instructions just to delete them immediately. But it's better than to
@@ -136,12 +139,25 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
   friend class SILInstructionVisitor<InstructionVisitor>;
   friend class CrossModuleOptimization;
 
+public:
+  /// This visitor is used for 2 passes, 1st pass that detects whether the instruction
+  /// visited can be serialized, and 2nd pass that does the serializing.
+  enum class VisitMode {
+    DetectSerializableInst,
+    SerializeInst
+  };
+
 private:
   CrossModuleOptimization &CMS;
+  VisitMode mode;
+  bool isInstSerializable = true;
+  // Tracks whether an instruction can be serialized by
+  // checking its contained types and fields.
+  llvm::DenseMap<SILInstruction *, bool> instToSerializeMap;
 
 public:
-  InstructionVisitor(SILFunction &F, CrossModuleOptimization &CMS) :
-    SILCloner(F), CMS(CMS) {}
+  InstructionVisitor(SILFunction &F, CrossModuleOptimization &CMS, VisitMode visitMode) :
+    SILCloner(F), CMS(CMS), mode(visitMode) {}
 
   SILType remapType(SILType Ty) {
     if (Ty.hasLocalArchetype()) {
@@ -151,7 +167,15 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
                     SubstFlags::SubstituteLocalArchetypes);
     }
 
-    CMS.makeTypeUsableFromInline(Ty.getASTType());
+    switch (mode) {
+      case VisitMode::DetectSerializableInst:
+        if (!CMS.canSerializeType(Ty))
+          isInstSerializable = false;
+        break;
+      case VisitMode::SerializeInst:
+        CMS.makeTypeUsableFromInline(Ty.getASTType());
+        break;
+    }
     return Ty;
   }
 
@@ -162,7 +186,15 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
                     SubstFlags::SubstituteLocalArchetypes)->getCanonicalType();
     }
 
-    CMS.makeTypeUsableFromInline(Ty);
+    switch (mode) {
+      case VisitMode::DetectSerializableInst:
+        if (!CMS.canSerializeType(Ty))
+          isInstSerializable = false;
+        break;
+      case VisitMode::SerializeInst:
+        CMS.makeTypeUsableFromInline(Ty);
+        break;
+    }
     return Ty;
   }
 
@@ -173,18 +205,57 @@ class InstructionVisitor : public SILCloner<InstructionVisitor> {
                         SubstFlags::SubstituteLocalArchetypes);
     }
 
-    CMS.makeSubstUsableFromInline(Subs);
+    for (Type replType : Subs.getReplacementTypes()) {
+      switch (mode) {
+        case VisitMode::DetectSerializableInst:
+          CMS.canSerializeType(replType->getCanonicalType());
+          break;
+        case VisitMode::SerializeInst:
+          /// Ensure that all replacement types of \p Subs are usable from serialized
+          /// functions.
+          CMS.makeTypeUsableFromInline(replType->getCanonicalType());
+          break;
+      }
+    }
+    for (ProtocolConformanceRef pref : Subs.getConformances()) {
+      if (pref.isConcrete()) {
+        ProtocolConformance *concrete = pref.getConcrete();
+        switch (mode) {
+          case VisitMode::DetectSerializableInst:
+            if (!CMS.canSerializeDecl(concrete->getProtocol()))
+              isInstSerializable = false;
+            break;
+          case VisitMode::SerializeInst:
+            CMS.makeDeclUsableFromInline(concrete->getProtocol());
+            break;
+        }
+      }
+    }
     return Subs;
   }
 
   void postProcess(SILInstruction *Orig, SILInstruction *Cloned) {
     SILCloner<InstructionVisitor>::postProcess(Orig, Cloned);
+    switch (mode) {
+      case VisitMode::DetectSerializableInst:
+        instToSerializeMap[Orig] = isInstSerializable;
+        break;
+      case VisitMode::SerializeInst:
+        break;
+    }
     Cloned->eraseFromParent();
   }
 
   SILValue getMappedValue(SILValue Value) { return Value; }
 
   SILBasicBlock *remapBasicBlock(SILBasicBlock *BB) { return BB; }
+
+  bool canSerializeTypesInInst(SILInstruction *inst) {
+    auto iter = instToSerializeMap.find(inst);
+    if (iter != instToSerializeMap.end())
+      return iter->getSecond();
+    return false;
+  }
 };
 
 static bool isPackageCMOEnabled(ModuleDecl *mod) {
@@ -434,32 +505,32 @@ bool CrossModuleOptimization::canSerializeFunction(
     return false;
 
   // Check if any instruction prevents serializing the function.
+  InstructionVisitor visitor(*function, *this, InstructionVisitor::VisitMode::DetectSerializableInst);
   for (SILBasicBlock &block : *function) {
     for (SILInstruction &inst : block) {
-      if (!canSerializeInstruction(&inst, canSerializeFlags, maxDepth)) {
+      visitor.getBuilder().setInsertionPoint(&inst);
+      // First, visit each instruction and see if its lowered types
+      // are serializalbe and track them in a map.
+      visitor.visit(&inst);
+      if (!visitor.canSerializeTypesInInst(&inst))
+        return false;
+      // If serializable, check if any fields in the instruction (that
+      // are not covered in the visitor) can be serialized.
+    if (!canSerializeFieldsByInstructionKind(&inst, canSerializeFlags, maxDepth))
         return false;
-      }
     }
   }
+  M.reclaimUnresolvedLocalArchetypeDefinitions();
+
   canSerializeFlags[function] = true;
   return true;
 }
 
-/// Returns true if \p inst can be serialized.
+/// Returns true if \p inst can be serialized by checking its fields per instruction kind.
 ///
 /// If \p inst is a function_ref, recursively visits the referenced function.
-bool CrossModuleOptimization::canSerializeInstruction(
+bool CrossModuleOptimization::canSerializeFieldsByInstructionKind(
     SILInstruction *inst, FunctionFlags &canSerializeFlags, int maxDepth) {
-  // First check if any result or operand types prevent serialization.
-  for (SILValue result : inst->getResults()) {
-    if (!canSerializeType(result->getType()))
-      return false;
-  }
-  for (Operand &op : inst->getAllOperands()) {
-    if (!canSerializeType(op.get()->getType()))
-      return false;
-  }
-
   if (auto *FRI = dyn_cast<FunctionRefBaseInst>(inst)) {
     SILFunction *callee = FRI->getReferencedFunctionOrNull();
     if (!callee)
@@ -551,6 +622,45 @@ bool CrossModuleOptimization::canSerializeInstruction(
   return true;
 }
 
+bool CrossModuleOptimization::canSerializeType(SILType type) {
+  return canSerializeType(type.getASTType());
+}
+
+bool CrossModuleOptimization::canSerializeType(CanType type) {
+  auto iter = canTypesChecked.find(type);
+  if (iter != canTypesChecked.end())
+    return iter->getSecond();
+
+  bool success = type.findIf(
+     [this](Type rawSubType) {
+       CanType subType = rawSubType->getCanonicalType();
+       if (auto nominal = subType->getNominalOrBoundGenericNominal()) {
+         return canSerializeDecl(nominal);
+       }
+       // If reached here, the type is a Builtin type or similar,
+       // e.g. Builtin.Int64, Builtin.Word, Builtin.NativeObject, etc.
+       return true;
+     });
+
+  canTypesChecked[type] = success;
+  return success;
+}
+
+bool CrossModuleOptimization::canSerializeDecl(NominalTypeDecl *decl) {
+  assert(decl && "Decl should not be null when checking if it can be serialized");
+
+  // In conservative mode we don't want to change the access level of types.
+  if (conservative && decl->getEffectiveAccess() < AccessLevel::Package) {
+    return false;
+  }
+  // Exclude types which are defined in an @_implementationOnly imported
+  // module. Such modules are not transitively available.
+  if (!canUseFromInline(decl)) {
+    return false;
+  }
+  return true;
+}
+
 bool CrossModuleOptimization::canSerializeGlobal(SILGlobalVariable *global) {
   // Check for referenced functions in the initializer.
   for (const SILInstruction &initInst : *global) {
@@ -572,32 +682,6 @@ bool CrossModuleOptimization::canSerializeGlobal(SILGlobalVariable *global) {
   return true;
 }
 
-bool CrossModuleOptimization::canSerializeType(SILType type) {
-  auto iter = typesChecked.find(type);
-  if (iter != typesChecked.end())
-    return iter->getSecond();
-
-  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;
-        }
-
-        // Exclude types which are defined in an @_implementationOnly imported
-        // module. Such modules are not transitively available.
-        if (!canUseFromInline(subNT)) {
-          return true;
-        }
-      }
-      return false;
-    });
-  typesChecked[type] = success;
-  return success;
-}
-
 /// Returns true if the function in \p funcCtxt could be linked statically to
 /// this module.
 static bool couldBeLinkedStatically(DeclContext *funcCtxt, SILModule &module) {
@@ -730,7 +814,7 @@ void CrossModuleOptimization::serializeFunction(SILFunction *function,
   }
   function->setSerializedKind(getRightSerializedKind(M));
 
-  InstructionVisitor visitor(*function, *this);
+  InstructionVisitor visitor(*function, *this, InstructionVisitor::VisitMode::SerializeInst);
   for (SILBasicBlock &block : *function) {
     for (SILInstruction &inst : block) {
       visitor.getBuilder().setInsertionPoint(&inst);
@@ -929,21 +1013,6 @@ void CrossModuleOptimization::makeTypeUsableFromInline(CanType type) {
   });
 }
 
-/// Ensure that all replacement types of \p substs are usable from serialized
-/// functions.
-void CrossModuleOptimization::makeSubstUsableFromInline(
-                                                const SubstitutionMap &substs) {
-  for (Type replType : substs.getReplacementTypes()) {
-    makeTypeUsableFromInline(replType->getCanonicalType());
-  }
-  for (ProtocolConformanceRef pref : substs.getConformances()) {
-    if (pref.isConcrete()) {
-      ProtocolConformance *concrete = pref.getConcrete();
-      makeDeclUsableFromInline(concrete->getProtocol());
-    }
-  }
-}
-
 class CrossModuleOptimizationPass: public SILModuleTransform {
   void run() override {
     auto &M = *getModule();

test/SILOptimizer/package-cmo-cast-internal-dst

@@ -0,0 +1,29 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend -emit-sil %t/Lib.swift -package-name pkg \
+// RUN:   -wmo -allow-non-resilient-access -package-cmo \
+// RUN:   -enable-library-evolution -swift-version 5 \
+// RUN:   -Xllvm -sil-print-function=topFunc -o %t/Lib.sil
+
+// RUN: %FileCheck %s < %t/Lib.sil
+
+/// Verify that `InternalKlass` is visited and the instruction containing it is not serialized.
+// CHECK: sil @$s3Lib7topFuncySiAA3PubCF : $@convention(thin) (@guaranteed Pub) -> Int {
+// CHECK:   checked_cast_br Pub in %0 : $Pub to InternalKlass
+
+
+//--- Lib.swift
+public class Pub {
+  public var pubVar: Int
+  public init(_ arg: Int) {
+    pubVar = arg
+  }
+}
+
+class InternalKlass: Pub {}
+
+public func topFunc(_ arg: Pub) -> Int {
+  let x = arg as? InternalKlass
+  return x != nil ? 1 : 0
+}