Merge pull request #80025 from eeckstein/optimize-array-cast

Simplification: remove redundant array casts

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift

@@ -25,6 +25,9 @@ extension ApplyInst : OnoneSimplifiable, SILCombineSimplifiable {
     if context.tryDevirtualize(apply: self, isMandatory: false) != nil {
       return
     }
+    if tryRemoveArrayCast(apply: self, context) {
+      return
+    }
     if !context.preserveDebugInfo {
       _ = tryReplaceExistentialArchetype(of: self, context)
     }
@@ -73,6 +76,40 @@ private func tryTransformThickToThinCallee(of apply: ApplyInst, _ context: Simpl
   return false
 }
 
+/// Removes casts between arrays of the same type.
+///
+///   %1 = function_ref @_arrayConditionalCast : (@guaranteed Array<Int>) -> @owned Optional<Array<Int>>
+///   %2 = apply %1(%0) : (@guaranteed Array<Int>) -> @owned Optional<Array<Int>>
+/// ->
+///   %1 = copy_value %0
+///   %2 = enum $Optional<Array<Int>>, #Optional.some!enumelt, %1
+///
+private func tryRemoveArrayCast(apply: ApplyInst, _ context: SimplifyContext) -> Bool {
+  guard let callee = apply.referencedFunction,
+        callee.hasSemanticsAttribute("array.conditional_cast"),
+        apply.parentFunction.hasOwnership,
+
+          // Check if the cast function has the expected calling convention
+        apply.arguments.count == 1,
+        apply.convention(of: apply.argumentOperands[0]) == .directGuaranteed,
+        apply.functionConvention.results[0].convention == .owned,
+        apply.type.isOptional,
+
+        // Check if the source and target type of the cast is identical.
+        // Note that we are checking the _formal_ element types and not the lowered types, because
+        // the element types are replacement type in the Array's substitution map and this is a formal type.
+        apply.arguments[0].type == apply.type.optionalPayloadType(in: apply.parentFunction)
+  else {
+    return false
+  }
+
+  let builder = Builder(after: apply, context)
+  let copiedArray = builder.createCopyValue(operand: apply.arguments[0])
+  let optional = builder.createEnum(caseIndex: 1, payload: copiedArray, enumType: apply.type)
+  apply.replace(with: optional, context)
+  return true
+}
+
 /// If the apply uses an existential archetype (`@opened("...")`) and the concrete type is known,
 /// replace the existential archetype with the concrete type
 ///   1. in the apply's substitution map
@@ -199,3 +236,10 @@ private extension FullApplySite {
     return SubstitutionMap(genericSignature: genSig, replacementTypes: newReplacementTypes)
   }
 }
+
+private extension Type {
+  func optionalPayloadType(in function: Function) -> Type {
+    let subs = contextSubstitutionMap
+    return subs.replacementTypes[0].loweredType(in: function)
+  }
+}

stdlib/public/core/ArrayCast.swift

@@ -71,6 +71,7 @@ internal func _arrayDownCastConditionalIndirect<SourceValue, TargetValue>(
 ///
 /// - Complexity: O(n), because each element must be checked.
 @inlinable //for performance reasons
+@_semantics("array.conditional_cast")
 public func _arrayConditionalCast<SourceElement, TargetElement>(
   _ source: [SourceElement]
 ) -> [TargetElement]? {

test/SILOptimizer/cast_folding.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -O -Xllvm -sil-print-types -emit-sil %s | %FileCheck %s
+// RUN: %target-swift-frontend -O -Xllvm -sil-print-types -Xllvm -sil-disable-pass=function-signature-opts -emit-sil %s | %FileCheck %s
 
 // We want to check two things here:
 // - Correctness
@@ -723,6 +723,19 @@ func test33() -> Bool {
     return cast33(A())
 }
 
+func castArray<T>(from: [T]) -> [Int]? {
+  return from as? [Int]
+}
+
+// CHECK-LABEL: sil @$s12cast_folding13callCastArray1aSaySiGSgAD_tF :
+// CHECK-NOT:     apply
+// CHECK:         [[E:%.*]] = enum $Optional<Array<Int>>, #Optional.some!enumelt, %0
+// CHECK-NOT:     apply
+// CHECK:         return [[E]]
+// CHECK:       } // end sil function '$s12cast_folding13callCastArray1aSaySiGSgAD_tF'
+public func callCastArray(a: [Int]) -> [Int]? {
+  return castArray(from: a)
+}
 
 protocol PP {
   func foo() -> Int

test/SILOptimizer/simplify_apply.sil

@@ -139,3 +139,29 @@ bb0(%0 : $Int):
   return %11
 }
 
+sil [_semantics "array.conditional_cast"] @cast_int_array : $@convention(thin) (@guaranteed Array<Int>) -> @owned Optional<Array<Int>>
+sil [_semantics "array.conditional_cast"] @cast_any_int_array : $@convention(thin) (@guaranteed Array<Any>) -> @owned Optional<Array<Int>>
+
+// CHECK-LABEL: sil [ossa] @remove_array_cast :
+// CHECK:         %1 = copy_value %0
+// CHECK:         %2 = enum $Optional<Array<Int>>, #Optional.some!enumelt, %1
+// CHECK:         return %2
+// CHECK:       } // end sil function 'remove_array_cast'
+sil [ossa] @remove_array_cast : $@convention(thin) (@guaranteed Array<Int>) -> @owned Optional<Array<Int>> {
+bb0(%0 : @guaranteed $Array<Int>):
+  %1 = function_ref @cast_int_array : $@convention(thin) (@guaranteed Array<Int>) -> @owned Optional<Array<Int>>
+  %2 = apply %1(%0) : $@convention(thin) (@guaranteed Array<Int>) -> @owned Optional<Array<Int>>
+  return %2
+}
+
+// CHECK-LABEL: sil [ossa] @dont_remove_array_cast :
+// CHECK:         %2 = apply
+// CHECK:         return %2
+// CHECK:       } // end sil function 'dont_remove_array_cast'
+sil [ossa] @dont_remove_array_cast : $@convention(thin) (@guaranteed Array<Any>) -> @owned Optional<Array<Int>> {
+bb0(%0 : @guaranteed $Array<Any>):
+  %1 = function_ref @cast_any_int_array : $@convention(thin) (@guaranteed Array<Any>) -> @owned Optional<Array<Int>>
+  %2 = apply %1(%0) : $@convention(thin) (@guaranteed Array<Any>) -> @owned Optional<Array<Int>>
+  return %2
+}
+