Optimizer: re-compute the re-borrow flags of phi arguments in updateAllBorrowArguments and updateBorrowArguments

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowArgumentsUpdater.swift

@@ -13,6 +13,18 @@
 import SIL
 import OptimizerBridging
 
+/// Updates the reborrow flags and the borrowed-from instructions for all guaranteed phis in `function`.
+func updateBorrowArguments(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 updateBorrowArguments(for 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() }
@@ -103,7 +156,7 @@ func registerBorrowArgumentsUpdater() {
     { (bridgedCtxt: BridgedPassContext, bridgedFunction: BridgedFunction) in
       let context = FunctionPassContext(_bridged: bridgedCtxt)
       let function = bridgedFunction.function;
-      updateBorrowedFrom(in: function, context)
+      updateBorrowArguments(in: function, context)
     },
     { (bridgedCtxt: BridgedPassContext, bridgedPhiArray: BridgedArrayRef) in
       let context = FunctionPassContext(_bridged: bridgedCtxt)
@@ -117,7 +170,7 @@ func registerBorrowArgumentsUpdater() {
           }
         }
       }
-      updateBorrowedFrom(for: guaranteedPhis, context)
+      updateBorrowArguments(for: guaranteedPhis, context)
     }
   )
 }

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

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