Merge pull request swiftlang#27904 from gottesmm/pr-624292f32572e0362c0566deebe886758e730510

[sil] Refactor out BeginApplyInst::getCoroutineEndPoints() impl from Inliner into method on the class itself.

Author: gottesmm

include/swift/SIL/SILInstruction.h

@@ -2226,6 +2226,9 @@ class BeginApplyResult final : public MultipleValueInstructionResult {
   }
 };
 
+class EndApplyInst;
+class AbortApplyInst;
+
 /// BeginApplyInst - Represents the beginning of the full application of
 /// a yield_once coroutine (up until the coroutine yields a value back).
 class BeginApplyInst final
@@ -2280,6 +2283,13 @@ class BeginApplyInst final
   bool isNonThrowing() const {
     return isNonThrowingApply();
   }
+
+  void getCoroutineEndPoints(
+      SmallVectorImpl<EndApplyInst *> &endApplyInsts,
+      SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const;
+
+  void getCoroutineEndPoints(SmallVectorImpl<Operand *> &endApplyInsts,
+                             SmallVectorImpl<Operand *> &abortApplyInsts) const;
 };
 
 inline BeginApplyInst *BeginApplyResult::getParent() {

lib/SIL/SILInstructions.cpp

@@ -499,6 +499,35 @@ BeginApplyInst::create(SILDebugLocation loc, SILValue callee,
                                       isNonThrowing, specializationInfo);
 }
 
+void BeginApplyInst::getCoroutineEndPoints(
+    SmallVectorImpl<EndApplyInst *> &endApplyInsts,
+    SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const {
+  for (auto *tokenUse : getTokenResult()->getUses()) {
+    auto *user = tokenUse->getUser();
+    if (auto *end = dyn_cast<EndApplyInst>(user)) {
+      endApplyInsts.push_back(end);
+      continue;
+    }
+
+    abortApplyInsts.push_back(cast<AbortApplyInst>(user));
+  }
+}
+
+void BeginApplyInst::getCoroutineEndPoints(
+    SmallVectorImpl<Operand *> &endApplyInsts,
+    SmallVectorImpl<Operand *> &abortApplyInsts) const {
+  for (auto *tokenUse : getTokenResult()->getUses()) {
+    auto *user = tokenUse->getUser();
+    if (isa<EndApplyInst>(user)) {
+      endApplyInsts.push_back(tokenUse);
+      continue;
+    }
+
+    assert(isa<AbortApplyInst>(user));
+    abortApplyInsts.push_back(tokenUse);
+  }
+}
+
 bool swift::doesApplyCalleeHaveSemantics(SILValue callee, StringRef semantics) {
   if (auto *FRI = dyn_cast<FunctionRefBaseInst>(callee))
     if (auto *F = FRI->getReferencedFunctionOrNull())

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -72,7 +72,8 @@ bool SILInliner::canInlineApplySite(FullApplySite apply) {
   return true;
 }
 
-namespace swift {
+namespace {
+
 /// Utility class for rewiring control-flow of inlined begin_apply functions.
 class BeginApplySite {
   SILLocation Loc;
@@ -102,14 +103,15 @@ class BeginApplySite {
   }
 
   void preprocess(SILBasicBlock *returnToBB) {
-    // Get the end_apply, abort_apply instructions.
-    auto Token = BeginApply->getTokenResult();
-    for (auto *TokenUse : Token->getUses()) {
-      if (auto End = dyn_cast<EndApplyInst>(TokenUse->getUser())) {
-        collectEndApply(End);
-      } else {
-        collectAbortApply(cast<AbortApplyInst>(TokenUse->getUser()));
-      }
+    SmallVector<EndApplyInst *, 1> endApplyInsts;
+    SmallVector<AbortApplyInst *, 1> abortApplyInsts;
+    BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts);
+    while (!endApplyInsts.empty()) {
+      auto *endApply = endApplyInsts.pop_back_val();
+      collectEndApply(endApply);
+    }
+    while (!abortApplyInsts.empty()) {
+      collectAbortApply(abortApplyInsts.pop_back_val());
     }
   }
 
@@ -228,7 +230,8 @@ class BeginApplySite {
     assert(!BeginApply->hasUsesOfAnyResult());
   }
 };
-} // namespace swift
+
+} // end anonymous namespace
 
 namespace swift {
 class SILInlineCloner