MandatoryPerformanceOptimizations: make sure to handle de-serialized vtable methods

When de-serializing a function and this function allocates a class, the methods of the de-serialized vtable must be handled, too.

Fixes an IRGen crash
rdar://152311945

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -68,7 +68,7 @@ private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
     // We need handle this case with a function signature optimization.
     removeMetatypeArgumentsInCallees(of: f, moduleContext)
 
-    worklist.addCallees(of: f)
+    worklist.addCallees(of: f, moduleContext)
   }
 }
 
@@ -521,29 +521,41 @@ fileprivate struct FunctionWorklist {
     return
   }
 
-  mutating func addCallees(of function: Function) {
+  mutating func addCallees(of function: Function, _ context: ModulePassContext) {
     for inst in function.instructions {
       switch inst {
-      case let apply as ApplySite:
-        if let callee = apply.referencedFunction {
-          pushIfNotVisited(callee)
+      case let fri as FunctionRefInst:
+        pushIfNotVisited(fri.referencedFunction)
+      case let alloc as AllocRefInst:
+        if context.options.enableEmbeddedSwift {
+          addVTableMethods(forClassType: alloc.type, context)
         }
-      case let bi as BuiltinInst:
-        switch bi.id {
-        case .Once, .OnceWithContext:
-          if let fri = bi.operands[1].value as? FunctionRefInst {
-            pushIfNotVisited(fri.referencedFunction)
+      case let metatype as MetatypeInst:
+        if context.options.enableEmbeddedSwift {
+          let instanceType = metatype.type.loweredInstanceTypeOfMetatype(in: function)
+          if instanceType.isClass {
+            addVTableMethods(forClassType: instanceType, context)
           }
-          break;
-        default:
-          break
         }
+
       default:
         break
       }
     }
   }
 
+  mutating func addVTableMethods(forClassType classType: Type, _ context: ModulePassContext) {
+    guard let vtable = classType.isGenericAtAnyLevel ?
+                        context.lookupSpecializedVTable(for: classType) :
+                        context.lookupVTable(for: classType.nominal!)
+    else {
+      return
+    }
+    for entry in vtable.entries where !entry.implementation.isGeneric {
+      pushIfNotVisited(entry.implementation)
+    }
+  }
+
   mutating func addWitnessMethods(of witnessTable: WitnessTable) {
     for entry in witnessTable.entries {
       if case .method(_, let witness) = entry,

test/SILOptimizer/mandatory_performance_optimizations.sil

@@ -215,7 +215,7 @@ bb0(%0 : $Int, %1 : @owned $Builtin.NativeObject):
   return %8 : $Builtin.NativeObject
 }
 
-// CHECK-LABEL: sil [signature_optimized_thunk] [ossa] @metatype_arg :
+// CHECK-LABEL: sil [signature_optimized_thunk] [perf_constraint] [ossa] @metatype_arg :
 sil [ossa] @metatype_arg : $@convention(thin) (Int, @thick Int.Type, @owned Builtin.NativeObject) -> @owned Builtin.NativeObject {
 bb0(%0 : $Int, %1 : $@thick Int.Type, %2 : @owned $Builtin.NativeObject):
   fix_lifetime %1 : $@thick Int.Type
@@ -243,7 +243,7 @@ bb2(%13 : @owned $any Error):
   throw %13 : $any Error
 }
 
-// CHECK-LABEL: sil [signature_optimized_thunk] [ossa] @metatype_arg_throws :
+// CHECK-LABEL: sil [signature_optimized_thunk] [perf_constraint] [ossa] @metatype_arg_throws :
 sil [ossa] @metatype_arg_throws : $@convention(thin) (Int, @thick Int.Type, @owned Builtin.NativeObject) -> (@owned Builtin.NativeObject, @error any Error) {
 bb0(%0 : $Int, %1 : $@thick Int.Type, %2 : @owned $Builtin.NativeObject):
   fix_lifetime %1 : $@thick Int.Type

test/embedded/classes-multi-module.swift

@@ -0,0 +1,58 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend -enable-experimental-feature Embedded -c -I%t -parse-as-library %t/MyModule.swift -o %t/MyModule.o -emit-module -emit-module-path %t/MyModule.swiftmodule -emit-empty-object-file
+// RUN: %target-swift-frontend -enable-experimental-feature Embedded -c -I%t %t/Main.swift -o %t/Main.o
+// RUN: %target-clang %t/Main.o %t/MyModule.o -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: swift_feature_Embedded
+
+//--- MyModule.swift
+
+public class C<T> {
+  public func foo() {
+    let x = X<Int>(x: 27)
+    x.bar()
+  }
+}
+
+class D<T>: C<T> {
+  override public func foo() {
+    print("D")
+  }
+}
+
+class X<T: BinaryInteger> {
+  var x: T
+
+  init(x: T) { self.x = x }
+  
+  func bar() {
+    print(x)
+  }
+}
+
+class Y<T: BinaryInteger>: X<T> {
+  override func bar() {
+  }
+}
+
+@inline(never)
+public func create<T>(_ t: T.Type) -> C<T> {
+  return C<T>()
+}
+
+//--- Main.swift
+
+import MyModule
+
+@inline(never)
+public func testit() {
+  let c = create(Int.self)
+  c.foo()
+}
+
+// CHECK: 27
+testit()