Pitch: Add `withTemporaryAllocation` using `Output(Raw)Span`

Hello, Swift community!

I'd like to pitch a few small additions to the standard library that you'd be interested in. One notable part is that these proposed withTemporaryAllocation functions can be back-deployed, which means you'd get safe temporary spans (potentially optimized to stack allocations) on versions of Darwin that don't have InlineArray available.

A partial draft implementation is available at swiftlang/swift#85866.

Introduction

This proposal introduces new top-level functions that provide a temporary buffer wrapped in an OutputSpan or OutputRawSpan. This enables safe initialization of temporary memory, leveraging the safety guarantees of these span types while utilizing the stack-allocation optimization of withUnsafeTemporaryAllocation.

Motivation

SE-0322 and SE-0437 introduced and refined withUnsafeTemporaryAllocation, a facility for allocating temporary storage that may be stack-allocated. This function yields an UnsafeMutableBufferPointer or UnsafeMutableRawBufferPointer, requiring the user to manually manage initialization and deinitialization of the elements. This is error-prone, as the user must ensure that all initialized elements are correctly deinitialized before the closure returns, even in the presence of errors.

SE-0485 introduced OutputSpan and OutputRawSpan, types that manage the initialization state of a contiguous region of memory. These types track the number of initialized elements and ensure that memory operations maintain initialization invariants.

By combining these two facilities, we can provide a high-level, safe API for temporary allocations. Users can use the append methods on the span types to initialize the temporary memory without dealing with raw pointers or manually tracking the initialized count for deinitialization.

Proposed solution

We propose adding new global functions that wrap withUnsafeTemporaryAllocation. Instead of yielding a raw buffer pointer, they yield an inout OutputSpan for typed allocations, and an inout OutputRawSpan for raw byte allocations.

Typed Allocation

try withTemporaryAllocation(of: Int.self, capacity: 10) { span in
  span.append(1)
  span.append(2)
  // `OutputSpan` passed to this closure deinitializes and deallocates
  // elements upon exit
}

Raw Bytes Allocation

try withTemporaryAllocation(byteCount: 16, alignment: 4) { rawSpan in
  rawSpan.append(UInt32(0xF00), as: UInt32.self)
  // `OutputRawSpan` deallocates raw bytes upon exit
}

These functions handle the creation of the span types and ensure that any initialized elements are correctly deallocated (and deinitialized in the case of OutputSpan) when the scope exits.

Detailed design

The proposal adds two functions:

Typed Allocation with OutputSpan

This function is for working with temporary allocations of a specific, homogenous type.


@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)
  }
}

Here's the implementation walkthrough:

  1. Allocation: It calls withUnsafeTemporaryAllocation(of:capacity:) to obtain a typed buffer of uninitialized memory.

  2. Span Creation: It creates an OutputSpan covering the buffer, with an initializedCount of 0.

  3. Execution: It yields the OutputSpan to the user's closure as an inout parameter.

  4. Cleanup: A defer block ensures that upon exit, finalize(for:) is called on the span to get the count of initialized elements, and then those elements are deinitialized via deinitialize().

Raw Byte Allocation with OutputRawSpan

This function is for working with temporary raw byte buffers.

@available(SwiftCompatibilitySpan 5.0, *)
@_alwaysEmitIntoClient @_transparent
public func withTemporaryAllocation<R: ~Copyable, E: Error>(
  byteCount: Int,
  alignment: Int,
  _ body: (inout OutputRawSpan) throws(E) -> R
) throws(E) -> R where R: ~Copyable {
  try withUnsafeTemporaryAllocation(byteCount: byteCount, alignment: alignment) { (buffer) throws(E) in
    var span = OutputRawSpan(buffer: buffer, initializedCount: 0)
    defer {
      _ = span.finalize(for: buffer)
      span = OutputRawSpan()
    }

    return try body(&span)
  }
}

The flow slightly differs from the OutputSpan version in the cleanup step 4, here's the full walkthrough for completeness:

  1. Allocation: It calls withUnsafeTemporaryAllocation(byteCount:alignment:) to obtain a raw byte buffer.

  2. Span Creation: It creates an OutputRawSpan covering the buffer, with an initializedCount of 0.

  3. Execution: It yields the OutputRawSpan to the user's closure as an inout parameter.

  4. Cleanup: A defer block ensures finalize(for:) is called to consume the span. Since OutputRawSpan deals with raw bytes (presumed to be BitwiseCopyable), no explicit deinitialization call is needed on the buffer itself. The temporary memory is automatically deallocated.

Source compatibility

This is an additive change and does not affect existing code.

ABI compatibility

The functions are marked @_alwaysEmitIntoClient and @_transparent. They will be emitted directly into the client's binary and do not constitute new ABI entry points in the standard library. They rely on existing ABI entry points.

Implications on adoption

