Fix the computation of the re-borrow flags for guaranteed phi arguments

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -628,6 +628,13 @@ extension BasicBlock {
   }
 }
 
+extension Argument {
+  func set(reborrow: Bool, _ context: some MutatingContext) {
+    context.notifyInstructionsChanged()
+    bridged.setReborrow(reborrow)
+  }
+}
+
 extension AllocRefInstBase {
   func setIsStackAllocatable(_ context: some MutatingContext) {
     context.notifyInstructionsChanged()

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -136,5 +136,5 @@ private func registerSwiftAnalyses() {
 
 private func registerUtilities() {
   registerVerifier()
-  registerBorrowedFromUpdater()
+  registerGuaranteedPhiUpdater()
 }

SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowUtils.swift

@@ -592,7 +592,7 @@ func gatherEnclosingValuesFromPredecessors(
     let incomingOperand = phi.incomingOperand(inPredecessor: predecessor)
 
     for predEV in incomingOperand.value.getEnclosingValues(context) {
-      let ev = predecessor.mapToPhiInSuccessor(incomingEnclosingValue: predEV)
+      let ev = predecessor.getEnclosingValueInSuccessor(ofIncoming: predEV)
       if alreadyAdded.insert(ev) {
         enclosingValues.push(ev)
       }
@@ -601,13 +601,31 @@ func gatherEnclosingValuesFromPredecessors(
 }
 
 extension BasicBlock {
-  func mapToPhiInSuccessor(incomingEnclosingValue: Value) -> Value {
+  // Returns either the `incomingEnclosingValue` or an adjacent phi in the successor block.
+  func getEnclosingValueInSuccessor(ofIncoming incomingEnclosingValue: Value) -> Value {
     let branch = terminator as! BranchInst
-    if let incomingEV = branch.operands.first(where: { $0.value.lookThroughBorrowedFrom == incomingEnclosingValue }) {
+    if let incomingEV = branch.operands.first(where: { branchOp in
+      // Only if the lifetime of `branchOp` ends at the branch (either because it's a reborrow or an owned value),
+      // the corresponding phi argument can be the adjacent phi for the incoming value.
+      //   bb1:
+      //     %incomingEnclosingValue = some_owned_value
+      //     %2 = begin_borrow %incomingEnclosingValue  // %incomingEnclosingValue = the enclosing value of %2 in bb1
+      //     br bb2(%incomingEnclosingValue, %2)        // lifetime of incomingEnclosingValue ends here
+      //   bb2(%4 : @owned, %5 : @guaranteed):          // -> %4 = the enclosing value of %5 in bb2
+      //
+      branchOp.endsLifetime &&
+      branchOp.value.lookThroughBorrowedFrom == incomingEnclosingValue
+    }) {
       return branch.getArgument(for: incomingEV)
     }
-    // No candidates phi are outer-adjacent phis. The incoming
-    // `predDef` must dominate the current guaranteed phi.
+    // No candidates phi are outer-adjacent phis. The incomingEnclosingValue must dominate the successor block.
+    //   bb1: // dominates bb3
+    //     %incomingEnclosingValue = some_owned_value
+    //   bb2:
+    //     %2 = begin_borrow %incomingEnclosingValue
+    //     br bb3(%2)
+    //   bb3(%5 : @guaranteed):          // -> %incomingEnclosingValue = the enclosing value of %5 in bb3
+    //
     return incomingEnclosingValue
   }
 }

SwiftCompilerSources/Sources/Optimizer/Utilities/CMakeLists.txt

@@ -9,7 +9,7 @@
 swift_compiler_sources(Optimizer
   AccessUtilsTest.swift
   AddressUtils.swift
-  BorrowedFromUpdater.swift
+  GuaranteedPhiUpdater.swift
   BorrowUtils.swift
   SpecializationCloner.swift
   DiagnosticEngine.swift

SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowedFromUpdater.swift → SwiftCompilerSources/Sources/Optimizer/Utilities/GuaranteedPhiUpdater.swift

@@ -1,4 +1,4 @@
-//===--- BorrowedFromUpdater.swift ----------------------------------------===//
+//===--- GuaranteedPhiUpdater.swift ---------------------------------------===//
 //
 // This source file is part of the Swift.org open source project
 //
@@ -13,6 +13,18 @@
 import SIL
 import OptimizerBridging
 
+/// Updates the reborrow flags and the borrowed-from instructions for all guaranteed phis in `function`.
+func updateGuaranteedPhis(in function: Function, _ context: some MutatingContext) {
+  updateReborrowFlags(in: function, context)
+  updateBorrowedFrom(in: function, context)
+}
+
+/// Updates the reborrow flags and the borrowed-from instructions for all `phis`.
+func updateGuaranteedPhis(phis: some Sequence<Phi>, _ context: some MutatingContext) {
+  updateReborrowFlags(for: phis, context)
+  updateBorrowedFrom(for: phis, context)
+}
+
 /// Update all borrowed-from instructions in the `function`
 func updateBorrowedFrom(in function: Function, _ context: some MutatingContext) {
   if !function.hasOwnership {
@@ -54,6 +66,47 @@ func updateBorrowedFrom(for phis: some Sequence<Phi>, _ context: some MutatingCo
   } while changed
 }
 
+/// Updates the reborrow flags for all guaranteed phis in `function`.
+func updateReborrowFlags(in function: Function, _ context: some MutatingContext) {
+  if !function.hasOwnership {
+    return
+  }
+  var guaranteedPhis = Stack<Phi>(context)
+  defer { guaranteedPhis.deinitialize() }
+
+  for block in function.blocks.reversed() {
+    for arg in block.arguments {
+      if let phi = Phi(arg), phi.value.ownership == .guaranteed {
+        guaranteedPhis.append(phi)
+      }
+    }
+  }
+  updateReborrowFlags(for: guaranteedPhis, context)
+}
+
+/// Updates the reborrow flags for all `phis`.
+func updateReborrowFlags(for phis: some Sequence<Phi>, _ context: some MutatingContext) {
+  // TODO: clear reborrow flags before re-computing when we have complete OSSA lifetimes.
+  // It would be cleaner to first clear all flags. But this is not possible because some end_borrow instructions
+  // might be missing in dead-end blocks. This will be fixed with complete OSSA lifetimes.
+
+  if let phi = phis.first(where: { phi in true }), !phi.value.parentFunction.hasOwnership {
+    return
+  }
+
+  var changed: Bool
+  repeat {
+    changed = false
+
+    for phi in phis where phi.value.ownership == .guaranteed {
+      if !phi.value.isReborrow && phi.hasBorrowEndingUse {
+        phi.value.set(reborrow: true, context)
+        changed = true
+      }
+    }
+  } while changed
+}
+
 private func updateBorrowedFrom(for phi: Phi, _ context: some MutatingContext) -> Bool {
   var computedEVs = Stack<Value>(context)
   defer { computedEVs.deinitialize() }
@@ -98,12 +151,12 @@ private func addEnclosingValues(
   return true
 }
 
-func registerBorrowedFromUpdater() {
-  BridgedUtilities.registerBorrowedFromUpdater(
+func registerGuaranteedPhiUpdater() {
+  BridgedUtilities.registerGuaranteedPhiUpdater(
     { (bridgedCtxt: BridgedPassContext, bridgedFunction: BridgedFunction) in
       let context = FunctionPassContext(_bridged: bridgedCtxt)
       let function = bridgedFunction.function;
-      updateBorrowedFrom(in: function, context)
+      updateGuaranteedPhis(in: function, context)
     },
     { (bridgedCtxt: BridgedPassContext, bridgedPhiArray: BridgedArrayRef) in
       let context = FunctionPassContext(_bridged: bridgedCtxt)
@@ -117,7 +170,7 @@ func registerBorrowedFromUpdater() {
           }
         }
       }
-      updateBorrowedFrom(for: guaranteedPhis, context)
+      updateGuaranteedPhis(phis: guaranteedPhis, context)
     }
   )
 }

SwiftCompilerSources/Sources/Optimizer/Utilities/Verifier.swift

@@ -57,18 +57,32 @@ private extension Instruction {
 private extension Argument {
   func verify(_ context: FunctionPassContext) {
     if let phi = Phi(self), phi.value.ownership == .guaranteed {
-      var forwardingBorrowedFromFound = false
-      for use in phi.value.uses {
-        require(use.instruction is BorrowedFromInst,
-                     "guaranteed phi: \(self)\n has non borrowed-from use: \(use)")
-        if use.index == 0 {
-          require(!forwardingBorrowedFromFound, "phi \(self) has multiple forwarding borrowed-from uses")
-          forwardingBorrowedFromFound = true
-        }
+
+      phi.verifyBorrowedFromUse()
+
+      require(phi.isReborrow == phi.hasBorrowEndingUse ||
+              // In a dead-end block an end_borrow might have been deleted.
+              // TODO: this check is not needed anymore once we have complete OSSA lifetimes.
+              (isReborrow && context.deadEndBlocks.isDeadEnd(parentBlock)),
+              "\(self) has stale reborrow flag");
+    }
+  }
+
+}
+
+private extension Phi {
+  func verifyBorrowedFromUse() {
+    var forwardingBorrowedFromFound = false
+    for use in value.uses {
+      require(use.instruction is BorrowedFromInst,
+              "guaranteed phi: \(self)\n has non borrowed-from use: \(use)")
+      if use.index == 0 {
+        require(!forwardingBorrowedFromFound, "phi \(self) has multiple forwarding borrowed-from uses")
+        forwardingBorrowedFromFound = true
       }
-      require (forwardingBorrowedFromFound,
-                   "missing forwarding borrowed-from user of guaranteed phi \(phi)")
     }
+    require(forwardingBorrowedFromFound,
+            "missing forwarding borrowed-from user of guaranteed phi \(self)")
   }
 }
 

SwiftCompilerSources/Sources/SIL/Argument.swift

@@ -24,8 +24,8 @@ public class Argument : Value, Hashable {
     return bridged.getParent().block
   }
 
-  var bridged: BridgedArgument { BridgedArgument(obj: SwiftObject(self)) }
-  
+  public var bridged: BridgedArgument { BridgedArgument(obj: SwiftObject(self)) }
+
   public var index: Int {
     return parentBlock.arguments.firstIndex(of: self)!
   }
@@ -165,6 +165,18 @@ public struct Phi {
     return nil
   }
 
+  // Returns true if the phi has an end_borrow or a re-borrowing branch as user.
+  // This should be consistent with the `isReborrow` flag.
+  public var hasBorrowEndingUse: Bool {
+    let phiValue: Value = borrowedFrom ?? value
+    return phiValue.uses.contains {
+      switch $0.ownership {
+        case .endBorrow, .reborrow: return true
+        default:                    return false
+      }
+    }
+  }
+
   public static func ==(lhs: Phi, rhs: Phi) -> Bool {
     lhs.value === rhs.value
   }

include/swift/SIL/OwnershipUtils.h

@@ -57,14 +57,6 @@ bool canOpcodeForwardInnerGuaranteedValues(SILValue value);
 /// the operation may be trivially rewritten with Guaranteed ownership.
 bool canOpcodeForwardInnerGuaranteedValues(Operand *use);
 
-bool computeIsScoped(SILArgument *arg);
-
-bool computeIsReborrow(SILArgument *arg);
-
-// This is the use-def equivalent of use->getOperandOwnership() ==
-// OperandOwnership::GuaranteedForwarding.
-bool computeIsGuaranteedForwarding(SILValue value);
-
 /// Is the opcode that produces \p value capable of forwarding owned values?
 ///
 /// This may be true even if the current instance of the instruction is not a
@@ -626,6 +618,12 @@ struct BorrowedValue {
   /// BorrowScopeIntroducingValue::isLocalScope().
   bool visitLocalScopeEndingUses(function_ref<bool(Operand *)> visitor) const;
 
+  /// Returns false if the value has no scope-ending uses because all control flow
+  /// paths end in dead-end blocks.
+  bool hasLocalScopeEndingUses() const {
+    return !visitLocalScopeEndingUses([](Operand *) { return false; });
+  }
+
   bool isLocalScope() const { return kind.isLocalScope(); }
 
   /// Add this scope's live blocks into the PrunedLiveness result. This
@@ -1404,6 +1402,10 @@ bool isNestedLexicalBeginBorrow(BeginBorrowInst *bbi);
 /// then the move_value is redundant.
 bool isRedundantMoveValue(MoveValueInst *mvi);
 
+/// Sets the reborrow flags for all transitively incoming phi-arguments of
+/// `forEndBorrowValue`, which is the operand value of an `end_borrow`.
+void updateReborrowFlags(SILValue forEndBorrowValue);
+
 } // namespace swift
 
 #endif

include/swift/SIL/SILBridging.h

@@ -864,6 +864,7 @@ struct BridgedArgument {
   BRIDGED_INLINE swift::SILArgument * _Nonnull getArgument() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedBasicBlock getParent() const;
   BRIDGED_INLINE bool isReborrow() const;
+  BRIDGED_INLINE void setReborrow(bool reborrow) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE OptionalBridgedDeclObj getVarDecl() const;
   BRIDGED_INLINE void copyFlags(BridgedArgument fromArgument) const;
 };

include/swift/SIL/SILBridgingImpl.h

@@ -672,6 +672,9 @@ BridgedBasicBlock BridgedArgument::getParent() const {
 }
 
 bool BridgedArgument::isReborrow() const { return getArgument()->isReborrow(); }
+void BridgedArgument::setReborrow(bool reborrow) const {
+  getArgument()->setReborrow(reborrow);
+}
 
 OptionalBridgedDeclObj BridgedArgument::getVarDecl() const {
   return {llvm::dyn_cast_or_null<swift::VarDecl>(getArgument()->getDecl())};

include/swift/SIL/SILBuilder.h

@@ -941,14 +941,7 @@ class SILBuilder {
     return lowering.emitStore(*this, Loc, Src, DestAddr, Qualifier);
   }
 
-  EndBorrowInst *createEndBorrow(SILLocation loc, SILValue borrowedValue) {
-    ASSERT(!SILArgument::isTerminatorResult(borrowedValue) &&
-               "terminator results do not have end_borrow");
-    ASSERT(!isa<SILFunctionArgument>(borrowedValue) &&
-           "Function arguments should never have an end_borrow");
-    return insert(new (getModule())
-                      EndBorrowInst(getSILDebugLocation(loc), borrowedValue));
-  }
+  EndBorrowInst *createEndBorrow(SILLocation loc, SILValue borrowedValue);
 
   BeginAccessInst *createBeginAccess(SILLocation loc, SILValue address,
                                      SILAccessKind accessKind,

include/swift/SIL/SILValue.h

@@ -1073,9 +1073,11 @@ class Operand {
     removeFromCurrent();
     TheValue = newValue;
     insertIntoCurrent();
+    updateReborrowFlags();
     verify();
   }
 
+  void updateReborrowFlags();
   void verify() const;
 
   /// Swap the given operand with the current one.

include/swift/SILOptimizer/OptimizerBridging.h

@@ -130,12 +130,12 @@ struct BridgedPostDomTree {
 
 struct BridgedUtilities {
   typedef void (* _Nonnull VerifyFunctionFn)(BridgedPassContext, BridgedFunction);
-  typedef void (* _Nonnull UpdateBorrowedFromFn)(BridgedPassContext, BridgedFunction);
-  typedef void (* _Nonnull UpdateBorrowedFromPhisFn)(BridgedPassContext, BridgedArrayRef);
+  typedef void (* _Nonnull UpdateFunctionFn)(BridgedPassContext, BridgedFunction);
+  typedef void (* _Nonnull UpdatePhisFn)(BridgedPassContext, BridgedArrayRef);
 
   static void registerVerifier(VerifyFunctionFn verifyFunctionFn);
-  static void registerBorrowedFromUpdater(UpdateBorrowedFromFn updateBorrowedFromFn,
-                                          UpdateBorrowedFromPhisFn updateBorrowedFromPhisFn);
+  static void registerGuaranteedPhiUpdater(UpdateFunctionFn updateBorrowedFromFn,
+                                           UpdatePhisFn updateBorrowedFromPhisFn);
 };
 
 struct BridgedBasicBlockSet {

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -357,9 +357,12 @@ bool extendStoreBorrow(StoreBorrowInst *sbi,
                        DeadEndBlocks *deadEndBlocks,
                        InstModCallbacks callbacks = InstModCallbacks());
 
-void updateBorrowedFrom(SILPassManager *pm, SILFunction *f);
+/// Updates the reborrow flags and the borrowed-from instructions for all
+/// guaranteed phis in function `f`.
+void updateAllGuaranteedPhis(SILPassManager *pm, SILFunction *f);
 
-void updateBorrowedFromPhis(SILPassManager *pm, ArrayRef<SILPhiArgument *> phis);
+/// Updates the reborrow flags and the borrowed-from instructions for all `phis`.
+void updateGuaranteedPhis(SILPassManager *pm, ArrayRef<SILPhiArgument *> phis);
 
 } // namespace swift
 

lib/SIL/IR/SILBuilder.cpp

@@ -13,6 +13,7 @@
 #include "swift/SIL/SILBuilder.h"
 #include "swift/AST/Expr.h"
 #include "swift/Basic/Assertions.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILGlobalVariable.h"
 
@@ -681,6 +682,17 @@ void SILBuilder::emitScopedBorrowOperation(SILLocation loc, SILValue original,
     createEndBorrow(loc, value);
 }
 
+EndBorrowInst *SILBuilder::createEndBorrow(SILLocation loc, SILValue borrowedValue) {
+  ASSERT(!SILArgument::isTerminatorResult(borrowedValue) &&
+             "terminator results do not have end_borrow");
+  ASSERT(!isa<SILFunctionArgument>(borrowedValue) &&
+         "Function arguments should never have an end_borrow");
+  updateReborrowFlags(borrowedValue);
+  return insert(new (getModule())
+                    EndBorrowInst(getSILDebugLocation(loc), borrowedValue));
+}
+
+
 SILPhiArgument *SILBuilder::createSwitchOptional(
                                 SILLocation loc, SILValue operand,
                                 SILBasicBlock *someBB, SILBasicBlock *noneBB,