- getBorrowIntroducers and getEnclosingValues with iterators

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/LifetimeDependenceDiagnostics.swift

@@ -236,24 +236,16 @@ private struct LifetimeVariable {
       self = Self(accessBase: value.accessBase, context)
       return
     }
-    if let firstIntroducer = getFirstBorrowIntroducer(of: value, context) {
+    // FUTURE: consider diagnosing multiple variable introducers. It's
+    // unclear how more than one can happen.
+    if let firstIntroducer = value.getBorrowIntroducers(context).first {
       self = Self(introducer: firstIntroducer)
       return
     }
     self.varDecl = nil
     self.sourceLoc = nil
   }
 
-  // FUTURE: consider diagnosing multiple variable introducers. It's
-  // unclear how more than one can happen.
-  private func getFirstBorrowIntroducer(of value: Value,
-                                        _ context: some Context)
-    -> Value? {
-    var introducers = Stack<Value>(context)
-    gatherBorrowIntroducers(for: value, in: &introducers, context)
-    return introducers.pop()
-  }
-
   private func getFirstLifetimeIntroducer(of value: Value,
                                           _ context: some Context)
     -> Value? {

SwiftCompilerSources/Sources/Optimizer/Utilities/BorrowUtils.swift

@@ -388,138 +388,172 @@ extension Value {
   }
 }
 
