Merge pull request #79872 from WeZZard/bugfix/unchecked-ref-casting-with-optional-any-object

[SILOptimizer] Fixes the ValueDefUseWalker to handle potential unchecked_ref_cast between optional and non-optional class types.

Author: eeckstein

SwiftCompilerSources/Sources/AST/Type.swift

@@ -93,6 +93,7 @@ extension TypeProperties {
   public var isEscapable: Bool { type.bridged.isEscapable() }
   public var isNoEscape: Bool { type.bridged.isNoEscape() }
   public var isInteger: Bool { type.bridged.isInteger() }
+  public var isOptional: Bool { type.bridged.isOptional() }
   public var isMetatypeType: Bool { type.bridged.isMetatypeType() }
   public var isExistentialMetatypeType: Bool { type.bridged.isExistentialMetatypeType() }
   public var representationOfMetatype: AST.`Type`.MetatypeRepresentation {

SwiftCompilerSources/Sources/Optimizer/TestPasses/EscapeInfoDumper.swift

@@ -14,7 +14,7 @@ import SIL
 
 /// Dumps the results of escape analysis.
 ///
-/// Dumps the EscapeInfo query results for all `alloc_stack` instructions in a function.
+/// Dumps the EscapeInfo query results for all `alloc_ref` instructions in a function.
 ///
 /// This pass is used for testing EscapeInfo.
 let escapeInfoDumper = FunctionPass(name: "dump-escape-info") {

SwiftCompilerSources/Sources/SIL/Type.swift

@@ -63,6 +63,7 @@ public struct Type : CustomStringConvertible, NoReflectionChildren {
   public var isFunction: Bool { bridged.isFunction() }
   public var isMetatype: Bool { bridged.isMetatype() }
   public var isClassExistential: Bool { bridged.isClassExistential() }
+  public var isOptional: Bool { astType.isOptional }
   public var isNoEscapeFunction: Bool { bridged.isNoEscapeFunction() }
   public var containsNoEscapeFunction: Bool { bridged.containsNoEscapeFunction() }
   public var isThickFunction: Bool { bridged.isThickFunction() }

SwiftCompilerSources/Sources/SIL/Utilities/WalkUtils.swift

@@ -370,7 +370,21 @@ extension ValueDefUseWalker {
         // We need to ignore this because otherwise the path wouldn't contain the right `existential` field kind.
         return leafUse(value: operand, path: path)
       }
-      return walkDownUses(ofValue: urc, path: path)
+      // The `unchecked_ref_cast` is designed to be able to cast between
+      // `Optional<ClassType>` and `ClassType`. We need to handle these
+      // cases by checking if the type is optional and adjust the path
+      // accordingly.
+      switch (urc.type.isOptional, urc.fromInstance.type.isOptional) {
+        case (true, false):
+          return walkDownUses(ofValue: urc, path: path.push(.enumCase, index: 1))
+        case (false, true):
+          if let path = path.popIfMatches(.enumCase, index: 1) {
+            return walkDownUses(ofValue: urc, path: path)
+          }
+          return unmatchedPath(value: operand, path: path)
+        default:
+          return walkDownUses(ofValue: urc, path: path)
+      }
     case let beginDealloc as BeginDeallocRefInst:
       if operand.index == 0 {
         return walkDownUses(ofValue: beginDealloc, path: path)
@@ -699,7 +713,21 @@ extension ValueUseDefWalker {
         // We need to ignore this because otherwise the path wouldn't contain the right `existential` field kind.
         return rootDef(value: urc, path: path)
       }
-      return walkUp(value: urc.fromInstance, path: path)
+      // The `unchecked_ref_cast` is designed to be able to cast between
+      // `Optional<ClassType>` and `ClassType`. We need to handle these
+      // cases by checking if the type is optional and adjust the path
+      // accordingly.
+      switch (urc.type.isOptional, urc.fromInstance.type.isOptional) {
+        case (true, false):
+          if let path = path.popIfMatches(.enumCase, index: 1) {
+            return walkUp(value: urc.fromInstance, path: path)
+          }
+          return unmatchedPath(value: urc.fromInstance, path: path)
+        case (false, true):
+          return walkUp(value: urc.fromInstance, path: path.push(.enumCase, index: 1))
+        default:
+          return walkUp(value: urc.fromInstance, path: path)
+      }
     case let arg as Argument:
       if let phi = Phi(arg) {
         for incoming in phi.incomingValues {

include/swift/AST/ASTBridging.h

@@ -3023,6 +3023,7 @@ struct BridgedASTType {
   BRIDGED_INLINE bool isInteger() const;
   BRIDGED_INLINE bool isMetatypeType() const;
   BRIDGED_INLINE bool isExistentialMetatypeType() const;
+  BRIDGED_INLINE bool isOptional() const;
   BRIDGED_INLINE TraitResult canBeClass() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getAnyNominal() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getInstanceTypeOfMetatype() const;

include/swift/AST/ASTBridgingImpl.h

@@ -430,6 +430,10 @@ bool BridgedASTType::isMetatypeType() const {
   return unbridged()->is<swift::AnyMetatypeType>();
 }
 
+bool BridgedASTType::isOptional() const {
+  return unbridged()->getCanonicalType()->isOptional();
+}
+
 bool BridgedASTType::isExistentialMetatypeType() const {
   return unbridged()->is<swift::ExistentialMetatypeType>();
 }

include/swift/SIL/SILBridging.h

@@ -263,6 +263,7 @@ struct BridgedType {
   BRIDGED_INLINE bool isFunction() const;
   BRIDGED_INLINE bool isMetatype() const;
   BRIDGED_INLINE bool isClassExistential() const;
+  BRIDGED_INLINE bool isOptional() const;
   BRIDGED_INLINE bool isNoEscapeFunction() const;
   BRIDGED_INLINE bool containsNoEscapeFunction() const;
   BRIDGED_INLINE bool isThickFunction() const;

include/swift/SIL/SILBridgingImpl.h

@@ -402,6 +402,11 @@ bool BridgedType::isClassExistential() const {
   return unbridged().isClassExistentialType();
 }
 
+bool BridgedType::isOptional() const {
+  swift::CanType astType = unbridged().getASTType();
+  return astType->isOptional();
+}
+
 bool BridgedType::isNoEscapeFunction() const {
   return unbridged().isNoEscapeFunction();
 }

test/SILOptimizer/escape_info.sil

@@ -1458,6 +1458,45 @@ bb0:
   return %2 : $X
 }
 
+// CHECK-LABEL: Escape information for test_unchecked_ref_cast_from_optional_to_non_optional:
+// CHECK: return[]: %0 = alloc_ref $Derived
+// CHECK:  -    :   %2 = alloc_ref $Derived
+// CHECK: End function test_unchecked_ref_cast_from_optional_to_non_optional
+sil @test_unchecked_ref_cast_from_optional_to_non_optional : $@convention(thin) () -> @owned X {
+bb0:
+  %0 = alloc_ref $Derived
+  %1 = enum $Optional<Derived>, #Optional.some!enumelt, %0 : $Derived
+  %2 = alloc_ref $Derived
+  %3 = enum $Optional<Derived>, #Optional.some!enumelt, %2 : $Derived
+  %4 = unchecked_ref_cast %1 : $Optional<Derived> to $X
+  %5 = unchecked_ref_cast %3 : $Optional<Derived> to $X
+  return %4 : $X
+}
+
+// CHECK-LABEL: Escape information for test_unchecked_ref_cast_from_optional_to_non_optional_2:
+// CHECK-NEXT: End function test_unchecked_ref_cast_from_optional_to_non_optional_2
+sil @test_unchecked_ref_cast_from_optional_to_non_optional_2 : $@convention(thin) () -> @owned X {
+bb0:
+  %0 = enum $Optional<Derived>, #Optional.none!enumelt
+  %1 = enum $Optional<Derived>, #Optional.none!enumelt
+  %2 = unchecked_ref_cast %0 : $Optional<Derived> to $X
+  %3 = unchecked_ref_cast %1 : $Optional<Derived> to $X
+  return %2 : $X
+}
+
+// CHECK-LABEL: Escape information for test_unchecked_ref_cast_from_non_optional_to_optional:
+// CHECK: return[e1]: %0 = alloc_ref $Derived
+// CHECK:  -    :   %1 = alloc_ref $Derived
+// CHECK: End function test_unchecked_ref_cast_from_non_optional_to_optional
+sil @test_unchecked_ref_cast_from_non_optional_to_optional : $@convention(thin) () -> @owned Optional<X> {
+bb0:
+  %0 = alloc_ref $Derived
+  %1 = alloc_ref $Derived
+  %2 = unchecked_ref_cast %0 : $Derived to $Optional<X>
+  %3 = unchecked_ref_cast %1 : $Derived to $Optional<X>
+  return %2 : $Optional<X>
+}
+
 // CHECK-LABEL: Escape information for test_unchecked_addr_cast:
 // CHECK: global:   %1 = alloc_ref $X
 // CHECK: End function test_unchecked_addr_cast

test/SILOptimizer/redundant_load_elim_with_casts.sil

@@ -52,6 +52,7 @@ sil @use : $@convention(thin) (Builtin.Int32) -> ()
 sil @use_a : $@convention(thin) (@in A) -> ()
 sil @escaped_a_ptr : $@convention(thin) () -> @out A
 sil @escaped_a : $@convention(thin) () -> Builtin.RawPointer
+sil @b_i_plus_one : $@convention(method) (@guaranteed B) -> ()
 
 // *NOTE* This does not handle raw pointer since raw pointer is only layout compatible with heap references.
 // CHECK-FUTURE: sil @store_to_load_forward_unchecked_addr_cast_struct : $@convention(thin) (Optional<A>) -> () {
@@ -547,3 +548,44 @@ bb0(%0 : $*(I32, I32, I32), %1 : $*((I32, I32), I32)):
   %63 = tuple (%60 : $(I32, I32), %61 : $(I32, I32), %62 : $(I32, I32, I32))
   return %63 : $((I32, I32), (I32, I32), (I32, I32, I32))
 }
+
+// Tests unchecked_ref_cast between Optional<ClassType1> and ClassType2.
+// E? -> B is safe
+//
+// CHECK-FUTURE: sil @unchecked_ref_cast_from_optional_class
+// CHECK: bb3(%6 : $Optional<AnyObject>):
+// CHECK: %8 = load %7 : $*Builtin.Int32
+// CHECK: %10 = apply %9(%5) : $@convention(method) (@guaranteed B) -> ()
+// CHECK: %12 = load %11 : $*Builtin.Int32
+// CHECK: return
+sil @unchecked_ref_cast_from_optional_class : $@convention(thin) (Optional<E>) -> () {
+bb0(%0 : $Optional<E>):
+  switch_enum %0 : $Optional<E>, case #Optional.some!enumelt: bb1, case #Optional.none!enumelt: bb2
+
+bb1(%1 : $E):
+  %2 = enum $Optional<E>, #Optional.some!enumelt, %1 : $E
+  br bb3(%2 : $Optional<E>)
+
+bb2:
+  %3 = enum $Optional<E>, #Optional.none!enumelt
+  br bb3(%3 : $Optional<E>)
+
+bb3(%4 : $Optional<E>):
+  %5 = unchecked_ref_cast %4 : $Optional<E> to $B
+
+  %6 = ref_element_addr %5 : $B, #B.i
+  %7 = begin_access [read] [dynamic] [no_nested_conflict] %6 : $* Builtin.Int32
+  %8 = load %7 : $*Builtin.Int32
+  end_access %7 : $*Builtin.Int32
+
+  %9 = function_ref @b_i_plus_one : $@convention(method) (@guaranteed B) -> ()
+  %10 = apply %9(%5) : $@convention(method) (@guaranteed B) -> ()
+
+  %11 = begin_access [read] [dynamic] [no_nested_conflict] %6 : $*Builtin.Int32
+  %12 = load %11 : $*Builtin.Int32
+  end_access %11 : $*Builtin.Int32
+
+  release_value %4 : $Optional<E>
+  %13 = tuple ()
+  return %13 : $()
+}