Merge pull request #78543 from eeckstein/fix-copy-to-borrow-optimization

Fix two problems in the copy-to-borrow optimization

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CopyToBorrowOptimization.swift

@@ -122,7 +122,11 @@ private struct Uses {
   // Exit blocks of the load/copy_value's liverange which don't have a destroy.
   // Those are successor blocks of terminators, like `switch_enum`, which do _not_ forward the value.
   // E.g. the none-case of a switch_enum of an Optional.
-  private(set) var nonDestroyingLiverangeExits: Stack<BasicBlock>
+  private(set) var nonDestroyingLiverangeExits: Stack<Instruction>
+
+  var allLifetimeEndingInstructions: [Instruction] {
+    Array(destroys.lazy.map { $0 }) + Array(nonDestroyingLiverangeExits)
+  }
 
   private(set) var usersInDeadEndBlocks: Stack<Instruction>
 
@@ -179,8 +183,13 @@ private struct Uses {
       // A terminator instruction can implicitly end the lifetime of its operand in a success block,
       // e.g. a `switch_enum` with a non-payload case block. Such success blocks need an `end_borrow`, though.
       for succ in termInst.successors where !succ.arguments.contains(where: {$0.ownership == .owned}) {
-        nonDestroyingLiverangeExits.append(succ)
+        nonDestroyingLiverangeExits.append(succ.instructions.first!)
       }
+    } else if !forwardingInst.forwardedResults.contains(where: { $0.ownership == .owned }) {
+      // The forwarding instruction has no owned result, which means it ends the lifetime of its owned operand.
+      // This can happen with an `unchecked_enum_data` which extracts a trivial payload out of a
+      // non-trivial enum.
+      nonDestroyingLiverangeExits.append(forwardingInst.next!)
     }
   }
 
@@ -268,7 +277,7 @@ private func remove(copy: CopyValueInst, collectedUses: Uses, liverange: Instruc
   context.erase(instructions: collectedUses.destroys)
 }
 
-// Handle the special case if the `load` or `copy_valuw` is immediately followed by a single `move_value`.
+// Handle the special case if the `load` or `copy_value` is immediately followed by a single `move_value`.
 // In this case we have to preserve the move's flags by inserting a `begin_borrow` with the same flags.
 // For example:
 //
@@ -314,15 +323,11 @@ private func createEndBorrows(for beginBorrow: Value, atEndOf liverange: Instruc
   //   destroy_value %2
   //   destroy_value %3  // The final destroy. Here we need to create the `end_borrow`(s)
   //
-  for destroy in collectedUses.destroys where !liverange.contains(destroy) {
-    let builder = Builder(before: destroy, context)
-    builder.createEndBorrow(of: beginBorrow)
-  }
-  for liverangeExitBlock in collectedUses.nonDestroyingLiverangeExits where
-      !liverange.blockRange.contains(liverangeExitBlock)
-  {
-    let builder = Builder(atBeginOf: liverangeExitBlock, context)
-    builder.createEndBorrow(of: beginBorrow)
+  for endInst in collectedUses.allLifetimeEndingInstructions {
+    if !liverange.contains(endInst) {
+      let builder = Builder(before: endInst, context)
+      builder.createEndBorrow(of: beginBorrow)
+    }
   }
 }
 

test/SILOptimizer/copy-to-borrow-optimization.sil

@@ -2119,3 +2119,43 @@ bb0(%0 : @owned $Klass):
   unreachable
 }
 
+// CHECK-LABEL: sil [ossa] @end_borrows_in_liferange_exit_block :
+// CHECK:         %1 = load_borrow %0
+// CHECK:         %2 = begin_borrow [lexical] %1
+// CHECK:       bb1({{.*}} : @guaranteed $C):
+// CHECK-NEXT:    end_borrow %2
+// CHECK-NEXT:    end_borrow %1
+// CHECK:       bb2:
+// CHECK-NEXT:    end_borrow %2
+// CHECK-NEXT:    end_borrow %1
+// CHECK:       } // end sil function 'end_borrows_in_liferange_exit_block'
+sil [ossa] @end_borrows_in_liferange_exit_block : $@convention(thin) (@in_guaranteed Optional<C>) -> () {
+bb0(%0 : $*Optional<C>):
+  %1 = load [copy] %0
+  %2 = move_value [lexical] %1
+  switch_enum %2, case #Optional.some!enumelt: bb1, case #Optional.none!enumelt: bb2
+
+bb1(%4 : @owned $C):
+  destroy_value %4
+  br bb3
+
+bb2:
+  br bb3
+
+bb3:
+  %r = tuple ()
+  return %r
+}
+
+// CHECK-LABEL: sil [ossa] @lifetime_ending_enum_data :
+// CHECK:         %1 = load_borrow %0
+// CHECK-NEXT:    %2 = unchecked_enum_data %1
+// CHECK-NEXT:    end_borrow %1
+// CHECK:       } // end sil function 'lifetime_ending_enum_data'
+sil [ossa] @lifetime_ending_enum_data : $@convention(method) (@in_guaranteed MultiPayload) -> MyInt {
+bb0(%0 : $*MultiPayload):
+  %1 = load [copy] %0
+  %2 = unchecked_enum_data %1, #MultiPayload.b!enumelt
+  return %2
+}
+