Fix a stack-discipline bug with inlining coroutines.

Author: rjmccall

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -20,6 +20,25 @@
 
 namespace swift {
 
+struct InstructionMatchResult {
+  /// The instruction matches.
+  bool matches;
+
+  /// The search should be halted.
+  bool halt;
+};
+
+using InstructionMatcher =
+  llvm::function_ref<InstructionMatchResult(SILInstruction *i)>;
+
+/// Given a predicate defining interesting instructions, check whether
+/// the stack is different at any of them from how it stood at the start
+/// of the search.
+///
+/// The matcher function must halt the search before any edges which would
+/// lead back to before the start point.
+bool hasStackDifferencesAt(SILInstruction *start, InstructionMatcher matcher);
+
 /// A utility to correct the nesting of stack allocating/deallocating
 /// instructions.
 ///

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -17,10 +17,54 @@
 #include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/Utils/CFG.h"
 #include "swift/SILOptimizer/Utils/SILOptFunctionBuilder.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 using namespace swift;
 
+/// Does the given coroutine make any stack allocations that are live across
+/// its yields?
+static bool allocatesStackAcrossYields(SILFunction *F) {
+  assert(F->getLoweredFunctionType()->isCoroutine());
+
+  return hasStackDifferencesAt(&F->getEntryBlock()->front(),
+                               [](SILInstruction *i) -> InstructionMatchResult {
+    if (isa<YieldInst>(i)) {
+      return {
+        /*matches*/ true,
+        /*halt*/ i->getFunction()->getLoweredFunctionType()->getCoroutineKind()
+                       == SILCoroutineKind::YieldOnce
+      };
+    }
+
+    // Otherwise, search until the end of the function.
+    return { false, false };
+  });
+}
+
+static bool isEndOfApply(SILInstruction *i, BeginApplyInst *beginApply) {
+  if (auto endApply = dyn_cast<EndApplyInst>(i)) {
+    return endApply->getBeginApply() == beginApply;
+  } else if (auto abortApply = dyn_cast<AbortApplyInst>(i)) {
+    return abortApply->getBeginApply() == beginApply;
+  } else {
+    return false;
+  }
+}
+
+/// Are there any stack differences from the given begin_apply to any
+/// corresponding end_apply/abort_apply?
+static bool hasStackDifferencesAtEnds(BeginApplyInst *apply) {
+  return hasStackDifferencesAt(apply, [apply](SILInstruction *i)
+                                                     -> InstructionMatchResult {
+    // Search for ends of the original apply.  We can stop searching
+    // at these points.
+    if (isEndOfApply(i, apply))
+      return { true, true };
+    return { false, false };
+  });
+}
+
 static bool canInlineBeginApply(BeginApplyInst *BA) {
   // Don't inline if we have multiple resumption sites (i.e. end_apply or
   // abort_apply instructions).  The current implementation clones a single
@@ -75,6 +119,7 @@ class BeginApplySite {
   SILBuilder *Builder;
   BeginApplyInst *BeginApply;
   bool HasYield = false;
+  bool NeedsStackCorrection;
 
   EndApplyInst *EndApply = nullptr;
   SILBasicBlock *EndApplyBB = nullptr;
@@ -86,15 +131,31 @@ class BeginApplySite {
 
 public:
   BeginApplySite(BeginApplyInst *BeginApply, SILLocation Loc,
-                 SILBuilder *Builder)
-      : Loc(Loc), Builder(Builder), BeginApply(BeginApply) {}
+                 SILBuilder *Builder, bool NeedsStackCorrection)
+      : Loc(Loc), Builder(Builder), BeginApply(BeginApply),
+        NeedsStackCorrection(NeedsStackCorrection) {}
 
   static Optional<BeginApplySite> get(FullApplySite AI, SILLocation Loc,
                                       SILBuilder *Builder) {
     auto *BeginApply = dyn_cast<BeginApplyInst>(AI);
     if (!BeginApply)
       return None;
-    return BeginApplySite(BeginApply, Loc, Builder);
+
+    // We need stack correction if there are both:
+    //   - stack allocations in the callee that are live across the yield and
+    //   - stack differences in the caller from the begin_apply to any
+    //     end_apply or abort_apply.
+    // In these cases, naive cloning will cause the allocations to become
+    // improperly nested.
+    //
+    // We need to compute this here before we do any splitting in the parent
+    // function.
+    bool NeedsStackCorrection = false;
+    if (allocatesStackAcrossYields(BeginApply->getReferencedFunction()) &&
+        hasStackDifferencesAtEnds(BeginApply))
+      NeedsStackCorrection = true;
+
+    return BeginApplySite(BeginApply, Loc, Builder, NeedsStackCorrection);
   }
 
   void preprocess(SILBasicBlock *returnToBB) {
@@ -222,6 +283,11 @@ class BeginApplySite {
       AbortApply->eraseFromParent();
 
     assert(!BeginApply->hasUsesOfAnyResult());
+
+    // Correct the stack if necessary.
+    if (NeedsStackCorrection) {
+      StackNesting().correctStackNesting(BeginApply->getFunction());
+    }
   }
 };
 } // namespace swift

lib/SILOptimizer/Utils/StackNesting.cpp

@@ -18,6 +18,75 @@
 
 using namespace swift;
 
+bool swift::hasStackDifferencesAt(SILInstruction *start,
+                                  InstructionMatcher matcher) {
+  struct StackStatus {
+    /// The number of currently-active stack allocations on this path.
+    /// Because of the stack discipline, we only have to maintain a count.
+    unsigned depth;
+
+    /// Whether we've popped anything off of the stack that was active
+    /// at the start of the search.
+    bool hasPops;
+
+    bool hasDifferences() const { return hasPops || depth != 0; }
+  };
+
+  SmallPtrSet<SILBasicBlock*, 8> visited;
+  SmallVector<std::pair<SILInstruction *, StackStatus>, 8> worklist;
+  worklist.push_back({start, {0, false}});
+
+  while (!worklist.empty()) {
+    // Pull a block and depth off the worklist.
+    // Visitation order doesn't matter.
+    auto pair = worklist.pop_back_val();
+    auto status = pair.second;
+    auto firstInst = pair.first;
+    auto block = firstInst->getParent();
+
+    for (SILBasicBlock::iterator i(firstInst), e = block->end(); i != e; ++i) {
+      auto match = matcher(&*i);
+
+      // If this is an interesting instruction, check for stack
+      // differences here.
+      if (match.matches && status.hasDifferences())
+        return true;
+
+      // If we should halt, just go to the next work-list item, bypassing
+      // visiting the successors.
+      if (match.halt)
+        goto nextWorkListItem; // a labelled continue would be nice
+
+      // Otherwise, do stack-depth bookkeeping.
+      if (i->isAllocatingStack()) {
+        status.depth++;
+      } else if (i->isDeallocatingStack()) {
+        // If the stack depth is already zero, we're deallocating a stack
+        // allocation that was live into the search.
+        if (status.depth > 0) {
+          status.depth--;
+        } else {
+          status.hasPops = true;
+        }
+      }
+    }
+
+    // Add the successors to the worklist.
+    for (auto &succ : block->getSuccessors()) {
+      auto succBlock = succ.getBB();
+      auto insertResult = visited.insert(succBlock);
+      if (insertResult.second) {
+        worklist.push_back({&succBlock->front(), status});
+      }
+    }
+
+  nextWorkListItem: ;
+  }
+
+  return false;
+}
+
+
 void StackNesting::setup(SILFunction *F) {
   SmallVector<BlockInfo *, 8> WorkList;
   llvm::DenseMap<SILBasicBlock *, BlockInfo *> BlockMapping;

test/SILOptimizer/inline_begin_apply.sil

@@ -436,3 +436,99 @@ entry:
   %ret = tuple ()
   return %ret : $()
 }
+
+sil [transparent] @stack_overlap : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32) {
+entry:
+  %marker = function_ref @marker : $@convention(thin) (Builtin.Int32) -> ()
+  %temp = alloc_stack $Builtin.Int32
+  %1000 = integer_literal $Builtin.Int32, 1000
+  store %1000 to [trivial] %temp : $*Builtin.Int32
+  %2000 = integer_literal $Builtin.Int32, 2000
+  apply %marker(%2000) : $@convention(thin) (Builtin.Int32) -> ()
+  yield %temp : $*Builtin.Int32, resume resume, unwind unwind
+
+resume:
+  %3000 = integer_literal $Builtin.Int32, 3000
+  apply %marker(%3000) : $@convention(thin) (Builtin.Int32) -> ()
+  dealloc_stack %temp : $*Builtin.Int32
+  %4000 = integer_literal $Builtin.Int32, 4000
+  apply %marker(%4000) : $@convention(thin) (Builtin.Int32) -> ()
+  %ret = tuple ()
+  return %ret : $()
+
+unwind:
+  dealloc_stack %temp : $*Builtin.Int32
+  unwind
+}
+
+// CHECK-LABEL: sil @test_stack_overlap_dealloc
+// CHECK:       bb0:
+// CHECK-NEXT:    [[A16:%.*]] = alloc_stack $Builtin.Int16
+// CHECK-NEXT:    // function_ref
+// CHECK-NEXT:    [[MARKER:%.*]] = function_ref @marker
+// CHECK-NEXT:    [[A32:%.*]] = alloc_stack $Builtin.Int32
+// CHECK-NEXT:    [[I1000:%.*]] = integer_literal $Builtin.Int32, 1000
+// CHECK-NEXT:    store [[I1000]] to [trivial] [[A32]]
+// CHECK-NEXT:    [[I2000:%.*]] = integer_literal $Builtin.Int32, 2000
+// CHECK-NEXT:    apply [[MARKER]]([[I2000]])
+// CHECK-NEXT:    br bb1
+// CHECK:       bb1:
+// CHECK-NEXT:    [[I3000:%.*]] = integer_literal $Builtin.Int32, 3000
+// CHECK-NEXT:    apply [[MARKER]]([[I3000]])
+// CHECK-NEXT:    dealloc_stack [[A32]] : $*Builtin.Int32
+//   Note that this has been delayed to follow stack discipline.
+// CHECK-NEXT:    dealloc_stack [[A16]] : $*Builtin.Int16
+// CHECK-NEXT:    [[I4000:%.*]] = integer_literal $Builtin.Int32, 4000
+// CHECK-NEXT:    apply [[MARKER]]([[I4000]])
+// CHECK-NEXT:    tuple ()
+// CHECK-NEXT:    br [[END:bb[0-9]+]]
+// CHECK:       [[END]]:
+// CHECK-NEXT:    [[RET:%.*]] = tuple ()
+// CHECK-NEXT:    return [[RET]] : $()
+// CHECK-LABEL: // end sil function 'test_stack_overlap_dealloc'
+sil @test_stack_overlap_dealloc : $() -> () {
+bb0:
+  %stack = alloc_stack $Builtin.Int16
+  %0 = function_ref @stack_overlap : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  (%value, %token) = begin_apply %0() : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  dealloc_stack %stack : $*Builtin.Int16
+  end_apply %token
+  %ret = tuple ()
+  return %ret : $()
+}
+
+// CHECK-LABEL: sil @test_stack_overlap_alloc
+// CHECK:       bb0:
+// CHECK-NEXT:    // function_ref
+// CHECK-NEXT:    [[MARKER:%.*]] = function_ref @marker
+// CHECK-NEXT:    [[A32:%.*]] = alloc_stack $Builtin.Int32
+// CHECK-NEXT:    [[I1000:%.*]] = integer_literal $Builtin.Int32, 1000
+// CHECK-NEXT:    store [[I1000]] to [trivial] [[A32]]
+// CHECK-NEXT:    [[I2000:%.*]] = integer_literal $Builtin.Int32, 2000
+// CHECK-NEXT:    apply [[MARKER]]([[I2000]])
+// CHECK-NEXT:    [[A16:%.*]] = alloc_stack $Builtin.Int16
+// CHECK-NEXT:    br bb1
+// CHECK:       bb1:
+// CHECK-NEXT:    [[I3000:%.*]] = integer_literal $Builtin.Int32, 3000
+// CHECK-NEXT:    apply [[MARKER]]([[I3000]])
+// CHECK-NEXT:    [[I4000:%.*]] = integer_literal $Builtin.Int32, 4000
+// CHECK-NEXT:    apply [[MARKER]]([[I4000]])
+// CHECK-NEXT:    tuple ()
+// CHECK-NEXT:    br [[END:bb[0-9]+]]
+// CHECK:       [[END]]:
+// CHECK-NEXT:    dealloc_stack [[A16]] : $*Builtin.Int16
+//   Note that this has been delayed to follow stack discipline.
+// CHECK-NEXT:    dealloc_stack [[A32]] : $*Builtin.Int32
+// CHECK-NEXT:    [[RET:%.*]] = tuple ()
+// CHECK-NEXT:    return [[RET]] : $()
+// CHECK-LABEL: // end sil function 'test_stack_overlap_alloc'
+sil @test_stack_overlap_alloc : $() -> () {
+bb0:
+  %0 = function_ref @stack_overlap : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  (%value, %token) = begin_apply %0() : $@convention(thin) @yield_once () -> (@yields @inout Builtin.Int32)
+  %stack = alloc_stack $Builtin.Int16
+  end_apply %token
+  dealloc_stack %stack : $*Builtin.Int16
+  %ret = tuple ()
+  return %ret : $()
+}