From 8181154f075c5785de392935162e8f2a65fe15dc Mon Sep 17 00:00:00 2001
From: Max Desiatov <m_desiatov@apple.com>
Date: Fri, 5 Dec 2025 15:19:31 +0000
Subject: [PATCH 1/3] stdlib: add `withTemporaryAllocation` using `OutputSpan`

---
 stdlib/public/core/TemporaryAllocation.swift | 22 ++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/stdlib/public/core/TemporaryAllocation.swift b/stdlib/public/core/TemporaryAllocation.swift
index 960eec28492f6..babbbe6016067 100644
--- a/stdlib/public/core/TemporaryAllocation.swift
+++ b/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,24 @@ 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)
+      buffer.extracting(..<initializedCount).deinitialize()
+    }
+
+    return try body(&span)
+  }
+}
+
 /// Provides scoped access to a raw buffer pointer with the specified byte count
 /// and alignment.
 ///

From 8c18646ad27bfdec68eb851d11d8c9c22e9947ac Mon Sep 17 00:00:00 2001
From: Max Desiatov <m_desiatov@apple.com>
Date: Fri, 5 Dec 2025 15:54:02 +0000
Subject: [PATCH 2/3] `withTemporaryAllocation`: fix use of `defer`, add tests

---
 stdlib/public/core/TemporaryAllocation.swift | 14 +++-
 test/stdlib/TemporaryAllocation.swift        | 68 ++++++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/stdlib/public/core/TemporaryAllocation.swift b/stdlib/public/core/TemporaryAllocation.swift
index babbbe6016067..869702a1b7277 100644
--- a/stdlib/public/core/TemporaryAllocation.swift
+++ b/stdlib/public/core/TemporaryAllocation.swift
@@ -278,12 +278,20 @@ public func withTemporaryAllocation<T: ~Copyable, R: ~Copyable, E: Error>(
 ) 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 {
+
+    do throws(E) {
+      let result = try body(&span)
+
       let initializedCount = span.finalize(for: buffer)
       buffer.extracting(..<initializedCount).deinitialize()
-    }
 
-    return try body(&span)
+      return result
+    } catch {
+      let initializedCount = span.finalize(for: buffer)
+      buffer.extracting(..<initializedCount).deinitialize()
+
+      throw error
+    }
   }
 }
 
diff --git a/test/stdlib/TemporaryAllocation.swift b/test/stdlib/TemporaryAllocation.swift
index 4d4eb31bb9128..f27772e14b2d6 100644
--- a/test/stdlib/TemporaryAllocation.swift
+++ b/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()

From 8b68d54904fae257ebae3e254527a6ee4333ec4d Mon Sep 17 00:00:00 2001
From: Max Desiatov <m_desiatov@apple.com>
Date: Fri, 5 Dec 2025 21:43:52 +0000
Subject: [PATCH 3/3] withTemporaryAllocation: revert back to using `defer`

---
 stdlib/public/core/TemporaryAllocation.swift | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/stdlib/public/core/TemporaryAllocation.swift b/stdlib/public/core/TemporaryAllocation.swift
index 869702a1b7277..b1e0dfbe5b35f 100644
--- a/stdlib/public/core/TemporaryAllocation.swift
+++ b/stdlib/public/core/TemporaryAllocation.swift
@@ -278,20 +278,13 @@ public func withTemporaryAllocation<T: ~Copyable, R: ~Copyable, E: Error>(
 ) 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)
-
-    do throws(E) {
-      let result = try body(&span)
-
+    defer {
       let initializedCount = span.finalize(for: buffer)
+      span = OutputSpan()
       buffer.extracting(..<initializedCount).deinitialize()
-
-      return result
-    } catch {
-      let initializedCount = span.finalize(for: buffer)
-      buffer.extracting(..<initializedCount).deinitialize()
-
-      throw error
     }
+
+    return try body(&span)
   }
 }