[sil-verifier] Make flow sensitive verification handle begin_apply token/allocation correctly.

Specifically:

1. We were tracking the stack allocation both in handleScopeInst and in the
Stack. We should only track it in one of them. I also used this as an
opportunity to make sure the code worked for non_nested code.

2. I made it so that we properly handle the end tracking part of the code so we
handle the token/stack part of begin_apply correctly. I generalized the code so
that we should handle non_nested stack allocations as well.

I included tests that validated that we now handle this correctly.

lib/SIL/Verifier/FlowSensitiveVerifier.cpp

@@ -210,6 +210,13 @@ struct BBState {
       verificationFailure("operation has already been ended", &i, nullptr);
     }
   }
+
+  void handleScopeEndingUse(SILValue value, SILInstruction *origUser) {
+    if (!ActiveOps.erase(value)) {
+      verificationFailure("operation has already been ended", origUser,
+                          nullptr);
+    }
+  }
 };
 
 struct DeadEndRegionState {
@@ -224,7 +231,7 @@ struct DeadEndRegionState {
 /// If this is a scope ending inst, return the result from the instruction
 /// that provides the scoped value whose lifetime must be ended by some other
 /// scope ending instruction.
-static SILValue isScopeInst(SILInstruction *i) {
+static SILValue getScopeEndingValueOfScopeInst(SILInstruction *i) {
   if (auto *bai = dyn_cast<BeginAccessInst>(i))
     return bai;
   if (auto *bai = dyn_cast<BeginApplyInst>(i))
@@ -325,15 +332,15 @@ void swift::silverifier::verifyFlowSensitiveRules(SILFunction *F) {
         // verify their joint post-dominance.
         if (i.isStackAllocationNested()) {
           state.Stack.push_back(i.getStackAllocation());
-
-          // Also track begin_apply as a scope instruction.
-          if (auto *bai = dyn_cast<BeginApplyInst>(&i)) {
-            state.handleScopeInst(bai->getStackAllocation());
-            state.handleScopeInst(bai->getTokenResult());
-          }
         } else {
           state.handleScopeInst(i.getStackAllocation());
         }
+
+        // Also track begin_apply's token as an ActiveOp so we can also verify
+        // its joint dominance.
+        if (auto *bai = dyn_cast<BeginApplyInst>(&i)) {
+          state.handleScopeInst(bai->getTokenResult());
+        }
         continue;
       }
 
@@ -343,28 +350,48 @@ void swift::silverifier::verifyFlowSensitiveRules(SILFunction *F) {
           op = mvi->getOperand();
         }
 
-        auto beginInst = op->getDefiningInstruction();
-        if (beginInst && (!beginInst->isAllocatingStack() ||
-                          !beginInst->isStackAllocationNested())) {
-          state.handleScopeEndingInst(i);
-        } else if (!state.Stack.empty() && op == state.Stack.back()) {
-          state.Stack.pop_back();
-          if (llvm::isa_and_present<BeginApplyInst>(beginInst))
-            state.handleScopeEndingInst(i);
-        } else {
-          verificationFailure(
-              "deallocating allocation that is not the top of the stack", &i,
-              [&](SILPrintContext &ctx) {
-                state.printStack(ctx, "Current stack state");
-                ctx.OS() << "Stack allocation:\n" << *op;
-                // The deallocation is printed out as the focus of the
-                // failure.
-              });
+        auto *definingInst = op->getDefiningInstruction();
+
+        // First see if we have a begin_inst. In such a case, we need to handle
+        // the token result first.
+        if (auto *beginInst =
+                llvm::dyn_cast_if_present<BeginApplyInst>(definingInst)) {
+          // Check if our value is the token result... in such a case, we need
+          // to handle scope ending inst and continue.
+          if (op == beginInst->getTokenResult()) {
+            state.handleScopeEndingUse(op, &i);
+            continue;
+          }
         }
+
+        // Ok, we have some sort of memory. Lets check if our definingInst is
+        // unnested stack memory. In such a case, we want to just verify post
+        // dominance and not validate that we are in stack order.
+        if (definingInst && definingInst->isAllocatingStack() &&
+            !definingInst->isStackAllocationNested()) {
+          state.handleScopeEndingUse(op, &i);
+          continue;
+        }
+
+        // Ok, we have nested stack memory. Lets try to pop the stack. If we
+        // fail due to an empty stack or a lack of a match, emit an error.
+        if (!state.Stack.empty() && op == state.Stack.back()) {
+          state.Stack.pop_back();
+          continue;
+        }
+
+        verificationFailure(
+            "deallocating allocation that is not the top of the stack", &i,
+            [&](SILPrintContext &ctx) {
+              state.printStack(ctx, "Current stack state");
+              ctx.OS() << "Stack allocation:\n" << *op;
+              // The deallocation is printed out as the focus of the
+              // failure.
+            });
         continue;
       }
 
-      if (auto scopeEndingValue = isScopeInst(&i)) {
+      if (auto scopeEndingValue = getScopeEndingValueOfScopeInst(&i)) {
         state.handleScopeInst(scopeEndingValue);
         continue;
       }

test/SIL/verifier_failures.sil

@@ -104,3 +104,37 @@ bb2:
   abort_apply %3
   unwind
 }
+
+// We shouldn't error here. Make sure that we properly handle the dealloc_stack
+// user vs the end_apply user.
+sil [ossa] @test_begin_apply_1 : $@yield_once @convention(method) (@guaranteed Builtin.Int32) -> @yields Builtin.Int32 {
+bb0(%0 : $Builtin.Int32):
+  (%3, %4, %5) = begin_apply undef(%0) : $@yield_once_2 @convention(method) (@guaranteed Builtin.Int32) -> @yields Builtin.Int32
+  %6 = end_apply %4 as $()
+  yield %3 : $Builtin.Int32, resume bb1, unwind bb2
+
+bb1:
+  dealloc_stack %5 : $*Builtin.SILToken
+  %9 = tuple ()
+  return %9 : $()
+
+bb2:
+  dealloc_stack %5 : $*Builtin.SILToken
+  unwind
+}
+
+sil [ossa] @test_begin_apply_2 : $@yield_once @convention(method) (@guaranteed Builtin.Int32) -> @yields Builtin.Int32 {
+bb0(%0 : $Builtin.Int32):
+  (%3, %4, %5) = begin_apply undef(%0) : $@yield_once_2 @convention(method) (@guaranteed Builtin.Int32) -> @yields Builtin.Int32
+  dealloc_stack %5 : $*Builtin.SILToken
+  yield %3 : $Builtin.Int32, resume bb1, unwind bb2
+
+bb1:
+  %6 = end_apply %4 as $()
+  %9 = tuple ()
+  return %9 : $()
+
+bb2:
+  %7 = end_apply %4 as $()
+  unwind
+}

test/SIL/verifier_loweredsil_failures.sil

@@ -1,4 +1,4 @@
-// RUN: %target-sil-opt -sil-verify-all=false -enable-sil-verify-all=false -emit-sorted-sil -verify-continue-on-failure -o /dev/null %s 2>&1 | %FileCheck %s
+// RUN: %target-sil-opt -sil-verify-all=false -enable-sil-verify-all=false -emit-sorted-sil -verify-continue-on-failure -o /dev/null %s 2>&1 | %FileCheck '--implicit-check-not=Begin Error' %s
 
 // REQUIRES: asserts
 
@@ -19,7 +19,7 @@ sil @alloc_pack_metadata_before_tuple : $@convention(thin) () -> () {
 }
 
 // CHECK-LABEL: Begin Error in function dealloc_pack_metadata_with_bad_operand
-// CHECK:       SIL verification failed: operation has already been ended
+// CHECK:       SIL verification failed: deallocating allocation that is not the top of the stack
 // CHECK-LABEL: End Error in function dealloc_pack_metadata_with_bad_operand
 // CHECK-LABEL: Begin Error in function dealloc_pack_metadata_with_bad_operand
 // CHECK:       SIL verification failed: return with stack allocs that haven't been deallocated