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

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -117,6 +117,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) {
@@ -259,6 +264,37 @@ private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlined
   return false
 }
 
+private func specializeWitnessTables(for initExRef: InitExistentialRefInst, _ context: ModulePassContext,
+                                     _ worklist: inout FunctionWorklist)
+{
+  for conformance in initExRef.conformances where conformance.isConcrete {
+    let origWitnessTable = context.lookupWitnessTable(for: conformance)
+    if conformance.isSpecialized {
+      if origWitnessTable == nil {
+        let wt = specializeWitnessTable(forConformance: conformance, errorLocation: initExRef.location, context)
+        worklist.addWitnessMethods(of: wt)
+      }
+    } else if let origWitnessTable {
+      checkForGenericMethods(in: origWitnessTable, errorLocation: initExRef.location, context)
+    }
+  }
+}
+
+private func checkForGenericMethods(in witnessTable: WitnessTable,
+                                    errorLocation: Location,
+                                    _ context: ModulePassContext)
+{
+  for entry in witnessTable.entries where entry.kind == .method {
+    if let method = entry.methodFunction,
+       method.isGeneric
+    {
+      context.diagnosticEngine.diagnose(errorLocation.sourceLoc, .cannot_specialize_witness_method,
+                                        entry.methodRequirement)
+      return
+    }
+  }
+}
+
 private func checkVTablesForGenericFunctions(_ context: ModulePassContext) {
   for vTable in context.vTables where !vTable.class.isGenericAtAnyLevel {
     for entry in vTable.entries where entry.implementation.isGeneric {
@@ -472,6 +508,18 @@ fileprivate struct FunctionWorklist {
     }
   }
 
+  mutating func addWitnessMethods(of witnessTable: WitnessTable) {
+    for entry in witnessTable.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
+      {
+        pushIfNotVisited(method)
+      }
+    }
+  }
+
   mutating func pushIfNotVisited(_ element: Function) {
     if pushedFunctions.insert(element).inserted {
       functions.append(element)

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -75,3 +75,38 @@ func specializeVTablesOfSuperclasses(_ moduleContext: ModulePassContext) {
     }
   }
 }
+
+func specializeWitnessTable(forConformance conformance: ProtocolConformance,
+                            errorLocation: Location,
+                            _ 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:
+      guard let origMethod = origEntry.methodFunction else {
+        return origEntry
+      }
+      let methodDecl = origEntry.methodRequirement
+      let methodSubs = conformance.specializedSubstitutions.getMethodSubstitutions(for: origMethod)
+
+      guard !methodSubs.conformances.contains(where: {!$0.isValid}),
+            let specializedMethod = context.specialize(function: origMethod, for: methodSubs) else
+      {
+        context.diagnosticEngine.diagnose(errorLocation.sourceLoc, .cannot_specialize_witness_method, methodDecl)
+        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,
+      "an existential type cannot contain a generic method %0 in embedded Swift", (DeclName))
 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()

test/embedded/existential-class-bound7.swift

@@ -0,0 +1,28 @@
+// RUN: %target-swift-emit-ir -parse-as-library -module-name main -verify %s -enable-experimental-feature Embedded -wmo
+
+// REQUIRES: swift_in_compiler
+// REQUIRES: optimized_stdlib
+// REQUIRES: OS=macosx || OS=linux-gnu
+
+public protocol P: AnyObject {
+  func foo<T>(t: T)
+}
+
+final public class Class: P {
+  public func foo<T>(t: T) {}
+}
+
+
+public func testClass() -> P {
+  return Class() // expected-error {{an existential type cannot contain a generic method 'foo(t:)' in embedded Swift}}
+}
+
+final public class GenClass<X>: P {
+  public func foo<T>(t: T) {}
+}
+
+
+public func testGenClass() -> P {
+  return GenClass<Int>() // expected-error {{an existential type cannot contain a generic method 'foo(t:)' in embedded Swift}}
+}
+