embedded: move the VTableSpecializer pass into MandatoryPerformanceOptimizations

MandatoryPerformanceOptimizations already did most of the vtable specialization work.
So it makes sense to remove the VTableSpecializerPass completely and do everything in MandatoryPerformanceOptimizations.

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -33,13 +33,23 @@ let mandatoryPerformanceOptimizations = ModulePass(name: "mandatory-performance-
   // For embedded Swift, optimize all the functions (there cannot be any
   // generics, type metadata, etc.)
   if moduleContext.options.enableEmbeddedSwift {
+    // We need to specialize all vtables which are referenced from non-generic contexts. Beside
+    // `alloc_ref`s of generic classes in non-generic functions, we also need to specialize generic
+    // superclasses of non-generic classes. E.g. `class Derived : Base<Int> {}`
+    specializeVTablesOfSuperclasses(moduleContext)
+
     worklist.addAllNonGenericFunctions(of: moduleContext)
   } else {
     worklist.addAllPerformanceAnnotatedFunctions(of: moduleContext)
     worklist.addAllAnnotatedGlobalInitOnceFunctions(of: moduleContext)
   }
 
   optimizeFunctionsTopDown(using: &worklist, moduleContext)
+
+  if moduleContext.options.enableEmbeddedSwift {
+    // Print errors for generic functions in vtables, which is not allowed in embedded Swift.
+    checkVTablesForGenericFunctions(moduleContext)
+  }
 }
 
 private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
@@ -92,11 +102,15 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
       // Embedded Swift specific transformations
       case let alloc as AllocRefInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: alloc.type, in: function, context, moduleContext, &worklist)
+          specializeVTableAndAddEntriesToWorklist(for: alloc.type, in: function,
+                                                  errorLocation: alloc.location,
+                                                  moduleContext, &worklist)
         }
       case let metatype as MetatypeInst:
         if context.options.enableEmbeddedSwift {
-          specializeVTableAndAddEntriesToWorklist(for: metatype.type, in: function, context, moduleContext, &worklist)
+          specializeVTableAndAddEntriesToWorklist(for: metatype.type, in: function,
+                                                  errorLocation: metatype.location,
+                                                  moduleContext, &worklist)
         }
       case let classMethod as ClassMethodInst:
         if context.options.enableEmbeddedSwift {
@@ -144,18 +158,23 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
 }
 
 private func specializeVTableAndAddEntriesToWorklist(for type: Type, in function: Function,
-                                                     _ context: FunctionPassContext, _ moduleContext: ModulePassContext,
+                                                     errorLocation: Location,
+                                                     _ moduleContext: ModulePassContext,
                                                      _ worklist: inout FunctionWorklist) {
   let vTablesCountBefore = moduleContext.vTables.count
 
-  guard context.specializeVTable(for: type, in: function) != nil else {
+  guard specializeVTable(forClassType: type, errorLocation: errorLocation, moduleContext) != nil else {
     return
   }
 
   // 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 {
+    for entry in vTables[i].entries
+      // A new vtable can still contain a generic function if the method couldn't be specialized for some reason
+      // and an error has been printed. Exclude generic functions to not run into an assert later.
+      where !entry.implementation.isGeneric
+    {
       worklist.pushIfNotVisited(entry.implementation)
     }
   }
@@ -240,6 +259,14 @@ private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlined
   return false
 }
 
