EscapeUtils: better handling of noescape closures and convert_function instructions

SwiftCompilerSources/Sources/Optimizer/Utilities/EscapeUtils.swift

@@ -361,10 +361,6 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
     case is ApplyInst, is TryApplyInst, is BeginApplyInst:
       return walkDownCallee(argOp: operand, apply: instruction as! FullApplySite, path: path)
     case let pai as PartialApplyInst:
-      // This is a non-stack closure.
-      // For `stack` closures, `hasRelevantType` in `walkDown` will return false
-      // stopping the walk since they don't escape.
-
       // Check whether the partially applied argument can escape in the body.
       if walkDownCallee(argOp: operand, apply: pai, path: path.with(knownType: nil)) == .abortWalk {
         return .abortWalk
@@ -376,8 +372,9 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
       // 2. something can escape in a destructor when the context is destroyed
       return walkDownUses(ofValue: pai, path: path.with(knownType: nil))
     case let pta as PointerToAddressInst:
-      assert(operand.index == 0)
       return walkDownUses(ofAddress: pta, path: path.with(knownType: nil))
+    case let cv as ConvertFunctionInst:
+      return walkDownUses(ofValue: cv, path: path.with(knownType: nil))
     case let bi as BuiltinInst:
       switch bi.id {
       case .DestroyArray:
@@ -561,6 +558,12 @@ fileprivate struct EscapeWalker<V: EscapeVisitor> : ValueDefUseWalker,
       return .continueWalk
     }
 
+    if argOp.value.type.isNoEscapeFunction {
+      // Per definition a `partial_apply [on_stack]` cannot escape the callee.
+      // Potential escapes of its captured values are already handled when visiting the `partial_apply`.
+      return .continueWalk
+    }
+
     // Argument effects do not consider any potential stores to the argument (or it's content).
     // Therefore, if we need to track stores, the argument effects do not correctly describe what we need.
     // For example, argument 0 in the following function is marked as not-escaping, although there

SwiftCompilerSources/Sources/SIL/Type.swift

@@ -41,6 +41,7 @@ public struct Type : CustomStringConvertible, NoReflectionChildren {
   public var isEnum: Bool { bridged.getEnumOrBoundGenericEnum() != nil }
   public var isFunction: Bool { bridged.isFunction() }
   public var isMetatype: Bool { bridged.isMetatype() }
+  public var isNoEscapeFunction: Bool { bridged.isNoEscapeFunction() }
 
   /// Can only be used if the type is in fact a nominal type (`isNominal` is true).
   public var nominal: NominalTypeDecl {

include/swift/SIL/SILType.h

@@ -471,6 +471,13 @@ class SILType {
     return false;
   }
 
+  bool isNoEscapeFunction() const {
+    if (auto *fTy = getASTType()->getAs<SILFunctionType>()) {
+      return fTy->isNoEscape();
+    }
+    return false;
+  }
+
   /// True if the type involves any archetypes.
   bool hasArchetype() const { return getASTType()->hasArchetype(); }
 

test/SILOptimizer/escape_info.sil

@@ -1022,6 +1022,9 @@ sil [ossa] @closure6 : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> @own
 [%1: escape c*.v** => %r.c*.v**]
 }
 
+sil @take_noescape_closure : $@convention(thin) (@guaranteed @noescape @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+sil @take_escaping_closure : $@convention(thin) (@owned @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+
 // CHECK-LABEL: Escape information for callClosure1:
 // CHECK:  -    :   %0 = alloc_ref $Y
 // CHECK: global:   %1 = alloc_ref $Y
@@ -1332,8 +1335,8 @@ sil @test_debug_value : $@convention(thin) () -> () {
 }
 
 // CHECK-LABEL: Escape information for test_walk_up_partial_apply_argument:
-// CHECK: global:   %0 = alloc_ref $Y                               // users: %3, %2
-// CHECK: global:   %6 = alloc_ref $X                               // user: %7
+// CHECK: global:   %0 = alloc_ref $Y
+// CHECK: global:   %6 = alloc_ref $X
 // CHECK: End function test_walk_up_partial_apply_argument
 sil @test_walk_up_partial_apply_argument : $@convention(thin) () -> () {
 bb0:
@@ -1350,6 +1353,36 @@ bb0:
   return %13 : $()
 }
 
+// CHECK-LABEL: Escape information for test_escaping_closure:
+// CHECK: global:   %0 = alloc_ref $Y
+// CHECK: End function test_escaping_closure
+sil @test_escaping_closure : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_ref $Y
+  %1 = function_ref @closure1 : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
+  %2 = partial_apply [callee_guaranteed] %1(%0) : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
+  %3 = function_ref @take_escaping_closure : $@convention(thin) (@owned @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+  %5 = apply %3(%2) : $@convention(thin) (@owned @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+  %13 = tuple ()
+  return %13 : $()
+}
+
+// CHECK-LABEL: Escape information for test_noescape_partial_apply_and_convert_function:
+// CHECK:  -    :   %0 = alloc_ref $Y
+// CHECK: End function test_noescape_partial_apply_and_convert_function
+sil @test_noescape_partial_apply_and_convert_function : $@convention(thin) () -> () {
+bb0:
+  %0 = alloc_ref $Y
+  %1 = function_ref @closure1 : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
+  %2 = partial_apply [callee_guaranteed] [on_stack] %1(%0) : $@convention(thin) (@guaranteed Y, @guaranteed Y) -> ()
+  %3 = convert_function %2 : $@noescape @callee_guaranteed (@guaranteed Y) -> () to $@noescape @callee_guaranteed (@guaranteed Y) -> ()
+  %4 = function_ref @take_noescape_closure : $@convention(thin) (@guaranteed @noescape @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+  %5 = apply %4(%3) : $@convention(thin) (@guaranteed @noescape @callee_guaranteed (@guaranteed Y) -> ()) -> ()
+  dealloc_stack %2 : $@noescape @callee_guaranteed (@guaranteed Y) -> ()
+  %13 = tuple ()
+  return %13 : $()
+}
+
 // CHECK-LABEL: Escape information for test_mismatching_existential_types:
 // CHECK:  -    :   %1 = alloc_ref $Y
 // CHECK: End function test_mismatching_existential_types