stdlib: add withTemporaryAllocation using OutputSpan

stdlib/public/core/TemporaryAllocation.swift

@@ -141,14 +141,14 @@ internal func _withUnsafeTemporaryAllocation<
   // Builtin.stackDealloc() will end up blowing it away (and the verifier will
   // notice and complain.)
   let result: R
-  
+
 #if compiler(>=5.5) && $BuiltinStackAlloc
   let stackAddress = Builtin.stackAlloc(
     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.
@@ -269,6 +269,25 @@ public func withUnsafeTemporaryAllocation<R: ~Copyable, E: Error>(
   return try result.get()
 }
 
+@available(SwiftCompatibilitySpan 5.0, *)
+@_alwaysEmitIntoClient @_transparent
+public func withTemporaryAllocation<T: ~Copyable, R: ~Copyable, E: Error>(
+  of type: T.Type,
+  capacity: Int,
+  _ body: (inout OutputSpan<T>) throws(E) -> R
+) throws(E) -> R where T : ~Copyable, R : ~Copyable {
+  try withUnsafeTemporaryAllocation(of: type, capacity: capacity) { (buffer) throws(E) in
+    var span = OutputSpan(buffer: buffer, initializedCount: 0)
+    defer {
+      let initializedCount = span.finalize(for: buffer)
+      span = OutputSpan()
+      buffer.extracting(..<initializedCount).deinitialize()
+    }
+
+    return try body(&span)
+  }
+}
+
 /// Provides scoped access to a raw buffer pointer with the specified byte count
 /// and alignment.
 ///

test/stdlib/TemporaryAllocation.swift

@@ -102,6 +102,15 @@ TemporaryAllocationTestSuite.test("typedAllocationOnStack") {
   }
 }
 
+TemporaryAllocationTestSuite.test("spanOnStack") {
+  withTemporaryAllocation(of: Int.self, capacity: 1) { span in
+    span.withUnsafeMutableBufferPointer { buffer, initializedCount in
+      expectStackAllocated(buffer.baseAddress!)
+      initializedCount = 0
+    }
+  }
+}
+
 TemporaryAllocationTestSuite.test("typedAllocationOnHeap") {
   // EXPECTATION: a very large allocated buffer is heap-allocated. (Note if
   // swift_stdlib_isStackAllocationSafe() gets fleshed out, this test may need
@@ -111,6 +120,18 @@ TemporaryAllocationTestSuite.test("typedAllocationOnHeap") {
   }
 }
 
+TemporaryAllocationTestSuite.test("spanOnHeap") {
+  // EXPECTATION: a very large allocated buffer is heap-allocated. (Note if
+  // swift_stdlib_isStackAllocationSafe() gets fleshed out, this test may need
+  // to be changed.)
+  withTemporaryAllocation(of: Int.self, capacity: 100_000) { span in
+    span.withUnsafeMutableBufferPointer { buffer, initializedCount in
+      expectNotStackAllocated(buffer.baseAddress!)
+      initializedCount = 0
+    }
+  }
+}
+
 TemporaryAllocationTestSuite.test("unprotectedTypedAllocationOnStack") {
   _withUnprotectedUnsafeTemporaryAllocation(of: Int.self, capacity: 1) { buffer in
       expectStackAllocated(buffer.baseAddress!)
@@ -132,9 +153,27 @@ TemporaryAllocationTestSuite.test("typedEmptyAllocationIsStackAllocated") {
   }
 }
 
+TemporaryAllocationTestSuite.test("emptySpanHasNoBuffer") {
+  withTemporaryAllocation(of: Int.self, capacity: 0) { span in
+    span.withUnsafeMutableBufferPointer { buffer, initializedCount in
+      expectTrue(buffer.baseAddress == nil)
+      initializedCount = 0
+    }
+  }
+}
+
 TemporaryAllocationTestSuite.test("voidAllocationIsStackAllocated") {
   withUnsafeTemporaryAllocation(of: Void.self, capacity: 1) { buffer in
+    expectStackAllocated(buffer.baseAddress!)
+  }
+}
+
+TemporaryAllocationTestSuite.test("voidSpanIsStackAllocated") {
+  withTemporaryAllocation(of: Void.self, capacity: 1) { span in
+    span.withUnsafeMutableBufferPointer { buffer, initializedCount in
       expectStackAllocated(buffer.baseAddress!)
+      initializedCount = 0
+    }
   }
 }
 
@@ -145,6 +184,13 @@ TemporaryAllocationTestSuite.test("crashOnNegativeValueCount") {
     withUnsafeTemporaryAllocation(of: Int.self, capacity: capacity) { _ in }
   }
 }
+
+TemporaryAllocationTestSuite.test("spanCrashOnNegativeValueCount") {
+  expectCrash {
+    let capacity = Int.random(in: -2 ..< -1)
+    withTemporaryAllocation(of: Int.self, capacity: capacity) { _ in }
+  }
+}
 #endif
 
 TemporaryAllocationTestSuite.test("typedAllocationIsAligned") {
@@ -155,6 +201,17 @@ TemporaryAllocationTestSuite.test("typedAllocationIsAligned") {
   }
 }
 
+TemporaryAllocationTestSuite.test("spanIsAligned") {
+  withTemporaryAllocation(of: Int.self, capacity: 1) { span in
+    span.withUnsafeMutableBufferPointer { buffer, initializedCount in
+      let pointerBits = Int(bitPattern: buffer.baseAddress!)
+      let alignmentMask = MemoryLayout<Int>.alignment - 1
+      expectEqual(pointerBits & alignmentMask, 0)
+      initializedCount = 0
+    }
+  }
+}
+
 // MARK: Typed throws
 enum HomeworkError: Error, Equatable {
 case dogAteIt
@@ -172,4 +229,15 @@ TemporaryAllocationTestSuite.test("typedAllocationWithThrow") {
   }
 }
 
+TemporaryAllocationTestSuite.test("spanWithThrow") {
+  do throws(HomeworkError) {
+    try withTemporaryAllocation(of: Int.self, capacity: 1) { (span) throws(HomeworkError) -> Void in
+      throw HomeworkError.forgot
+    }
+    fatalError("did not throw!?!")
+  } catch {
+    expectEqual(error, .forgot)
+  }
+}
+
 runAllTests()