MandatoryPerformanceOptimizations: fix adding witness methods to the function worklist

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -138,9 +138,8 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
       case let initExRef as InitExistentialRefInst:
         if context.options.enableEmbeddedSwift {
           for c in initExRef.conformances where c.isConcrete {
-            specializeWitnessTable(for: c, moduleContext) {
-              worklist.addWitnessMethods(of: $0)
-            }
+            specializeWitnessTable(for: c, moduleContext)
+            worklist.addWitnessMethods(of: c, moduleContext)
           }
         }
 
@@ -149,9 +148,9 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
         case .BuildOrdinaryTaskExecutorRef,
              .BuildOrdinarySerialExecutorRef,
              .BuildComplexEqualitySerialExecutorRef:
-          specializeWitnessTable(for: bi.substitutionMap.conformances[0], moduleContext) {
-            worklist.addWitnessMethods(of: $0)
-          }
+          let conformance = bi.substitutionMap.conformances[0]
+          specializeWitnessTable(for: conformance, moduleContext)
+          worklist.addWitnessMethods(of: conformance, moduleContext)
 
         default:
           break
@@ -515,15 +514,44 @@ extension FunctionWorklist {
     }
   }
 
-  mutating func addWitnessMethods(of witnessTable: WitnessTable) {
+  mutating func addWitnessMethods(of conformance: Conformance, _ context: ModulePassContext) {
+    var visited = Set<Conformance>()
+    addWitnessMethodsRecursively(of: conformance, visited: &visited, context)
+  }
+
+  private mutating func addWitnessMethodsRecursively(of conformance: Conformance,
+                                                     visited: inout Set<Conformance>,
+                                                     _ context: ModulePassContext)
+  {
+    guard conformance.isConcrete,
+          visited.insert(conformance).inserted
+    else {
+      return
+    }
+    let witnessTable: WitnessTable
+    if let wt = context.lookupWitnessTable(for: conformance) {
+      witnessTable = wt
+    } else if let wt = context.lookupWitnessTable(for: conformance.rootConformance) {
+      witnessTable = wt
+    } else {
+      return
+    }
     for entry in witnessTable.entries {
-      if case .method(_, let witness) = entry,
-         let method = witness,
-         // 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)
+      switch entry {
+      case .invalid, .associatedType:
+        break
+      case .method(_, let witness):
+        if let method = witness,
+           // 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)
+        }
+      case .baseProtocol(_, let baseConf):
+        addWitnessMethodsRecursively(of: baseConf, visited: &visited, context)
+      case .associatedConformance(_, let assocConf):
+        addWitnessMethodsRecursively(of: assocConf, visited: &visited, context)
       }
     }
   }

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -106,10 +106,7 @@ private struct VTableSpecializer {
 }
 
 /// Specializes a witness table of `conformance` for the concrete type of the conformance.
-func specializeWitnessTable(for conformance: Conformance,
-                            _ context: ModulePassContext,
-                            _ notifyNewWitnessTable: (WitnessTable) -> ())
-{
+func specializeWitnessTable(for conformance: Conformance, _ context: ModulePassContext) {
   if let existingSpecialization = context.lookupWitnessTable(for: conformance),
          existingSpecialization.isSpecialized
   {
@@ -125,7 +122,7 @@ func specializeWitnessTable(for conformance: Conformance,
   let baseConf = conformance.isInherited ? conformance.inheritedConformance: conformance
   if !baseConf.isSpecialized {
     var visited = Set<Conformance>()
-    specializeDefaultMethods(for: conformance, visited: &visited, context, notifyNewWitnessTable)
+    specializeDefaultMethods(for: conformance, visited: &visited, context)
     return
   }
 
@@ -158,7 +155,7 @@ func specializeWitnessTable(for conformance: Conformance,
       let baseConf = context.getSpecializedConformance(of: witness,
                                                        for: conformance.type,
                                                        substitutions: conformance.specializedSubstitutions)
-      specializeWitnessTable(for: baseConf, context, notifyNewWitnessTable)
+      specializeWitnessTable(for: baseConf, context)
       return .baseProtocol(requirement: requirement, witness: baseConf)
     case .associatedType(let requirement, let witness):
       let substType = witness.subst(with: conformance.specializedSubstitutions)
@@ -169,15 +166,14 @@ func specializeWitnessTable(for conformance: Conformance,
       let concreteAssociateConf = conformance.getAssociatedConformance(ofAssociatedType: requirement.rawType,
                                                                        to: assocConf.protocol)
       if concreteAssociateConf.isSpecialized {
-        specializeWitnessTable(for: concreteAssociateConf, context, notifyNewWitnessTable)
+        specializeWitnessTable(for: concreteAssociateConf, context)
       }
       return .associatedConformance(requirement: requirement,
                                     witness: concreteAssociateConf)
     }
   }
-  let newWT = context.createSpecializedWitnessTable(entries: newEntries,conformance: conformance,
-                                                    linkage: .shared, serialized: false)
-  notifyNewWitnessTable(newWT)
+  context.createSpecializedWitnessTable(entries: newEntries,conformance: conformance,
+                                        linkage: .shared, serialized: false)
 }
 
 /// Specializes the default methods of a non-generic witness table.
@@ -186,8 +182,7 @@ func specializeWitnessTable(for conformance: Conformance,
 /// specialize inherited conformances so that the concrete self type is correct, even for derived classes.
 private func specializeDefaultMethods(for conformance: Conformance,
                                       visited: inout Set<Conformance>,
-                                      _ context: ModulePassContext,
-                                      _ notifyNewWitnessTable: (WitnessTable) -> ())
+                                      _ context: ModulePassContext)
 {
   // Avoid infinite recursion, which may happen if an associated conformance is the conformance itself.
   guard visited.insert(conformance).inserted,
@@ -224,21 +219,20 @@ private func specializeDefaultMethods(for conformance: Conformance,
       specialized = true
       return .method(requirement: requirement, witness: specializedMethod)
     case .baseProtocol(_, let witness):
-      specializeDefaultMethods(for: witness, visited: &visited, context, notifyNewWitnessTable)
+      specializeDefaultMethods(for: witness, visited: &visited, context)
       return origEntry
     case .associatedType:
       return origEntry
     case .associatedConformance(_, let assocConf):
-      specializeDefaultMethods(for: assocConf, visited: &visited, context, notifyNewWitnessTable)
+      specializeDefaultMethods(for: assocConf, visited: &visited, context)
       return origEntry
     }
   }
   // If the witness table does not contain any default methods, there is no need to create a
   // specialized witness table.
   if specialized {
-    let newWT = context.createSpecializedWitnessTable(entries: newEntries,conformance: conformance,
-                                                      linkage: .shared, serialized: false)
-    notifyNewWitnessTable(newWT)
+    context.createSpecializedWitnessTable(entries: newEntries,conformance: conformance,
+                                          linkage: .shared, serialized: false)
   }
 }