Fix optimization of checked_cast_addr_br to checked_cast_br

Author: meg-gupta

include/swift/SIL/DynamicCasts.h

@@ -97,6 +97,12 @@ bool canSILUseScalarCheckedCastInstructions(SILModule &M,
                                             CanType sourceType,
                                             CanType targetType);
 
+/// Can the given cast be performed by the scalar checked-cast instructions in
+/// the current SIL stage, or do we need to use the indirect instructions?
+bool canOptimizeToScalarCheckedCastInstructions(
+    SILFunction *func, CanType sourceType, CanType targetType,
+    CastConsumptionKind consumption);
+
 /// Can the given cast be performed by the scalar checked-cast
 /// instructions at IRGen, or do we need to use the indirect instructions?
 bool canIRGenUseScalarCheckedCastInstructions(SILModule &M,

include/swift/SIL/TypeSubstCloner.h

@@ -244,8 +244,9 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
     auto FalseCount = inst->getFalseBBCount();
 
     // Try to use the scalar cast instruction.
-    if (canSILUseScalarCheckedCastInstructions(B.getModule(),
-                                               sourceType, targetType)) {
+    if (canOptimizeToScalarCheckedCastInstructions(
+            &B.getFunction(), sourceType, targetType,
+            inst->getConsumptionKind())) {
       emitIndirectConditionalCastWithScalar(
           B, SwiftMod, loc, inst->getCheckedCastOptions(),
           inst->getConsumptionKind(), src, sourceType, dest,

lib/SIL/Utils/DynamicCasts.cpp

@@ -1355,6 +1355,27 @@ bool swift::canSILUseScalarCheckedCastInstructions(SILModule &M,
                                                   targetFormalType);
 }
 
+bool swift::canOptimizeToScalarCheckedCastInstructions(
+    SILFunction *func, CanType sourceType, CanType targetType,
+    CastConsumptionKind consumption) {
+  if (!canSILUseScalarCheckedCastInstructions(func->getModule(), sourceType,
+                                              targetType)) {
+    return false;
+  }
+
+  if (consumption == CastConsumptionKind::CopyOnSuccess) {
+    // If it's a copy-on-success cast, check whether the cast preserves
+    // ownership in ossa. This is needed because the optimization creates a
+    // scalar cast which is a guaranteed forwarding instruction and is valid
+    // only when ownership can be preserved.
+    return !func->hasOwnership() ||
+           doesCastPreserveOwnershipForTypes(func->getModule(), sourceType,
+                                             targetType);
+  }
+
+  return true;
+}
+
 /// Can the given cast be performed by the scalar checked-cast
 /// instructions?
 bool swift::canIRGenUseScalarCheckedCastInstructions(SILModule &M,

lib/SILOptimizer/Utils/CastOptimizer.cpp

@@ -1064,8 +1064,7 @@ CastOptimizer::simplifyCheckedCastBranchInst(CheckedCastBranchInst *Inst) {
       // The unconditional_cast can be skipped, if the result of a cast
       // is not used afterwards.
       if (!ResultNotUsed) {
-        if (!dynamicCast.canSILUseScalarCheckedCastInstructions())
-          return nullptr;
+        ASSERT(dynamicCast.canSILUseScalarCheckedCastInstructions());
 
         CastedValue =
             emitSuccessfulScalarUnconditionalCast(Builder, Loc, dynamicCast);
@@ -1160,9 +1159,9 @@ SILInstruction *CastOptimizer::optimizeCheckedCastAddrBranchInst(
 
       if (MI) {
         if (SuccessBB->getSinglePredecessorBlock() &&
-            canSILUseScalarCheckedCastInstructions(
-                Inst->getModule(), MI->getType().getASTType(),
-                Inst->getTargetFormalType())) {
+            canOptimizeToScalarCheckedCastInstructions(
+                Inst->getFunction(), MI->getType().getASTType(),
+                Inst->getTargetFormalType(), Inst->getConsumptionKind())) {
           SILBuilderWithScope B(Inst, builderContext);
           auto NewI = B.createCheckedCastBranch(
               Loc, false /*isExact*/,

test/SILOptimizer/specialize_casts.sil

@@ -0,0 +1,62 @@
+// RUN: %target-sil-opt -enable-sil-verify-all -generic-specializer %s | %FileCheck %s
+
+// REQUIRES: objc_interop
+
+sil_stage canonical
+
+import Builtin
+import Swift
+import SwiftShims
+
+import Foundation
+
+public enum ResultType : Equatable {
+  case object(NSObject)
+  case undefined
+  internal init<T>(metatypeOf value: T)
+}
+
+// CHECK-LABEL: // specialized callee_anyobject
+// CHECK-LABEL: sil shared [ossa] @$s16callee_anyobjectyXl_Tt0g5 :
+// CHECK: checked_cast_addr_br 
+// CHECK-LABEL: } // end sil function '$s16callee_anyobjectyXl_Tt0g5'
+sil hidden [ossa] @callee_anyobject : $@convention(method) <T> (@in T, @thin ResultType.Type) -> @owned ResultType {
+bb0(%0 : $*T, %1 : $@thin ResultType.Type):
+  %3 = alloc_stack $T
+  copy_addr %0 to [init] %3
+  %5 = alloc_stack $NSObject
+  checked_cast_addr_br copy_on_success T in %3 to NSObject in %5, bb1, bb2
+
+bb1:
+  destroy_addr %0
+  %8 = enum $ResultType, #ResultType.undefined!enumelt
+  destroy_addr %5
+  dealloc_stack %5
+  br bb5(%8)
+
+bb2:
+  destroy_addr %0
+  dealloc_stack %5
+  %28 = enum $ResultType, #ResultType.undefined!enumelt
+  br bb5(%28)
+
+bb5(%30 : @owned $ResultType):
+  destroy_addr %3
+  dealloc_stack %3
+  %34 = move_value [lexical] %30
+  return %34
+}
+
+sil hidden [ossa] @caller_anyobject : $@convention(thin) (@owned AnyObject) -> () {
+bb0(%0 : @owned $AnyObject):
+  %1 = metatype $@thin ResultType.Type
+  %2 = alloc_stack $AnyObject
+  store %0 to [init] %2
+  %4 = function_ref @callee_anyobject : $@convention(method) <τ_0_0> (@in τ_0_0, @thin ResultType.Type) -> @owned ResultType
+  %5 = apply %4<AnyObject>(%2, %1) : $@convention(method) <τ_0_0> (@in τ_0_0, @thin ResultType.Type) -> @owned ResultType
+  dealloc_stack %2
+  destroy_value %5
+  %8 = tuple ()
+  return %8
+}
+