These functions make temporary allocations significantly safer and easier to use. They lower the barrier to entry for using stack-allocated temporary memory, as users no longer need to be comfortable with "unsafe" pointer APIs.

Future directions

This proposal covers the primary safe wrappers for temporary allocation. Future work could consider specialized versions, like async overloads.

Alternatives considered

  • Do nothing: Users would continue to use the withUnsafe... variants and manually wrap them in OutputSpan or OutputRawSpan, replicating the boilerplate code proposed here.

  • Member of OutputSpan/OutputRawSpan: We could make these static methods on their respective span types. However, top-level functions better match the existing withUnsafeTemporaryAllocation and withExtendedLifetime patterns.

15 Likes

Some of the most common use cases for withUnsafeTemporaryAllocation() involve working in an unsafe buffer of count 1, then finally move()ing the initialized value out of the buffer into the result and returning it to the caller, allowing construction of values that cannot be cleanly/correctly constructed in Swift the usual way. Is there an equivalent mechanism I can use with OutputSpan?

3 Likes

Yes, it’s OutputSpan's func removeLast() → Element. Once we have a non-copyable container in the standard library, the multiple-element version of removeLast() will be able to return multiple items as well.

7 Likes

Super excited to see this! I’ve been reaching for this in Foundation as we’ve been moving from unsafe primitives to safe versions using span types. I feel this is a fairly straightforward adding that brings the existing APIs in line with the latest span APIs added and it will be a huge improvement for developers reaching for safer APIs. Foundation can definitely adopt these new additions in a number of places. Thanks for working on this!

1 Like

Did you consider exposing new TemporaryAllocation/RawTemporaryAllocation types instead? This would allow you to do async work with the allocation:

let ints = TemporaryAllocation<Int>(capacity: 10)
ints.span.append(1)
await something()
ints.span.append(2)
// [deinit destroys the buffer]
2 Likes

It did come up. At least currently, it's not possible to constrain the lifetime of ints (in your example) to its scope alone. It must be constrained to the lifetime of another value. So it would be possible for ints to escape the context in which its inner stack allocation might be valid.

2 Likes

I consider all async-supporting variants out of scope of this pitch. There are unobvious interactions with async that deserve a separate in-depth discussion. This pitch is laser-focused on making two existing non-async allocation functions safer with Output(Raw)Span and nothing else.

6 Likes

Would you mind clarifying "unobvious interactions"? Because this specific API has a very clear semantic in async-land IMHO, quite a miss that we didn't introduce this immediately. It's fine if you just didn't want to pitch it in this proposal but this made me wonder if there's something I'm missing here?

An async implementation would, instead of the thread stack, use the task local stack allocator. It has the exact semantics we need here and is accessible by the runtime, by being an async function it is guaranteed to have a current task available as well.

I do agree though that all those with bytes APIs are missing async versions, and this is better handled separately. This specific API is kinda special though, because it has a very concrete allocator we'd use here.

2 Likes

Can we also increase the guaranteed stack allocation size from 1024 bytes to 4096 bytes? That’s still less than 1% of the 512KiB stack given to Dispatch and Swift Concurrency threads.

As far as I'm aware, withTemporaryAllocation provides no such guarantees, at least in cross-platform way. Libcs on different platforms have different stack sizes that can also sometimes be overridden by linker flags. withTemporaryAllocation has no upfront knowledge of libc that you'll use in your code and/or what linker flags you'll pass. Without that knowledge it can't appropriately adjust stack allocation optimizations that may or may not be applied.

In summary, I consider that stack size allocation size discussion out of scope of this proposal. These new proposed functions are explicitly specified as implemented in terms of withUnsafeTemporaryAllocation without proposing a change in the existing behavior. If you disagree with how withUnsafeTemporaryAllocation should work under the hood, you're welcome to discuss that in a separate thread.

3 Likes

This limit is imposed by the Swift compiler for all stack promotions (see StackPromotionSizeLimit in the source). withUnsafeTemporaryAllocation() should be using StackPromotionSizeLimit rather than hard-coding 1024, but I ran out of time during the initial implementation to plumb that through. Oops.

In any event, changing the limit here is (as Max has said) out of scope for this particular pitch.

2 Likes

Yes, by "withTemporaryAllocation provides no such guarantees" I'm referring to guarantees in the proposal that introduced it and its public documentation. Even if its implementation may use a hardcoded value at some point, that's still an implementation detail and not a public API contract.

2 Likes

Indeed, to quote:

This function is useful for cheaply allocating storage for a sequence of values for a brief duration. Storage may be allocated on the heap or on the stack, depending on the required size and alignment.

The guarantee here is that with[Unsafe]TemporaryAllocation will somehow allocate memory, and that it will try to do so "cheaply". It would be totally valid for it to never use the stack (though possibly not ideal, depending on the implementation; e.g. calling malloc() even for small blocks would probably be bad, but on a system with less stack space it might choose to use a bump allocator or something).

2 Likes