Merge pull request #67913 from nate-chandler/opaque-values/20230814/partial-apply-loadable

[AddressLowering] Rewrite indirect partial_apply.

Author: nate-chandler

include/swift/SIL/ApplySite.h

@@ -541,9 +541,9 @@ class ApplySite {
 
   void dump() const LLVM_ATTRIBUTE_USED { getInstruction()->dump(); }
 
-  /// Attempt to cast this apply site to a full apply site, returning None on
-  /// failure.
-  llvm::Optional<FullApplySite> asFullApplySite() const;
+  /// Form a FullApplySite.  Note that it will be null if this apply site is not
+  /// in fact a FullApplySite.
+  FullApplySite asFullApplySite() const;
 };
 
 //===----------------------------------------------------------------------===//
@@ -818,15 +818,15 @@ template <> struct DenseMapInfo<::swift::FullApplySite> {
 
 namespace swift {
 
-inline llvm::Optional<FullApplySite> ApplySite::asFullApplySite() const {
+inline FullApplySite ApplySite::asFullApplySite() const {
   return FullApplySite::isa(getInstruction());
 }
 
 inline bool ApplySite::isIndirectResultOperand(const Operand &op) const {
   auto fas = asFullApplySite();
   if (!fas)
     return false;
-  return fas->isIndirectResultOperand(op);
+  return fas.isIndirectResultOperand(op);
 }
 
 } // namespace swift

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -509,7 +509,7 @@ struct AddressLoweringState {
   // Applies with indirect results are removed as they are rewritten. Applies
   // with only indirect arguments are rewritten in a post-pass, only after all
   // parameters are rewritten.
-  SmallBlotSetVector<FullApplySite, 16> indirectApplies;
+  SmallBlotSetVector<ApplySite, 16> indirectApplies;
 
   // unconditional_checked_cast instructions of loadable type which need to be
   // rewritten.
@@ -732,7 +732,7 @@ class OpaqueValueVisitor {
   void mapValueStorage();
 
 protected:
-  void checkForIndirectApply(FullApplySite applySite);
+  void checkForIndirectApply(ApplySite applySite);
   void visitValue(SILValue value);
   void canonicalizeReturnValues();
 };
@@ -755,7 +755,7 @@ void OpaqueValueVisitor::mapValueStorage() {
       }
     }
     for (auto &inst : *block) {
-      if (auto apply = FullApplySite::isa(&inst))
+      if (auto apply = ApplySite::isa(&inst))
         checkForIndirectApply(apply);
 
       if (auto *yieldInst = dyn_cast<YieldInst>(&inst)) {
@@ -774,7 +774,7 @@ void OpaqueValueVisitor::mapValueStorage() {
 }
 
 /// Populate `indirectApplies`.
-void OpaqueValueVisitor::checkForIndirectApply(FullApplySite applySite) {
+void OpaqueValueVisitor::checkForIndirectApply(ApplySite applySite) {
   auto calleeConv = applySite.getSubstCalleeConv();
   unsigned calleeArgIdx = applySite.getCalleeArgIndexOfFirstAppliedArg();
   for (Operand &operand : applySite.getArgumentOperands()) {
@@ -2228,7 +2228,7 @@ void CallArgRewriter::rewriteIndirectArgument(Operand *operand) {
   // Allocate temporary storage for a loadable operand.
   AllocStackInst *allocInst =
       argBuilder.createAllocStack(callLoc, argValue->getType());
-  if (apply.getArgumentConvention(*operand).isOwnedConvention()) {
+  if (apply.getCaptureConvention(*operand).isOwnedConvention()) {
     argBuilder.createTrivialStoreOr(apply.getLoc(), argValue, allocInst,
                                     StoreOwnershipQualifier::Init);
     apply.insertAfterApplication([&](SILBuilder &callBuilder) {
@@ -4139,10 +4139,17 @@ class DefRewriter : SILInstructionVisitor<DefRewriter> {
 
 // Rewrite applies with indirect parameters or results of loadable types which
 // were not visited during opaque value rewriting.
-static void rewriteIndirectApply(FullApplySite apply,
+static void rewriteIndirectApply(ApplySite anyApply,
                                  AddressLoweringState &pass) {
   // If all indirect args were loadable, then they still need to be rewritten.
-  CallArgRewriter(apply, pass).rewriteArguments();
+  CallArgRewriter(anyApply, pass).rewriteArguments();
+
+  auto apply = anyApply.asFullApplySite();
+  if (!apply) {
+    // The "result" of the partial_apply isn't rewritten.
+    assert(anyApply.getKind() == ApplySiteKind::PartialApplyInst);
+    return;
+  }
 
   switch (apply.getKind()) {
   case FullApplySiteKind::ApplyInst:

test/SILOptimizer/address_lowering.sil

@@ -93,6 +93,9 @@ public protocol Comparable {
 
 sil [ossa] @unknown : $@convention(thin) () -> ()
 sil [ossa] @getT : $@convention(thin) <T> () -> @out T
+sil [ossa] @getKlass : $@convention(thin) () -> (@owned Klass)
+sil [ossa] @borrowKlass : $(@in_guaranteed Klass) -> ()
+sil [ossa] @borrowT : $@convention(thin) <T> (@in_guaranteed T) -> ()
 sil [ossa] @getPair : $@convention(thin) <T> () -> @out Pair<T>
 sil [ossa] @getOwned : $@convention(thin) <T : AnyObject> () -> (@owned T)
 sil [ossa] @takeGuaranteedObject : $@convention(thin) (@guaranteed AnyObject) -> ()
@@ -2521,6 +2524,101 @@ bb0:
   return %retval : $()
 }
 
+// CHECK-LABEL: sil [ossa] @test_partial_apply_1_addronly_heap : {{.*}} {
+// CHECK:         [[STACK:%[^,]+]] = alloc_stack $T
+// CHECK:         [[GET:%[^,]+]] = function_ref @getT
+// CHECK:         apply [[GET]]<T>([[STACK]])
+// CHECK:         [[BORROW:%[^,]+]] = function_ref @borrowT
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [[BORROW]]<T>([[STACK]])
+// CHECK:         apply [[CLOSURE]]()
+// CHECK:         destroy_value [[CLOSURE]]
+// CHECK:         dealloc_stack [[STACK]]
+// CHECK-LABEL: } // end sil function 'test_partial_apply_1_addronly_heap'
+sil [ossa] @test_partial_apply_1_addronly_heap : $<T> () -> () {
+  %get = function_ref @getT : $@convention(thin) <T> () -> (@out T)
+  %instance = apply %get<T>() : $@convention(thin) <T> () -> (@out T)
+  %callee = function_ref @borrowT : $@convention(thin) <T> (@in_guaranteed T) -> ()
+  %closure = partial_apply [callee_guaranteed] %callee<T>(%instance) : $@convention(thin) <T> (@in_guaranteed T) -> ()
+  apply %closure() : $@callee_guaranteed () -> ()
+  destroy_value %closure : $@callee_guaranteed () -> ()
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_partial_apply_2_addronly_stack : {{.*}} {
+// CHECK:         [[STACK:%[^,]+]] = alloc_stack $T
+// CHECK:         [[GET:%[^,]+]] = function_ref @getT
+// CHECK:         apply [[GET]]<T>([[STACK]])
+// CHECK:         [[BORROW:%[^,]+]] = function_ref @borrowT
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[BORROW]]<T>([[STACK]])
+// CHECK:         apply [[CLOSURE]]()
+// CHECK:         destroy_value [[CLOSURE]]
+// CHECK:         destroy_addr [[STACK]]
+// CHECK:         dealloc_stack [[STACK]]
+// CHECK-LABEL: } // end sil function 'test_partial_apply_2_addronly_stack'
+sil [ossa] @test_partial_apply_2_addronly_stack : $<T> () -> () {
+  %get = function_ref @getT : $@convention(thin) <T> () -> (@out T)
+  %instance = apply %get<T>() : $@convention(thin) <T> () -> (@out T)
+  %callee = function_ref @borrowT : $@convention(thin) <T> (@in_guaranteed T) -> ()
+  %closure = partial_apply [on_stack] [callee_guaranteed] %callee<T>(%instance) : $@convention(thin) <T> (@in_guaranteed T) -> ()
+  apply %closure() : $@noescape @callee_guaranteed () -> ()
+  destroy_value %closure : $@noescape @callee_guaranteed () -> ()
+  destroy_value %instance : $T
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_partial_apply_3_loadable_heap : {{.*}} {
+// CHECK:         [[GET:%[^,]+]] = function_ref @getKlass
+// CHECK:         [[INSTANCE:%[^,]+]] = apply [[GET]]()
+// CHECK:         [[BORROW:%[^,]+]] = function_ref @borrowKlass
+// CHECK:         [[ADDR:%[^,]+]] = alloc_stack $Klass
+// CHECK:         store [[INSTANCE]] to [init] [[ADDR]]
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [[BORROW]]([[ADDR]])
+// CHECK:         dealloc_stack [[ADDR]]
+// CHECK:         apply [[CLOSURE]]()
+// CHECK:         destroy_value [[CLOSURE]]
+// CHECK-LABEL: } // end sil function 'test_partial_apply_3_loadable_heap'
+sil [ossa] @test_partial_apply_3_loadable_heap : $() -> () {
+  %get = function_ref @getKlass : $@convention(thin) () -> (@owned Klass)
+  %instance = apply %get() : $@convention(thin) () -> (@owned Klass)
+  %callee = function_ref @borrowKlass : $@convention(thin) (@in_guaranteed Klass) -> ()
+  %closure = partial_apply [callee_guaranteed] %callee(%instance) : $@convention(thin) (@in_guaranteed Klass) -> ()
+  apply %closure() : $@callee_guaranteed () -> ()
+  destroy_value %closure : $@callee_guaranteed () -> ()
+  // destroy_value %instance : $Klass
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_partial_apply_4_loadable_stack : {{.*}} {
+// CHECK:         [[GET:%[^,]+]] = function_ref @getKlass
+// CHECK:         [[INSTANCE:%[^,]+]] = apply [[GET]]()
+// CHECK:         [[BORROW:%[^,]+]] = function_ref @borrowKlass
+// CHECK:         [[STACK:%[^,]+]] = alloc_stack $Klass
+// CHECK:         [[LIFETIME:%[^,]+]] = begin_borrow [[INSTANCE]]
+// CHECK:         [[STORAGE_LIFETIME:%[^,]+]] = store_borrow [[LIFETIME]] to [[STACK]]
+// CHECK:         [[CLOSURE:%[^,]+]] = partial_apply [callee_guaranteed] [on_stack] [[BORROW]]([[STORAGE_LIFETIME]])
+// FIXME: These end_borrows must be _after_ the apply of the closure!
+// CHECK:         end_borrow [[STORAGE_LIFETIME]]
+// CHECK:         end_borrow [[LIFETIME]]
+// CHECK:         dealloc_stack [[STACK]]
+// CHECK:         apply [[CLOSURE]]() : $@noescape @callee_guaranteed () -> ()
+// CHECK:         destroy_value [[CLOSURE]] : $@noescape @callee_guaranteed () -> ()
+// CHECK:         destroy_value [[INSTANCE]]
+// CHECK-LABEL: } // end sil function 'test_partial_apply_4_loadable_stack'
+sil [ossa] @test_partial_apply_4_loadable_stack : $() -> () {
+  %get = function_ref @getKlass : $@convention(thin) () -> (@owned Klass)
+  %instance = apply %get() : $@convention(thin) () -> (@owned Klass)
+  %callee = function_ref @borrowKlass : $@convention(thin) (@in_guaranteed Klass) -> ()
+  %closure = partial_apply [on_stack] [callee_guaranteed] %callee(%instance) : $@convention(thin) (@in_guaranteed Klass) -> ()
+  apply %closure() : $@noescape @callee_guaranteed () -> ()
+  destroy_value %closure : $@noescape @callee_guaranteed () -> ()
+  destroy_value %instance : $Klass
+  %retval = tuple ()
+  return %retval : $()
+}
+
 // CHECK-LABEL: sil hidden [ossa] @test_store_1 : {{.*}} {
 // CHECK:         [[MAYBE_ADDR:%[^,]+]] = alloc_stack $Optional<Self>
 // CHECK:         [[LOAD_ADDR:%[^,]+]] = alloc_stack $Self