+private func checkVTablesForGenericFunctions(_ context: ModulePassContext) {
+  for vTable in context.vTables where !vTable.class.isGenericAtAnyLevel {
+    for entry in vTable.entries where entry.implementation.isGeneric {
+      context.diagnosticEngine.diagnose(entry.methodDecl.location.sourceLoc, .non_final_generic_class_function)
+    }
+  }
+}
+
 private extension FullApplySite {
   func resultIsUsedInGlobalInitialization() -> SmallProjectionPath? {
     guard parentFunction.isGlobalInitOnceFunction,

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -339,13 +339,6 @@ struct FunctionPassContext : MutatingContext {
     return false
   }
 
-  func specializeVTable(for type: Type, in function: Function) -> VTable? {
-    guard let vtablePtr = _bridged.specializeVTableForType(type.bridged, function.bridged) else {
-      return nil
-    }
-    return VTable(bridged: BridgedVTable(vTable: vtablePtr))
-  }
-
   func specializeClassMethodInst(_ cm: ClassMethodInst) -> Bool {
     if _bridged.specializeClassMethodInst(cm.bridged) {
       notifyInstructionsChanged()

SwiftCompilerSources/Sources/Optimizer/Utilities/CMakeLists.txt

@@ -17,6 +17,7 @@ swift_compiler_sources(Optimizer
   EscapeUtils.swift
   ForwardingUtils.swift
   FunctionSignatureTransforms.swift
+  GenericSpecialization.swift
   LifetimeDependenceUtils.swift
   LocalVariableUtils.swift
   OptUtils.swift

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -0,0 +1,77 @@
+//===--- GenericSpecialization.swift ---------------------------------------==//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+
+@discardableResult
+func specializeVTable(forClassType classType: Type,
+                      errorLocation: Location,
+                      _ context: ModulePassContext) -> VTable?
+{
+  if !classType.isClass {
+    return nil
+  }
+  if !classType.isGenericAtAnyLevel {
+    return nil
+  }
+  let classDecl = classType.nominal
+
+  if context.lookupSpecializedVTable(for: classType) != nil {
+    return nil
+  }
+
+  guard let origVTable = context.lookupVTable(for: classDecl) else {
+    context.diagnosticEngine.diagnose(errorLocation.sourceLoc, .cannot_specialize_class, classType)
+    return nil
+  }
+
+  let classContextSubs = classType.contextSubstitutionMap
+
+  let newEntries = origVTable.entries.map { origEntry in
+    if !origEntry.implementation.isGeneric {
+      return origEntry
+    }
+    let methodSubs = classContextSubs.getMethodSubstitutions(for: origEntry.implementation)
+
+    guard !methodSubs.conformances.contains(where: {!$0.isValid}),
+          let specializedMethod = context.specialize(function: origEntry.implementation, for: methodSubs) else
+    {
+      context.diagnosticEngine.diagnose(origEntry.methodDecl.location.sourceLoc, .non_final_generic_class_function)
+      return origEntry
+    }
+
+    context.deserializeAllCallees(of: specializedMethod, mode: .allFunctions)
+    specializedMethod.set(linkage: .public, context)
+    specializedMethod.set(isSerialized: false, context)
+
+    return VTable.Entry(kind: origEntry.kind, isNonOverridden: origEntry.isNonOverridden,
+                        methodDecl: origEntry.methodDecl, implementation: specializedMethod)
+  }
+
+  let specializedVTable = context.createSpecializedVTable(entries: newEntries, for: classType, isSerialized: false)
+  if let superClassTy = classType.superClassType {
+    specializeVTable(forClassType: superClassTy, errorLocation: classDecl.location, context)
+  }
+  return specializedVTable
+}
+
+func specializeVTablesOfSuperclasses(_ moduleContext: ModulePassContext) {
+  for vtable in moduleContext.vTables {
+    if !vtable.isSpecialized,
+       !vtable.class.isGenericAtAnyLevel,
+       let superClassTy = vtable.class.superClassType,
+       superClassTy.nominal.isGenericAtAnyLevel
+    {
+      specializeVTable(forClassType: superClassTy, errorLocation: vtable.class.location, moduleContext)
+    }
+  }
+}

include/swift/SILOptimizer/OptimizerBridging.h

@@ -252,8 +252,6 @@ struct BridgedPassContext {
   SWIFT_IMPORT_UNSAFE DevirtResult tryDevirtualizeApply(BridgedInstruction apply, bool isMandatory) const;
   bool tryOptimizeKeypath(BridgedInstruction apply) const;
   SWIFT_IMPORT_UNSAFE OptionalBridgedValue constantFoldBuiltin(BridgedInstruction builtin) const;
-  SWIFT_IMPORT_UNSAFE swift::SILVTable * _Nullable specializeVTableForType(BridgedType type,
-                                                                           BridgedFunction function) const;
   SWIFT_IMPORT_UNSAFE OptionalBridgedFunction specializeFunction(BridgedFunction function,
                                                                  BridgedSubstitutionMap substitutions) const;
   void deserializeAllCallees(BridgedFunction function, bool deserializeAll) const;

include/swift/SILOptimizer/PassManager/Passes.def

@@ -258,8 +258,6 @@ SWIFT_FUNCTION_PASS(DeadStoreElimination, "dead-store-elimination",
      "Dead Store Elimination")
 PASS(GenericSpecializer, "generic-specializer",
      "Generic Function Specialization on Static Types")
-PASS(VTableSpecializer, "vtable-specializer",
-     "Generic Specialization on VTables")
 PASS(ExistentialSpecializer, "existential-specializer",
      "Existential Specializer")
 PASS(SILSkippingChecker, "check-sil-skipping",

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -606,8 +606,6 @@ bool optimizeMemoryAccesses(SILFunction *fn, DominanceInfo *domInfo);
 /// See the PredictableDeadAllocationElimination pass.
 bool eliminateDeadAllocations(SILFunction *fn, DominanceInfo *domInfo);
 
-SILVTable *specializeVTableForType(SILType type, SILModule &mod, SILTransform *transform);
-
 bool specializeClassMethodInst(ClassMethodInst *cm);
 
 bool specializeAppliesInFunction(SILFunction &F,

lib/SILOptimizer/Mandatory/PerformanceDiagnostics.cpp

@@ -914,7 +914,7 @@ class PerformanceDiagnosticsPass : public SILModuleTransform {
       for (SILFunction *function : constructorsAndDestructors) {
         diagnoser.visitFunctionEmbeddedSwift(function);
       }
-    }    
+    }
   }
 };
 

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1735,12 +1735,6 @@ bool BridgedPassContext::canMakeStaticObjectReadOnly(BridgedType type) const {
   return false;
 }
 
-swift::SILVTable * BridgedPassContext::specializeVTableForType(BridgedType type, BridgedFunction function) const {
-  return ::specializeVTableForType(type.unbridged(),
-                                   function.getFunction()->getModule(),
-                                   invocation->getTransform());
-}
-
 OptionalBridgedFunction BridgedPassContext::specializeFunction(BridgedFunction function,
                                                                BridgedSubstitutionMap substitutions) const {
   swift::SILModule *mod = invocation->getPassManager()->getModule();

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -255,9 +255,6 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   // Canonical swift requires all non cond_br critical edges to be split.
   P.addSplitNonCondBrCriticalEdges();
 
-  // For embedded Swift: Specialize generic class vtables.
-  P.addVTableSpecializer();
-
   P.addMandatoryPerformanceOptimizations();
   P.addOnoneSimplification();
   P.addAllocVectorLowering();