ConstantCapturePropagation: don't propagate keypaths with multiple uses in non-OSSA

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ConstantCapturePropagation.swift

@@ -66,6 +66,10 @@ let constantCapturePropagation = FunctionPass(name: "constant-capture-propagatio
       continue
     }
 
+    if !context.continueWithNextSubpassRun(for: partialApply) {
+      return
+    }
+
     optimizeClosureWithDeadCaptures(of: partialApply, context)
 
     if partialApply.isDeleted {
@@ -301,9 +305,15 @@ private extension PartialApplyInst {
     var nonConstArgs = [Operand]()
     var hasKeypath = false
     for argOp in argumentOperands {
-      if argOp.value.isConstant(hasKeypath: &hasKeypath) {
+      // In non-OSSA we don't know where to insert the compensating release for a propagated keypath.
+      // Therefore bail if a keypath has multiple uses.
+      switch argOp.value.isConstant(requireSingleUse: !parentFunction.hasOwnership && !isOnStack) {
+      case .constant:
+        constArgs.append(argOp)
+      case .constantWithKeypath:
         constArgs.append(argOp)
-      } else {
+        hasKeypath = true
+      case .notConstant:
         nonConstArgs.append(argOp)
       }
     }
@@ -333,33 +343,49 @@ private extension FullApplySite {
   }
 }
 
+private enum ConstantKind {
+  case notConstant
+  case constant
+  case constantWithKeypath
+
+  func merge(with other: ConstantKind) -> ConstantKind {
+    switch (self, other) {
+      case (.notConstant, _):      return .notConstant
+      case (_, .notConstant):      return .notConstant
+      case (.constant, .constant): return .constant
+      default:                     return .constantWithKeypath
+    }
+  }
+}
+
 private extension Value {
-  func isConstant(hasKeypath: inout Bool) -> Bool {
+  func isConstant(requireSingleUse: Bool) -> ConstantKind {
     // All instructions handled here must also be handled in
     // `FunctionSignatureSpecializationMangler::mangleConstantProp`.
+    let result: ConstantKind
     switch self {
     case let si as StructInst:
-      for op in si.operands {
-        if !op.value.isConstant(hasKeypath: &hasKeypath) {
-          return false
-        }
-      }
-      return true
+      result = si.operands.reduce(.constant, {
+        $0.merge(with: $1.value.isConstant(requireSingleUse: requireSingleUse))
+      })
     case is ThinToThickFunctionInst, is ConvertFunctionInst, is UpcastInst, is OpenExistentialRefInst:
-      return (self as! UnaryInstruction).operand.value.isConstant(hasKeypath: &hasKeypath)
+      result = (self as! UnaryInstruction).operand.value.isConstant(requireSingleUse: requireSingleUse)
     case is StringLiteralInst, is IntegerLiteralInst, is FloatLiteralInst, is FunctionRefInst, is GlobalAddrInst:
-      return true
+      result = .constant
     case let keyPath as KeyPathInst:
-      hasKeypath = true
       guard keyPath.operands.isEmpty,
             keyPath.hasPattern,
             !keyPath.substitutionMap.hasAnySubstitutableParams
       else {
-        return false
+        return .notConstant
       }
-      return true
+      result = .constantWithKeypath
     default:
-      return false
+      return .notConstant
+    }
+    if requireSingleUse, result == .constantWithKeypath, !uses.ignoreDebugUses.isSingleUse {
+      return .notConstant
     }
+    return result
   }
 }

test/SILOptimizer/capture_propagation.sil

@@ -558,6 +558,29 @@ bb0:
   return %7 : $()
 }
 
+// CHECK-LABEL: sil @dontPropagateMulitUseKeyPathInNonOSSA :
+// CHECK:         %0 = keypath
+// CHECK:         [[U:%.*]] = upcast %0
+// CHECK:         [[C:%.*]] = function_ref @closureWithKeypath
+// CHECK:         partial_apply [callee_guaranteed] [[C]]([[U]])
+// CHECK:       } // end sil function 'dontPropagateMulitUseKeyPathInNonOSSA'
+sil @dontPropagateMulitUseKeyPathInNonOSSA : $@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
+  strong_retain %0
+  %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
+  strong_release %0
+  %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
@@ -824,6 +847,46 @@ bb0:
   return %12 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @testNonConstStruct1 :
+// CHECK:         [[S:%.*]] = struct $S
+// CHECK:         partial_apply [callee_guaranteed] {{%[0-9]+}}([[S]])
+// CHECK:       } // end sil function 'testNonConstStruct1'
+sil [ossa] @testNonConstStruct1 : $@convention(thin) (Int32) -> () {
+bb0(%0 : $Int32):
+  %2 = integer_literal $Builtin.Int1, 0
+  %3 = struct $Bool (%2)
+  %4 = struct $S (%0, %3)
+  %5 = function_ref @closureWithStruct : $@convention(thin) (Str, S) -> Builtin.Int32
+  %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (Str, S) -> Builtin.Int32
+  %7 = convert_escape_to_noescape %6 to $@noescape @callee_guaranteed (Str) -> Builtin.Int32
+  %8 = function_ref @useIntClosure : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Builtin.Int32) -> ()
+  %9 = apply %8(%7) : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Builtin.Int32) -> ()
+  destroy_value %7
+  destroy_value %6
+  %12 = tuple ()
+  return %12 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @testNonConstStruct2 :
+// CHECK:         [[S:%.*]] = struct $S
+// CHECK:         partial_apply [callee_guaranteed] {{%[0-9]+}}([[S]])
+// CHECK:       } // end sil function 'testNonConstStruct2'
+sil [ossa] @testNonConstStruct2 : $@convention(thin) (Bool) -> () {
+bb0(%0 : $Bool):
+  %1 = integer_literal $Builtin.Int32, 3
+  %2 = struct $Int32 (%1)
+  %4 = struct $S (%2, %0)
+  %5 = function_ref @closureWithStruct : $@convention(thin) (Str, S) -> Builtin.Int32
+  %6 = partial_apply [callee_guaranteed] %5(%4) : $@convention(thin) (Str, S) -> Builtin.Int32
+  %7 = convert_escape_to_noescape %6 to $@noescape @callee_guaranteed (Str) -> Builtin.Int32
+  %8 = function_ref @useIntClosure : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Builtin.Int32) -> ()
+  %9 = apply %8(%7) : $@convention(thin) (@noescape @callee_guaranteed (Str) -> Builtin.Int32) -> ()
+  destroy_value %7
+  destroy_value %6
+  %12 = tuple ()
+  return %12 : $()
+}
+
 // CHECK-LABEL: sil shared [ossa] @$s17closureWithStruct4main1SVs5Int32VSbTf3npSSi3Si0_n : $@convention(thin) (Str) -> Builtin.Int32 {
 // CHECK:       bb0(%0 : $Str):
 // CHECK:         %1 = integer_literal $Builtin.Int32, 3