Merge pull request swiftlang#83921 from DougGregor/with-unsafe-temporary-alloc-typed-throws

Adopt typed throws in withUnsafeTemporaryAllocation

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,22 +245,28 @@ 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)
+    }
   }
+
+  return try result.get()
 }
 
 /// Provides scoped access to a raw buffer pointer with the specified byte count
@@ -283,22 +275,28 @@ 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)
+    }
   }
+
+  return try result.get()
 }
 
 /// Provides scoped access to a buffer pointer to memory of the specified type
@@ -334,24 +332,31 @@ 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)
+    }
   }
+
+  return try result.get()
 }
 
 /// Provides scoped access to a buffer pointer to memory of the specified type
@@ -361,22 +366,29 @@ 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)
+    }
   }
+
+  return try result.get()
 }

test/IRGen/temporary_allocation/codegen.swift

@@ -1,5 +1,5 @@
 // RUN: %target-swift-frontend -primary-file %s -O -emit-ir | %FileCheck %s
-// REQUIRES: swift_stdlib_no_asserts,optimized_stdlib
+// REQUIRES: optimized_stdlib
 
 @_silgen_name("blackHole")
 func blackHole(_ value: UnsafeMutableRawPointer?) -> Void
@@ -16,10 +16,10 @@ do {
 // CHECK: [[ONE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca i8, align 1
 // CHECK: [[FIVE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca [5 x i8], align 1
 // CHECK: [[ONE_KB_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 8
-// CHECK: [[ONE_KB_RAND_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 16
 // CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
 // CHECK: [[INT_PTR_RAW2:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
 // CHECK: [[VOID_PTR_RAW:%temp_alloc[0-9]*]] = alloca [2 x i8], align 1
+// CHECK: [[ONE_KB_RAND_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 16
 
 // CHECK: ptrtoint ptr {{.*}} to [[WORD:i[0-9]+]]
 
@@ -49,14 +49,6 @@ withUnsafeTemporaryAllocation(byteCount: 1024, alignment: 8) { buffer in
 // CHECK: [[ONE_KB_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_KB_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_KB_PTR]])
 
-// MARK: Alignment unknown at compile-time
-
-withUnsafeTemporaryAllocation(byteCount: 1024, alignment: Int.random(in: 0 ..< 16)) { buffer in
-  blackHole(buffer.baseAddress)
-}
-// CHECK: [[ONE_KB_RAND_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_KB_RAND_PTR_RAW]] to [[WORD]]
-// CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_KB_RAND_PTR]])
-
 // MARK: Typed buffers
 
 withUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
@@ -77,3 +69,10 @@ withUnsafeTemporaryAllocation(of: Void.self, capacity: 2) { buffer in
 // CHECK: [[VOID_PTR:%[0-9]+]] = ptrtoint ptr [[VOID_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[VOID_PTR]])
 
+// MARK: Alignment unknown at compile-time
+
+withUnsafeTemporaryAllocation(byteCount: 1024, alignment: Int.random(in: 0 ..< 16)) { buffer in
+  blackHole(buffer.baseAddress)
+}
+// CHECK: [[ONE_KB_RAND_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_KB_RAND_PTR_RAW]] to [[WORD]]
+// CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_KB_RAND_PTR]])

test/stdlib/TemporaryAllocation.swift

@@ -155,4 +155,21 @@ TemporaryAllocationTestSuite.test("typedAllocationIsAligned") {
   }
 }
 
+// MARK: Typed throws
+enum HomeworkError: Error, Equatable {
+case dogAteIt
+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()