Fix DestroyHoisting: InteriorLiveness for noescape closures

SwiftCompilerSources/Sources/Optimizer/Utilities/OwnershipLiveness.swift

@@ -715,8 +715,17 @@ extension InteriorUseWalker: AddressUseVisitor {
 
   mutating func dependentAddressUse(of operand: Operand, into value: Value)
     -> WalkResult {
-    walkDownUses(of: value)
-  }    
+    // For Escapable values, simply continue the walk.
+    if value.mayEscape {
+      return walkDownUses(of: value)
+    }
+    // TODO: Handle non-escapable values by walking through copies as done by LifetimeDependenceDefUseWalker or
+    // NonEscapingClosureDefUseWalker. But this code path also handles non-escaping closures that have not been promoted
+    // to [on_stack] (and still have an escapable function type). Such closures may be incorrectly destroyed after their
+    // captures. To avoid this problem, either rewrite ClosureLifetimeFixup to produce correct OSSA lifetimes, or check
+    // for that special case and continue to bailout here.
+    return escapingAddressUse(of: operand)
+  }
 
   mutating func escapingAddressUse(of operand: Operand) -> WalkResult {
     pointerStatus.setEscaping(operand: operand)
@@ -830,12 +839,14 @@ let linearLivenessTest = FunctionTest("linear_liveness_swift") {
 let interiorLivenessTest = FunctionTest("interior_liveness_swift") {
   function, arguments, context in
   let value = arguments.takeValue()
-  print("Interior liveness: \(value)")
+  let visitInnerUses = arguments.hasUntaken ? arguments.takeBool() : false
+
+  print("Interior liveness\(visitInnerUses ? " with inner uses" : ""): \(value)")
 
   var range = InstructionRange(for: value, context)
   defer { range.deinitialize() }
 
-  var visitor = InteriorUseWalker(definingValue: value, ignoreEscape: true, visitInnerUses: false, context) {
+  var visitor = InteriorUseWalker(definingValue: value, ignoreEscape: true, visitInnerUses: visitInnerUses, context) {
     range.insert($0.instruction)
     return .continueWalk
   }

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1182,10 +1182,15 @@ class ClassifyBridgeObjectInst : SingleValueInstruction, UnaryInstruction {}
 final public class PartialApplyInst : SingleValueInstruction, ApplySite {
   public var numArguments: Int { bridged.PartialApplyInst_numArguments() }
 
-  /// WARNING: isOnStack incorrectly returns false for all closures prior to ClosureLifetimeFixup, even if they need to
-  /// be diagnosed as on-stack closures. Use has mayEscape instead.
+  /// Warning: isOnStack returns false for all closures prior to ClosureLifetimeFixup, even if they capture on-stack
+  /// addresses and need to be diagnosed as non-escaping closures. Use mayEscape to determine whether a closure is
+  /// non-escaping prior to ClosureLifetimeFixup.
   public var isOnStack: Bool { bridged.PartialApplyInst_isOnStack() }
 
+  // Warning: ClosureLifetimeFixup does not promote all non-escaping closures to on-stack. When that promotion fails, it
+  // creates a fake destroy of the closure after the captured values that the closure depends on. This is invalid OSSA,
+  // so OSSA utilities need to bail-out when isOnStack is false even if hasNoescapeCapture is true to avoid encoutering
+  // an invalid nested lifetime.
   public var mayEscape: Bool { !isOnStack && !hasNoescapeCapture }
 
   /// True if this closure captures anything nonescaping.

SwiftCompilerSources/Sources/SIL/Type.swift

@@ -73,11 +73,19 @@ public struct Type : CustomStringConvertible, NoReflectionChildren {
 
   public var isMoveOnly: Bool { bridged.isMoveOnly() }
 
-  // Note that invalid types are not considered Escapable. This makes it difficult to make any assumptions about
-  // nonescapable types.
+  /// Return true if this type conforms to Escapable.
+  ///
+  /// Note: invalid types are non-Escapable, so take care not to make assumptions about non-Escapable types.
   public func isEscapable(in function: Function) -> Bool {
     bridged.isEscapable(function.bridged)
   }
+
+  /// Return true if this type conforms to Escapable and is not a noescape function type.
+  ///
+  /// Warning: may return true for (source-level) non-escaping closures. After SILGen, all partial_apply instructions
+  /// have an escaping function type. ClosureLifetimeFixup only changes them to noescape function type if they are
+  /// promoted to [on_stack]. But regardless of stack promotion, a non-escaping closure value's lifetime is constrained
+  /// by the lifetime of its captures. Use Value.mayEscape instead to check for this case.
   public func mayEscape(in function: Function) -> Bool {
     !isNoEscapeFunction && isEscapable(in: function)
   }
@@ -204,16 +212,6 @@ public struct Type : CustomStringConvertible, NoReflectionChildren {
   }
 }
 
-extension Value {
-  public var isEscapable: Bool {
-    type.objectType.isEscapable(in: parentFunction)
-  }
-
-  public var mayEscape: Bool {
-    type.objectType.mayEscape(in: parentFunction)
-  }
-}
-
 extension Type: Equatable {
   public static func ==(lhs: Type, rhs: Type) -> Bool { 
     lhs.bridged.opaqueValue == rhs.bridged.opaqueValue

SwiftCompilerSources/Sources/SIL/Value.swift

@@ -140,6 +140,31 @@ extension Value {
     }
   }
 
+  /// Return true if the object type conforms to Escapable.
+  ///
+  /// Note: noescape function types conform to Escapable, use mayEscape instead to exclude them.
+  public var isEscapable: Bool {
+    type.objectType.isEscapable(in: parentFunction)
+  }
+
+  /// Return true only if this value's lifetime is unconstrained by an outer lifetime. Requires all of the following:
+  /// - the object type conforms to Escapable
+  /// - the type is not a noescape function
+  /// - the value is not the direct result of a partial_apply with a noescape (inout_aliasable) capture.
+  public var mayEscape: Bool {
+    if !type.objectType.mayEscape(in: parentFunction) {
+      return false
+    }
+    // A noescape partial_apply has an escaping function type if it has not been promoted to on_stack, but it's value
+    // still cannot "escape" its captures.
+    //
+    // TODO: This would be much more robust if pai.hasNoescapeCapture simply implied !pai.type.isEscapable
+    if let pai = self as? PartialApplyInst {
+      return pai.mayEscape
+    }
+    return true
+  }
+
   public var definingInstructionOrTerminator: Instruction? {
     if let def = definingInstruction {
       return def

lib/SIL/Verifier/SILVerifier.cpp

@@ -2202,6 +2202,9 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     require(resultInfo->getExtInfo().hasContext(),
             "result of closure cannot have a thin function type");
 
+    require(PAI->getType().isNoEscapeFunction() == PAI->isOnStack(),
+            "closure must be on stack to have a noescape function type");
+
     // We rely on all indirect captures to be in the argument list.
     require(PAI->getCallee()->getType().isObject(),
             "Closure callee must not be an address type.");

test/SILOptimizer/ownership_liveness_unit.sil

@@ -43,6 +43,9 @@ struct NE: ~Escapable {
   init(object: borrowing C) { self.object = copy object }
 }
 
+sil [ossa] @useCAddress : $@convention(thin) (@inout_aliasable C) -> ()
+
+
 // =============================================================================
 // LinearLiveness
 // =============================================================================
@@ -1948,6 +1951,54 @@ bb1(%reborrow2 : @reborrow $D):
   return %99 : $()
 }
 
+// CHECK-LABEL: partial_apply_noescape_onheap: interior_liveness_swift with: %box, true
+// CHECK: Interior liveness with inner uses:   %{{.*}} = alloc_box ${ var C }
+// CHECK-NEXT: Inner scope:   %{{.*}} = begin_borrow [lexical] %{{.*}} : ${ var C }
+// CHECK-NEXT: Pointer escape:   %{{.*}} = partial_apply [callee_guaranteed] %{{.*}}(%{{.*}}) : $@convention(thin) (@inout_aliasable C) -> ()
+// CHECK-NEXT: begin:      %{{.*}} = alloc_box ${ var C }
+// CHECK-NEXT: ends:       destroy_value %{{.*}} : ${ var C }
+// CHECK-NEXT: exits:
+// CHECK-NEXT: interiors:  end_borrow %{{.*}} : ${ var C }
+// CHECK-NEXT:             %{{.*}} = partial_apply [callee_guaranteed] %{{.*}}(%{{.*}}) : $@convention(thin) (@inout_aliasable C) -> ()
+// CHECK-NEXT:             %{{.*}} = project_box %{{.*}} : ${ var C }, 0
+// CHECK-NEXT:             %{{.*}} = begin_borrow [lexical] %3 : ${ var C }
+// CHECK-NEXT: last user:   destroy_value %3 : ${ var C }
+// CHECK-NEXT: partial_apply_noescape_onheap: interior_liveness_swift with: %box, true
+sil [ossa] @partial_apply_noescape_onheap : $@convention(thin) () -> () {
+bb0:
+  cond_br undef, bb2, bb1
+
+bb1:
+  %1 = enum $Optional<@callee_guaranteed () -> ()>, #Optional.none!enumelt
+  br bb4(%1)
+
+bb2:
+  %box = alloc_box ${ var C }
+  %borrow = begin_borrow [lexical] %box
+  specify_test "interior_liveness_swift %box true"
+  %adr = project_box %borrow : ${ var C }, 0
+  br bb3
+
+bb3:
+  %f = function_ref @useCAddress : $@convention(thin) (@inout_aliasable C) -> ()
+  %pa1 = partial_apply [callee_guaranteed] %f(%adr) : $@convention(thin) (@inout_aliasable C) -> ()
+  %pa2 = copy_value %pa1
+  %pane = convert_escape_to_noescape %pa2 : $@callee_guaranteed () -> () to $@noescape @callee_guaranteed () -> ()
+  %enum = enum $Optional<@callee_guaranteed () -> ()>, #Optional.some!enumelt, %pa2
+  end_borrow %borrow
+  destroy_value %box
+  destroy_value %pa1
+  destroy_value %pane
+  br bb4(%enum)
+
+bb4(%phi : @owned $Optional<@callee_guaranteed () -> ()>):
+  // ClosureLifetimeFixup destroys closures after destroying its addressable captures:
+  // this is bad OSSA but still accepted as valid until ClosureLifetimeFixup is rewritten.
+  destroy_value %phi
+  %r = tuple ()
+  return %r
+}
+
 // =============================================================================
 // ExtendedLiveness
 // =============================================================================