Merge pull request swiftlang#78160 from eeckstein/keypath-in-cpature-propagation

CapturePropagation: handle keypaths which are upcast.

Author: eeckstein

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -68,6 +68,9 @@ static SILInstruction *getConstant(SILValue V) {
   if (auto *lit = dyn_cast<LiteralInst>(V))
     return lit;
 
+  if (auto *uc = dyn_cast<UpcastInst>(V))
+    V = uc->getOperand();
+
   if (auto *kp = dyn_cast<KeyPathInst>(V)) {
     // We could support operands, if they are constants, to enable propagation
     // of subscript keypaths. This would require to add the operands in the
@@ -567,7 +570,11 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
       // keypath instruction in this pass, but let dead-object-elimination clean
       // it up later.
       if (!PAI->isOnStack()) {
-        if (getSingleNonDebugUser(kp) != PAI)
+        SILInstruction *user = getSingleNonDebugUser(kp);
+        if (auto *uc = dyn_cast_or_null<UpcastInst>(user))
+          user = getSingleNonDebugUser(uc);
+
+        if (user != PAI)
           return false;
         toDelete.push_back(kp);
       }

test/SILOptimizer/capture_propagate_keypath.swift

@@ -28,7 +28,14 @@ struct GenStr<T> {
   }
 }
 
+struct GetID<C: RandomAccessCollection, ID> {
+  var getID: (C.Element) -> ID
 
+  @inline(__always)
+  init(id: KeyPath<C.Element, ID>) {
+    getID = { $0[keyPath: id] }
+  }
+}
 
 // CHECK-LABEL: sil {{.*}} @$s4test0A6SimpleyyySiAA3StrVXEXEF :
 // CHECK-NOT:     keypath
@@ -50,6 +57,14 @@ func testGenStr(_ mymap: ((GenStr<Int>)->Int) -> ()) {
   mymap(\.c)
 }
 
+// CHECK-LABEL: sil {{.*}} @$s4test0A22GenericEscapingClosureAA5GetIDVySayAA3StrVGSiGyF :
+// CHECK-NOT:     keypath
+// CHECK-LABEL:  } // end sil function '$s4test0A22GenericEscapingClosureAA5GetIDVySayAA3StrVGSiGyF'
+@inline(never)
+func testGenericEscapingClosure() -> GetID<[Str], Int> {
+    GetID(id: \.i)
+} 
+
 // CHECK-LABEL: sil {{.*}} @$s4test0A7GenericyyyxAA6GenStrVyxGXEXElF :
 // CHECK:         keypath
 // CHECK:         keypath
@@ -103,6 +118,11 @@ func calltests() {
     print(c(s))
   }
 
+  // CHECK-OUTPUT-LABEL: testGenericEscapingClosure:
+  print("testGenericEscapingClosure:")
+
+  // CHECK-OUTPUT-NEXT:  27
+  print(testGenericEscapingClosure().getID(Str()))
 }
 
 calltests()

test/SILOptimizer/capture_propagation.sil

@@ -530,6 +530,23 @@ bb0:
   return %7 : $()
 }
 
+// CHECK-LABEL: sil @testCastKeypath :
+// CHECK-NOT:     keypath
+// CHECK:       } // end sil function 'testCastKeypath'
+sil @testCastKeypath : $@convention(thin) () -> () {
+bb0:
+  %0 = keypath $WritableKeyPath<Str, Int>, (root $Str; stored_property #Str.a : $Int)
+  %c = upcast %0 to $KeyPath<Str, Int>
+  %1 = function_ref @closureWithKeypath : $@convention(thin) (Str, @guaranteed KeyPath<Str, Int>) -> Int
+  %2 = partial_apply [callee_guaranteed] %1(%c) : $@convention(thin) (Str, @guaranteed KeyPath<Str, Int>) -> Int
+  %3 = convert_escape_to_noescape %2 : $@callee_guaranteed (Str) -> Int to $@noescape @callee_guaranteed (Str) -> Int
+  %4 = function_ref @calleeWithKeypath : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Int) -> ()
+  %5 = apply %4(%3) : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Int) -> ()
+  strong_release %2 : $@callee_guaranteed (Str) -> Int
+  %7 = tuple ()
+  return %7 : $()
+}
+
 // CHECK-LABEL: sil shared @$s18closureWithKeypath{{.*}}main3StrVSiTf3npk_n : $@convention(thin) (Str) -> Int {
 // CHECK:         [[K:%[0-9]+]] = keypath
 // CHECK:         [[F:%[0-9]+]] = function_ref @swift_getAtKeyPath
@@ -566,6 +583,24 @@ bb0:
   return %7 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @testCastKeypathOSSA :
+// CHECK-NOT:     keypath
+// CHECK:       } // end sil function 'testCastKeypathOSSA'
+sil [ossa] @testCastKeypathOSSA : $@convention(thin) () -> () {
+bb0:
+  %0 = keypath $WritableKeyPath<Str, Int>, (root $Str; stored_property #Str.a : $Int)
+  %c = upcast %0 to $KeyPath<Str, Int>
+  %1 = function_ref @closureWithKeypathOSSA : $@convention(thin) (Str, @guaranteed KeyPath<Str, Int>) -> Int
+  %2 = partial_apply [callee_guaranteed] %1(%c) : $@convention(thin) (Str, @guaranteed KeyPath<Str, Int>) -> Int
+  %3 = convert_escape_to_noescape %2 : $@callee_guaranteed (Str) -> Int to $@noescape @callee_guaranteed (Str) -> Int
+  %4 = function_ref @calleeWithKeypath : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Int) -> ()
+  %5 = apply %4(%3) : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Int) -> ()
+  destroy_value %3 : $@noescape @callee_guaranteed (Str) -> Int
+  destroy_value %2 : $@callee_guaranteed (Str) -> Int
+  %7 = tuple ()
+  return %7 : $()
+}
+
 // CHECK-LABEL: sil [ossa] @testKeypathNoescapeOSSA
 // CHECK:         [[K:%[0-9]+]] = keypath
 // CHECK:         [[C:%[0-9]+]] = function_ref @$s22closureWithKeypathOSSA{{.*}}main3StrVSiTf3npk_n