MandatoryPerformanceOptimizations: specialize witness tables to support class existentials with generic classes

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -112,6 +112,11 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
           _ = context.specializeClassMethodInst(classMethod)
         }
 
+      case let initExRef as InitExistentialRefInst:
+        if context.options.enableEmbeddedSwift {
+          specializeWitnessTables(for: initExRef, moduleContext, &worklist)
+        }
+
       // We need to de-virtualize deinits of non-copyable types to be able to specialize the deinitializers.
       case let destroyValue as DestroyValueInst:
         if !devirtualizeDeinits(of: destroyValue, simplifyCtxt) {
@@ -254,6 +259,31 @@ private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlined
   return false
 }
 
+private func specializeWitnessTables(for initExRef: InitExistentialRefInst, _ context: ModulePassContext,
+                                     _ worklist: inout FunctionWorklist)
+{
+  let conformingType = initExRef.instance.type
+  assert(conformingType.isClass)
+
+  for conformance in initExRef.conformances {
+    if conformance.isConcrete,
+       conformance.isSpecialized,
+       context.lookupWitnessTable(for: conformance) == nil
+    {
+      let wt = specializeWitnessTable(forConformance: conformance, context)
+      for entry in wt.entries where entry.kind == .method {
+        if let method = entry.methodFunction,
+           // A new witness table 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.
+           !method.isGeneric
+        {
+          worklist.pushIfNotVisited(method)
+        }
+      }
+    }
+  }
+}
+
 private extension FullApplySite {
   func resultIsUsedInGlobalInitialization() -> SmallProjectionPath? {
     guard parentFunction.isGlobalInitOnceFunction,

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -75,3 +75,31 @@ func specializeVTablesOfSuperclasses(_ moduleContext: ModulePassContext) {
     }
   }
 }
+
+func specializeWitnessTable(forConformance conformance: ProtocolConformance,
+                            _ context: ModulePassContext) -> WitnessTable
+{
+  let genericConformance = conformance.genericConformance
+  guard let witnessTable = context.lookupWitnessTable(for: genericConformance) else {
+    fatalError("no witness table found")
+  }
+  assert(witnessTable.isDefinition, "No witness table available")
+
+  let newEntries = witnessTable.entries.map { origEntry in
+    switch origEntry.kind {
+    case .method:
+      let methodDecl = origEntry.methodRequirement
+      guard let specializedMethod = context.specialize(function: origEntry.methodFunction!,
+                                                       for: conformance.specializedSubstitutions) else {
+        context.diagnosticEngine.diagnose(methodDecl.location.sourceLoc, .cannot_specialize_witness_method)
+        return origEntry
+      }
+      return WitnessTable.Entry(methodRequirement: methodDecl, methodFunction: specializedMethod)
+    default:
+      // TODO: handle other witness table entry kinds
+      fatalError("unsupported witness table etnry")
+    }
+    return origEntry
+  }
+  return context.createWitnessTable(entries: newEntries, conformance: conformance, linkage: .shared, serialized: false)
+}

include/swift/AST/DiagnosticsSIL.def

@@ -377,6 +377,8 @@ ERROR(global_must_be_compile_time_const,none,
       "global variable must be a compile-time constant", ())
 ERROR(non_final_generic_class_function,none,
       "classes cannot have non-final generic fuctions in embedded Swift", ())
+ERROR(cannot_specialize_witness_method,none,
+      "generic method cannot be used in existential type", ())
 ERROR(cannot_specialize_class,none,
       "cannot specialize %0 because class definition is not available (make sure to build with -wmo)", (Type))
 ERROR(embedded_swift_existential_type,none,

test/embedded/existential-class-bound4.swift

@@ -5,9 +5,6 @@
 // REQUIRES: optimized_stdlib
 // REQUIRES: OS=macosx || OS=linux-gnu
 
-// Generic classes don't work yet.
-// XFAIL: *
-
 protocol ClassBound: AnyObject {
     func foo()
     func bar()