Optimizer: support partial_apply of thunks in the copy_block simplification

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyCopyBlock.swift

@@ -29,7 +29,7 @@ extension CopyBlockInst : Simplifiable, SILCombineSimplifiable {
   ///
   func simplify(_ context: SimplifyContext) {
     if hasValidUses(block: self) {
-      replaceBlock( self, with: operand.value, context)
+      replaceBlock(self, with: operand.value, context)
       context.erase(instruction: self)
     }
   }
@@ -44,6 +44,13 @@ private func hasValidUses(block: Value) -> Bool {
       }
     case let apply as FullApplySite where apply.isCallee(operand: use):
       break
+    case let partialApply as PartialApplyInst:
+      // If the block is passed to another function - either as closure argument or as closure capture -
+      // it's "converted" to a swift closure with the help of a thunk. The thunk just calls the block.
+      // If this is a non-escaping partial_apply and it's such a thunk, the block does not escape.
+      if partialApply.canClosureArgumentEscape(closure: use) {
+        return false
+      }
     case is EndBorrowInst, is DestroyValueInst:
       break
     default:
@@ -61,10 +68,40 @@ private func replaceBlock(_ block: Value, with original: Value, _ context: Simpl
       context.erase(instruction: beginBorrow)
     case is FullApplySite:
       use.set(to: original, context)
+    case let partialApply as PartialApplyInst:
+      if original.ownership == .unowned {
+        let builder = Builder(before: partialApply, context)
+        let conv = builder.createUncheckedOwnershipConversion(operand: original, resultOwnership: .guaranteed)
+        use.set(to: conv, context)
+      } else {
+        use.set(to: original, context)
+      }
     case is EndBorrowInst, is DestroyValueInst:
       context.erase(instruction: use.instruction)
     default:
       fatalError("unhandled use")
     }
   }
 }
+
+private extension PartialApplyInst {
+  func canClosureArgumentEscape(closure: Operand) -> Bool {
+    guard isOnStack,
+          let callee = referencedFunction,
+          callee.isDefinition,
+          let argIdx = calleeArgumentIndex(of: closure),
+          // If the callee only _calls_ the closure argument, it does not escape.
+          callee.arguments[argIdx].uses.allSatisfy(isCalleeOperandOfApply)
+    else {
+      return true
+    }
+    return false
+  }
+}
+
+private func isCalleeOperandOfApply(_ operand: Operand) -> Bool {
+  if let apply = operand.instruction as? FullApplySite, apply.isCallee(operand: operand) {
+    return true
+  }
+  return false
+}

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -215,6 +215,13 @@ public struct Builder {
     return notifyNew(cast.getAs(UnconditionalCheckedCastAddrInst.self))
   }
 
+  public func createUncheckedOwnershipConversion(
+    operand: Value, resultOwnership: Ownership
+  ) -> UncheckedOwnershipConversionInst {
+    let uoc = bridged.createUncheckedOwnershipConversion(operand.bridged, resultOwnership._bridged)
+    return notifyNew(uoc.getAs(UncheckedOwnershipConversionInst.self))
+  }
+
   public func createLoad(fromAddress: Value, ownership: LoadInst.LoadOwnership) -> LoadInst {
     let load = bridged.createLoad(fromAddress.bridged, ownership.rawValue)
     return notifyNew(load.getAs(LoadInst.self))

include/swift/SIL/SILBridging.h

@@ -1178,6 +1178,8 @@ struct BridgedBuilder{
         BridgedInstruction::CastingIsolatedConformances isolatedConformances,
         BridgedValue source, BridgedCanType sourceFormalType,
         BridgedValue destination, BridgedCanType targetFormalType) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUncheckedOwnershipConversion(
+        BridgedValue op, BridgedValue::Ownership ownership) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createLoad(BridgedValue op, SwiftInt ownership) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createLoadBorrow(BridgedValue op) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createBeginDeallocRef(BridgedValue reference,

include/swift/SIL/SILBridgingImpl.h

@@ -2205,6 +2205,13 @@ BridgedInstruction BridgedBuilder::createLoad(BridgedValue op, SwiftInt ownershi
                                  (swift::LoadOwnershipQualifier)ownership)};
 }
 
