stdlib: add an underscored _withUnprotectedUnsafeTemporaryAllocation

This function is similar to `withUnsafeTemporaryAllocation`, except that it doesn't trigger stack protection for the stack allocated memory.

Author: eeckstein

stdlib/public/core/TemporaryAllocation.swift

@@ -129,16 +129,7 @@ internal func _withUnsafeTemporaryAllocation<T, R>(
   let byteCount = _byteCountForTemporaryAllocation(of: type, capacity: capacity)
 
   guard _isStackAllocationSafe(byteCount: byteCount, alignment: alignment) else {
-    // Fall back to the heap. This may still be optimizable if escape analysis
-    // shows that the allocated pointer does not escape.
-    let buffer = UnsafeMutableRawPointer.allocate(
-      byteCount: byteCount,
-      alignment: alignment
-    )
-    defer {
-      buffer.deallocate()
-    }
-    return try body(buffer._rawValue)
+    return try _fallBackToHeapAllocation(byteCount: byteCount, alignment: alignment, body)
   }
 
   // This declaration must come BEFORE Builtin.stackAlloc() or
@@ -170,6 +161,63 @@ internal func _withUnsafeTemporaryAllocation<T, R>(
 #endif
 }
 
+#if $BuiltinUnprotectedStackAlloc
+@_alwaysEmitIntoClient @_transparent
+internal func _withUnprotectedUnsafeTemporaryAllocation<T, R>(
+  of type: T.Type,
+  capacity: Int,
+  alignment: Int,
+  _ body: (Builtin.RawPointer) throws -> R
+) rethrows -> 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)
+  }
+
+  // This declaration must come BEFORE Builtin.unprotectedStackAlloc() or
+  // Builtin.stackDealloc() will end up blowing it away (and the verifier will
+  // notice and complain.)
+  let result: R
+
+  let stackAddress = Builtin.unprotectedStackAlloc(
+    capacity._builtinWordValue,
+    MemoryLayout<T>.stride._builtinWordValue,
+    alignment._builtinWordValue
+  )
+
+  // 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
+  }
+}
+#endif
+
+@_alwaysEmitIntoClient @_transparent
+internal func _fallBackToHeapAllocation<R>(
+  byteCount: Int,
+  alignment: Int,
+  _ body: (Builtin.RawPointer) throws -> R
+) rethrows -> R {
+  let buffer = UnsafeMutableRawPointer.allocate(
+    byteCount: byteCount,
+    alignment: alignment
+  )
+  defer {
+    buffer.deallocate()
+  }
+  return try body(buffer._rawValue)
+}
+
 // MARK: - Public interface
 
 /// Provides scoped access to a raw buffer pointer with the specified byte count
@@ -222,6 +270,34 @@ public func withUnsafeTemporaryAllocation<R>(
   }
 }
 
+/// Provides scoped access to a raw buffer pointer with the specified byte count
+/// and alignment.
+///
+/// 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>(
+  byteCount: Int,
+  alignment: Int,
+  _ body: (UnsafeMutableRawBufferPointer) throws -> R
+) rethrows -> R {
+  return try _withUnsafeTemporaryAllocation(
+    of: Int8.self,
+    capacity: byteCount,
+    alignment: alignment
+  ) { pointer in
+#if $BuiltinUnprotectedStackAlloc
+    let buffer = UnsafeMutableRawBufferPointer(
+      start: .init(pointer),
+      count: byteCount
+    )
+    return try body(buffer)
+#else
+    withUnsafeTemporaryAllocation(byteCount: byteCount, alignment: alignment, body)
+#endif
+  }
+}
+
 /// Provides scoped access to a buffer pointer to memory of the specified type
 /// and with the specified capacity.
 ///
@@ -272,3 +348,32 @@ public func withUnsafeTemporaryAllocation<T, R>(
     return try body(buffer)
   }
 }
