SILBuilder: add an API to insert _after_ an instruction.

... and use that API in FullApplySite::insertAfterInvocation.

Also change FullApplySite::insertAfterInvocation/insertAfterFullEvaluation to directly pass a SILBuilder instead of just an insertion point to the callback.
This makes more sense (given the function names) and simplifies the usages.

It's a NFC.

include/swift/SIL/ApplySite.h

@@ -585,62 +585,28 @@ class FullApplySite : public ApplySite {
     llvm_unreachable("Covered switch isn't covered?!");
   }
 
-  /// If this is a terminator apply site, then pass the first instruction of
-  /// each successor to fun. Otherwise, pass std::next(Inst).
+  /// 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 return std::next() for begin_apply. If one wishes to insert code
+  /// 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(SILBasicBlock::iterator)> func) const {
-    switch (getKind()) {
-    case FullApplySiteKind::ApplyInst:
-    case FullApplySiteKind::BeginApplyInst:
-      return func(std::next(getInstruction()->getIterator()));
-    case FullApplySiteKind::TryApplyInst:
-      for (auto *succBlock :
-           cast<TermInst>(getInstruction())->getSuccessorBlocks()) {
-        func(succBlock->begin());
-      }
-      return;
-    }
-    llvm_unreachable("Covered switch isn't covered");
-  }
+  void insertAfterInvocation(function_ref<void(SILBuilder &)> func) const;
 
-  /// Pass to func insertion points that are guaranteed to be immediately after
-  /// this full apply site has completely finished executing.
+  /// 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(SILBasicBlock::iterator)> func) const {
-    switch (getKind()) {
-    case FullApplySiteKind::ApplyInst:
-    case FullApplySiteKind::TryApplyInst:
-      return insertAfterInvocation(func);
-    case FullApplySiteKind::BeginApplyInst:
-      SmallVector<EndApplyInst *, 2> endApplies;
-      SmallVector<AbortApplyInst *, 2> abortApplies;
-      auto *bai = cast<BeginApplyInst>(getInstruction());
-      bai->getCoroutineEndPoints(endApplies, abortApplies);
-      for (auto *eai : endApplies) {
-        func(std::next(eai->getIterator()));
-      }
-      for (auto *aai : abortApplies) {
-        func(std::next(aai->getIterator()));
-      }
-      return;
-    }
-
-    llvm_unreachable("covered switch isn't covered");
-  }
+  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.

include/swift/SIL/SILBuilder.h

@@ -2499,6 +2499,20 @@ class SILBuilderWithScope : public SILBuilder {
     assert(DS && "Instruction without debug scope associated!");
     setCurrentDebugScope(DS);
   }
+
+  /// If \p inst 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 \p inst is a terminator when inserting instructions
+  /// after \p inst.
+  ///
+  /// Precondition: It's the responsibility of the caller to ensure that if
+  /// \p inst is a terminator, all successor blocks have only a single
+  /// predecessor block: the parent of \p inst.
+  static void insertAfter(SILInstruction *inst,
+                          function_ref<void(SILBuilder &)> func);
 };
 
 class SavedInsertionPointRAII {

lib/SIL/IR/ApplySite.cpp

@@ -0,0 +1,46 @@
+//===--- ApplySite.cpp - Wrapper around apply instructions ----------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/SIL/ApplySite.h"
+#include "swift/SIL/SILBuilder.h"
+
+
+using namespace swift;
+
+void FullApplySite::insertAfterInvocation(function_ref<void(SILBuilder &)> func) const {
+  SILBuilderWithScope::insertAfter(getInstruction(), func);
+}
+
+void FullApplySite::insertAfterFullEvaluation(
+    function_ref<void(SILBuilder &)> func) const {
+  switch (getKind()) {
+  case FullApplySiteKind::ApplyInst:
+  case FullApplySiteKind::TryApplyInst:
+    return insertAfterInvocation(func);
+  case FullApplySiteKind::BeginApplyInst:
+    SmallVector<EndApplyInst *, 2> endApplies;
+    SmallVector<AbortApplyInst *, 2> abortApplies;
+    auto *bai = cast<BeginApplyInst>(getInstruction());
+    bai->getCoroutineEndPoints(endApplies, abortApplies);
+    for (auto *eai : endApplies) {
+      SILBuilderWithScope builder(std::next(eai->getIterator()));
+      func(builder);
+    }
+    for (auto *aai : abortApplies) {
+      SILBuilderWithScope builder(std::next(aai->getIterator()));
+      func(builder);
+    }
+    return;
+  }
+  llvm_unreachable("covered switch isn't covered");
+}
+

lib/SIL/IR/CMakeLists.txt

@@ -1,5 +1,6 @@
 target_sources(swiftSIL PRIVATE
   AbstractionPattern.cpp
+  ApplySite.cpp
   Bridging.cpp
   Linker.cpp
   Notifications.cpp

lib/SIL/IR/SILBuilder.cpp

@@ -661,3 +661,19 @@ CheckedCastBranchInst *SILBuilder::createCheckedCastBranch(
       destLoweredTy, destFormalTy, successBB, failureBB,
       getFunction(), C.OpenedArchetypes, target1Count, target2Count));
 }
