[embedded] Specialize superclasses in VTableSpecializer as part of MandatoryPerformanceOptimizations

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -56,7 +56,7 @@ private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
         f.set(isPerformanceConstraint: true, context)
       }
 
-      optimize(function: f, context, &worklist)
+      optimize(function: f, context, moduleContext, &worklist)
     }
 
     // Generic specialization takes care of removing metatype arguments of generic functions.
@@ -73,7 +73,7 @@ fileprivate struct PathFunctionTuple: Hashable {
   var function: Function
 }
 
-private func optimize(function: Function, _ context: FunctionPassContext, _ worklist: inout FunctionWorklist) {
+private func optimize(function: Function, _ context: FunctionPassContext, _ moduleContext: ModulePassContext, _ worklist: inout FunctionWorklist) {
   var alreadyInlinedFunctions: Set<PathFunctionTuple> = Set()
 
   var changed = true
@@ -92,11 +92,11 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ work
       // Embedded Swift specific transformations
       case let alloc as AllocRefInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: alloc.type, in: function, context, &worklist)
+          specializeVTableAndAddEntriesToWorklist(for: alloc.type, in: function, context, moduleContext, &worklist)
         }
       case let metatype as MetatypeInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: metatype.type, in: function, context, &worklist)
+          specializeVTableAndAddEntriesToWorklist(for: metatype.type, in: function, context, moduleContext, &worklist)
         }
       case let classMethod as ClassMethodInst:
         if context.options.enableEmbeddedSwift {
@@ -133,13 +133,21 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ work
   }
 }
 
-private func specializeVTableAndAddEntriesToWorklist(for type: Type, in function: Function, _ context: FunctionPassContext, _ worklist: inout FunctionWorklist) {
-  guard let vtable = context.specializeVTable(for: type, in: function) else {
+private func specializeVTableAndAddEntriesToWorklist(for type: Type, in function: Function,
+                                                     _ context: FunctionPassContext, _ moduleContext: ModulePassContext,
+                                                     _ worklist: inout FunctionWorklist) {
+  let vTablesCountBefore = moduleContext.vTables.count
+
+  guard context.specializeVTable(for: type, in: function) != nil else {
     return
   }
 
-  for entry in vtable.entries {
-    worklist.pushIfNotVisited(entry.function)
+  // More than one new vtable might have been created (superclasses), process them all
+  let vTables = moduleContext.vTables
+  for i in vTablesCountBefore ..< vTables.count {
+    for entry in vTables[i].entries {
+      worklist.pushIfNotVisited(entry.function)
+    }
   }
 }
 

lib/SILOptimizer/Transforms/VTableSpecializer.cpp

@@ -38,7 +38,6 @@ namespace {
 
 class VTableSpecializer : public SILModuleTransform {
   bool specializeVTables(SILModule &module);
-  bool specializeVTablesOfSuperclasses(SILModule &module);
 
   /// The entry point to the transformation.
   void run() override {
@@ -59,6 +58,8 @@ static SILFunction *specializeVTableMethod(SILFunction *origMethod,
                                            SILModule &module,
                                            SILTransform *transform);
 
+static bool specializeVTablesOfSuperclasses(SILModule &module, SILTransform *transform);
+
 static bool specializeVTablesInFunction(SILFunction &func, SILModule &module,
                                         SILTransform *transform) {
   bool changed = false;
@@ -89,7 +90,7 @@ bool VTableSpecializer::specializeVTables(SILModule &module) {
     changed |= specializeVTablesInFunction(func, module, this);
   }
 
-  changed |= specializeVTablesOfSuperclasses(module);
+  changed |= specializeVTablesOfSuperclasses(module, this);
 
   for (SILVTable *vtable : module.getVTables()) {
     if (vtable->getClass()->isGenericContext()) continue;
@@ -115,24 +116,36 @@ bool VTableSpecializer::specializeVTables(SILModule &module) {
   return changed;
 }
 
-bool VTableSpecializer::specializeVTablesOfSuperclasses(SILModule &module) {
+static bool specializeVTablesOfSuperclasses(SILVTable *vtable,
+                                            SILModule &module,
+                                            SILTransform *transform) {
+  if (vtable->getClass()->isGenericContext() && !vtable->getClassType())
+    return false;
+
+  SILType superClassTy;
+  if (SILType classTy = vtable->getClassType()) {
+    superClassTy = classTy.getSuperclass();
+  } else {
+    if (Type superTy = vtable->getClass()->getSuperclass())
+      superClassTy =
+          SILType::getPrimitiveObjectType(superTy->getCanonicalType());
+  }
+  if (superClassTy) {
+    return (specializeVTableForType(superClassTy, module, transform) !=
+            nullptr);
+  }
+
+  return false;
+}
+
+static bool specializeVTablesOfSuperclasses(SILModule &module,
+                                            SILTransform *transform) {
   bool changed = false;
-  // The module's vtable table can grow while we are specializing superclass vtables.
+  // The module's vtable table can grow while we are specializing superclass
+  // vtables.
   for (unsigned i = 0; i < module.getVTables().size(); ++i) {
     SILVTable *vtable = module.getVTables()[i];
-    if (vtable->getClass()->isGenericContext() && !vtable->getClassType())
-      continue;
-
-    SILType superClassTy;
-    if (SILType classTy = vtable->getClassType()) {
-      superClassTy = classTy.getSuperclass();
-    } else {
-      if (Type superTy = vtable->getClass()->getSuperclass())
-        superClassTy = SILType::getPrimitiveObjectType(superTy->getCanonicalType());
-    }
-    if (superClassTy) {
-      changed |= (specializeVTableForType(superClassTy, module, this) != nullptr);
-    }
+    specializeVTablesOfSuperclasses(vtable, module, transform);
   }
   return changed;
 }
@@ -173,6 +186,9 @@ SILVTable *swift::specializeVTableForType(SILType classTy, SILModule &module,
 
   SILVTable *vtable = SILVTable::create(module, classDecl, classTy,
                                         IsNotSerialized, newEntries);
+
+  specializeVTablesOfSuperclasses(vtable, module, transform);
+
   return vtable;
 }
 

test/embedded/generic-classes2.swift

@@ -0,0 +1,31 @@
+// RUN: %target-run-simple-swift(-enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) | %FileCheck %s
+// RUN: %target-run-simple-swift(-Osize -enable-experimental-feature Embedded -parse-as-library -runtime-compatibility-version none -wmo -Xfrontend -disable-objc-interop) | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+// REQUIRES: executable_test
+// REQUIRES: optimized_stdlib
+// REQUIRES: OS=macosx || OS=linux-gnu
+
+class B<T> {
+}
+
+class D<T>: B<T> {
+}
+
+func callee<T>(_ t: T.Type) {
+  _ = D<T>()
+}
+
+public func test() {
+  callee(Int.self)
+}
+
+@main
+struct Main {
+  static func main() {
+    test()
+    print("OK!")
+  }
+}
+
+// CHECK: OK!