+
+BridgedInstruction BridgedBuilder::createUncheckedOwnershipConversion(BridgedValue op,
+                                                                      BridgedValue::Ownership ownership) const {
+  return {unbridged().createUncheckedOwnershipConversion(regularLoc(), op.getSILValue(),
+                                                         BridgedValue::unbridge(ownership))};
+}
+
 BridgedInstruction BridgedBuilder::createLoadBorrow(BridgedValue op) const {
   return {unbridged().createLoadBorrow(regularLoc(), op.getSILValue())};
 }

lib/SIL/Verifier/SILOwnershipVerifier.cpp

@@ -658,6 +658,15 @@ bool SILValueOwnershipChecker::checkValueWithoutLifetimeEndingUses(
     }
   }
 
+  if (auto *uoc = dyn_cast<UncheckedOwnershipConversionInst>(value)) {
+    // When converting from `unowned`, which can happen for ObjectiveC blocks,
+    // we don't require and scope-ending instruction. This falls in the category:
+    // "trust me, I know what I'm doing".
+    if (uoc->getOperand()->getOwnershipKind() == OwnershipKind::Unowned) {
+      return true;
+    }
+  }
+
   // If we have an unowned value, then again there is nothing left to do.
   if (value->getOwnershipKind() == OwnershipKind::Unowned)
     return true;

test/SILOptimizer/optimize_copy_block.swift

