Merge pull request #67384 from nate-chandler/opaque-values/20230718/1/emit_end_borrows/destructure_and_extract

[AddressLowering] End borrow scopes for extracts/destructures at enclosing guaranteed boundary.

Author: nate-chandler

lib/SIL/Verifier/SILVerifier.cpp

@@ -1371,7 +1371,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     // guaranteed address-only values are forwarded with the same
     // representation. Non-destructive projection is allowed. Aggregation and
     // destructive disaggregation is not allowed. See SIL.rst, Forwarding
-    // Addres-Only Values.
+    // Address-Only Values.
     if (ownership == OwnershipKind::Guaranteed && fwdOp.isAddressOnly()) {
       require(fwdOp.hasSameRepresentation(),
               "Forwarding a guaranteed address-only value requires the same "

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -150,6 +150,7 @@
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILValue.h"
 #include "swift/SIL/SILVisitor.h"
+#include "swift/SIL/StackList.h"
 #include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/Analysis/PostOrderAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
@@ -2259,6 +2260,13 @@ void ApplyRewriter::rewriteApply(ArrayRef<SILValue> newCallArgs) {
 /// ending uses.
 static void emitEndBorrows(SILValue value, AddressLoweringState &pass);
 
+/// Emit end_borrows at the lifetime ends of the guaranteed value which
+/// encloses \p enclosingValue.
+static void
+emitEndBorrowsAtEnclosingGuaranteedBoundary(SILValue lifetimeToEnd,
+                                            SILValue enclosingValue,
+                                            AddressLoweringState &pass);
+
 void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
   // Avoid revisiting this apply.
   bool erased = pass.indirectApplies.erase(apply);
@@ -3268,7 +3276,8 @@ void UseRewriter::rewriteDestructure(SILInstruction *destructure) {
       }
       result->replaceAllUsesWith(loadElement);
       if (guaranteed) {
-        emitEndBorrows(loadElement, pass);
+        emitEndBorrowsAtEnclosingGuaranteedBoundary(
+            loadElement, destructure->getOperand(0), pass);
       }
     }
   }
@@ -3395,8 +3404,56 @@ static void emitEndBorrows(SILValue value, AddressLoweringState &pass) {
       });
 }
 
