Merge pull request #5306 from swiftix/fixes-for-witness-method-devirt

Fixes for witness method devirtualization

Author: swiftix

include/swift/SILOptimizer/Utils/Devirtualize.h

@@ -55,6 +55,8 @@ DevirtualizationResult devirtualizeClassMethod(FullApplySite AI,
 DevirtualizationResult tryDevirtualizeClassMethod(FullApplySite AI,
                                                   SILValue ClassInstance);
 DevirtualizationResult tryDevirtualizeWitnessMethod(ApplySite AI); 
+/// Check if an upcast is legal.
+bool isLegalUpcast(SILType FromTy, SILType ToTy);
 }
 
 #endif

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -871,12 +871,17 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite AI,
       return nullptr;
   }
 
-  // Obtain the protocol whose which should be used by the conformance.
+  // The lookup type is not a opened existential type, 
+  // thus it cannot be made more concrete.
+  if (!WMI->getLookupType()->isOpenedExistential())
+    return nullptr;
+
+  // Obtain the protocol which should be used by the conformance.
   auto *PD = WMI->getLookupProtocol();
 
   // Propagate the concrete type into a callee-operand, which is a
   // witness_method instruction.
-  auto PropagateIntoOperand = [this, &WMI](CanType ConcreteType,
+  auto PropagateIntoOperand = [this, &WMI, &AI](CanType ConcreteType,
                                            ProtocolConformanceRef Conformance) {
     // Keep around the dependence on the open instruction unless we've
     // actually eliminated the use.
@@ -885,8 +890,15 @@ SILCombiner::propagateConcreteTypeOfInitExistential(FullApplySite AI,
                                                 Conformance, WMI->getMember(),
                                                 WMI->getType(),
                                                 WMI->isVolatile());
-    replaceInstUsesWith(*WMI, NewWMI);
-    eraseInstFromFunction(*WMI);
+    // Replace only uses of the witness_method in the apply that is going to
+    // be changed.
+    MutableArrayRef<Operand> Operands = AI.getInstruction()->getAllOperands();
+    for (auto &Op : Operands) {
+      if (Op.get() == WMI)
+         Op.set(NewWMI); 
+    }
+    if (WMI->use_empty())
+      eraseInstFromFunction(*WMI);
   };
 
   // Try to perform the propagation.

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -866,6 +866,73 @@ static void getWitnessMethodSubstitutions(ApplySite AI, SILFunction *F,
                                 origSubs, NewSubs);
 }
 
+/// Check if an upcast is legal.
+/// THe logic in this function is heavily based on the checks in
+/// the SILVerifier.
+bool swift::isLegalUpcast(SILType FromTy, SILType ToTy) {
+  if (ToTy.is<MetatypeType>()) {
+    CanType InstTy(ToTy.castTo<MetatypeType>()->getInstanceType());
+    if (!FromTy.is<MetatypeType>())
+      return false;
+    CanType OpInstTy(FromTy.castTo<MetatypeType>()->getInstanceType());
+    auto InstClass = InstTy->getClassOrBoundGenericClass();
+    if (!InstClass)
+      return false;
+
+    bool CanBeUpcasted =
+        InstClass->usesObjCGenericsModel()
+            ? InstClass->getDeclaredTypeInContext()->isBindableToSuperclassOf(
+                  OpInstTy, nullptr)
+            : InstTy->isExactSuperclassOf(OpInstTy, nullptr);
+
+    return CanBeUpcasted;
+  }
+
+  // Upcast from Optional<B> to Optional<A> is legal as long as B is a
+  // subclass of A.
+  if (ToTy.getSwiftRValueType().getAnyOptionalObjectType() &&
+      FromTy.getSwiftRValueType().getAnyOptionalObjectType()) {
+    ToTy = SILType::getPrimitiveObjectType(
+        ToTy.getSwiftRValueType().getAnyOptionalObjectType());
+    FromTy = SILType::getPrimitiveObjectType(
+        FromTy.getSwiftRValueType().getAnyOptionalObjectType());
+  }
+
+  auto ToClass = ToTy.getClassOrBoundGenericClass();
+  if (!ToClass)
+    return false;
+  bool CanBeUpcasted =
+      ToClass->usesObjCGenericsModel()
+          ? ToClass->getDeclaredTypeInContext()->isBindableToSuperclassOf(
+                FromTy.getSwiftRValueType(), nullptr)
+          : ToTy.isExactSuperclassOf(FromTy);
+
+  return CanBeUpcasted;
+}
+
+/// Check if we can pass/convert all arguments of the original apply
+/// as required by the found devirtualized method.
+static bool
+canPassOrConvertAllArguments(ApplySite AI,
+                             CanSILFunctionType SubstCalleeCanType) {
+  for (unsigned ArgN = 0, ArgE = AI.getNumArguments(); ArgN != ArgE; ++ArgN) {
+    SILValue A = AI.getArgument(ArgN);
+    auto ParamType = SubstCalleeCanType->getSILArgumentType(
+      SubstCalleeCanType->getNumSILArguments() - AI.getNumArguments() + ArgN);
+    // Check if we can cast the provided argument into the required
+    // parameter type.
+    auto FromTy = A->getType();
+    auto ToTy = ParamType;
+    // If types are the same, no conversion will be required.
+    if (FromTy == ToTy)
+      continue;
+    // Otherwise, it should be possible to upcast the arguments.
+    if (!isLegalUpcast(FromTy, ToTy))
+      return false;
+  }
+  return true;
+}
+
 /// 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.
