LifetimeDependenceScopeFixup: crash handling dead-end coroutine

When extending a coroutine, handle the end_borrow instruction used to end a
coroutine lifetime at a dead-end block.

Fixes rdar://153479358 (Compiler crash when force-unwrapping optional ~Copyable type)

(cherry picked from commit 5b5f370ce1)

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceScopeFixup.swift

@@ -801,7 +801,7 @@ extension ExtendableScope {
   func canExtend(beginApply: BeginApplyInst, over range: inout InstructionRange, _ context: some Context) -> Bool {
     let canEndAtBoundary = { (boundaryInst: Instruction) in
       switch beginApply.endReaches(block: boundaryInst.parentBlock, context) {
-      case .abortReaches, .endReaches:
+      case .abortReaches, .endReaches, .deadEndReaches:
         return true
       case .none:
         return false
@@ -929,64 +929,62 @@ private extension BeginApplyInst {
       return builder.createEndApply(beginApply: self)
     case .abortReaches:
       return builder.createAbortApply(beginApply: self)
+    case .deadEndReaches:
+      return builder.createEndBorrow(of: self.token)
     }
   }
 
   enum EndReaches {
     case endReaches
     case abortReaches
+    case deadEndReaches
   }
 
   /// Return the single kind of coroutine termination that reaches 'reachableBlock' or nil.
   func endReaches(block reachableBlock: BasicBlock, _ context: some Context) -> EndReaches? {
-    var endBlocks = BasicBlockSet(context)
-    var abortBlocks = BasicBlockSet(context)
+    // TODO: use InlineArray<3> once bootstrapping is fixed.
+    var endingBlockMap: [(EndReaches, BasicBlockSet)] = [
+      (.endReaches, BasicBlockSet(context)),
+      (.abortReaches, BasicBlockSet(context)),
+      (.deadEndReaches, BasicBlockSet(context))
+    ]
     defer {
-      endBlocks.deinitialize()
-      abortBlocks.deinitialize()
+      for index in endingBlockMap.indices {
+        endingBlockMap[index].1.deinitialize()
+      }
     }
     for endInst in endInstructions {
+      let endKind: EndReaches
       switch endInst {
       case let endApply as EndApplyInst:
         // Cannot extend the scope of a coroutine when the resume produces a value.
         if !endApply.type.isEmpty(in: parentFunction) {
           return nil
         }
-        endBlocks.insert(endInst.parentBlock)
+        endKind = .endReaches
       case is AbortApplyInst:
-        abortBlocks.insert(endInst.parentBlock)
+        endKind = .abortReaches
+      case is EndBorrowInst:
+        endKind = .deadEndReaches
       default:
         fatalError("invalid begin_apply ending instruction")
       }
+      let endingBlocksIndex = endingBlockMap.firstIndex(where: { $0.0 == endKind })!
+      endingBlockMap[endingBlocksIndex].1.insert(endInst.parentBlock)
     }
     var endReaches: EndReaches?
     var backwardWalk = BasicBlockWorklist(context)
     defer { backwardWalk.deinitialize() }
 
     let backwardVisit = { (block: BasicBlock) -> WalkResult in
-      if endBlocks.contains(block) {
-        switch endReaches {
-        case .none:
-          endReaches = .endReaches
-          break
-        case .endReaches:
-          break
-        case .abortReaches:
-          return .abortWalk
+      for (endKind, endingBlocks) in endingBlockMap {
+        if endingBlocks.contains(block) {
+          if let endReaches = endReaches, endReaches != endKind {
+            return .abortWalk
+          }
+          endReaches = endKind
+          return .continueWalk
         }
-        return .continueWalk
-      }
-      if abortBlocks.contains(block) {
-        switch endReaches {
-        case .none:
-          endReaches = .abortReaches
-          break
-        case .abortReaches:
-          break
-        case .endReaches:
-          return .abortWalk
-        }
-        return .continueWalk
       }
       if block == self.parentBlock {
         // the insertion point is not dominated by the coroutine

docs/SIL/Instructions.md

@@ -2615,11 +2615,11 @@ except:
     instead of its normal results.
 
 The final (in the case of `@yield_once`) or penultimate (in the case of
-`@yield_once_2`) result of a `begin_apply` is a "token", a special
-value which can only be used as the operand of an `end_apply` or
-`abort_apply` instruction. Before this second instruction is executed,
-the coroutine is said to be "suspended", and the token represents a
-reference to its suspended activation record.
+`@yield_once_2`) result of a `begin_apply` is a "token", a special value which
+can only be used as the operand of an `end_apply`, `abort_apply`, or
+`end_borrow` instruction. Before this second instruction is executed, the
+coroutine is said to be "suspended", and the token represents a reference to its
+suspended activation record.
 
 If the coroutine's kind `yield_once_2`, its final result is an address
 of a "token", representing the allocation done by the callee

test/SILOptimizer/lifetime_dependence/scope_fixup.sil

@@ -39,6 +39,13 @@ struct NCContainer : ~Copyable {
   var wrapper: Wrapper { get } // _read
 }
 
+struct NCWrapper : ~Copyable, ~Escapable {
+  @_hasStorage let a: NE { get }
+  deinit
+}
+
+sil @NCWrapper_getNE : $@convention(method) (@guaranteed NCWrapper) -> @lifetime(borrow 0) @owned NE
+
 struct TrivialHolder {
   var pointer: UnsafeRawPointer
 }
@@ -86,6 +93,12 @@ sil @readAccess : $@yield_once @convention(method) (@guaranteed Holder) -> @life
 sil @yieldInoutHolder : $@yield_once @convention(method) (@inout Holder) -> @yields @inout Holder
 sil @yieldInoutNE : $@yield_once @convention(method) (@inout Holder) -> @lifetime(borrow 0) @owned NE
 
+class C {
+  @_hasStorage @_hasInitialValue private var nc: NCWrapper? { get set }
+}
+
+sil @C_read : $@yield_once @convention(method) (@guaranteed C) -> @yields @guaranteed Optional<NCWrapper>
+
 // NCContainer.wrapper._read:
 //   var wrapper: Wrapper {
 //     _read {
@@ -613,3 +626,56 @@ bb0(%0 : @owned $NCContainer):
   %31 = tuple ()
   return %31
 }
+
+// rdar://153479358 (Compiler crash when force-unwrapping optional ~Copyable type)
+//
+// Handle dead end coroutines: begin_apply -> end_borrow
+// CHECK-LABEL: sil hidden [ossa] @testReadDeadEnd : $@convention(method) (@guaranteed C) -> () {
+// CHECK: bb0(%0 : @guaranteed $C):
+// CHECK:   ({{.*}}, [[TOKEN:%[0-9]+]]) = begin_apply %{{.*}}(%0) : $@yield_once @convention(method) (@guaranteed C) -> @yields @guaranteed Optional<NCWrapper>
+// CHECK:   switch_enum %{{.*}}, case #Optional.some!enumelt: bb2, case #Optional.none!enumelt: bb1
+// CHECK: bb1:
+// CHECK:   destroy_value [dead_end]
+// CHECK:   end_borrow [[TOKEN]]
+// CHECK:   unreachable
+// CHECK: bb2(%{{.*}} : @guaranteed $NCWrapper):
+// CHECK:   mark_dependence [unresolved]
+// CHECK:   destroy_value
+// CHECK:   destroy_value
+// CHECK:   end_apply [[TOKEN]] as $()
+// CHECK-LABEL: } // end sil function 'testReadDeadEnd'
+sil hidden [ossa] @testReadDeadEnd : $@convention(method) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  // FIXME: I don't know how to print a lifetime-dependent class method in SIL.
+  // %2 = class_method %0, #C.nc!read : (C) -> () -> (), $@yield_once @convention(method) (@guaranteed C) -> @yields @guaranteed Optional<NCWrapper>
+  %2 = function_ref @C_read : $@yield_once @convention(method) (@guaranteed C) -> @yields @guaranteed Optional<NCWrapper>
+  (%3, %4) = begin_apply %2(%0) : $@yield_once @convention(method) (@guaranteed C) -> @yields @guaranteed Optional<NCWrapper>
+  %5 = copy_value %3
+  %6 = mark_unresolved_non_copyable_value [no_consume_or_assign] %5
+  %7 = alloc_stack $Optional<NCWrapper>
+  %8 = mark_unresolved_non_copyable_value [no_consume_or_assign] %7
+  %9 = store_borrow %6 to %8
+  %10 = load_borrow [unchecked] %9
+  switch_enum %10, case #Optional.some!enumelt: bb2, case #Optional.none!enumelt: bb1
+
+bb1:
+  end_borrow %10
+  end_borrow %9
+  destroy_value [dead_end] %6
+  end_borrow %4
+  unreachable
+
+bb2(%25 : @guaranteed $NCWrapper):
+  %26 = function_ref @NCWrapper_getNE : $@convention(method) (@guaranteed NCWrapper) -> @lifetime(borrow 0) @owned NE
+  %27 = apply %26(%25) : $@convention(method) (@guaranteed NCWrapper) -> @lifetime(borrow 0) @owned NE
+  %28 = mark_dependence [unresolved] %27 on %3
+  end_borrow %10
+  end_borrow %9
+  destroy_value %6
+  %32 = end_apply %4 as $()
+  dealloc_stack %7
+  %34 = move_value [var_decl] %28
+  destroy_value %34
+  %37 = tuple ()
+  return %37
+}

test/SILOptimizer/lifetime_dependence/verify_diagnostics.swift

@@ -90,6 +90,24 @@ func getImmutableSpan(_ array: inout [Int]) -> Span<Int> {
  return array.span
 }
 
+struct NCInt: ~Copyable {
+  var i: Int
+
+  @_lifetime(borrow self)
+  func getNE() -> NEInt {
+    NEInt(owner: self)
+  }
+}
+
+public struct NEInt: ~Escapable {
+  var i: Int
+
+  @_lifetime(borrow owner)
+  init(owner: borrowing NCInt) {
+    self.i = owner.i
+  }
+}
+
 struct TestDeinitCallsAddressor: ~Copyable, ~Escapable {
   let a: Borrow<A>
 
@@ -190,3 +208,23 @@ func testIndirectClosureResult<T>(f: () -> GNE<T>) -> GNE<T> {
   // expected-note  @-3{{it depends on the lifetime of argument '$return_value'}}
   // expected-note  @-3{{this use causes the lifetime-dependent value to escape}}
 }
+
+// =============================================================================
+// Coroutines
+// =============================================================================
+
+// Test _read of a noncopyable type with a dead-end (end_borrow)
+//
+// rdar://153479358 (Compiler crash when force-unwrapping optional ~Copyable type)
+class ClassStorage {
+  private var nc: NCInt?
+
+  init(nc: consuming NCInt?) {
+    self.nc = nc
+  }
+
+  func readNoncopyable() {
+    let ne = self.nc!.getNE()
+    _ = ne
+  }
+}