[AddressLowering] Rewrite indirect loadable yields

nate-chandler · nate-chandler · commit b64510db0d1d · 2022-11-10T18:11:29.000-08:00
While visiting the function, record not only applies which have indirect
formal _results_ but also applies which have indirect formal _yields_.

When rewriting indirect applies, check the apply site kind and rewrite
according to it.

Taught ApplyRewriter to handle rewriting indirect loadable yields.
diff --git a/lib/SILOptimizer/Mandatory/AddressLowering.cpp b/lib/SILOptimizer/Mandatory/AddressLowering.cpp
@@ -195,6 +195,8 @@ static SILFunctionConventions getLoweredCallConv(ApplySite call) {
 /// If \p pseudoResult represents multiple results and at least one result is
 /// used, then return the destructure.
 static DestructureTupleInst *getCallDestructure(FullApplySite apply) {
+  if (apply.getKind() == FullApplySiteKind::BeginApplyInst)
+    return nullptr;
   if (apply.getSubstCalleeConv().getNumDirectSILResults() == 1)
     return nullptr;
 
@@ -640,7 +642,8 @@ void OpaqueValueVisitor::checkForIndirectApply(FullApplySite applySite) {
     ++calleeArgIdx;
   }
 
-  if (applySite.getSubstCalleeType()->hasIndirectFormalResults()) {
+  if (applySite.getSubstCalleeType()->hasIndirectFormalResults() ||
+      applySite.getSubstCalleeType()->hasIndirectFormalYields()) {
     pass.indirectApplies.insert(applySite);
   }
 }
@@ -2125,6 +2128,11 @@ void ApplyRewriter::rewriteApply(ArrayRef<SILValue> newCallArgs) {
 static void emitEndBorrows(SILValue value, AddressLoweringState &pass);
 
 void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
+  // Avoid revisiting this apply.
+  bool erased = pass.indirectApplies.erase(apply);
+  assert(erased && "all yields should be rewritten at the same time");
+  (void)erased;
+
   auto *origCall = cast<BeginApplyInst>(apply.getInstruction());
   SmallVector<SILValue, 4> opValues;
 
@@ -2144,11 +2152,33 @@ void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
   auto newResults = newCall->getAllResultsBuffer();
   assert(oldResults.size() == newResults.size());
   for (auto i : indices(oldResults)) {
-    if (oldResults[i].getType().isAddressOnly(*pass.function)) {
-      pass.valueStorageMap.setStorageAddress(&oldResults[i], &newResults[i]);
-      pass.valueStorageMap.getStorage(&oldResults[i]).markRewritten();
+    auto &oldResult = oldResults[i];
+    auto &newResult = newResults[i];
+    if (oldResult.getType().isAddressOnly(*pass.function)) {
+      if (!oldResult.getType().isAddress()) {
+        pass.valueStorageMap.setStorageAddress(&oldResult, &newResult);
+        pass.valueStorageMap.getStorage(&oldResult).markRewritten();
+      } else {
+        oldResult.replaceAllUsesWith(&newResult);
+      }
     } else {
-      oldResults[i].replaceAllUsesWith(&newResults[i]);
+      if (oldResult.getType().isObject() && newResult.getType().isAddress()) {
+        if (oldResult.getOwnershipKind() == OwnershipKind::Guaranteed) {
+          SILValue load =
+              resultBuilder.emitLoadBorrowOperation(callLoc, &newResult);
+          oldResult.replaceAllUsesWith(load);
+          emitEndBorrows(load, pass);
+        } else {
+          auto *load = resultBuilder.createLoad(
+              callLoc, &newResult,
+              newResult.getType().isTrivial(*pass.function)
+                  ? LoadOwnershipQualifier::Trivial
+                  : LoadOwnershipQualifier::Take);
+          oldResult.replaceAllUsesWith(load);
+        }
+      } else {
+        oldResult.replaceAllUsesWith(&newResult);
+      }
     }
   }
 }
@@ -3369,14 +3399,27 @@ static void rewriteIndirectApply(FullApplySite apply,
   // If all indirect args were loadable, then they still need to be rewritten.
   CallArgRewriter(apply, pass).rewriteArguments();
 
-  if (!apply.getSubstCalleeType()->hasIndirectFormalResults()) {
-    return;
+  switch (apply.getKind()) {
+  case FullApplySiteKind::ApplyInst:
+  case FullApplySiteKind::TryApplyInst: {
+    if (!apply.getSubstCalleeType()->hasIndirectFormalResults()) {
+      return;
+    }
+    // If the call has indirect results and wasn't already rewritten, rewrite it
+    // now. This handles try_apply, which is not rewritten when DefRewriter
+    // visits block arguments. It also handles apply with loadable indirect
+    // results.
+    ApplyRewriter(apply, pass).convertApplyWithIndirectResults();
+    break;
+  }
+  case FullApplySiteKind::BeginApplyInst: {
+    if (!apply.getSubstCalleeType()->hasIndirectFormalYields()) {
+      return;
+    }
+    ApplyRewriter(apply, pass).convertBeginApplyWithOpaqueYield();
+    break;
+  }
   }
-
-  // If the call has indirect results and wasn't already rewritten, rewrite it
-  // now. This handles try_apply, which is not rewritten when DefRewriter visits
-  // block arguments. It also handles apply with loadable indirect results.
-  ApplyRewriter(apply, pass).convertApplyWithIndirectResults();
 
   if (!apply.getInstruction()->isDeleted()) {
     assert(!getCallDestructure(apply)
diff --git a/test/SILOptimizer/address_lowering.sil b/test/SILOptimizer/address_lowering.sil
@@ -29,6 +29,14 @@ struct I {}
 
 class Klass {}
 
+struct LoadableTrivialExtractable {
+    var i: I
+}
+
+struct LoadableNontrivial {
+    var x: Klass
+}
+
 struct SI<T> {
   var element: T
   var index: I
@@ -1468,6 +1476,202 @@ bb0(%0 : @guaranteed $TestGeneric<T>):
   return %10 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @testBeginApply4YieldLoadableTrivial : {{.*}} {
+// CHECK:         ([[ADDR:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE:%[^,]+]] = load [trivial] [[ADDR]]
+// CHECK:         struct_extract [[INSTANCE]]
+// CHECK-LABEL: } // end sil function 'testBeginApply4YieldLoadableTrivial'
+sil [ossa] @testBeginApply4YieldLoadableTrivial : $@convention(thin) () -> I {
+entry:
+  (%instance, %token) = begin_apply undef<LoadableTrivialExtractable>() : $@yield_once @convention(thin) <U> () -> @yields @in_guaranteed U
+  end_apply %token
+  %extract = struct_extract %instance : $LoadableTrivialExtractable, #LoadableTrivialExtractable.i
+  return %extract : $I
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApply5Yield2LoadableTrivial : {{.*}} {
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE_1:%[^,]+]] = load [trivial] [[ADDR_1]]
+// CHECK:         [[INSTANCE_2:%[^,]+]] = load [trivial] [[ADDR_2]]
+// CHECK:         struct_extract [[INSTANCE_1]]
+// CHECK:         struct_extract [[INSTANCE_2]]
+// CHECK-LABEL: } // end sil function 'testBeginApply5Yield2LoadableTrivial'
+sil [ossa] @testBeginApply5Yield2LoadableTrivial : $@convention(thin) () -> (I, I) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableTrivialExtractable>() : $@yield_once @convention(thin) <U> () -> (@yields @in_guaranteed U, @yields @in_guaranteed U)
+  end_apply %token
+  %i_1 = struct_extract %instance_1 : $LoadableTrivialExtractable, #LoadableTrivialExtractable.i
+  %i_2 = struct_extract %instance_2 : $LoadableTrivialExtractable, #LoadableTrivialExtractable.i
+  %tuple = tuple (%i_1 : $I, %i_2 : $I)
+  return %tuple : $(I, I)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApply6Yield1LoadableTrivial1OpaqueGuaranteed : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[TRIV_OUT:%[^,]+]] : $*I, [[OPAQUE_OUT:%[^,]+]] : $*T):
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE_1:%[^,]+]] = load [trivial] [[ADDR_1]]
+// CHECK:         copy_addr [[ADDR_2]] to [init] [[OPAQUE_OUT]]
+// CHECK:         [[EXTRACT:%[^,]+]]  = struct_extract [[INSTANCE_1]]
+// CHECK:         store [[EXTRACT]] to [trivial] [[TRIV_OUT]]
+// CHECK-LABEL: } // end sil function 'testBeginApply6Yield1LoadableTrivial1OpaqueGuaranteed'
+sil [ossa] @testBeginApply6Yield1LoadableTrivial1OpaqueGuaranteed : $@convention(thin) <T> () -> (@out I, @out T) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableTrivialExtractable, T>() : $@yield_once @convention(thin) <U, T> () -> (@yields @in_guaranteed U, @yields @in_guaranteed T)
+  %copy = copy_value %instance_2 : $T
+  end_apply %token
+  %i = struct_extract %instance_1 : $LoadableTrivialExtractable, #LoadableTrivialExtractable.i
+  %tuple = tuple (%i : $I, %copy : $T)
+  return %tuple : $(I, T)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApply7YieldLoadableNontrivialGuaranteed : {{.*}} {
+// CHECK:         ([[ADDR:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[BORROW:%[^,]+]] = load_borrow [[ADDR]]
+// CHECK:         end_borrow [[BORROW]]
+// CHECK-LABEL: } // end sil function 'testBeginApply7YieldLoadableNontrivialGuaranteed'
+sil [ossa] @testBeginApply7YieldLoadableNontrivialGuaranteed : $@convention(thin) () -> @owned Klass {
+entry:
+  (%instance, %token) = begin_apply undef<LoadableNontrivial>() : $@yield_once @convention(thin) <U> () -> @yields @in_guaranteed U
+  %extract = struct_extract %instance : $LoadableNontrivial, #LoadableNontrivial.x
+  %retval = copy_value %extract : $Klass
+  end_apply %token
+  return %retval : $Klass
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApply8Yield2LoadableNontrivialGuaranteed : {{.*}} {
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE_1:%[^,]+]] = load_borrow [[ADDR_1]]
+// CHECK:         [[INSTANCE_2:%[^,]+]] = load_borrow [[ADDR_2]]
+// CHECK:         struct_extract [[INSTANCE_1]]
+// CHECK:         struct_extract [[INSTANCE_2]]
+// CHECK:         end_borrow [[INSTANCE_1]]
+// CHECK:         end_borrow [[INSTANCE_2]]
+// CHECK-LABEL: } // end sil function 'testBeginApply8Yield2LoadableNontrivialGuaranteed'
+sil [ossa] @testBeginApply8Yield2LoadableNontrivialGuaranteed : $@convention(thin) () -> @owned (Klass, Klass) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableNontrivial>() : $@yield_once @convention(thin) <U> () -> (@yields @in_guaranteed U, @yields @in_guaranteed U)
+  %extract_1 = struct_extract %instance_1 : $LoadableNontrivial, #LoadableNontrivial.x
+  %extract_2 = struct_extract %instance_2 : $LoadableNontrivial, #LoadableNontrivial.x
+  %copy_1 = copy_value %extract_1 : $Klass
+  %copy_2 = copy_value %extract_2 : $Klass
+  end_apply %token
+  %retval = tuple (%copy_1 : $Klass, %copy_2 : $Klass)
+  return %retval : $(Klass, Klass)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApply9Yield1LoadableNontrivialGuaranteed1OpaqueGuaranteed : {{.*}} {
+// CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Klass, [[OPAQUE_OUT:%[^,]+]] : $*T):
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[BORROW_1:%[^,]+]] = load_borrow [[ADDR_1]] : $*LoadableNontrivial
+// CHECK:         struct_extract [[BORROW_1]] : $LoadableNontrivial, #LoadableNontrivial.x
+// CHECK:         end_borrow [[BORROW_1]] : $LoadableNontrivial
+// CHECK:         copy_addr [[ADDR_2]] to [init] [[OPAQUE_OUT]] : $*T
+// CHECK-LABEL: } // end sil function 'testBeginApply9Yield1LoadableNontrivialGuaranteed1OpaqueGuaranteed'
+sil [ossa] @testBeginApply9Yield1LoadableNontrivialGuaranteed1OpaqueGuaranteed : $@convention(thin) <T> () -> (@out Klass, @out T) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableNontrivial, T>() : $@yield_once @convention(thin) <U, T> () -> (@yields @in_guaranteed U, @yields @in_guaranteed T)
+  %extract_1 = struct_extract %instance_1 : $LoadableNontrivial, #LoadableNontrivial.x
+  %copy_1 = copy_value %extract_1 : $Klass
+  %copy_2 = copy_value %instance_2 : $T
+  end_apply %token
+  %retval = tuple (%copy_1 : $Klass, %copy_2 : $T)
+  return %retval : $(Klass, T)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApplyAYieldLoadableNontrivialOwned : {{.*}} {
+// CHECK:         ([[ADDR:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE:%[^,]+]] = load [take] [[ADDR]] : $*LoadableNontrivial
+// CHECK:         [[BORROW:%[^,]+]] = begin_borrow [[INSTANCE]] : $LoadableNontrivial
+// CHECK:         struct_extract [[BORROW]] : $LoadableNontrivial, #LoadableNontrivial.x
+// CHECK:         end_borrow [[BORROW]] : $LoadableNontrivial
+// CHECK-LABEL: } // end sil function 'testBeginApplyAYieldLoadableNontrivialOwned'
+sil [ossa] @testBeginApplyAYieldLoadableNontrivialOwned : $@convention(thin) () -> @owned Klass {
+entry:
+  (%instance, %token) = begin_apply undef<LoadableNontrivial>() : $@yield_once @convention(thin) <U> () -> @yields @in U
+  end_apply %token
+  %borrow = begin_borrow %instance : $LoadableNontrivial
+  %extract = struct_extract %borrow : $LoadableNontrivial, #LoadableNontrivial.x
+  %retval = copy_value %extract : $Klass
+  end_borrow %borrow : $LoadableNontrivial
+  destroy_value %instance : $LoadableNontrivial
+  return %retval : $Klass
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApplyBYield2LoadableNontrivialOwned : {{.*}} {
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE_1:%[^,]+]] = load [take] [[ADDR_1]]
+// CHECK:         [[INSTANCE_2:%[^,]+]] = load [take] [[ADDR_2]]
+// CHECK:         begin_borrow [[INSTANCE_1]]
+// CHECK:         begin_borrow [[INSTANCE_2]]
+// CHECK-LABEL: } // end sil function 'testBeginApplyBYield2LoadableNontrivialOwned'
+sil [ossa] @testBeginApplyBYield2LoadableNontrivialOwned : $@convention(thin) () -> @owned (Klass, Klass) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableNontrivial>() : $@yield_once @convention(thin) <U> () -> (@yields @in U, @yields @in U)
+  end_apply %token
+  %borrow_1 = begin_borrow %instance_1 : $LoadableNontrivial
+  %extract_1 = struct_extract %borrow_1 : $LoadableNontrivial, #LoadableNontrivial.x
+  %copy_1 = copy_value %extract_1 : $Klass
+  end_borrow %borrow_1 : $LoadableNontrivial
+  destroy_value %instance_1 : $LoadableNontrivial
+  %borrow_2 = begin_borrow %instance_2 : $LoadableNontrivial
+  %extract_2 = struct_extract %borrow_2 : $LoadableNontrivial, #LoadableNontrivial.x
+  %copy_2 = copy_value %extract_2 : $Klass
+  end_borrow %borrow_2 : $LoadableNontrivial
+  destroy_value %instance_2 : $LoadableNontrivial
+  %retval = tuple (%copy_1 : $Klass, %copy_2 : $Klass)
+  return %retval : $(Klass, Klass)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApplyCYield1LoadableNontrivialOwned1OpaqueOwned : $@convention(thin) <T> () -> (@out Klass, @out T) {
+// CHECK:       {{bb[0-9]+}}({{%[^,]+}} : $*Klass, [[OPAQUE_OUT:%[^,]+]] : $*T):
+// CHECK:         ([[ADDR_1:%[^,]+]], [[ADDR_2:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         [[INSTANCE_1:%[^,]+]] = load [take] [[ADDR_1]]
+// CHECK:         begin_borrow [[INSTANCE_1]]
+// CHECK:         copy_addr [take] [[ADDR_2]] to [init] [[OPAQUE_OUT]]
+// CHECK-LABEL: } // end sil function 'testBeginApplyCYield1LoadableNontrivialOwned1OpaqueOwned'
+sil [ossa] @testBeginApplyCYield1LoadableNontrivialOwned1OpaqueOwned : $@convention(thin) <T> () -> (@out Klass, @out T) {
+entry:
+  (%instance_1, %instance_2, %token) = begin_apply undef<LoadableNontrivial, T>() : $@yield_once @convention(thin) <U, T> () -> (@yields @in U, @yields @in T)
+  end_apply %token
+  %borrow_1 = begin_borrow %instance_1 : $LoadableNontrivial
+  %extract_1 = struct_extract %borrow_1 : $LoadableNontrivial, #LoadableNontrivial.x
+  %copy_1 = copy_value %extract_1 : $Klass
+  end_borrow %borrow_1 : $LoadableNontrivial
+  destroy_value %instance_1 : $LoadableNontrivial
+  %retval = tuple (%copy_1 : $Klass, %instance_2 : $T)
+  return %retval : $(Klass, T)
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApplyDYieldOpaqueInout : {{.*}} {
+// CHECK:         begin_apply undef<T>() : $@yield_once @convention(method) <τ_0_0> () -> @yields @inout τ_0_0
+// CHECK-LABEL: } // end sil function 'testBeginApplyDYieldOpaqueInout'
+sil [ossa] @testBeginApplyDYieldOpaqueInout : $@yield_once @convention(thin) <T> () -> () {
+entry:
+  (%addr, %token) = begin_apply undef<T>() : $@yield_once @convention(method) <U> () -> @yields @inout U
+  end_apply %token
+  %8 = tuple ()
+  return %8 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @testBeginApplyEYieldOpaqueInoutAndYield : {{.*}}{
+// CHECK:         ([[REGISTER_0:%[^,]+]], {{%[^,]+}}) = begin_apply
+// CHECK:         yield [[REGISTER_0]] : $*T
+// CHECK-LABEL: } // end sil function 'testBeginApplyEYieldOpaqueInoutAndYield'
+sil [ossa] @testBeginApplyEYieldOpaqueInoutAndYield : $@yield_once @convention(thin) <T> () -> @yields @inout T {
+entry:
+  (%addr, %token) = begin_apply undef<T>() : $@yield_once @convention(method) <U> () -> @yields @inout U
+  yield %addr : $*T, resume bb1, unwind bb2
+
+bb1:
+  end_apply %token
+  %8 = tuple ()
+  return %8 : $()
+
+bb2:
+  abort_apply %token
+  unwind
+}
+
 // CHECK-LABEL: sil hidden [ossa] @testOpaqueYield :
 // CHECK: bb0(%0 : @guaranteed $TestGeneric<T>):
 // CHECK:   [[REF:%.*]] = ref_element_addr %0 : $TestGeneric<T>, #TestGeneric.borrowedGeneric