@@ -890,6 +957,11 @@ static ApplySite devirtualizeWitnessMethod(ApplySite AI, SILFunction *F,
   auto SubstCalleeCanType = CalleeCanType->substGenericArgs(
     Module, Module.getSwiftModule(), NewSubs);
 
+  // Bail if some of the arguments cannot be converted into
+  // types required by the found devirtualized method.
+  if (!canPassOrConvertAllArguments(AI, SubstCalleeCanType))
+    return ApplySite();
+
   // Collect arguments from the apply instruction.
   auto Arguments = SmallVector<SILValue, 4>();
 

test/SILOptimizer/sil_combine.sil

@@ -3002,3 +3002,69 @@ bb0(%0 : $VV):
   return %26 : $()
 }
 
+protocol PPP : class {
+  func foo()
+}
+
+class BBB: PPP {
+  @inline(never)
+  func foo()
+}
+
+final class XXX : BBB {
+  @inline(never)
+  override func foo()
+}
+
+// Check that sil-combine does not crash on this example and does not generate a wrong
+// upcast.
+// CHECK-LABEL: sil @silcombine_dont_generate_wrong_upcasts_during_devirt
+// CHECK-NOT: upcast
+// CHECK: witness_method $XXX, #PPP.foo!1 : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+// CHECK-NOT: upcast
+// CHECK: return
+sil @silcombine_dont_generate_wrong_upcasts_during_devirt: $@convention(thin) (@owned BBB) -> () {
+bb0(%0 : $BBB):
+  strong_retain %0 : $BBB
+  %3 = init_existential_ref %0 : $BBB : $BBB, $PPP
+  %5 = open_existential_ref %3 : $PPP to $@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP
+  %6 = witness_method $XXX, #PPP.foo!1, %5 : $@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  %7 = apply %6<@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP>(%5) : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  %8 = apply %6<@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP>(%5) : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  strong_release %3 : $PPP
+  %9 = tuple ()
+  strong_release %0 : $BBB
+  %11 = tuple ()
+  return %11 : $()
+}
+
+// Check that both applies can be devirtualized by means of propagating the concrete
+// type of the existential into witness_method and apply instructions.
+// CHECK-LABEL: sil @silcombine_devirt_both_applies_of_witness_method
+// CHECK-NOT: open_existential_ref
+// CHECK-NOT: witness_method
+// CHECK: [[FR1:%.*]] = function_ref @_TTWC4nix23XXXS_3PPPS_FS1_3foofT_T_
+// CHECK: apply [[FR1]](%0)
+// CHECK: [[FR2:%.*]] = function_ref @_TTWC4nix23XXXS_3PPPS_FS1_3foofT_T_
+// CHECK: apply [[FR2]](%0)
+// CHECK: return
+sil @silcombine_devirt_both_applies_of_witness_method : $@convention(thin) (@owned XXX) -> () {
+bb0(%0 : $XXX):
+  strong_retain %0 : $XXX
+  %3 = init_existential_ref %0 : $XXX : $XXX, $PPP
+  %5 = open_existential_ref %3 : $PPP to $@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP
+  %6 = witness_method $@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP, #PPP.foo!1, %5 : $@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  %7 = apply %6<@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP>(%5) : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  %8 = apply %6<@opened("0AC9A62E-926E-11E6-8BF5-685B35C48C83") PPP>(%5) : $@convention(witness_method) <τ_0_0 where τ_0_0 : PPP> (@guaranteed τ_0_0) -> ()
+  strong_release %3 : $PPP
+  %9 = tuple ()
+  strong_release %0 : $XXX
+  %11 = tuple ()
+  return %11 : $()
+}
+
+sil @_TTWC4nix23XXXS_3PPPS_FS1_3foofT_T_ : $@convention(witness_method) (@guaranteed XXX) -> ()
+
+sil_witness_table XXX: PPP module nix2 {
+  method #PPP.foo!1: @_TTWC4nix23XXXS_3PPPS_FS1_3foofT_T_	// protocol witness for PPP.foo() -> () in conformance XXX
+}