+
+void SILBuilderWithScope::insertAfter(SILInstruction *inst,
+                                      function_ref<void(SILBuilder &)> func) {
+  if (isa<TermInst>(inst)) {
+    for (const SILSuccessor &succ : inst->getParent()->getSuccessors()) {
+      SILBasicBlock *succBlock = succ;
+      assert(succBlock->getSinglePredecessorBlock() == inst->getParent() &&
+             "the terminator instruction must not have critical successors");
+      SILBuilderWithScope builder(succBlock->begin());
+      func(builder);
+    }
+  } else {
+    SILBuilderWithScope builder(std::next(inst->getIterator()));
+    func(builder);
+  }
+}

lib/SILOptimizer/Mandatory/MandatoryInlining.cpp

@@ -123,9 +123,9 @@ static  bool fixupReferenceCounts(
           });
 
       if (!consumedInLoop) {
-        applySite.insertAfterInvocation([&](SILBasicBlock::iterator iter) {
-          SILBuilderWithScope(iter).createDestroyAddr(loc, stackLoc);
-          SILBuilderWithScope(iter).createDeallocStack(loc, stackLoc);
+        applySite.insertAfterInvocation([&](SILBuilder &builder) {
+          builder.createDestroyAddr(loc, stackLoc);
+          builder.createDeallocStack(loc, stackLoc);
         });
       }
       v = stackLoc;
@@ -176,11 +176,11 @@ static  bool fixupReferenceCounts(
       // uses in the top of a diamond and need to insert a destroy after the
       // apply since the leak will just cover the other path.
       if (!consumedInLoop) {
-        applySite.insertAfterInvocation([&](SILBasicBlock::iterator iter) {
+        applySite.insertAfterInvocation([&](SILBuilder &builder) {
           if (hasOwnership) {
-            SILBuilderWithScope(iter).createEndBorrow(loc, argument);
+            builder.createEndBorrow(loc, argument);
           }
-          SILBuilderWithScope(iter).emitDestroyValueOperation(loc, copy);
+          builder.emitDestroyValueOperation(loc, copy);
         });
       }
       v = argument;
@@ -217,8 +217,8 @@ static  bool fixupReferenceCounts(
 
       // Then insert destroys after the apply site since our value is not being
       // consumed as part of the actual apply.
-      applySite.insertAfterInvocation([&](SILBasicBlock::iterator iter) {
-        SILBuilderWithScope(iter).emitDestroyValueOperation(loc, v);
+      applySite.insertAfterInvocation([&](SILBuilder &builder) {
+        builder.emitDestroyValueOperation(loc, v);
       });
       break;
     }
@@ -263,8 +263,8 @@ static  bool fixupReferenceCounts(
   // Destroy the callee as the apply would have done if our function is not
   // callee guaranteed.
   if (!isCalleeGuaranteed) {
-    applySite.insertAfterInvocation([&](SILBasicBlock::iterator iter) {
-      SILBuilderWithScope(iter).emitDestroyValueOperation(loc, calleeValue);
+    applySite.insertAfterInvocation([&](SILBuilder &builder) {
+      builder.emitDestroyValueOperation(loc, calleeValue);
     });
   }
   return invalidatedStackNesting;

lib/SILOptimizer/Utils/Generics.cpp

@@ -2027,8 +2027,7 @@ static ApplySite replaceWithSpecializedCallee(ApplySite applySite,
     assert(resultBlock->getSinglePredecessorBlock() == tai->getParent());
     // First insert the cleanups for our arguments int he appropriate spot.
     FullApplySite(tai).insertAfterFullEvaluation(
-        [&](SILBasicBlock::iterator insertPt) {
-          SILBuilderWithScope argBuilder(insertPt);
+        [&](SILBuilder &argBuilder) {
           cleanupCallArguments(argBuilder, loc, arguments,
                                argsNeedingEndBorrow);
         });
@@ -2052,8 +2051,7 @@ static ApplySite replaceWithSpecializedCallee(ApplySite applySite,
   case ApplySiteKind::ApplyInst: {
     auto *ai = cast<ApplyInst>(applySite);
     FullApplySite(ai).insertAfterFullEvaluation(
-        [&](SILBasicBlock::iterator insertPt) {
-          SILBuilderWithScope argBuilder(insertPt);
+        [&](SILBuilder &argBuilder) {
           cleanupCallArguments(argBuilder, loc, arguments,
                                argsNeedingEndBorrow);
         });
@@ -2082,8 +2080,7 @@ static ApplySite replaceWithSpecializedCallee(ApplySite applySite,
     auto *bai = cast<BeginApplyInst>(applySite);
     assert(!resultOut);
     FullApplySite(bai).insertAfterFullEvaluation(
-        [&](SILBasicBlock::iterator insertPt) {
-          SILBuilderWithScope argBuilder(insertPt);
+        [&](SILBuilder &argBuilder) {
           cleanupCallArguments(argBuilder, loc, arguments,
                                argsNeedingEndBorrow);
         });
@@ -2227,8 +2224,7 @@ SILFunction *ReabstractionThunkGenerator::createThunk() {
 
   // Now that we have finished constructing our CFG (note the return above),
   // insert any compensating end borrows that we need.
-  ApplySite.insertAfterFullEvaluation([&](SILBasicBlock::iterator insertPt) {
-    SILBuilderWithScope argBuilder(insertPt);
+  ApplySite.insertAfterFullEvaluation([&](SILBuilder &argBuilder) {
     cleanupCallArguments(argBuilder, Loc, Arguments, ArgsThatNeedEndBorrow);
   });
 

lib/SILOptimizer/Utils/PartialApplyCombiner.cpp

@@ -153,8 +153,7 @@ void PartialApplyCombiner::processSingleApply(FullApplySite paiAI) {
         auto *ASI = builder.createAllocStack(pai->getLoc(), arg->getType());
         builder.createCopyAddr(pai->getLoc(), arg, ASI, IsTake_t::IsNotTake,
                                IsInitialization_t::IsInitialization);
-        paiAI.insertAfterFullEvaluation([&](SILBasicBlock::iterator insertPt) {
-          SILBuilderWithScope builder(insertPt);
+        paiAI.insertAfterFullEvaluation([&](SILBuilder &builder) {
           builder.createDeallocStack(destroyloc, ASI);
         });
         arg = ASI;
@@ -184,8 +183,7 @@ void PartialApplyCombiner::processSingleApply(FullApplySite paiAI) {
   // We also need to destroy the partial_apply instruction itself because it is
   // consumed by the apply_instruction.
   if (!pai->hasCalleeGuaranteedContext()) {
-    paiAI.insertAfterFullEvaluation([&](SILBasicBlock::iterator insertPt) {
-      SILBuilderWithScope builder(insertPt);
+    paiAI.insertAfterFullEvaluation([&](SILBuilder &builder) {
       builder.emitDestroyValueOperation(destroyloc, pai);
     });
   }