SILGen: let withoutActuallyEscaping also catch an escaping closure in the throwing case

withoutActuallyEscaping checks that the passed closure has not beed escaped by ensuring that its reference count is exactly 1 at the end of the code block.
So far this was only done in the regular return case. But if the code block throws, the check was not done.

Fixes an undetected undefined behavior.

Author: eeckstein

lib/SILGen/SILGenExpr.cpp

@@ -6041,6 +6041,34 @@ RValue RValueEmitter::visitOpenExistentialExpr(OpenExistentialExpr *E,
                                              });
 }
 
+namespace {
+class DestroyNotEscapedClosureCleanup : public Cleanup {
+  SILValue v;
+public:
+  DestroyNotEscapedClosureCleanup(SILValue v) : v(v) {}
+
+  void emit(SILGenFunction &SGF, CleanupLocation l,
+            ForUnwind_t forUnwind) override {
+    // Now create the verification of the withoutActuallyEscaping operand.
+    // Either we fail the uniqueness check (which means the closure has escaped)
+    // and abort or we continue and destroy the ultimate reference.
+    auto isEscaping = SGF.B.createIsEscapingClosure(
+        l, v, IsEscapingClosureInst::WithoutActuallyEscaping);
+    SGF.B.createCondFail(l, isEscaping, "non-escaping closure has escaped");
+    SGF.B.createDestroyValue(l, v);
+  }
+
+  void dump(SILGenFunction &) const override {
+#ifndef NDEBUG
+    llvm::errs() << "DestroyNotEscapedClosureCleanup\n"
+                 << "State:" << getState() << "\n"
+                 << "Value:" << v << "\n";
+#endif
+  }
+};
+} // end anonymous namespace
+
+
 RValue RValueEmitter::visitMakeTemporarilyEscapableExpr(
     MakeTemporarilyEscapableExpr *E, SGFContext C) {
   // Emit the non-escaping function value.
@@ -6075,19 +6103,15 @@ RValue RValueEmitter::visitMakeTemporarilyEscapableExpr(
   }
 
   // Convert it to an escaping function value.
-  auto escapingClosure =
+  auto ec =
       SGF.createWithoutActuallyEscapingClosure(E, functionValue, escapingFnTy);
+  SILValue ecv = ec.forward(SGF);
+  SGF.Cleanups.pushCleanup<DestroyNotEscapedClosureCleanup>(ecv);
+  auto escapingClosure = ManagedValue::forOwnedRValue(ecv, SGF.Cleanups.getTopCleanup());
   auto loc = SILLocation(E);
   auto borrowedClosure = escapingClosure.borrow(SGF, loc);
   RValue rvalue = visitSubExpr(borrowedClosure, false /* isClosureConsumable */);
 
-  // Now create the verification of the withoutActuallyEscaping operand.
-  // Either we fail the uniqueness check (which means the closure has escaped)
-  // and abort or we continue and destroy the ultimate reference.
-  auto isEscaping = SGF.B.createIsEscapingClosure(
-      loc, borrowedClosure.getValue(),
-      IsEscapingClosureInst::WithoutActuallyEscaping);
-  SGF.B.createCondFail(loc, isEscaping, "non-escaping closure has escaped");
   return rvalue;
 }
 

test/Interpreter/withoutActuallyEscaping.swift

@@ -21,6 +21,7 @@ var testShouldThrow = false
 
 struct MyError : Error {}
 
+@inline(never)
 func letEscapeThrowing(f: () -> ()) throws -> () -> () {
   return try withoutActuallyEscaping(f) {
     if testShouldThrow {
@@ -30,6 +31,15 @@ func letEscapeThrowing(f: () -> ()) throws -> () -> () {
   }
 }
 
+@inline(never)
+func letEscapeOnThrowingPath(f: () throws -> ()) throws {
+  try withoutActuallyEscaping(f) {
+    sink = $0
+    throw MyError()
+  }
+}
+
+
 WithoutEscapingSuite.test("ExpectNoCrash") {
   dontEscape(f: { print("foo") })
 }
@@ -53,7 +63,7 @@ WithoutEscapingSuite.test("ExpectCrash") {
   sink = letEscape(f: { print("Context: \(context.a) \(context.b)") })
 }
 
-WithoutEscapingSuite.test("ExpectThrowingCrash") {
+WithoutEscapingSuite.test("ExpectNonThrowingCrash") {
   expectCrashLater()
   let context = Context()
   var testDidThrow = false
@@ -78,4 +88,16 @@ WithoutEscapingSuite.test("ExpectThrowingNoCrash") {
   expectTrue(testDidThrow)
 }
 
+WithoutEscapingSuite.test("ExpectThrowingCrash") {
+  expectCrashLater()
+  let context = Context()
+  var testDidThrow = false
+  do {
+    try letEscapeOnThrowingPath(f: { print("Context: \(context.a) \(context.b)") })
+  } catch {
+    testDidThrow = true
+  }
+  expectTrue(testDidThrow)
+}
+
 runAllTests()

test/SILGen/without_actually_escaping.swift

@@ -35,14 +35,16 @@ func letEscape(f: () -> ()) -> () -> () {
 // CHECK:  try_apply [[USER]]([[BORROW]]) : {{.*}}, normal bb1, error bb2
 //
 // CHECK: bb1([[RES:%.*]] : @owned $@callee_guaranteed () -> ()):
-// CHECK:   [[ESCAPED:%.*]] = is_escaping_closure [[BORROW]]
-// CHECK:   cond_fail [[ESCAPED]] : $Builtin.Int1
 // CHECK:   end_borrow [[BORROW]]
+// CHECK:   [[ESCAPED:%.*]] = is_escaping_closure [[MD]]
+// CHECK:   cond_fail [[ESCAPED]] : $Builtin.Int1
 // CHECK:   destroy_value [[MD]]
 // CHECK:   return [[RES]]
 //
 // CHECK: bb2([[ERR:%.*]] : @owned $any Error):
 // CHECK:   end_borrow [[BORROW]]
+// CHECK:   [[ESCAPED:%.*]] = is_escaping_closure [[MD]]
+// CHECK:   cond_fail [[ESCAPED]] : $Builtin.Int1
 // CHECK:   destroy_value [[MD]]
 // CHECK:   throw [[ERR]] : $any Error
 // CHECK: }