[sil-devirtualizer] Fix devirtualization of partial_apply on generic witness_methods

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -846,8 +846,9 @@ static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
 /// Generate a new apply of a function_ref to replace an apply of a
 /// witness_method when we've determined the actual function we'll end
 /// up calling.
-static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
-                                           ProtocolConformanceRef C) {
+static DevirtualizationResult
+devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
+                          ProtocolConformanceRef C) {
   // We know the witness thunk and the corresponding set of substitutions
   // required to invoke the protocol method at this point.
   auto &Module = AI.getModule();
@@ -893,20 +894,38 @@ static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
   auto ResultSILType = substConv.getSILResultType();
   ApplySite SAI;
 
-  if (auto *A = dyn_cast<ApplyInst>(AI))
-    SAI = Builder.createApply(Loc, FRI, SubstCalleeSILType,
-                              ResultSILType, NewSubs, Arguments,
-                              A->isNonThrowing());
+  SILValue ResultValue;
+  if (auto *A = dyn_cast<ApplyInst>(AI)) {
+    auto *NewAI =
+        Builder.createApply(Loc, FRI, SubstCalleeSILType, ResultSILType,
+                            NewSubs, Arguments, A->isNonThrowing());
+    // Check if any casting is required for the return value.
+    ResultValue = castValueToABICompatibleType(&Builder, Loc, NewAI,
+                                               NewAI->getType(), AI.getType());
+    SAI = ApplySite::isa(NewAI);
+  }
   if (auto *TAI = dyn_cast<TryApplyInst>(AI))
     SAI = Builder.createTryApply(Loc, FRI, SubstCalleeSILType,
                                  NewSubs, Arguments,
                                  TAI->getNormalBB(), TAI->getErrorBB());
-  if (auto *PAI = dyn_cast<PartialApplyInst>(AI))
-    SAI = Builder.createPartialApply(Loc, FRI, SubstCalleeSILType,
-                                     NewSubs, Arguments, PAI->getType());
+  if (auto *PAI = dyn_cast<PartialApplyInst>(AI)) {
+    auto PartialApplyConvention = PAI->getType()
+                                      .getSwiftRValueType()
+                                      ->getAs<SILFunctionType>()
+                                      ->getCalleeConvention();
+    auto PAIResultType = SILBuilder::getPartialApplyResultType(
+        SubstCalleeSILType, Arguments.size(), Module, {},
+        PartialApplyConvention);
+    auto *NewPAI = Builder.createPartialApply(
+        Loc, FRI, SubstCalleeSILType, NewSubs, Arguments, PAIResultType);
+    // Check if any casting is required for the return value.
+    ResultValue = castValueToABICompatibleType(
+        &Builder, Loc, NewPAI, NewPAI->getType(), PAI->getType());
+    SAI = ApplySite::isa(NewPAI);
+  }
 
   NumWitnessDevirt++;
-  return SAI;
+  return std::make_pair(ResultValue, SAI);
 }
 
 static bool canDevirtualizeWitnessMethod(ApplySite AI) {
@@ -948,8 +967,7 @@ DevirtualizationResult swift::tryDevirtualizeWitnessMethod(ApplySite AI) {
     AI.getModule().lookUpFunctionInWitnessTable(WMI->getConformance(),
                                                 WMI->getMember());
 
-  auto Result = devirtualizeWitnessMethod(AI, F, WMI->getConformance());
-  return std::make_pair(Result.getInstruction(), Result);
+  return devirtualizeWitnessMethod(AI, F, WMI->getConformance());
 }
 
 //===----------------------------------------------------------------------===//

test/SILOptimizer/devirt_generic_witness_method.sil

@@ -0,0 +1,53 @@
+// RUN: %target-sil-opt -assume-parsing-unqualified-ownership-sil -enable-sil-verify-all %s -devirtualizer | %FileCheck %s
+
+sil_stage canonical
+
+import Builtin
+import Swift
+import SwiftShims
+
+public protocol Decodable {
+  init(json: [String : Any])
+}
+
+public class Item : Decodable {
+  required public init(json: [String : Any])
+  deinit
+}
+
+public class ItemSubclass : Item {
+  deinit
+  required public init(json: [String : Any])
+}
+
+sil @witness_method_impl : $@convention(witness_method) (@owned Dictionary<String, Any>, @thick Item.Type) -> @out Item {
+bb0(%0 : $*Item, %1 : $Dictionary<String, Any>, %2 : $@thick Item.Type):
+  %3 = class_method %2 : $@thick Item.Type, #Item.init!allocator.1 : (Item.Type) -> ([String : Any]) -> Item, $@convention(method) (@owned Dictionary<String, Any>, @thick Item.Type) -> @owned Item
+  %4 = apply %3(%1, %2) : $@convention(method) (@owned Dictionary<String, Any>, @thick Item.Type) -> @owned Item
+  store %4 to %0 : $*Item
+  %6 = tuple ()
+  return %6 : $()
+}
+
+// Check that it is possible to devirtualize a partial_appy of a generic witness_method.
+// Since it is a derived class that invokes an implementation from a base class,
+// make sure that the resulting closure is properly converted into a required type.
+
+// CHECK-LABEL: sil @test_generic_witness_method_partial_apply_devirt : $@convention(thin) (@thick ItemSubclass.Type) -> @owned @callee_owned (@owned Dictionary<String, Any>) -> @out ItemSubclass
+// CHECK: %[[UPCASTED_ARG:[0-9]+]] = upcast %0 : $@thick ItemSubclass.Type to $@thick Item.Type
+// CHECK: %[[FUNCTION_REF:[0-9]+]] = function_ref @witness_method_impl
+// CHECK: %[[PA:[0-9]+]] = partial_apply %[[FUNCTION_REF]](%[[UPCASTED_ARG]]) : $@convention(witness_method) (@owned Dictionary<String, Any>, @thick Item.Type) -> @out Item
+// CHECK: %[[RETURN_VALUE:[0-9]+]] = convert_function %[[PA]] : $@callee_owned (@owned Dictionary<String, Any>) -> @out Item to $@callee_owned (@owned Dictionary<String, Any>) -> @out ItemSubclass
+// CHECK: return %[[RETURN_VALUE]] : $@callee_owned (@owned Dictionary<String, Any>) -> @out ItemSubclass
+
+sil @test_generic_witness_method_partial_apply_devirt : $@convention(thin) (@thick ItemSubclass.Type) -> @owned @callee_owned (@owned Dictionary<String, Any>) -> @out ItemSubclass {
+bb0(%0 : $@thick ItemSubclass.Type):
+  %1 = witness_method $ItemSubclass, #Decodable.init!allocator.1 : <Self where Self : Decodable> (Self.Type) -> ([String : Any]) -> Self : $@convention(witness_method) <τ_0_0 where τ_0_0 : Decodable> (@owned Dictionary<String, Any>, @thick τ_0_0.Type) -> @out τ_0_0
+  %2 = partial_apply %1<ItemSubclass>(%0) : $@convention(witness_method) <τ_0_0 where τ_0_0 : Decodable> (@owned Dictionary<String, Any>, @thick τ_0_0.Type) -> @out τ_0_0
+  return %2 : $@callee_owned (@owned Dictionary<String, Any>) -> @out ItemSubclass
+}
+
+sil_witness_table [serialized] Item: Decodable module main {
+  method #Decodable.init!allocator.1: <Self where Self : Decodable> (Self.Type) -> ([String : Any]) -> Self : @witness_method_impl
+}
+