Optimizer: add simplification for begin_borrow

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/NamedReturnValueOptimization.swift

@@ -61,7 +61,7 @@ let namedReturnValueOptimization = FunctionPass(name: "named-return-value-optimi
 /// Returns a copy_addr which copies from an alloc_stack to the `outArg` at the end of the function.
 ///
 private func findCopyForNRVO(for outArg: FunctionArgument) -> CopyAddrInst? {
-  guard let singleArgUse = outArg.uses.singleNonDebugUse,
+  guard let singleArgUse = outArg.uses.ignoreDebugUses.singleUse,
         let copyToArg = singleArgUse.instruction as? CopyAddrInst else {
     return nil
   }
@@ -106,7 +106,7 @@ private func findCopyForNRVO(for outArg: FunctionArgument) -> CopyAddrInst? {
 }
 
 private func performNRVO(with copy: CopyAddrInst, _ context: FunctionPassContext) {
-  copy.source.uses.replaceAllExceptDealloc(with: copy.destination, context)
+  copy.source.replaceAllUsesExceptDealloc(with: copy.destination, context)
   assert(copy.source == copy.destination)
   context.erase(instruction: copy)
 }
@@ -122,10 +122,8 @@ private func isAnyInstructionWritingToMemory(after: Instruction) -> Bool {
   return false
 }
 
-private extension UseList {
-  func replaceAllExceptDealloc(with replacement: Value, _ context: some MutatingContext) {
-    for use in self where !(use.instruction is Deallocation)  {
-      use.set(to: replacement, context)
-    }
+private extension Value {
+  func replaceAllUsesExceptDealloc(with replacement: Value, _ context: some MutatingContext) {
+    uses.lazy.filter{!($0.instruction is Deallocation)}.replaceAll(with: replacement, context)
   }
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -8,6 +8,7 @@
 
 swift_compiler_sources(Optimizer
   SimplifyApply.swift
+  SimplifyBeginBorrow.swift
   SimplifyBeginCOWMutation.swift
   SimplifyBranch.swift
   SimplifyBuiltin.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginBorrow.swift

@@ -0,0 +1,211 @@
+//===--- SimplifyBeginBorrow.swift ----------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+
+extension BeginBorrowInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    if borrowedValue.ownership == .owned,
+       // We need to keep lexical lifetimes in place.
+       !isLexical
+    {
+      tryReplaceBorrowWithOwnedOperand(beginBorrow: self, context)
+    }
+  }
+}
+
+private func tryReplaceBorrowWithOwnedOperand(beginBorrow: BeginBorrowInst, _ context: SimplifyContext) {
+  // The last value of a (potentially empty) forwarding chain, beginning at the `begin_borrow`.
+  let forwardedValue = beginBorrow.lookThroughSingleForwardingUses()
+  if forwardedValue.allUsesCanBeConvertedToOwned {
+    if tryReplaceCopy(of: forwardedValue, withCopiedOperandOf: beginBorrow, context) {
+      return
+    }
+    if beginBorrow.borrowedValue.isDestroyed(after: beginBorrow) {
+      convertAllUsesToOwned(of: beginBorrow, context)
+    }
+  }
+}
+
+/// Replace
+/// ```
+///   %1 = begin_borrow %0
+///   %2 = struct_extract %1     // a chain of forwarding instructions
+///   %3 = copy_value %1
+///   // ... uses of %3
+///   end_borrow %1
+/// ```
+/// with
+/// ```
+///   %1 = copy_value %0
+///   %3 = destructure_struct %0 // owned version of the forwarding instructions
+///   // ... uses of %3
+/// ```
+private func tryReplaceCopy(
+  of forwardedValue: Value,
+  withCopiedOperandOf beginBorrow: BeginBorrowInst,
+  _ context: SimplifyContext
+) -> Bool {
+  guard let singleUser = forwardedValue.uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse?.instruction,
+        let copy = singleUser as? CopyValueInst,
+        copy.parentBlock == beginBorrow.parentBlock else {
+    return false
+  }
+  let builder = Builder(before: beginBorrow, context)
+  let copiedOperand = builder.createCopyValue(operand: beginBorrow.borrowedValue)
+  let forwardedOwnedValue = replace(guaranteedValue: beginBorrow, withOwnedValue: copiedOperand, context)
+  copy.uses.replaceAll(with: forwardedOwnedValue, context)
+  context.erase(instruction: copy)
+  context.erase(instructionIncludingAllUsers: beginBorrow)
+  return true
+}
+
+/// Replace
+/// ```
+///   %1 = begin_borrow %0
+///   %2 = struct_extract %1     // a chain of forwarding instructions
+///   // ... uses of %2
+///   end_borrow %1
+///   destroy_value %1           // the only other use of %0 beside begin_borrow
+/// ```
+/// with
+/// ```
+///   %2 = destructure_struct %0 // owned version of the forwarding instructions
+///   // ... uses of %2
+///   destroy_value %2
+/// ```
+private func convertAllUsesToOwned(of beginBorrow: BeginBorrowInst, _ context: SimplifyContext) {
+  let forwardedOwnedValue = replace(guaranteedValue: beginBorrow, withOwnedValue: beginBorrow.borrowedValue, context)
+  beginBorrow.borrowedValue.replaceAllDestroys(with: forwardedOwnedValue, context)
+  context.erase(instructionIncludingAllUsers: beginBorrow)
+}
+
+private extension Value {
+  /// Returns the last value of a (potentially empty) forwarding chain.
+  /// For example, returns %3 for the following def-use chain:
+  /// ```
+  ///   %1 = struct_extract %self, #someField
+  ///   %2 = tuple_extract %1, 0
+  ///   %3 = struct $S(%2)   // %3 has no forwarding users
+  /// ```
+  /// Returns self if this value has no uses which are ForwardingInstructions.
+  func lookThroughSingleForwardingUses() -> Value {
+    if let singleUse = uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse,
+       let fwdInst = singleUse.instruction as? (SingleValueInstruction & ForwardingInstruction),
+       fwdInst.canConvertToOwned,
+       fwdInst.isSingleForwardedOperand(singleUse),
+       fwdInst.parentBlock == parentBlock
+    {
+      return fwdInst.lookThroughSingleForwardingUses()
+    }
+    return self
+  }
+
+  var allUsesCanBeConvertedToOwned: Bool {
+    let relevantUses = uses.ignoreUsers(ofType: EndBorrowInst.self)
+    return relevantUses.allSatisfy { $0.canAccept(ownership: .owned) }
+  }
+
+  func isDestroyed(after nonDestroyUser: Instruction) -> Bool {
+    uses.getSingleUser(notOfType: DestroyValueInst.self) == nonDestroyUser
+  }
+
+  func replaceAllDestroys(with replacement: Value, _ context: SimplifyContext) {
+    uses.filterUsers(ofType: DestroyValueInst.self).replaceAll(with: replacement, context)
+  }
+}
+
+private extension ForwardingInstruction {
+  func isSingleForwardedOperand(_ operand: Operand) -> Bool {
+    switch self {
+    case is StructInst, is TupleInst:
+      // TODO: we could move that logic to StructInst/TupleInst.singleForwardedOperand.
+      return operands.lazy.map({ $0.value.type }).hasSingleNonTrivialElement(at: operand.index, in: parentFunction)
+    default:
+      if let sfo = singleForwardedOperand {
+        return sfo == operand
+      }
+      return false
+    }
+  }
+}
+
+/// Replaces a guaranteed value with an owned value.
+/// 
+/// If the `guaranteedValue`'s use is a ForwardingInstruction (or forwarding instruction chain),
+/// it is converted to an owned version of the forwarding instruction (or instruction chain).
+///
+/// Returns the last owned value in a forwarding-chain or `ownedValue` if `guaranteedValue` has
+/// no forwarding uses.
+private func replace(guaranteedValue: Value, withOwnedValue ownedValue: Value, _ context: SimplifyContext) -> Value {
+  var result = ownedValue
+  var numForwardingUses = 0
+  for use in guaranteedValue.uses {
+
+    switch use.instruction {
+    case let tei as TupleExtractInst:
+      numForwardingUses += 1
+      let dti = Builder(before: tei, context).createDestructureTuple(tuple: ownedValue)
+      result = replace(guaranteedValue: tei, withOwnedValue: dti.results[tei.fieldIndex], context)
+      context.erase(instruction: tei)
+    case let sei as StructExtractInst:
+      numForwardingUses += 1
+      let dsi = Builder(before: sei, context).createDestructureStruct(struct: ownedValue)
+      result = replace(guaranteedValue: sei, withOwnedValue: dsi.results[sei.fieldIndex], context)
+      context.erase(instruction: sei)
+    case let fwdInst as (SingleValueInstruction & ForwardingInstruction) where
+         fwdInst.isSingleForwardedOperand(use):
+      // Other forwarding instructions beside tuple_extract and struct_extract
+      numForwardingUses += 1
+      use.set(to: ownedValue, context)
+      fwdInst.setForwardingOwnership(to: .owned, context)
+      result = replace(guaranteedValue: fwdInst, withOwnedValue: fwdInst, context)
+    case is EndBorrowInst:
+      break
+    default:
+      precondition(use.canAccept(ownership: .owned))
+      use.set(to: ownedValue, context)
+    }
+  }
+  precondition(numForwardingUses <= 1, "guaranteed value must not have multiple forwarding uses")
+  return result
+}
+
+private extension ForwardingInstruction {
+  var canConvertToOwned: Bool {
+    switch self {
+    case let si as StructExtractInst:
+      if si.struct.type.isMoveOnly {
+        // We cannot easily convert a struct_extract to a destructure_struct of a move-only type, because
+        // the deinit would get lost.
+        return false
+      }
+      let structFields = si.struct.type.getNominalFields(in: parentFunction)
+      return structFields.hasSingleNonTrivialElement(at: si.fieldIndex, in: parentFunction)
+    case let ti as TupleExtractInst:
+      return ti.tuple.type.tupleElements.hasSingleNonTrivialElement(at: ti.fieldIndex, in: parentFunction)
+    default:
+      return canForwardOwnedValues
+    }
+  }
+}
+
+private extension Collection where Element == Type {
+  func hasSingleNonTrivialElement(at nonTrivialElementIndex: Int, in function: Function) -> Bool {
+    for (elementIdx, elementTy) in self.enumerated() {
+      if elementTy.isTrivial(in: function) != (elementIdx != nonTrivialElementIndex) {
+        return false
+      }
+    }
+    return true
+  }
+}

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginCOWMutation.swift

@@ -59,7 +59,7 @@ private extension BeginCOWMutationInst {
     if !isEmptyCOWSingleton(instance) {
       return
     }
-    if uniquenessResult.nonDebugUses.isEmpty {
+    if uniquenessResult.uses.ignoreDebugUses.isEmpty {
       /// Don't create an integer_literal which would be dead. This would result
       /// in an infinite loop in SILCombine.
       return
@@ -70,15 +70,15 @@ private extension BeginCOWMutationInst {
   }
 
   func optimizeEmptyBeginEndPair(_ context: SimplifyContext) -> Bool {
-    if !uniquenessResult.nonDebugUses.isEmpty {
+    if !uniquenessResult.uses.ignoreDebugUses.isEmpty {
       return false
     }
     let buffer = instanceResult
-    if buffer.nonDebugUses.contains(where: { !($0.instruction is EndCOWMutationInst) }) {
+    if buffer.uses.ignoreDebugUses.contains(where: { !($0.instruction is EndCOWMutationInst) }) {
       return false
     }
 
-    for use in buffer.nonDebugUses {
+    for use in buffer.uses.ignoreDebugUses {
       let endCOW = use.instruction as! EndCOWMutationInst
       endCOW.uses.replaceAll(with: instance, context)
       context.erase(instruction: endCOW)
@@ -88,13 +88,13 @@ private extension BeginCOWMutationInst {
   }
 
   func optimizeEmptyEndBeginPair(_ context: SimplifyContext) -> Bool {
-    if !uniquenessResult.nonDebugUses.isEmpty {
+    if !uniquenessResult.uses.ignoreDebugUses.isEmpty {
       return false
     }
     guard let endCOW = instance as? EndCOWMutationInst else {
       return false
     }
-    if endCOW.nonDebugUses.contains(where: { $0.instruction != self }) {
+    if endCOW.uses.ignoreDebugUses.contains(where: { $0.instruction != self }) {
       return false
     }
 

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyDestructure.swift

@@ -87,7 +87,7 @@ extension DestructureStructInst : OnoneSimplifyable {
 private func tryReplaceConstructDestructPair(construct: SingleValueInstruction,
                                              destruct: MultipleValueInstruction,
                                              _ context: SimplifyContext) {
-  let singleUse = context.preserveDebugInfo ? construct.uses.singleUse : construct.uses.singleNonDebugUse
+  let singleUse = context.preserveDebugInfo ? construct.uses.singleUse : construct.uses.ignoreDebugUses.singleUse
   let canEraseFirst = singleUse?.instruction == destruct
 
   if !canEraseFirst && construct.parentFunction.hasOwnership && construct.ownership == .owned {

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyLoad.swift

@@ -197,7 +197,7 @@ extension LoadInst : OnoneSimplifyable, SILCombineSimplifyable {
     if context.preserveDebugInfo {
       return !uses.contains { !($0.instruction is DestroyValueInst) }
     } else {
-      return !nonDebugUses.contains { !($0.instruction is DestroyValueInst) }
+      return !uses.ignoreDebugUses.contains { !($0.instruction is DestroyValueInst) }
     }
   }
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyRefCasts.swift

@@ -65,7 +65,7 @@ private extension UnaryInstruction {
   }
 
   func tryReplaceSource(withOperandOf inst: SingleValueInstruction, _ context: SimplifyContext) -> Bool {
-    let singleUse = context.preserveDebugInfo ? inst.uses.singleUse : inst.uses.singleNonDebugUse
+    let singleUse = context.preserveDebugInfo ? inst.uses.singleUse : inst.uses.ignoreDebugUses.singleUse
     let canEraseInst = singleUse?.instruction == self
     let replacement = inst.operands[0].value
 

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -265,7 +265,14 @@ private extension Value {
     var singleUseValue: any Value = self
     var path = SmallProjectionPath()
     while true {
-      guard let use = singleUseValue.uses.singleRelevantUse else {
+      // The initializer value of a global can contain access instructions if it references another
+      // global variable by address, e.g.
+      //   var p = Point(x: 10, y: 20)
+      //   let o = UnsafePointer(&p)
+      // Therefore ignore the `end_access` use of a `begin_access`.
+      let relevantUses = singleUseValue.uses.ignoreDebugUses.ignoreUsers(ofType: EndAccessInst.self)
+
+      guard let use = relevantUses.singleUse else {
         return nil
       }
 
@@ -383,27 +390,3 @@ fileprivate struct FunctionWorklist {
     }
   }
 }
-
-private extension UseList {
-  var singleRelevantUse: Operand? {
-    var singleUse: Operand?
-    for use in self {
-      switch use.instruction {
-      case is DebugValueInst,
-           // The initializer value of a global can contain access instructions if it references another
-           // global variable by address, e.g.
-           //   var p = Point(x: 10, y: 20)
-           //   let o = UnsafePointer(&p)
-           // Therefore ignore the `end_access` use of a `begin_access`.
-           is EndAccessInst:
-        continue
-      default:
-        if singleUse != nil {
-          return nil
-        }
-        singleUse = use
-      }
-    }
-    return singleUse
-  }
-}

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -365,9 +365,7 @@ private struct StackProtectionOptimization {
     let builder = Builder(after: beginAccess, location: beginAccess.location.autoGenerated, context)
     let temporary = builder.createAllocStack(beginAccess.type)
 
-    for use in beginAccess.uses where !(use.instruction is EndAccessInst) {
-      use.instruction.setOperand(at: use.index, to: temporary, context)
-    }
+    beginAccess.uses.ignoreUsers(ofType: EndAccessInst.self).replaceAll(with: temporary, context)
 
     for endAccess in beginAccess.endInstructions {
       let endBuilder = Builder(before: endAccess, location: endAccess.location.autoGenerated, context)