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

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.

Author: WeZZard

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 { bridged.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,28 @@ 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):
+          if walkDownUses(ofValue: urc, path: path.push(.enumCase, index: 0)) == .abortWalk {
+            return .abortWalk
+          }
+          return walkDownUses(ofValue: urc, path: path.push(.enumCase, index: 1))
+        case (false, true):
+          if let path = path.popIfMatches(.enumCase, index: 0) {
+            if walkDownUses(ofValue: urc, path: path) == .abortWalk {
+              return .abortWalk
+            } else if let path = path.popIfMatches(.enumCase, index: 1) {
+              return walkDownUses(ofValue: urc, path: path)
+            }
+          }
+          return .abortWalk
+        default:
+          return walkDownUses(ofValue: urc, path: path)
+      }
     case let beginDealloc as BeginDeallocRefInst:
       if operand.index == 0 {
         return walkDownUses(ofValue: beginDealloc, path: path)
@@ -699,7 +720,29 @@ 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: 0) {
+            if walkUp(value: urc.fromInstance, path: path) == .abortWalk {
+              return .abortWalk
+            } else if let path = path.popIfMatches(.enumCase, index: 1) {
+              return walkUp(value: urc.fromInstance, path: path)
+            }
+          }
+          return .abortWalk
+        case (false, true):
+          if walkUp(value: urc.fromInstance, path: path.push(.enumCase, index: 0)) == .abortWalk {
+            return .abortWalk
+          } else {
+            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/Types.h

@@ -795,6 +795,9 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Determine whether the type is an opened existential type with Error inside
   bool isOpenedExistentialWithError();
 
+  /// Determine whether the type is an Optional type.
+  bool isOptional() const;
+
   /// Retrieve the set of type parameter packs that occur within this type.
   void getTypeParameterPacks(SmallVectorImpl<Type> &rootParameterPacks);
 
@@ -7926,6 +7929,10 @@ inline bool TypeBase::isClassExistentialType() {
   return false;
 }
 
+inline bool TypeBase::isOptional() const {
+  return getCanonicalType()->isOptional();
+}
+
 inline bool TypeBase::canDynamicallyBeOptionalType(bool includeExistential) {
   CanType T = getCanonicalType();
   auto isArchetypeOrExistential = isa<ArchetypeType>(T) ||

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,10 @@ bool BridgedType::isClassExistential() const {
   return unbridged().isClassExistentialType();
 }
 
+bool BridgedType::isOptional() const {
+  return unbridged().isOptional();
+}
+
 bool BridgedType::isNoEscapeFunction() const {
   return unbridged().isNoEscapeFunction();
 }

include/swift/SIL/SILType.h

@@ -458,6 +458,11 @@ class SILType {
     return getASTType()->hasOpenedExistential();
   }
 
+  /// Returns true if the referenced type is an Optional type.
+  bool isOptional() const {
+    return getASTType()->isOptional();
+  }
+
   TypeTraitResult canBeClass() const {
     return getASTType()->canBeClass();
   }

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 : $()
+}