Devirtualizer: be less restrictive when checking for correct types to enable witness method devirtualization

rdar://114202760

Author: eeckstein

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -1210,6 +1210,17 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite, bool isMandatory)
   if (!interfaceTy->hasTypeParameter())
     return true;
 
+  auto subs = getWitnessMethodSubstitutions(f->getModule(), applySite,
+                                            f, wmi->getConformance());
+  CanSILFunctionType substCalleTy = f->getLoweredFunctionType()->substGenericArgs(
+      f->getModule(), subs,
+      applySite.getFunction()->getTypeExpansionContext());
+  CanSILFunctionType applySubstCalleeTy = applySite.getSubstCalleeType();
+
+  // If the function types match, there is no problem.
+  if (substCalleTy == applySubstCalleeTy)
+    return true;
+
   auto selfGP = wmi->getLookupProtocol()->getSelfInterfaceType();
   auto isSelfRootedTypeParameter = [selfGP](Type T) -> bool {
     if (!T->hasTypeParameter())

test/SILOptimizer/devirt_protocol_method_invocations.swift

@@ -344,3 +344,33 @@ public func testNoDevirt() {
   p.assoc3 { _ in }
   p.assoc4 { _ in }
 }
+
+protocol MyProtocol {
+    associatedtype Element
+    var array: [Element] { get }
+    var foo: Bool { get }
+}
+
+extension Array {
+    var isThisACoolArray: Bool {
+        return true
+    }
+}
+
+extension MyProtocol {
+    var foo: Bool { array.isThisACoolArray }
+}
+
+public struct MyStruct {
+    var array: [Int] = []
+}
+
+extension MyStruct: MyProtocol {}
+
+// CHECK-LABEL: sil @$s34devirt_protocol_method_invocations15testArrayReturn1xSbAA8MyStructVz_tF :
+// CHECK-NOT:     witness_method
+// CHECK:       } // end sil function '$s34devirt_protocol_method_invocations15testArrayReturn1xSbAA8MyStructVz_tF'
+public func testArrayReturn(x: inout MyStruct) -> Bool {
+    return x.foo
+}
+