-/// Find the borrow introducers for `value`. This gives you a set of
-/// OSSA lifetimes that directly include `value`. If `value` is owned,
-/// or introduces a borrow scope, then `value` is the single
-/// introducer for itself.
-///
-/// If `value` is an address or any trivial type, then it has no introducers.
-///
-/// Example:                                       // introducers:
-///                                                // ~~~~~~~~~~~~
-///   bb0(%0 : @owned $Class,                      // %0
-///       %1 : @guaranteed $Class):                // %1
-///     %borrow0 = begin_borrow %0                 // %borrow0
-///     %pair = struct $Pair(%borrow0, %1)         // %borrow0, %1
-///     %first = struct_extract %pair              // %borrow0, %1
-///     %field = ref_element_addr %first           // (none)
-///     %load = load_borrow %field : $*C           // %load
-func gatherBorrowIntroducers(for initialValue: Value,
-                             in borrowIntroducers: inout Stack<Value>,
-                             _ context: some Context)
-{
-  var worklist = ValueWorklist(context)
-  defer { worklist.deinitialize() }
-
-  worklist.pushIfNotVisited(initialValue)
-  while let value = worklist.pop() {
-    switch value.ownership {
-    case .none, .unowned:
-      break
-
-    case .owned:
-      borrowIntroducers.append(value)
-
-    case .guaranteed:
-      if BeginBorrowValue(value) != nil {
-        borrowIntroducers.append(value)
-      } else if let bfi = value as? BorrowedFromInst,
-                let phi = Phi(bfi.borrowedValue)
-      {
-        if phi.isReborrow {
-          worklist.pushIfNotVisited(phi.value)
+struct BorrowIntroducers<Ctxt: Context> : CollectionLikeSequence {
+  let initialValue: Value
+  let context: Ctxt
+
+  func makeIterator() -> EnclosingValueIterator {
+    EnclosingValueIterator(forBorrowIntroducers: initialValue, context)
+  }
+}
+
+struct EnclosingValues<Ctxt: Context> : CollectionLikeSequence {
+  let initialValue: Value
+  let context: Ctxt
+
+  func makeIterator() -> EnclosingValueIterator {
+    EnclosingValueIterator(forEnclosingValues: initialValue, context)
+  }
+}
+
+// This iterator must be a class because we need a deinit.
+// It shouldn't be a performance problem because the optimizer should always be able to stack promote the iterator.
+// TODO: Make it a struct once this is possible with non-copyable types.
+final class EnclosingValueIterator : IteratorProtocol {
+  var worklist: ValueWorklist
+
+  init(forBorrowIntroducers value: Value, _ context: some Context) {
+    self.worklist = ValueWorklist(context)
+    self.worklist.pushIfNotVisited(value)
+  }
+
+  init(forEnclosingValues value: Value, _ context: some Context) {
+    self.worklist = ValueWorklist(context)
+    if value is Undef || value.ownership != .guaranteed {
+      return
+    }
+    if let beginBorrow = BeginBorrowValue(value.lookThroughBorrowedFrom) {
+      switch beginBorrow {
+      case let .beginBorrow(bbi):
+        // Gather the outer enclosing borrow scope.
+        worklist.pushIfNotVisited(bbi.borrowedValue)
+      case .loadBorrow, .beginApply, .functionArgument:
+        // There is no enclosing value on this path.
+        break
+      case .reborrow(let phi):
+        worklist.pushIfNotVisited(contentsOf: phi.borrowedFrom!.enclosingValues)
+      }
+    } else {
+      // Handle forwarded guaranteed values.
+      worklist.pushIfNotVisited(value)
+    }
+  }
+
+  deinit {
+    worklist.deinitialize()
+  }
+
+  func next() -> Value? {
+    while let value = worklist.pop() {
+      switch value.ownership {
+      case .none, .unowned:
+        break
+
+      case .owned:
+        return value
+
+      case .guaranteed:
+        if BeginBorrowValue(value) != nil {
+          return value
+        } else if let bfi = value as? BorrowedFromInst {
+          if bfi.borrowedPhi.isReborrow {
+            worklist.pushIfNotVisited(bfi.borrowedValue)
+          } else {
+            worklist.pushIfNotVisited(contentsOf: bfi.enclosingValues)
+          }
+        } else if let forwardingInst = value.forwardingInstruction {
+          // Recurse through guaranteed forwarding non-phi instructions.
+          let ops = forwardingInst.forwardedOperands
+          worklist.pushIfNotVisited(contentsOf: ops.lazy.map { $0.value })
         } else {
-          worklist.pushIfNotVisited(contentsOf: bfi.enclosingValues)
+          fatalError("cannot get borrow introducers for unknown guaranteed value")
         }
-      } else if let forwardingInst = value.forwardingInstruction {
-        // Recurse through guaranteed forwarding non-phi instructions.
-        let ops = forwardingInst.forwardedOperands
-        worklist.pushIfNotVisited(contentsOf: ops.lazy.map { $0.value })
-      } else {
-        fatalError("cannot get borrow introducers for unknown guaranteed value")
       }
     }
+    return nil
   }
 }
 
-/// Find each "enclosing value" whose OSSA lifetime immediately
-/// encloses a guaranteed value. The guaranteed `value` being enclosed
-/// effectively keeps these enclosing values alive. This lets you walk
-/// up the levels of nested OSSA lifetimes to determine all the
-/// lifetimes that are kept alive by a given SILValue. In particular,
-/// it discovers "outer-adjacent phis": phis that are kept alive by
-/// uses of another phi in the same block.
-///
-/// If `value` is a forwarded guaranteed value, then this finds the
-/// introducers of the current borrow scope, which is never an empty
-/// set.
-///
-/// If `value` introduces a borrow scope, then this finds the
-/// introducers of the outer enclosing borrow scope that contains this
-/// inner scope.
-///
-/// If `value` is a `begin_borrow`, then this returns its operand.
-///
-/// If `value` is an owned value, a function argument, or a
-/// load_borrow, then this is an empty set.
-///
-/// If `value` is a reborrow, then this either returns a dominating
-/// enclosing value or an outer adjacent phi.
-///
-/// Example:                                       // enclosing value:
-///                                                // ~~~~~~~~~~~~
-///   bb0(%0 : @owned $Class,                      // (none)
-///       %1 : @guaranteed $Class):                // (none)
-///     %borrow0 = begin_borrow %0                 // %0
-///     %pair = struct $Pair(%borrow0, %1)         // %borrow0, %1
-///     %first = struct_extract %pair              // %borrow0, %1
-///     %field = ref_element_addr %first           // (none)
-///     %load = load_borrow %field : $*C           // %load
-///
-/// Example:                                       // enclosing value:
-///                                                // ~~~~~~~~~~~~
-///     %outerBorrow = begin_borrow %0             // %0
-///     %innerBorrow = begin_borrow %outerBorrow   // %outerBorrow
-///     br bb1(%outerBorrow, %innerBorrow)
-///   bb1(%outerReborrow : @reborrow,              // %0
-///       %innerReborrow : @reborrow)              // %outerReborrow
-///
-func gatherEnclosingValues(for value: Value,
-                           in enclosingValues: inout Stack<Value>,
-                           _ context: some Context)
-{
-  if value is Undef || value.ownership != .guaranteed {
-    return
+
+extension Value {
+  /// Get the borrow introducers for `value`. This gives you a set of
+  /// OSSA lifetimes that directly include `value`. If `value` is owned,
+  /// or introduces a borrow scope, then `value` is the single
+  /// introducer for itself.
+  ///
+  /// If `value` is an address or any trivial type, then it has no introducers.
+  ///
+  /// Example:                                       // introducers:
+  ///                                                // ~~~~~~~~~~~~
+  ///   bb0(%0 : @owned $Class,                      // %0
+  ///       %1 : @guaranteed $Class):                // %1
+  ///     %borrow0 = begin_borrow %0                 // %borrow0
+  ///     %pair = struct $Pair(%borrow0, %1)         // %borrow0, %1
+  ///     %first = struct_extract %pair              // %borrow0, %1
+  ///     %field = ref_element_addr %first           // (none)
+  ///     %load = load_borrow %field : $*C           // %load
+  func getBorrowIntroducers<Ctxt: Context>(_ context: Ctxt) -> BorrowIntroducers<Ctxt> {
+    BorrowIntroducers(initialValue: self, context: context)
   }
-  if let beginBorrow = BeginBorrowValue(value.lookThroughBorrowedFrom) {
-    switch beginBorrow {
-    case let .beginBorrow(bbi):
-      // Gather the outer enclosing borrow scope.
-      gatherBorrowIntroducers(for: bbi.borrowedValue, in: &enclosingValues, context)
-    case .loadBorrow, .beginApply, .functionArgument:
-      // There is no enclosing value on this path.
-      break
-    case .reborrow(let phi):
-      enclosingValues.append(contentsOf: phi.borrowedFrom!.enclosingValues)
-    }
-  } else {
-    // Handle forwarded guaranteed values.
-    gatherBorrowIntroducers(for: value, in: &enclosingValues, context)
+
+  /// Get "enclosing values" whose OSSA lifetime immediately
+  /// encloses a guaranteed value. The guaranteed `value` being enclosed
+  /// effectively keeps these enclosing values alive. This lets you walk
+  /// up the levels of nested OSSA lifetimes to determine all the
+  /// lifetimes that are kept alive by a given SILValue. In particular,
+  /// it discovers "outer-adjacent phis": phis that are kept alive by
+  /// uses of another phi in the same block.
+  ///
+  /// If `value` is a forwarded guaranteed value, then this finds the
+  /// introducers of the current borrow scope, which is never an empty
+  /// set.
+  ///
+  /// If `value` introduces a borrow scope, then this finds the
+  /// introducers of the outer enclosing borrow scope that contains this
+  /// inner scope.
+  ///
+  /// If `value` is a `begin_borrow`, then this returns its operand.
+  ///
+  /// If `value` is an owned value, a function argument, or a
+  /// load_borrow, then this is an empty set.
+  ///
+  /// If `value` is a reborrow, then this either returns a dominating
+  /// enclosing value or an outer adjacent phi.
+  ///
+  /// Example:                                       // enclosing value:
+  ///                                                // ~~~~~~~~~~~~
+  ///   bb0(%0 : @owned $Class,                      // (none)
+  ///       %1 : @guaranteed $Class):                // (none)
+  ///     %borrow0 = begin_borrow %0                 // %0
+  ///     %pair = struct $Pair(%borrow0, %1)         // %borrow0, %1
+  ///     %first = struct_extract %pair              // %borrow0, %1
+  ///     %field = ref_element_addr %first           // (none)
+  ///     %load = load_borrow %field : $*C           // %load
+  ///
+  /// Example:                                       // enclosing value:
+  ///                                                // ~~~~~~~~~~~~
+  ///     %outerBorrow = begin_borrow %0             // %0
+  ///     %innerBorrow = begin_borrow %outerBorrow   // %outerBorrow
+  ///     br bb1(%outerBorrow, %innerBorrow)
+  ///   bb1(%outerReborrow : @reborrow,              // %0
+  ///       %innerReborrow : @reborrow)              // %outerReborrow
+  ///
+  func getEnclosingValues<Ctxt: Context>(_ context: Ctxt) -> EnclosingValues<Ctxt> {
+    EnclosingValues(initialValue: self, context: context)
   }
 }
 
-/// Find inner adjacent phis in the same block as `enclosingPhi`.
-/// These keep the enclosing (outer adjacent) phi alive.
-func gatherInnerAdjacentPhis(for enclosingPhi: Phi,
-                             in innerAdjacentPhis: inout Stack<Phi>,
-                             _ context: Context)
-{
-  for use in enclosingPhi.value.uses {
-    if let bfi = use.instruction as? BorrowedFromInst,
-       use.index != 0,
-       let innerPhi = Phi(bfi.borrowedValue)
-    {
-      innerAdjacentPhis.append(innerPhi)
+extension Phi {
+  /// Inner adjacent phis in the same block as `enclosingPhi`.
+  /// These keep the enclosing (outer adjacent) phi alive.
+  var innerAdjacentPhis: some Sequence<Phi> {
+    value.uses.lazy.compactMap { use in
+      if let bfi = use.instruction as? BorrowedFromInst,
+         use.index != 0
+      {
+         return Phi(bfi.borrowedValue)
+      }
+      return nil
     }
   }
 }
@@ -535,11 +569,7 @@ func gatherEnclosingValuesFromPredecessors(
   for predecessor in phi.predecessors {
     let incomingOperand = phi.incomingOperand(inPredecessor: predecessor)
 
-    var predecessorEVs = Stack<Value>(context)
-    defer { predecessorEVs.deinitialize() }
-    gatherEnclosingValues(for: incomingOperand.value, in: &predecessorEVs, context)
-
-    for predEV in predecessorEVs {
+    for predEV in incomingOperand.value.getEnclosingValues(context) {
       let ev = predecessor.mapToPhiInSuccessor(incomingEnclosingValue: predEV)
       if alreadyAdded.insert(ev) {
         enclosingValues.push(ev)
@@ -569,8 +599,9 @@ let borrowIntroducersTest = FunctionTest("borrow_introducers") {
   defer {
     introducers.deinitialize()
   }
-  gatherBorrowIntroducers(for: value.lookThroughBorrowedFromUser, in: &introducers, context)
-  introducers.forEach { print($0) }
+  for bi in value.lookThroughBorrowedFromUser.getBorrowIntroducers(context) {
+    print(bi)
+  }
 }
 
 let enclosingValuesTest = FunctionTest("enclosing_values") {
@@ -582,8 +613,9 @@ let enclosingValuesTest = FunctionTest("enclosing_values") {
   defer {
     enclosing.deinitialize()
   }
-  gatherEnclosingValues(for: value.lookThroughBorrowedFromUser, in: &enclosing, context)
-  enclosing.forEach { print($0) }
+  for ev in value.lookThroughBorrowedFromUser.getEnclosingValues(context) {
+    print(ev)
+  }
 }
 
 extension Value {

SwiftCompilerSources/Sources/Optimizer/Utilities/LifetimeDependenceUtils.swift

@@ -364,15 +364,14 @@ extension LifetimeDependence.Scope {
   }
 
   private init?(guaranteed base: Value, _ context: some Context) {
-    var introducers = Stack<Value>(context)
-    gatherBorrowIntroducers(for: base, in: &introducers, context)
     // If introducers is empty, then the dependence is on a trivial value, so
     // there is no dependence scope.
     //
     // TODO: Add a SIL verifier check that a mark_dependence [nonescaping]
     // base is never a guaranteed phi.
-    guard let introducer = introducers.pop() else { return nil }
-    assert(introducers.isEmpty,
+    let iter = base.getBorrowIntroducers(context).makeIterator()
+    guard let introducer = iter.next() else { return nil }
+    assert(iter.next() == nil,
            "guaranteed phis not allowed when diagnosing lifetime dependence")
     guard let beginBorrow = BeginBorrowValue(introducer) else {
       fatalError("unknown borrow introducer")