Merge pull request #61846 from nate-chandler/opaque-values/1/20221027

[OpaqueValues] Handle PartialApplyInsts and @out tuples.

Author: nate-chandler

include/swift/SIL/ApplySite.h

@@ -474,6 +474,29 @@ class ApplySite {
     llvm_unreachable("covered switch");
   }
 
+  /// If this is a terminator apply site, then pass a builder to insert at the
+  /// first instruction of each successor to \p func. Otherwise, pass a builder
+  /// to insert at std::next(Inst).
+  ///
+  /// The intention is that this abstraction will enable the compiler writer to
+  /// ignore whether or not an apply site is a terminator when inserting
+  /// instructions after an apply site. This results in eliminating unnecessary
+  /// if-else code otherwise required to handle such situations.
+  ///
+  /// NOTE: We pass std::next() for begin_apply. If one wishes to insert code
+  /// /after/ the end_apply/abort_apply, please use instead
+  /// insertAfterFullEvaluation.
+  void insertAfterInvocation(function_ref<void(SILBuilder &)> func) const;
+
+  /// Pass a builder with insertion points that are guaranteed to be immediately
+  /// after this full apply site has completely finished executing.
+  ///
+  /// This is just like insertAfterInvocation except that if the full apply site
+  /// is a begin_apply, we pass the insertion points after the end_apply,
+  /// abort_apply rather than an insertion point right after the
+  /// begin_apply. For such functionality, please invoke insertAfterInvocation.
+  void insertAfterFullEvaluation(function_ref<void(SILBuilder &)> func) const;
+
   /// Return whether the given apply is of a formally-throwing function
   /// which is statically known not to throw.
   bool isNonThrowing() const {
@@ -660,29 +683,6 @@ class FullApplySite : public ApplySite {
     llvm_unreachable("Covered switch isn't covered?!");
   }
 
-  /// If this is a terminator apply site, then pass a builder to insert at the
-  /// first instruction of each successor to \p func. Otherwise, pass a builder
-  /// to insert at std::next(Inst).
-  ///
-  /// The intention is that this abstraction will enable the compiler writer to
-  /// ignore whether or not an apply site is a terminator when inserting
-  /// instructions after an apply site. This results in eliminating unnecessary
-  /// if-else code otherwise required to handle such situations.
-  ///
-  /// NOTE: We pass std::next() for begin_apply. If one wishes to insert code
-  /// /after/ the end_apply/abort_apply, please use instead
-  /// insertAfterFullEvaluation.
-  void insertAfterInvocation(function_ref<void(SILBuilder &)> func) const;
-
-  /// Pass a builder with insertion points that are guaranteed to be immediately
-  /// after this full apply site has completely finished executing.
-  ///
-  /// This is just like insertAfterInvocation except that if the full apply site
-  /// is a begin_apply, we pass the insertion points after the end_apply,
-  /// abort_apply rather than an insertion point right after the
-  /// begin_apply. For such functionality, please invoke insertAfterInvocation.
-  void insertAfterFullEvaluation(function_ref<void(SILBuilder &)> func) const;
-
   /// Returns true if \p op is an operand that passes an indirect
   /// result argument to the apply site.
   bool isIndirectResultOperand(const Operand &op) const {

include/swift/SIL/SILFunctionConventions.h

@@ -256,6 +256,11 @@ class SILFunctionConventions {
                                     : funcTy->getNumResults();
   }
 
+  /// Like getNumDirectSILResults but @out tuples, which are not flattened in
+  /// the return type, are recursively flattened in the count.
+  template <bool _ = false>
+  unsigned getNumExpandedDirectSILResults(TypeExpansionContext context) const;
+
   struct DirectSILResultFilter {
     bool loweredAddresses;
     DirectSILResultFilter(bool loweredAddresses)
@@ -429,6 +434,34 @@ SILFunctionConventions::getDirectSILResultTypes(
                          SILResultTypeFunc(*this, context));
 }
 
+template <bool _>
+unsigned SILFunctionConventions::getNumExpandedDirectSILResults(
+    TypeExpansionContext context) const {
+  if (silConv.loweredAddresses)
+    return funcTy->getNumDirectFormalResults();
+  unsigned retval = 0;
+  // Worklist of elements to flatten or count.
+  SmallVector<SILType, 4> flattenedElements;
+  // Seed the worklist with the direct results.
+  for (auto result : getDirectSILResultTypes(context)) {
+    flattenedElements.push_back(result);
+  }
+  // Pop elements from the worklist and either increment the count or, if the
+  // type is a tuple, add its elements to the worklist.
+  while (!flattenedElements.empty()) {
+    auto ty = flattenedElements.pop_back_val();
+    if (auto tupleType = ty.getASTType()->getAs<TupleType>()) {
+      for (auto index :
+           llvm::reverse(indices(tupleType->getElementTypes()))) {
+        flattenedElements.push_back(ty.getTupleElementType(index));
+      }
+    } else {
+      retval += 1;
+    }
+  }
+  return retval;
+}
+
 struct SILFunctionConventions::SILParameterTypeFunc {
   SILFunctionConventions silConv;
   TypeExpansionContext context;

lib/SIL/IR/ApplySite.cpp

@@ -16,17 +16,18 @@
 
 using namespace swift;
 
-void FullApplySite::insertAfterInvocation(function_ref<void(SILBuilder &)> func) const {
+void ApplySite::insertAfterInvocation(function_ref<void(SILBuilder &)> func) const {
   SILBuilderWithScope::insertAfter(getInstruction(), func);
 }
 
-void FullApplySite::insertAfterFullEvaluation(
+void ApplySite::insertAfterFullEvaluation(
     function_ref<void(SILBuilder &)> func) const {
   switch (getKind()) {
-  case FullApplySiteKind::ApplyInst:
-  case FullApplySiteKind::TryApplyInst:
+  case ApplySiteKind::ApplyInst:
+  case ApplySiteKind::TryApplyInst:
+  case ApplySiteKind::PartialApplyInst:
     return insertAfterInvocation(func);
-  case FullApplySiteKind::BeginApplyInst:
+  case ApplySiteKind::BeginApplyInst:
     SmallVector<EndApplyInst *, 2> endApplies;
     SmallVector<AbortApplyInst *, 2> abortApplies;
     auto *bai = cast<BeginApplyInst>(getInstruction());

lib/SILGen/SILGenEpilog.cpp

@@ -10,10 +10,12 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "ASTVisitor.h"
 #include "SILGen.h"
 #include "SILGenFunction.h"
-#include "ASTVisitor.h"
+#include "swift/AST/Types.h"
 #include "swift/SIL/SILArgument.h"
+#include "llvm/ADT/STLExtras.h"
 
 using namespace swift;
 using namespace Lowering;
@@ -33,7 +35,28 @@ void SILGenFunction::prepareEpilog(Optional<Type> directResultType,
       for (auto directResult : fnConv.getDirectSILResults()) {
         SILType resultType = F.getLoweredType(F.mapTypeIntoContext(
             fnConv.getSILType(directResult, getTypeExpansionContext())));
-        epilogBB->createPhiArgument(resultType, OwnershipKind::Owned);
+        // @out tuples do not get flattened in the function's return type, but
+        // the epilog block expects (recursively) flattened arguments. Flatten
+        // the type now.
+        SmallVector<SILType, 4> worklist;
+        worklist.push_back(resultType);
+        while (!worklist.empty()) {
+          auto ty = worklist.pop_back_val();
+          if (auto tupleType = ty.getASTType()->getAs<TupleType>()) {
+            assert(!fnConv.useLoweredAddresses() &&
+                   "expanding tuple in non-opaque-values exit block?!");
+            // Push tuple elements in reverse order (resulting in later tuple
+            // elements appearing earlier in worklist) so that as the worklist
+            // is drained by popping the back, arguments are created for the
+            // earlier types first.
+            for (auto index :
+                 llvm::reverse(indices(tupleType->getElementTypes()))) {
+              worklist.push_back(ty.getTupleElementType(index));
+            }
+          } else {
+            epilogBB->createPhiArgument(ty, OwnershipKind::Owned);
+          }
+        }
       }
     }
   }
@@ -61,15 +84,84 @@ void SILGenFunction::prepareCoroutineUnwindEpilog(CleanupLocation cleanupLoc) {
   CoroutineUnwindDest = JumpDest(unwindBB, getCleanupsDepth(), cleanupLoc);
 }
 
+/// View a given SILType as a type-tree under the operation of tupling and visit
+/// its nodes (tuple elements) in post-order.
+///
+/// For convenience, the index of the type in the flattened tuple is passed to
+/// the visitor.
+template <typename Visit>
+void visitTupleTypeTreeInPostOrder(SILType root, Visit visit) {
+  struct Node {
+    SILType ty;
+    unsigned index;
+  };
+  SmallVector<std::pair<Node, unsigned>, 32> stack;
+  auto tupleElementCount = [](SILType ty) -> unsigned {
+    if (auto tupleType = ty.getASTType()->getAs<TupleType>())
+      return tupleType->getNumElements();
+    return 0;
+  };
+  auto tupleElement = [](SILType ty, unsigned index) -> SILType {
+    return ty.getTupleElementType(index);
+  };
+  unsigned flattenedIndex = 0;
+  stack.push_back({{root, flattenedIndex}, 0});
+  while (!stack.empty()) {
+    while (stack.back().second != tupleElementCount(stack.back().first.ty)) {
+      auto index = stack.back().second;
+      stack.back().second++;
+      stack.push_back(
+          {{tupleElement(stack.back().first.ty, index), flattenedIndex}, 0});
+    }
+    auto node = stack.pop_back_val().first;
+    visit(node.ty, node.index);
+    if (!node.ty.getASTType()->template is<TupleType>())
+      flattenedIndex += 1;
+  }
+}
+
 /// Given a list of direct results, form the direct result value.
 ///
 /// Note that this intentionally loses any tuple sub-structure of the
-/// formal result type.
+/// formal result type, except in the case of @out tuples where it must be
+/// preserved.
 static SILValue buildReturnValue(SILGenFunction &SGF, SILLocation loc,
                                  ArrayRef<SILValue> directResults) {
   if (directResults.size() == 1)
     return directResults[0];
 
+  auto fnConv = SGF.F.getConventions();
+  if (!fnConv.useLoweredAddresses()) {
+    // In opaque-values code, nested @out enums are not flattened.  Reconstruct
+    // nested tuples.
+    auto resultType = SGF.F.getLoweredType(SGF.F.mapTypeIntoContext(
+        fnConv.getSILResultType(SGF.getTypeExpansionContext())));
+    SmallVector<Optional<SILValue>, 4> mutableDirectResult;
+    for (auto result : directResults) {
+      mutableDirectResult.push_back({result});
+    }
+    visitTupleTypeTreeInPostOrder(resultType, [&](SILType ty, unsigned index) {
+      if (auto tupleTy = ty.getASTType()->getAs<TupleType>()) {
+        SmallVector<SILValue, 4> elements;
+        unsigned offset = 0;
+        auto elementCount = tupleTy->getNumElements();
+        while (elements.size() < elementCount) {
+          if (mutableDirectResult[index + offset].hasValue()) {
+            auto val = mutableDirectResult[index + offset].getValue();
+            elements.push_back(val);
+            mutableDirectResult[index + offset].reset();
+          }
+          ++offset;
+        }
+        assert(!mutableDirectResult[index].hasValue());
+        auto tuple = SGF.B.createTuple(loc, ty, elements);
+        mutableDirectResult[index] = tuple;
+      }
+    });
+    assert(mutableDirectResult[0].hasValue());
+    return mutableDirectResult[0].getValue();
+  }
+
   SmallVector<TupleTypeElt, 4> eltTypes;
   for (auto elt : directResults)
     eltTypes.push_back(elt->getType().getASTType());
@@ -195,7 +287,9 @@ SILGenFunction::emitEpilogBB(SILLocation topLevel) {
   // block.
   SILValue returnValue;
   if (!directResults.empty()) {
-    assert(directResults.size() == F.getConventions().getNumDirectSILResults());
+    assert(directResults.size() ==
+           F.getConventions().getNumExpandedDirectSILResults(
+               getTypeExpansionContext()));
     returnValue = buildReturnValue(*this, topLevel, directResults);
   }
 

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -1714,13 +1714,13 @@ namespace {
 /// object arguments with address-type arguments.
 class CallArgRewriter {
   AddressLoweringState &pass;
-  FullApplySite apply;
+  ApplySite apply;
   SILLocation callLoc;
   SILBuilder argBuilder;
   AddressMaterialization addrMat;
 
 public:
-  CallArgRewriter(FullApplySite apply, AddressLoweringState &pass)
+  CallArgRewriter(ApplySite apply, AddressLoweringState &pass)
       : pass(pass), apply(apply), callLoc(apply.getLoc()),
         argBuilder(pass.getBuilder(apply.getInstruction()->getIterator())),
         addrMat(pass, argBuilder) {}
@@ -2599,6 +2599,10 @@ class UseRewriter : SILInstructionVisitor<UseRewriter> {
     CallArgRewriter(applyInst, pass).rewriteIndirectArgument(use);
   }
 
+  void visitPartialApplyInst(PartialApplyInst *pai) {
+    CallArgRewriter(pai, pass).rewriteIndirectArgument(use);
+  }
+
   void visitBeginApplyInst(BeginApplyInst *bai) {
     CallArgRewriter(bai, pass).rewriteIndirectArgument(use);
   }

test/SILGen/opaque_values_silgen.swift

@@ -478,3 +478,24 @@ enum TestEnum<T> {
   }
 }
 
+// Verify exit block arguments are ordered correctly.
+// 
+// CHECK-LABEL: sil private [ossa] @$s20opaque_values_silgen19duplicate_with_int49condition5valueSi_S2ix_xttSb_xtlFSi_S2ix_xttyXEfU_ : {{.*}} {
+// CHECK:         br [[EPILOG:bb[0-9]+]]({{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : $Value, {{%[^,]+}} : $Value)
+// CHECK:         br [[EPILOG]]({{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : $Value, {{%[^,]+}} : $Value)
+// CHECK:       [[EPILOG]]({{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : $Int, {{%[^,]+}} : @owned $Value, {{%[^,]+}} : @owned $Value):
+// CHECK-LABEL: } // end sil function '$s20opaque_values_silgen19duplicate_with_int49condition5valueSi_S2ix_xttSb_xtlFSi_S2ix_xttyXEfU_'
+func doit<T>(_ f: () -> T) -> T {
+  f()
+}
+@_silgen_name("duplicate_with_int4")
+func duplicate_with_int4<Value>(condition: Bool, value: Value) -> (Int, Int, Int, (Value, Value)) {
+  doit {
+    if condition {
+      return (Int(), Int(), Int(), (value, value))
+    } else {
+      return (Int(), Int(), Int(), (value, value))
+    }
+  }
+}
+