+
+/// Provides scoped access to a buffer pointer to memory of the specified type
+/// and with the specified capacity.
+///
+/// This function is similar to `withUnsafeTemporaryAllocation`, except that it
+/// doesn't trigger stack protection for the stack allocated memory.
+@_alwaysEmitIntoClient @_transparent
+public func _withUnprotectedUnsafeTemporaryAllocation<T, R>(
+  of type: T.Type,
+  capacity: Int,
+  _ body: (UnsafeMutableBufferPointer<T>) throws -> R
+) rethrows -> R {
+  return try _withUnprotectedUnsafeTemporaryAllocation(
+    of: type,
+    capacity: capacity,
+    alignment: MemoryLayout<T>.alignment
+  ) { pointer in
+#if $BuiltinUnprotectedStackAlloc
+    Builtin.bindMemory(pointer, capacity._builtinWordValue, type)
+    let buffer = UnsafeMutableBufferPointer<T>(
+      start: .init(pointer),
+      count: capacity
+    )
+    return try body(buffer)
+#else
+    return try withUnsafeTemporaryAllocation(of: type, capacity: capacity, body)
+#endif
+  }
+}

test/IRGen/temporary_allocation/codegen.swift

@@ -59,6 +59,13 @@ withUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
 // CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint [16 x i8]* [[INT_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
 
+_withUnprotectedUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
+  blackHole(buffer.baseAddress)
+}
+// CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
+// CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint [16 x i8]* [[INT_PTR_RAW]] to [[WORD]]
+// CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
+
 withUnsafeTemporaryAllocation(of: Void.self, capacity: 2) { buffer in
   blackHole(buffer.baseAddress)
 }

test/stdlib/TemporaryAllocation.swift

@@ -49,6 +49,21 @@ TemporaryAllocationTestSuite.test("untypedAllocationOnHeap") {
   }
 }
 
+TemporaryAllocationTestSuite.test("unprotectedUntypedAllocationOnStack") {
+  _withUnprotectedUnsafeTemporaryAllocation(byteCount: 8, alignment: 1) { buffer in
+      expectStackAllocated(buffer.baseAddress!)
+  }
+}
+
+TemporaryAllocationTestSuite.test("unprotectedUntypedAllocationOnHeap") {
+  // EXPECTATION: a very large allocated buffer is heap-allocated. (Note if
+  // swift_stdlib_isStackAllocationSafe() gets fleshed out, this test may need
+  // to be changed.)
+  _withUnprotectedUnsafeTemporaryAllocation(byteCount: 100_000, alignment: 1) { buffer in
+      expectNotStackAllocated(buffer.baseAddress!)
+  }
+}
+
 TemporaryAllocationTestSuite.test("untypedEmptyAllocationIsStackAllocated") {
   withUnsafeTemporaryAllocation(byteCount: 0, alignment: 1) { buffer in
       expectStackAllocated(buffer.baseAddress!)
@@ -94,6 +109,21 @@ TemporaryAllocationTestSuite.test("typedAllocationOnHeap") {
   }
 }
 
+TemporaryAllocationTestSuite.test("unprotectedTypedAllocationOnStack") {
+  _withUnprotectedUnsafeTemporaryAllocation(of: Int.self, capacity: 1) { buffer in
+      expectStackAllocated(buffer.baseAddress!)
+  }
+}
+
+TemporaryAllocationTestSuite.test("unprotectedTypedAllocationOnHeap") {
+  // EXPECTATION: a very large allocated buffer is heap-allocated. (Note if
+  // swift_stdlib_isStackAllocationSafe() gets fleshed out, this test may need
+  // to be changed.)
+  _withUnprotectedUnsafeTemporaryAllocation(of: Int.self, capacity: 100_000) { buffer in
+      expectNotStackAllocated(buffer.baseAddress!)
+  }
+}
+
 TemporaryAllocationTestSuite.test("typedEmptyAllocationIsStackAllocated") {
   withUnsafeTemporaryAllocation(of: Int.self, capacity: 0) { buffer in
       expectStackAllocated(buffer.baseAddress!)