Fix speculative devirtualization to correctly use casted value

include/swift/SIL/SILBuilder.h

@@ -1783,12 +1783,6 @@ class SILBuilder {
   // Unchecked cast helpers
   //===--------------------------------------------------------------------===//
 
-  // Create an UncheckedRefCast if the source and dest types are legal,
-  // otherwise return null.
-  // Unwrap or wrap optional types as needed.
-  SingleValueInstruction *tryCreateUncheckedRefCast(SILLocation Loc, SILValue Op,
-                                                    SILType ResultTy);
-
   // Create the appropriate cast instruction based on result type.
   SingleValueInstruction *createUncheckedBitCast(SILLocation Loc, SILValue Op,
                                                  SILType Ty);

lib/SIL/SILBuilder.cpp

@@ -110,19 +110,6 @@ ProjectBoxInst *SILBuilder::createProjectBox(SILLocation Loc,
       getSILDebugLocation(Loc), boxOperand, index, fieldTy));
 }
 
-// If legal, create an unchecked_ref_cast from the given operand and result
-// type, otherwise return null.
-SingleValueInstruction *
-SILBuilder::tryCreateUncheckedRefCast(SILLocation Loc, SILValue Op,
-                                      SILType ResultTy) {
-  if (!SILType::canRefCast(Op->getType(), ResultTy, getModule()))
-    return nullptr;
-
-  return insert(UncheckedRefCastInst::create(getSILDebugLocation(Loc), Op,
-                                             ResultTy, getFunction(),
-                                             C.OpenedArchetypes));
-}
-
 ClassifyBridgeObjectInst *
 SILBuilder::createClassifyBridgeObject(SILLocation Loc, SILValue value) {
   auto &ctx = getASTContext();
@@ -142,8 +129,8 @@ SILBuilder::createUncheckedBitCast(SILLocation Loc, SILValue Op, SILType Ty) {
     return insert(UncheckedTrivialBitCastInst::create(
         getSILDebugLocation(Loc), Op, Ty, getFunction(), C.OpenedArchetypes));
 
-  if (auto refCast = tryCreateUncheckedRefCast(Loc, Op, Ty))
-    return refCast;
+  if (SILType::canRefCast(Op->getType(), Ty, getModule()))
+    return createUncheckedRefCast(Loc, Op, Ty);
 
   // The destination type is nontrivial, and may be smaller than the source
   // type, so RC identity cannot be assumed.

lib/SILGen/SILGenBuilder.cpp

@@ -804,17 +804,6 @@ ManagedValue SILGenBuilder::createUncheckedAddrCast(SILLocation loc, ManagedValu
   return cloner.clone(cast);
 }
 
-ManagedValue SILGenBuilder::tryCreateUncheckedRefCast(SILLocation loc,
-                                                      ManagedValue original,
-                                                      SILType type) {
-  CleanupCloner cloner(*this, original);
-  SILValue result = tryCreateUncheckedRefCast(loc, original.getValue(), type);
-  if (!result)
-    return ManagedValue();
-  original.forward(SGF);
-  return cloner.clone(result);
-}
-
 ManagedValue SILGenBuilder::createUncheckedTrivialBitCast(SILLocation loc,
                                                           ManagedValue original,
                                                           SILType type) {

lib/SILGen/SILGenBuilder.h

@@ -256,11 +256,6 @@ class SILGenBuilder : public SILBuilder {
   ManagedValue createUpcast(SILLocation loc, ManagedValue original,
                             SILType type);
 
-  using SILBuilder::tryCreateUncheckedRefCast;
-  ManagedValue tryCreateUncheckedRefCast(SILLocation loc,
-                                         ManagedValue op,
-                                         SILType type);
-
   using SILBuilder::createUncheckedTrivialBitCast;
   ManagedValue createUncheckedTrivialBitCast(SILLocation loc,
                                              ManagedValue original,

lib/SILGen/SILGenBuiltin.cpp

@@ -649,13 +649,14 @@ emitBuiltinCastReference(SILGenFunction &SGF,
   auto &toTL = SGF.getTypeLowering(toTy);
   assert(!fromTL.isTrivial() && !toTL.isTrivial() && "expected ref type");
 
+  auto arg = args[0];
+
   // TODO: Fix this API.
   if (!fromTL.isAddress() || !toTL.isAddress()) {
-    if (auto refCast = SGF.B.tryCreateUncheckedRefCast(loc, args[0],
-                                                       toTL.getLoweredType())) {
+    if (SILType::canRefCast(arg.getType(), toTL.getLoweredType(), SGF.SGM.M)) {
       // Create a reference cast, forwarding the cleanup.
       // The cast takes the source reference.
-      return refCast;
+      return SGF.B.createUncheckedRefCast(loc, arg, toTL.getLoweredType());
     }
   }
 
@@ -670,7 +671,7 @@ emitBuiltinCastReference(SILGenFunction &SGF,
   // TODO: For now, we leave invalid casts in address form so that the runtime
   // will trap. We could emit a noreturn call here instead which would provide
   // more information to the optimizer.
-  SILValue srcVal = args[0].ensurePlusOne(SGF, loc).forward(SGF);
+  SILValue srcVal = arg.ensurePlusOne(SGF, loc).forward(SGF);
   SILValue fromAddr;
   if (!fromTL.isAddress()) {
     // Move the loadable value into a "source temp".  Since the source and
@@ -745,17 +746,9 @@ static ManagedValue emitBuiltinReinterpretCast(SILGenFunction &SGF,
   }
   // Create the appropriate bitcast based on the source and dest types.
   ManagedValue in = args[0];
-  SILType resultTy = toTL.getLoweredType();
-  if (resultTy.isTrivial(SGF.F))
-    return SGF.B.createUncheckedTrivialBitCast(loc, in, resultTy);
 
-  // If we can perform a ref cast, just return.
-  if (auto refCast = SGF.B.tryCreateUncheckedRefCast(loc, in, resultTy))
-    return refCast;
-
-  // Otherwise leave the original cleanup and retain the cast value.
-  SILValue out = SGF.B.createUncheckedBitwiseCast(loc, in.getValue(), resultTy);
-  return SGF.emitManagedRetain(loc, out, toTL);
+  SILType resultTy = toTL.getLoweredType();
+  return SGF.B.createUncheckedBitCast(loc, in, resultTy);
 }
 
 /// Specialized emitter for Builtin.castToBridgeObject.

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -286,9 +286,12 @@ SILCombiner::visitUncheckedRefCastAddrInst(UncheckedRefCastAddrInst *URCI) {
   Builder.setCurrentDebugScope(URCI->getDebugScope());
   LoadInst *load = Builder.createLoad(Loc, URCI->getSrc(),
                                       LoadOwnershipQualifier::Unqualified);
-  auto *cast = Builder.tryCreateUncheckedRefCast(Loc, load,
-                                                 DestTy.getObjectType());
-  assert(cast && "SILBuilder cannot handle reference-castable types");
+
+  assert(SILType::canRefCast(load->getType(), DestTy.getObjectType(),
+                             Builder.getModule()) &&
+         "SILBuilder cannot handle reference-castable types");
+  auto *cast = Builder.createUncheckedRefCast(Loc, load,
+                                              DestTy.getObjectType());
   Builder.createStore(Loc, cast, URCI->getDest(),
                       StoreOwnershipQualifier::Unqualified);
 
@@ -393,11 +396,12 @@ visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
                                                  UBCI->getOperand(),
                                                  UBCI->getType());
 
-  if (auto refCast = Builder.tryCreateUncheckedRefCast(
-        UBCI->getLoc(), UBCI->getOperand(), UBCI->getType()))
-    return refCast;
+  if (!SILType::canRefCast(UBCI->getOperand()->getType(), UBCI->getType(),
+                           Builder.getModule()))
+    return nullptr;
 
-  return nullptr;
+  return Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
+                                        UBCI->getType());
 }
 
 SILInstruction *

lib/SILOptimizer/Transforms/SpeculativeDevirtualizer.cpp

@@ -81,14 +81,20 @@ static SILBasicBlock *cloneEdge(TermInst *TI, unsigned SuccIndex) {
 }
 
 // A utility function for cloning the apply instruction.
-static FullApplySite CloneApply(FullApplySite AI, SILBuilder &Builder) {
+static FullApplySite CloneApply(FullApplySite AI, SILValue SelfArg,
+                                SILBuilder &Builder) {
   // Clone the Apply.
   Builder.setCurrentDebugScope(AI.getDebugScope());
   Builder.addOpenedArchetypeOperands(AI.getInstruction());
   auto Args = AI.getArguments();
   SmallVector<SILValue, 8> Ret(Args.size());
-  for (unsigned i = 0, e = Args.size(); i != e; ++i)
-    Ret[i] = Args[i];
+  for (unsigned i = 0, e = Args.size(); i != e; ++i) {
+    if (i == e - 1 && SelfArg) {
+      Ret[i] = SelfArg;
+    } else {
+      Ret[i] = Args[i];
+    }
+  }
 
   FullApplySite NAI;
 
@@ -169,8 +175,8 @@ static FullApplySite speculateMonomorphicTarget(FullApplySite AI,
   SILValue DownCastedClassInstance = Iden->getArgument(0);
 
   // Copy the two apply instructions into the two blocks.
-  FullApplySite IdenAI = CloneApply(AI, IdenBuilder);
-  FullApplySite VirtAI = CloneApply(AI, VirtBuilder);
+  FullApplySite IdenAI = CloneApply(AI, DownCastedClassInstance, IdenBuilder);
+  FullApplySite VirtAI = CloneApply(AI, SILValue(), VirtBuilder);
 
   // See if Continue has a release on self as the instruction right after the
   // apply. If it exists, move it into position in the diamond.

test/SILOptimizer/devirt_speculate.swift

@@ -10,44 +10,96 @@ public class Base {
   public init() {}
   public func foo() {}
 }
+
+@_optimize(none)
+func blackHole<T : AnyObject>(_: T) {}
+
 class Sub1 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub2 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub3 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub4 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub5 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub6 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 class Sub7 : Base {
-  override func foo() {}
+  override func foo() { blackHole(self) }
 }
 // CHECK: @$s16devirt_speculate28testMaxNumSpeculativeTargetsyyAA4BaseCF
-// CHECK: checked_cast_br [exact] %0 : $Base to Base
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub1
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub2
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub3
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub4
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub5
-// CHECK: checked_cast_br [exact] %0 : $Base to Sub6
+// CHECK: bb0(%0 : $Base):
+// CHECK:   checked_cast_br [exact] %0 : $Base to Base, bb2, bb3
+
+// CHECK: bb2([[CASTED:%.*]]):
+// CHECK:   br bb1
+
+// CHECK: bb3:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub1, bb4, bb5
+
+// CHECK: bb4([[CASTED:%.*]] : $Sub1):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub1>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb5:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub2, bb6, bb7
+
+// CHECK: bb6([[CASTED:%.*]] : $Sub2):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub2>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb7:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub3, bb8, bb9
+
+// CHECK: bb8([[CASTED:%.*]] : $Sub3):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub3>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb9:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub4, bb10, bb11
+
+// CHECK: bb10([[CASTED:%.*]] : $Sub4):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub4>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb11:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub5, bb12, bb13
+
+// CHECK: bb12([[CASTED:%.*]] : $Sub5):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub5>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb13:
+// CHECK:   checked_cast_br [exact] %0 : $Base to Sub6, bb14, bb15
+
+// CHECK: bb14([[CASTED:%.*]] : $Sub6):
+// CHECK:   [[FN:%.*]] = function_ref @$s16devirt_speculate9blackHoleyyxRlzClF : $@convention(thin) <τ_0_0 where τ_0_0 : AnyObject> (@guaranteed τ_0_0) -> ()
+// CHECK:   apply [[FN]]<Sub6>([[CASTED]])
+// CHECK:   br bb1
+
+// CHECK: bb15:
 // CHECK-NOT: checked_cast_br
-// CHECK: %[[CM:[0-9]+]] = class_method %0 : $Base, #Base.foo!1 : (Base) -> () -> (), $@convention(method) (@guaranteed Base) -> ()
-// CHECK: apply %[[CM]](%0) : $@convention(method) (@guaranteed Base) -> ()
+// CHECK:   %[[CM:[0-9]+]] = class_method %0 : $Base, #Base.foo!1 : (Base) -> () -> (), $@convention(method) (@guaranteed Base) -> ()
+// CHECK:   apply %[[CM]](%0) : $@convention(method) (@guaranteed Base) -> ()
 
 // YAML:      Pass:            sil-speculative-devirtualizer
 // YAML-NEXT: Name:            sil.PartialSpecDevirt
 // YAML-NEXT: DebugLoc:
 // YAML-NEXT:   File:            {{.*}}/devirt_speculate.swift
-// YAML-NEXT:   Line:            66
+// YAML-NEXT:   Line:            118
 // YAML-NEXT:   Column:          5
 // YAML-NEXT: Function:        'testMaxNumSpeculativeTargets(_:)'
 // YAML-NEXT: Args:

test/SILOptimizer/devirt_speculative.sil

@@ -145,7 +145,6 @@ bb0(%0: $Base2):
 // CHECK-NEXT: unreachable
 // CHECK: checked_cast_br
 // CHECK: function_ref @Sub_exit
-// CHECK-NEXT: unchecked_ref_cast
 // CHECK-NEXT: apply
 // CHECK-NEXT: unreachable
 sil hidden [noinline] @test_devirt_of_noreturn_function : $@convention(thin) (@owned Base) -> () {