Adopt typed throws in withUnsafeTemporaryAllocation

The interior wrapping in Result<T, U> is a little unfortunate, but is
currently necessary because we end up with extra stack allocations
along the error-handling path in a do..catch that aren't there when
using untyped throws. We can simplify the implementation if we can
eliminate the extraneous stack allocation.

Fixes rdar://134973620.

Author: DougGregor

stdlib/public/core/TemporaryAllocation.swift

@@ -119,8 +119,6 @@ internal func _isStackAllocationSafe(byteCount: Int, alignment: Int) -> Bool {
 ///
 /// - Returns: Whatever is returned by `body`.
 ///
-/// - Throws: Whatever is thrown by `body`.
-///
 /// This function encapsulates the various calls to builtins required by
 /// `withUnsafeTemporaryAllocation()`.
 @_alwaysEmitIntoClient @_transparent
@@ -130,13 +128,13 @@ internal func _withUnsafeTemporaryAllocation<
   of type: T.Type,
   capacity: Int,
   alignment: Int,
-  _ body: (Builtin.RawPointer) throws -> R
-) rethrows -> R {
+  _ body: (Builtin.RawPointer) -> R
+) -> R {
   // How many bytes do we need to allocate?
   let byteCount = _byteCountForTemporaryAllocation(of: type, capacity: capacity)
 
   guard _isStackAllocationSafe(byteCount: byteCount, alignment: alignment) else {
-    return try _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
+    return _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
   }
 
   // This declaration must come BEFORE Builtin.stackAlloc() or
@@ -154,15 +152,9 @@ internal func _withUnsafeTemporaryAllocation<
   // The multiple calls to Builtin.stackDealloc() are because defer { } produces
   // a child function at the SIL layer and that conflicts with the verifier's
   // idea of a stack allocation's lifetime.
-  do {
-    result = try body(stackAddress)
-    Builtin.stackDealloc(stackAddress)
-    return result
-
-  } catch {
-    Builtin.stackDealloc(stackAddress)
-    throw error
-  }
+  result = body(stackAddress)
+  Builtin.stackDealloc(stackAddress)
+  return result
 #else
   fatalError("unsupported compiler")
 #endif
@@ -175,13 +167,13 @@ internal func _withUnprotectedUnsafeTemporaryAllocation<
   of type: T.Type,
   capacity: Int,
   alignment: Int,
-  _ body: (Builtin.RawPointer) throws -> R
-) rethrows -> R {
+  _ body: (Builtin.RawPointer) -> R
+) -> R {
   // How many bytes do we need to allocate?
   let byteCount = _byteCountForTemporaryAllocation(of: type, capacity: capacity)
 
   guard _isStackAllocationSafe(byteCount: byteCount, alignment: alignment) else {
-    return try _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
+    return _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
   }
 
   // This declaration must come BEFORE Builtin.unprotectedStackAlloc() or
@@ -198,23 +190,17 @@ internal func _withUnprotectedUnsafeTemporaryAllocation<
   // The multiple calls to Builtin.stackDealloc() are because defer { } produces
   // a child function at the SIL layer and that conflicts with the verifier's
   // idea of a stack allocation's lifetime.
-  do {
-    result = try body(stackAddress)
-    Builtin.stackDealloc(stackAddress)
-    return result
-
-  } catch {
-    Builtin.stackDealloc(stackAddress)
-    throw error
-  }
+  result = body(stackAddress)
+  Builtin.stackDealloc(stackAddress)
+  return result
 }
 
 @_alwaysEmitIntoClient @_transparent
