[closure-spec] Do not try to process begin_apply. It is not supported now.

include/swift/SIL/SILInstruction.h

@@ -7554,6 +7554,10 @@ class ApplySite {
     }
   }
 
+  ApplySiteKind getKind() const {
+    return ApplySiteKind(Inst->getKind());
+  }
+
   explicit operator bool() const {
     return Inst != nullptr;
   }
@@ -7897,6 +7901,10 @@ class FullApplySite : public ApplySite {
     }
   }
 
+  FullApplySiteKind getKind() const {
+    return FullApplySiteKind(getInstruction()->getKind());
+  }
+
   bool hasIndirectSILResults() const {
     return getSubstCalleeConv().hasIndirectSILResults();
   }

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -410,21 +410,29 @@ static void rewriteApplyInst(const CallSiteDescriptor &CSDesc,
 
   Builder.setInsertionPoint(AI.getInstruction());
   FullApplySite NewAI;
-  if (auto *TAI = dyn_cast<TryApplyInst>(AI)) {
+  switch (AI.getKind()) {
+  case FullApplySiteKind::TryApplyInst: {
+    auto *TAI = cast<TryApplyInst>(AI);
     NewAI = Builder.createTryApply(AI.getLoc(), FRI,
                                    SubstitutionMap(), NewArgs,
                                    TAI->getNormalBB(), TAI->getErrorBB());
     // If we passed in the original closure as @owned, then insert a release
     // right after NewAI. This is to balance the +1 from being an @owned
     // argument to AI.
-    if (CSDesc.isClosureConsumed() && CSDesc.closureHasRefSemanticContext()) {
-      Builder.setInsertionPoint(TAI->getNormalBB()->begin());
-      Builder.createReleaseValue(Closure->getLoc(), Closure, Builder.getDefaultAtomicity());
-      Builder.setInsertionPoint(TAI->getErrorBB()->begin());
-      Builder.createReleaseValue(Closure->getLoc(), Closure, Builder.getDefaultAtomicity());
-      Builder.setInsertionPoint(AI.getInstruction());
+    if (!CSDesc.isClosureConsumed() || !CSDesc.closureHasRefSemanticContext()) {
+      break;
     }
-  } else {
+
+    Builder.setInsertionPoint(TAI->getNormalBB()->begin());
+    Builder.createReleaseValue(Closure->getLoc(), Closure,
+                               Builder.getDefaultAtomicity());
+    Builder.setInsertionPoint(TAI->getErrorBB()->begin());
+    Builder.createReleaseValue(Closure->getLoc(), Closure,
+                               Builder.getDefaultAtomicity());
+    Builder.setInsertionPoint(AI.getInstruction());
+    break;
+  }
+  case FullApplySiteKind::ApplyInst: {
     auto oldApply = cast<ApplyInst>(AI);
     auto newApply = Builder.createApply(oldApply->getLoc(), FRI,
                                         SubstitutionMap(),
@@ -438,6 +446,10 @@ static void rewriteApplyInst(const CallSiteDescriptor &CSDesc,
 
     // Replace all uses of the old apply with the new apply.
     oldApply->replaceAllUsesWith(newApply);
+    break;
+  }
+  case FullApplySiteKind::BeginApplyInst:
+    llvm_unreachable("Unhandled case");
   }
 
   // Erase the old apply.
@@ -937,6 +949,16 @@ static bool isClosureAppliedIn(SILFunction *Callee, unsigned closureArgIdx,
   return false;
 }
 
+static bool canSpecializeFullApplySite(FullApplySiteKind kind) {
+  switch (kind) {
+  case FullApplySiteKind::TryApplyInst:
+  case FullApplySiteKind::ApplyInst:
+    return true;
+  case FullApplySiteKind::BeginApplyInst:
+    return false;
+  }
+}
+
 bool SILClosureSpecializerTransform::gatherCallSites(
     SILFunction *Caller,
     llvm::SmallVectorImpl<std::unique_ptr<ClosureInfo>> &ClosureCandidates,
@@ -986,7 +1008,7 @@ bool SILClosureSpecializerTransform::gatherCallSites(
           Uses.append(CFI->getUses().begin(), CFI->getUses().end());
           continue;
         }
-        if (auto *Cvt= dyn_cast<ConvertEscapeToNoEscapeInst>(Use->getUser())) {
+        if (auto *Cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(Use->getUser())) {
           Uses.append(Cvt->getUses().begin(), Cvt->getUses().end());
           continue;
         }
@@ -1010,11 +1032,12 @@ bool SILClosureSpecializerTransform::gatherCallSites(
           continue;
         }
 
-        // If this use is not an apply inst or an apply inst with
-        // substitutions, there is nothing interesting for us to do, so
-        // continue...
+        // If this use is not a full apply site that we can process or an apply
+        // inst with substitutions, there is nothing interesting for us to do,
+        // so continue...
         auto AI = FullApplySite::isa(Use->getUser());
-        if (!AI || AI.hasSubstitutions())
+        if (!AI || AI.hasSubstitutions() ||
+            !canSpecializeFullApplySite(AI.getKind()))
           continue;
 
         // Check if we have already associated this apply inst with a closure to

test/SILOptimizer/closure_specialize.sil

@@ -732,3 +732,39 @@ bb0(%0 : $Int):
   %empty = tuple ()
   return %empty : $()
 }
+
+//////////////////////
+// Begin Apply Test //
+//////////////////////
+
+sil @coroutine_user : $@yield_once @convention(thin) (@noescape @callee_guaranteed () -> Int) -> @yields Int {
+bb0(%0 : $@noescape @callee_guaranteed () -> Int):
+  %1 = apply %0() : $@noescape @callee_guaranteed () -> Int
+  unreachable
+}
+
+// CHECK-LABEL: sil @test_coroutine_user : $@convention(thin) (Int) -> Int {
+// CHECK: [[COROUTINE_USER:%.*]] = function_ref @coroutine_user
+// CHECK: begin_apply [[COROUTINE_USER]](
+// CHECK: } // end sil function 'test_coroutine_user'
+sil @test_coroutine_user : $@convention(thin) (Int) -> Int {
+bb0(%0 : $Int):
+  %1 = function_ref @testClosureConvertHelper2 : $@convention(thin) (Int) -> Int
+  %2 = partial_apply [callee_guaranteed] %1(%0) : $@convention(thin) (Int) -> Int
+  %3 = convert_escape_to_noescape %2 : $@callee_guaranteed () -> Int to $@noescape @callee_guaranteed () -> Int
+  %4 = function_ref @coroutine_user : $@yield_once @convention(thin) (@noescape @callee_guaranteed () -> Int) -> @yields Int
+  (%value, %token) = begin_apply %4(%3) : $@yield_once @convention(thin) (@noescape @callee_guaranteed () -> Int) -> @yields Int
+  cond_br undef, bb1, bb2
+
+bb1:
+  end_apply %token
+  br bb3
+
+bb2:
+  abort_apply %token
+  br bb3
+
+bb3:
+  release_value %2 : $@callee_guaranteed () -> Int
+  return %value : $Int
+}