+/// Emit end_borrows at the lifetime ends of the guaranteed value which
+/// encloses \p enclosingValue.
+///
+/// precondition: \p enclosingValue must be of opaque type
+static void
+emitEndBorrowsAtEnclosingGuaranteedBoundary(SILValue lifetimeToEnd,
+                                            SILValue enclosingValue,
+                                            AddressLoweringState &pass) {
+  /// It's necessary that enclosingValue be of opaque type.  In practice, this
+  /// is the only situation in which AddressLowering needs to create a borrow
+  /// scope corresponding to some guaranteed value: the enclosingValue is the
+  /// opaque value that AddressLowering is rewriting.
+  ///
+  /// It's required in order to ensure that no introducer of enclosingValue has
+  /// a phi scope-ending use.  That enables just visiting the local scope ending
+  /// uses.
+  ///
+  /// Given: enclosingValue is opaque.
+  ///
+  /// Then every introducer is opaque.  In fact every introducer's type is some
+  /// aggregate in which enclosingValue's type appears.  The reason is that
+  /// representation-changing forwarding guaranteed operations are not allowed
+  /// on opaque values.  (See SIL.rst, Forwarding Address-Only Values).
+  ///
+  /// Thus no introducer is reborrowed.  For suppose that some introducer were
+  /// reborrowed.  Then it would appear as an operand of some phi with
+  /// guaranteed ownership. And that phi would be of opaque type. But that's
+  /// invalid! (See SILVerifier::visitSILPhiArgument.)
+  assert(enclosingValue->getType().isAddressOnly(*pass.function));
+  TinyPtrVector<SILValue> introducers;
+  visitBorrowIntroducers(enclosingValue, [&](SILValue introducer) {
+    introducers.push_back(introducer);
+    return true;
+  });
+  assert(introducers.size() == 1);
+  for (auto introducer : introducers) {
+    assert(introducer->getType().isAddressOnly(*pass.function));
+    auto borrow = BorrowedValue(introducer);
+    borrow.visitLocalScopeEndingUses([&](Operand *use) {
+      assert(!PhiOperand(use));
+      pass.getBuilder(use->getUser()->getIterator())
+          .createEndBorrow(pass.genLoc(), lifetimeToEnd);
+      return true;
+    });
+  }
+}
+
 // Extract from an opaque struct or tuple.
 void UseRewriter::emitExtract(SingleValueInstruction *extractInst) {
+  auto source = extractInst->getOperand(0);
   SILValue extractAddr = addrMat.materializeDefProjection(extractInst);
 
   if (extractInst->getType().isAddressOnly(*pass.function)) {
@@ -3427,7 +3484,7 @@ void UseRewriter::emitExtract(SingleValueInstruction *extractInst) {
   SILValue loadElement =
       builder.emitLoadBorrowOperation(extractInst->getLoc(), extractAddr);
   replaceUsesWithLoad(extractInst, loadElement);
-  emitEndBorrows(loadElement, pass);
+  emitEndBorrowsAtEnclosingGuaranteedBoundary(loadElement, source, pass);
 }
 
 void UseRewriter::visitStructExtractInst(StructExtractInst *extractInst) {
@@ -3808,12 +3865,22 @@ static void rewriteFunction(AddressLoweringState &pass) {
     if (valueDef->getType().isAddress())
       continue;
 
+    // Collect end_borrows and rewrite them last so that when rewriting other
+    // users the lifetime ends of a borrow introducer can be used.
+    StackList<EndBorrowInst *> endBorrows(pass.function);
     SmallPtrSet<Operand *, 8> originalUses;
     SmallVector<Operand *, 8> uses(valueDef->getUses());
     for (auto *oper : uses) {
       originalUses.insert(oper);
+      if (auto *ebi = dyn_cast<EndBorrowInst>(oper->getUser())) {
+        endBorrows.push_back(ebi);
+        continue;
+      }
       UseRewriter::rewriteUse(oper, pass);
     }
+    for (auto *ebi : endBorrows) {
+      UseRewriter::rewriteUse(&ebi->getOperandRef(), pass);
+    }
     // Rewrite every new use that was added.
     uses.clear();
     for (auto *use : valueDef->getUses()) {

test/SILOptimizer/address_lowering.sil

@@ -60,7 +60,12 @@ enum Mixed<T> {
   case i(Int)
   case t(T)
   case o(AnyObject)
-};
+}
+
+indirect enum RE<T> {
+   case recursive(RE, T) 
+   case other(AnyObject)
+}
 
 class TestGeneric<T> {
   init()
@@ -756,9 +761,9 @@ bb0(%0 : @owned $SRef<T>):
 // CHECK:   [[L:%.*]] = load_borrow [[E0]] : $*AnyObject
 // CHECK:   apply %{{.*}}([[L]]) : $@convention(thin) (@guaranteed AnyObject) -> ()
 // CHECK:   [[C:%.*]] = copy_value [[L]] : $AnyObject
-// CHECK:   end_borrow [[L]] : $AnyObject
 // CHECK:   [[E1:%.*]] = struct_element_addr %2 : $*SRef<T>, #SRef.element
 // CHECK:   copy_addr [[E1]] to [init] %1 : $*T
+// CHECK:   end_borrow [[L]] : $AnyObject
 // CHECK:   destroy_addr %2 : $*SRef<T>
 // CHECK:   store [[C]] to [init] %0 : $*AnyObject
 // CHECK-NOT: dealloc_stack
@@ -808,9 +813,9 @@ bb0(%0 : @owned $(AnyObject, T)):
 // CHECK:   [[L:%.*]] = load_borrow %3 : $*AnyObject
 // CHECK:   apply %{{.*}}([[L]]) : $@convention(thin) (@guaranteed AnyObject) -> ()
 // CHECK:   [[C:%.*]] = copy_value [[L]] : $AnyObject
-// CHECK:   end_borrow [[L]] : $AnyObject
 // CHECK:   [[E1:%.*]] = tuple_element_addr %2 : $*(AnyObject, T), 1
 // CHECK:   copy_addr [[E1]] to [init] %1 : $*T
+// CHECK:   end_borrow [[L]] : $AnyObject
 // CHECK:   destroy_addr %2 : $*(AnyObject, T)
 // CHECK:   store [[C]] to [init] %0 : $*AnyObject
 // CHECK-NOT: dealloc_stack
@@ -2224,9 +2229,9 @@ exit:
 // CHECK:         [[INSTANCE:%[^,]+]] = unchecked_take_enum_data_addr
 // CHECK:         [[KLASS_ADDR:%[^,]+]] = tuple_element_addr [[INSTANCE]] {{.*}}, 0
 // CHECK:         [[KLASS:%[^,]+]] = load_borrow [[KLASS_ADDR]]
-// CHECK:         end_borrow [[KLASS]]
 // CHECK:         [[ANY_ADDR:%[^,]+]] = tuple_element_addr [[INSTANCE]] {{.*}}, 1
 // CHECK:         apply undef<Any>([[ANY_ADDR]])
+// CHECK:         end_borrow [[KLASS]]
 // CHECK:         destroy_addr [[INSTANCE]]
 // CHECK-LABEL: } // end sil function 'test_destructure_tuple_1_guaranteed'
 sil [ossa] @test_destructure_tuple_1_guaranteed : $@convention(thin) () -> () {
@@ -2249,6 +2254,62 @@ exit:
   return %505 : $()
 }
 
+// Check that the end_borrows created on behalf of the destructure_tuple end
+// after the use of the load_borrow of the projection of the terminator result
+// of the switch_enum that is guaranteed by that load_borrow.
+//
+// CHECK-LABEL: sil [ossa] @test_destructure_tuple_load_borrow_projected_box : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[OUTER:%[^,]+]] :
+// CHECK:         switch_enum [[OUTER]] : $RE<T>, case #RE.recursive!enumelt: [[OUTER_RECURSIVE:bb[0-9]+]], case #RE.other!enumelt: [[OUTER_OTHER:bb[0-9]+]]
+// CHECK:       [[OUTER_OTHER]]({{%[^,]+}} :
+// CHECK:         br [[EXIT:bb[0-9]+]]
+// CHECK:       [[OUTER_RECURSIVE]]([[TUPLE_BOX:%[^,]+]] :
+// CHECK-NEXT:    [[TUPLE_ADDR:%[^,]+]] = project_box [[TUPLE_BOX]]
+// CHECK:         [[INNER_ADDR:%[^,]+]] = tuple_element_addr [[TUPLE_ADDR]]
+// CHECK:         [[INNER:%[^,]+]] = load_borrow [[INNER_ADDR]]
+// CHECK:         switch_enum [[INNER]] : $RE<T>, case #RE.recursive!enumelt: [[INNER_RECURSIVE:bb[0-9]+]], case #RE.other!enumelt: [[INNER_OTHER:bb[0-9]+]]
+// CHECK:       [[INNER_OTHER]]([[OTHER_BOX:%[^,]+]] :
+// CHECK:         [[OTHER_ADDR:%[^,]+]] = project_box [[OTHER_BOX]]
+// CHECK:         [[OTHER:%[^,]+]] = load_borrow [[OTHER_ADDR]]
+// CHECK:         [[OTHER_COPY:%[^,]+]] = copy_value [[OTHER]]
+// CHECK:         end_borrow [[OTHER]]
+// CHECK:         end_borrow [[INNER]]
+// CHECK:         destroy_value [[OTHER_COPY]]
+// CHECK:         br [[EXIT]]
+// CHECK:       [[INNER_RECURSIVE]]
+// CHECK:         end_borrow [[INNER]]
+// CHECK:         br [[EXIT]]
+// CHECK-LABEL: } // end sil function 'test_destructure_tuple_load_borrow_projected_box'
+sil [ossa] @test_destructure_tuple_load_borrow_projected_box : $@convention(thin) <T> (@guaranteed RE<T>) -> () {
+entry(%o : @guaranteed $RE<T>):
+  switch_enum %o : $RE<T>, case #RE.recursive!enumelt: outer_recursive, case #RE.other!enumelt: outer_other
+
+outer_other(%other : @guaranteed $<τ_0_0> { var AnyObject } <T>):
+  br exit
+
+outer_recursive(%tuple_box : @guaranteed $<Tau> { var (RE<Tau>, Tau) } <T>):
+  %tuple_addr = project_box %tuple_box : $<Tau> { var (RE<Tau>, Tau) } <T>, 0
+  %tuple = load_borrow %tuple_addr : $*(RE<T>, T)
+  (%i, %c) = destructure_tuple %tuple : $(RE<T>, T)
+  switch_enum %i : $RE<T>, case #RE.recursive!enumelt: inner_recursive, case #RE.other!enumelt: inner_other
+
+inner_recursive(%in_tuple_box : @guaranteed $<Tau> { var (RE<Tau>, Tau) } <T>):
+  end_borrow %tuple : $(RE<T>, T)
+  br exit
+
+inner_other(%s_box : @guaranteed $<Tau> { var AnyObject } <T>):
+  %s_addr = project_box %s_box : $<Tau> { var AnyObject } <T>, 0
+  %s = load_borrow %s_addr : $*AnyObject
+  %sc = copy_value %s : $AnyObject
+  end_borrow %s : $AnyObject
+  end_borrow %tuple : $(RE<T>, T)
+  destroy_value %sc : $AnyObject
+  br exit
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
 
 // CHECK-LABEL: sil hidden [ossa] @test_store_1 : {{.*}} {
 // CHECK:         [[MAYBE_ADDR:%[^,]+]] = alloc_stack $Optional<Self>