-internal func _fallBackToHeapAllocation<R: ~Copyable>(
+internal func _fallBackToHeapAllocation<R: ~Copyable, E: Error>(
   byteCount: Int,
   alignment: Int,
-  _ body: (Builtin.RawPointer) throws -> R
-) rethrows -> R {
+  _ body: (Builtin.RawPointer) throws(E) -> R
+) throws(E) -> R {
   let buffer = UnsafeMutableRawPointer.allocate(
     byteCount: byteCount,
     alignment: alignment
@@ -259,21 +245,30 @@ internal func _fallBackToHeapAllocation<R: ~Copyable>(
 /// the buffer) must not escape. It will be deallocated when `body` returns and
 /// cannot be used afterward.
 @_alwaysEmitIntoClient @_transparent
-public func withUnsafeTemporaryAllocation<R: ~Copyable>(
+public func withUnsafeTemporaryAllocation<R: ~Copyable, E: Error>(
   byteCount: Int,
   alignment: Int,
-  _ body: (UnsafeMutableRawBufferPointer) throws -> R
-) rethrows -> R {
-  return try _withUnsafeTemporaryAllocation(
+  _ body: (UnsafeMutableRawBufferPointer) throws(E) -> R
+) throws(E) -> R {
+  let result: Result<R, E> = _withUnsafeTemporaryAllocation(
     of: Int8.self,
     capacity: byteCount,
     alignment: alignment
   ) { pointer in
-    let buffer = unsafe UnsafeMutableRawBufferPointer(
-      start: .init(pointer),
-      count: byteCount
-    )
-    return try unsafe body(buffer)
+    do throws(E) {
+      let buffer = unsafe UnsafeMutableRawBufferPointer(
+        start: .init(pointer),
+        count: byteCount
+      )
+      return .success(try unsafe body(buffer))
+    } catch {
+      return .failure(error)
+    }
+  }
+
+  switch consume result {
+  case .success(let resultValue): return resultValue
+  case .failure(let error): throw error
   }
 }
 
@@ -283,21 +278,30 @@ public func withUnsafeTemporaryAllocation<R: ~Copyable>(
 /// This function is similar to `withUnsafeTemporaryAllocation`, except that it
 /// doesn't trigger stack protection for the stack allocated memory.
 @_alwaysEmitIntoClient @_transparent
-public func _withUnprotectedUnsafeTemporaryAllocation<R: ~Copyable>(
+public func _withUnprotectedUnsafeTemporaryAllocation<R: ~Copyable, E: Error>(
   byteCount: Int,
   alignment: Int,
-  _ body: (UnsafeMutableRawBufferPointer) throws -> R
-) rethrows -> R {
-  return try _withUnprotectedUnsafeTemporaryAllocation(
+  _ body: (UnsafeMutableRawBufferPointer) throws(E) -> R
+) throws(E) -> R {
+  let result: Result<R, E> = _withUnprotectedUnsafeTemporaryAllocation(
     of: Int8.self,
     capacity: byteCount,
     alignment: alignment
   ) { pointer in
-    let buffer = unsafe UnsafeMutableRawBufferPointer(
-      start: .init(pointer),
-      count: byteCount
-    )
-    return try unsafe body(buffer)
+    do throws(E) {
+      let buffer = unsafe UnsafeMutableRawBufferPointer(
+        start: .init(pointer),
+        count: byteCount
+      )
+      return try unsafe .success(body(buffer))
+    } catch {
+      return .failure(error)
+    }
+  }
+
+  switch consume result {
+  case .success(let resultValue): return resultValue
+  case .failure(let error): throw error
   }
 }
 
@@ -334,23 +338,33 @@ public func _withUnprotectedUnsafeTemporaryAllocation<R: ~Copyable>(
 /// cannot be used afterward.
 @_alwaysEmitIntoClient @_transparent
 public func withUnsafeTemporaryAllocation<
-  T: ~Copyable,R: ~Copyable
+  T: ~Copyable,R: ~Copyable,
+  E: Error
 >(
   of type: T.Type,
   capacity: Int,
-  _ body: (UnsafeMutableBufferPointer<T>) throws -> R
-) rethrows -> R {
-  return try _withUnsafeTemporaryAllocation(
+  _ body: (UnsafeMutableBufferPointer<T>) throws(E) -> R
+) throws(E) -> R {
+  let result: Result<R, E> = _withUnsafeTemporaryAllocation(
     of: type,
     capacity: capacity,
     alignment: MemoryLayout<T>.alignment
   ) { pointer in
-    Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
-    let buffer = unsafe UnsafeMutableBufferPointer<T>(
-      start: .init(pointer),
-      count: capacity
-    )
-    return try unsafe body(buffer)
+    do throws(E) {
+      Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
+      let buffer = unsafe UnsafeMutableBufferPointer<T>(
+        start: .init(pointer),
+        count: capacity
+      )
+      return try unsafe .success(body(buffer))
+    } catch {
+      return .failure(error)
+    }
+  }
+
+  switch consume result {
+  case .success(let resultValue): return resultValue
+  case .failure(let error): throw error
   }
 }
 
@@ -361,22 +375,32 @@ public func withUnsafeTemporaryAllocation<
 /// doesn't trigger stack protection for the stack allocated memory.
 @_alwaysEmitIntoClient @_transparent
 public func _withUnprotectedUnsafeTemporaryAllocation<
-  T: ~Copyable, R: ~Copyable
+  T: ~Copyable, R: ~Copyable,
+  E: Error
 >(
   of type: T.Type,
   capacity: Int,
-  _ body: (UnsafeMutableBufferPointer<T>) throws -> R
-) rethrows -> R {
-  return try _withUnprotectedUnsafeTemporaryAllocation(
+  _ body: (UnsafeMutableBufferPointer<T>) throws(E) -> R
+) throws(E) -> R {
+  let result: Result<R, E> = _withUnprotectedUnsafeTemporaryAllocation(
     of: type,
     capacity: capacity,
     alignment: MemoryLayout<T>.alignment
   ) { pointer in
-    Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
-    let buffer = unsafe UnsafeMutableBufferPointer<T>(
-      start: .init(pointer),
-      count: capacity
-    )
-    return try unsafe body(buffer)
+    do throws(E) {
+      Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
+      let buffer = unsafe UnsafeMutableBufferPointer<T>(
+        start: .init(pointer),
+        count: capacity
+      )
+      return try unsafe .success(body(buffer))
+    } catch {
+      return .failure(error)
+    }
+  }
+
+  switch consume result {
+  case .success(let resultValue): return resultValue
+  case .failure(let error): throw error
   }
 }

test/stdlib/TemporaryAllocation.swift

@@ -155,4 +155,21 @@ TemporaryAllocationTestSuite.test("typedAllocationIsAligned") {
   }
 }
 
+// MARK: Typed throws
+enum HomeworkError: Error, Equatable {
+case dogAtIt
+case forgot
+}
+
+TemporaryAllocationTestSuite.test("typedAllocationWithThrow") {
+  do throws(HomeworkError) {
+    try withUnsafeTemporaryAllocation(of: Int.self, capacity: 1) { (buffer) throws(HomeworkError) -> Void in
+      throw HomeworkError.forgot
+    }
+    fatalError("did not throw!?!")
+  } catch {
+    expectEqual(error, .forgot)
+  }
+}
+
 runAllTests()