@@ -17,7 +17,9 @@ module CModule {
 @import Foundation;
 
 @interface TestClass : NSObject
-- (void)callHandlerInline: (NS_NOESCAPE _Nonnull dispatch_block_t)block;
+- (void)callHandler1: (NS_NOESCAPE _Nonnull dispatch_block_t)block;
+- (void)callHandler2: (NS_NOESCAPE _Nonnull dispatch_block_t)block;
+- (void)callHandler3: (NS_NOESCAPE _Nonnull dispatch_block_t)block;
 @end
 
 
@@ -27,13 +29,34 @@ import CModule
 
 @objc @implementation
 extension TestClass {
-  // CHECK-LABEL: sil private [thunk] @$sSo9TestClassC4testE17callHandlerInlineyyyyXEFTo :
+  // CHECK-LABEL: sil private [thunk] @$sSo9TestClassC4testE12callHandler1yyyyXEFTo :
   // CHECK-NOT:     copy_block
   // CHECK:         apply %0
   // CHECK-NOT:     destroy_value
-  // CHECK:       } // end sil function '$sSo9TestClassC4testE17callHandlerInlineyyyyXEFTo' 
-  func callHandlerInline(_ handler: () -> Void) {
+  // CHECK:       } // end sil function '$sSo9TestClassC4testE12callHandler1yyyyXEFTo'
+  func callHandler1(_ handler: () -> Void) {
     handler()
   }
+
+  // CHECK-LABEL: sil private [thunk] @$sSo9TestClassC4testE12callHandler2yyyyXEFTo :
+  // CHECK-NOT:     copy_block
+  // CHECK:       } // end sil function '$sSo9TestClassC4testE12callHandler2yyyyXEFTo' 
+  func callHandler2(_ handler: () -> Void) {
+    callee(handler)
+  }
+
+  // CHECK-LABEL: sil private [thunk] @$sSo9TestClassC4testE12callHandler3yyyyXEFTo :
+  // CHECK-NOT:     copy_block
+  // CHECK:       } // end sil function '$sSo9TestClassC4testE12callHandler3yyyyXEFTo' 
+  func callHandler3(_ handler: () -> Void) {
+    callClosure {
+      handler()
+    }
+  }
 }
 
+@_silgen_name("callee")
+func callee(_ handler: () -> Void)
+
+@_silgen_name("callClosure")
+func callClosure(_ closure: () -> ())

test/SILOptimizer/simplify_copy_block.sil

@@ -24,6 +24,119 @@ bb0(%0 : @unowned $@convention(block) @noescape () -> ()):
   return %7
 }
 
+// CHECK-LABEL: sil [ossa] @remove_copy_block_with_partial_apply :
+// CHECK-NOT:     copy_block
+// CHECK:         [[B:%.*]] = unchecked_ownership_conversion %0, @unowned to @guaranteed
+// CHECK:         [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] {{%[0-9]+}}([[B]])
+// CHECK:         destroy_value [[PA]]
+// CHECK-NOT:     destroy_value
+// CHECK:       } // end sil function 'remove_copy_block_with_partial_apply'
+sil [ossa] @remove_copy_block_with_partial_apply : $@convention(thin) (@convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @unowned $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] [on_stack] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  destroy_value %4
+  destroy_value %2
+  %7 = tuple ()
+  return %7
+}
+
+// CHECK-LABEL: sil [ossa] @remove_copy_block_with_partial_apply_guaranteed :
+// CHECK-NOT:     copy_block
+// CHECK:         [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] {{%[0-9]+}}(%0)
+// CHECK:         destroy_value [[PA]]
+// CHECK-NOT:     destroy_value
+// CHECK:       } // end sil function 'remove_copy_block_with_partial_apply_guaranteed'
+sil [ossa] @remove_copy_block_with_partial_apply_guaranteed : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @guaranteed $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] [on_stack] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  destroy_value %4
+  destroy_value %2
+  %7 = tuple ()
+  return %7
+}
+
+// CHECK-LABEL: sil [ossa] @remove_copy_block_with_partial_apply_owned :
+// CHECK-NOT:     copy_block
+// CHECK:         [[PA:%.*]] = partial_apply [callee_guaranteed] [on_stack] {{%[0-9]+}}(%0)
+// CHECK:       } // end sil function 'remove_copy_block_with_partial_apply_owned'
+sil [ossa] @remove_copy_block_with_partial_apply_owned : $@convention(thin) (@owned @convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @owned $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] [on_stack] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  destroy_value %4
+  destroy_value %2
+  destroy_value %0
+  %7 = tuple ()
+  return %7
+}
+
+sil [ossa] @thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @guaranteed $@convention(block) @noescape () -> ()):
+  %1 = apply %0() : $@convention(block) @noescape () -> ()
+  %2 = tuple ()
+  return %2
+}
+
+sil @use_closure : $@convention(thin) (@guaranteed @callee_guaranteed () -> ()) -> ()
+
+// CHECK-LABEL: sil [ossa] @dont_remove_copy_block_with_escaping_partial_apply :
+// CHECK:         copy_block
+// CHECK:       } // end sil function 'dont_remove_copy_block_with_escaping_partial_apply'
+sil [ossa] @dont_remove_copy_block_with_escaping_partial_apply : $@convention(thin) (@convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @unowned $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %5 = function_ref @use_closure : $@convention(thin) (@guaranteed @callee_guaranteed () -> ()) -> ()
+  apply %5(%4) : $@convention(thin) (@guaranteed @callee_guaranteed () -> ()) -> ()
+  destroy_value %4
+  %7 = tuple ()
+  return %7
+}
+
+sil [ossa] @unknown_thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+
+// CHECK-LABEL: sil [ossa] @dont_remove_copy_block_with_unknown_partial_apply :
+// CHECK:         copy_block
+// CHECK:       } // end sil function 'dont_remove_copy_block_with_unknown_partial_apply'
+sil [ossa] @dont_remove_copy_block_with_unknown_partial_apply : $@convention(thin) (@convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @unowned $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @unknown_thunk : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] [on_stack] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  destroy_value %4
+  destroy_value %2
+  %7 = tuple ()
+  return %7
+}
+
+// CHECK-LABEL: sil [ossa] @dont_remove_copy_block_with_partial_apply_escaping :
+// CHECK:         copy_block
+// CHECK:       } // end sil function 'dont_remove_copy_block_with_partial_apply_escaping'
+sil [ossa] @dont_remove_copy_block_with_partial_apply_escaping : $@convention(thin) (@convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @unowned $@convention(block) @noescape () -> ()):
+  %2 = copy_block %0
+  %3 = function_ref @thunk_escaping : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  %4 = partial_apply [callee_guaranteed] [on_stack] %3(%2) : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> ()
+  destroy_value %4
+  destroy_value %2
+  %7 = tuple ()
+  return %7
+}
+
+sil [ossa] @thunk_escaping : $@convention(thin) (@guaranteed @convention(block) @noescape () -> ()) -> () {
+bb0(%0 : @guaranteed $@convention(block) @noescape () -> ()):
+  %1 = function_ref @use_block : $@convention(thin) (@convention(block) @noescape () -> ()) -> ()
+  %2 = apply %1(%0) : $@convention(thin) (@convention(block) @noescape () -> ()) -> ()
+  %3 = tuple ()
+  return %3
+}
+
 // CHECK-LABEL: sil [ossa] @dont_remove_copied_block :
 // CHECK:         copy_block
 // CHECK:       } // end sil function 'dont_remove_copied_block'