TemporaryArray: a stack-seeded, heap-spilling array

[dnadoba]

TemporaryArray

A dynamically self-resizing, ~Copyable and ~Escapable array of
potentially noncopyable elements. Its initial storage is a borrowed buffer,
most usefully a stack allocation vended by withTemporaryArray, and it
transparently spills over into freshly allocated heap storage the moment it
grows beyond that initial buffer. As long as the element count stays within a
seed buffer that fit on the stack (currently up to 1KB), the array incurs no
heap traffic at all.

Because it can hold a dependency on borrowed (stack) memory, TemporaryArray
is non-escapable: instances cannot outlive the scope that provides their
initial buffer. To keep the contents past that scope, move them into an owning
container with take().

Its API mirrors UniqueArray (SE-0527), minus the pieces that don't apply to a
scratch/non-escapable type plus a small set of members
unique to its borrow-then-spill design.

Entry point

This is how this type is mainly intended to be initialized. It allocates capacity
on the stack if the storage required is below 1KB, otherwise on the heap.

/// Provides a dynamically-resizing array that is initially backed by a stack
/// allocation of the requested capacity, spilling over to the heap only if it
/// grows beyond it.
///
/// This is the primary way to create a `TemporaryArray`. The array passed to
/// `body` starts empty with room for `capacity` elements. As long as the
/// array's element count stays at or below that capacity, no heap allocation
/// occurs, provided the requested storage fits within the stack budget.
///
/// To keep latency predictable, the initial buffer is placed on the stack only
/// if it occupies at most 1024 bytes; a larger initial `capacity` is heap
/// allocated up front instead. Either way the array grows on the heap once it
/// exceeds its initial capacity.
///
/// The array cannot escape `body` (it is non-escapable), so when nothing is
/// moved out, small inputs never touch the heap:
///
///     let sum = withTemporaryArray(of: Int.self, capacity: 64) { scratch in
///       for x in numbers where isHot(x) {
///         scratch.append(x * x)
///       }
///       var total = 0
///       for i in scratch.indices { total += scratch[i] }
///       return total
///     }
///
/// To keep the elements themselves, move them into an owning container with
/// `take()`.
///
/// - Parameters:
///   - type: The element type of the array.
///   - capacity: The number of elements to reserve up front.
///   - body: A closure that receives the freshly created, empty array.
/// - Returns: The result of `body`.
func withTemporaryArray<Element: ~Copyable, E: Error, R: ~Copyable>(
  of type: Element.Type,
  capacity: Int,
  _ body: (inout TemporaryArray<Element>) throws(E) -> R
) throws(E) -> R

We could add more overloads similar to UniqueArray.init's that already copy elements into the array.

API new or different from UniqueArray

struct TemporaryArray<Element: ~Copyable>: ~Copyable, ~Escapable {

  /// Move the contents of this array into a newly created `UniqueArray`,
  /// leaving this array empty.
  ///
  /// Use this to hand the accumulated elements back out of the scope that owns
  /// the array's seed buffer: the result is an ordinary heap-backed, escapable
  /// container. If the array had already spilled to the heap, its existing
  /// storage is transferred without copying; otherwise the elements are moved
  /// out of the borrowed buffer into a fresh allocation.
  ///
  /// - Complexity: O(1) if the array has already spilled to the heap;
  ///    O(`count`) otherwise.
  mutating func take() -> UniqueArray<Element>

}

extension TemporaryArray /* where Element: Copyable */ {
  /// Appends `count` copies of `repeatedValue` to the end of the array, growing
  /// (and, if needed, spilling to the heap) to make room.
  ///
  /// This is `TemporaryArray`'s analogue of `UniqueArray`'s
  /// `init(repeating:count:)`: it's an append rather than an initializer,
  /// because a `TemporaryArray` is meant to be seeded (typically on the stack)
  /// via `withTemporaryArray` and then filled.
  ///
  /// - Complexity: O(`count`), amortized.
  mutating func append(repeating: Element, count: Int)
}

API shared with UniqueArray

The largest part of the API is identical to UniqueArray.
Similarly, the implementation is largely identical, the main difference being
that deinit checks whether the storage is on the stack or the heap and only
deallocates it if it is on the heap. Reallocation performs a similar check.

extension TemporaryArray where Element: ~Copyable {

  // Initializers
  @_lifetime(immortal) init()
  @_lifetime(immortal) init(capacity: Int)
  @_lifetime(immortal) init<E>(capacity: Int, initializingWith: (inout OutputSpan<Element>) throws(E) -> Void) throws(E)

  // Basics
  var capacity: Int
  var freeCapacity: Int
  @_lifetime(borrow self) var span: Span<Element>
  @_lifetime(&self) var mutableSpan: MutableSpan<Element>
  func isTriviallyIdentical(to: borrowing Self) -> Bool
  mutating func edit<E, R>(_: (inout OutputSpan<Element>) throws(E) -> R) throws(E) -> R
  mutating func reallocate(capacity: Int)
  mutating func reserveCapacity(_: Int)

  // Collection
  typealias Index = Int
  var isEmpty: Bool
  var count: Int
  var startIndex: Int
  var endIndex: Int
  var indices: Range<Int>
  subscript(position: Int) -> Element
  mutating func swapAt(_: Int, _: Int)
  func index(after: Int) -> Int
  func index(before: Int) -> Int
  func formIndex(after: inout Int)
  func formIndex(before: inout Int)
  func index(_: Int, offsetBy: Int) -> Int
  func distance(from: Int, to: Int) -> Int
  func formIndex(_: inout Int, offsetBy: inout Int, limitedBy: Int)

  // Appends
  mutating func append(_: consuming Element)
  mutating func append<E>(addingCount: Int, initializingWith: (inout OutputSpan<Element>) throws(E) -> Void) throws(E)
  mutating func append(moving: UnsafeMutableBufferPointer<Element>)
  mutating func append(moving: inout OutputSpan<Element>)
  mutating func append(copying: UnsafeBufferPointer<Element>)
  mutating func append(copying: UnsafeMutableBufferPointer<Element>)
  mutating func append(copying: Span<Element>)
  mutating func append(copying: some Sequence<Element>)

  // Insertions
  mutating func insert(_: consuming Element, at: Int)
  mutating func insert<E>(addingCount: Int, at: Int, initializingWith: (inout OutputSpan<Element>) throws(E) -> Void) throws(E)
  mutating func insert(moving: UnsafeMutableBufferPointer<Element>, at: Int)
  mutating func insert(moving: inout OutputSpan<Element>, at: Int)
  mutating func insert(copying: UnsafeBufferPointer<Element>, at: Int)
  mutating func insert(copying: UnsafeMutableBufferPointer<Element>, at: Int)
  mutating func insert(copying: Span<Element>, at: Int)
  mutating func insert(copying: some Collection<Element>, at: Int)

  // Removals
  mutating func popLast() -> Element?
  mutating func removeLast() -> Element
  mutating func removeLast(_ k: Int)
  mutating func remove(at: Int) -> Element
  mutating func removeSubrange(_: Range<Index>)
  mutating func removeSubrange(_: some RangeExpression<Index>)
  mutating func removeAll(keepingCapacity: Bool = false)

  // Replacements
  mutating func replaceSubrange<E>(_: Range<Int>, addingCount: Int, initializingWith: (inout OutputSpan<Element>) throws(E) -> Void) throws(E)
  mutating func replaceSubrange(_: Range<Int>, moving: UnsafeMutableBufferPointer<Element>)
  mutating func replaceSubrange(_: Range<Int>, moving: inout OutputSpan<Element>)
  mutating func replaceSubrange(_: Range<Int>, copying: UnsafeBufferPointer<Element>)
  mutating func replaceSubrange(_: Range<Int>, copying: UnsafeMutableBufferPointer<Element>)
  mutating func replaceSubrange(_: Range<Int>, copying: Span<Element>)
  mutating func replaceSubrange(_: Range<Int>, copying: some Collection<Element>)

  // Consumption (closure-based; needs no stored back-reference)
  mutating func consume(_: Range<Index>, consumingWith: (inout InputSpan<Element>) -> Void)
  mutating func replace<E>(removing: Range<Int>, consumingWith: (inout InputSpan<Element>) -> Void, addingCount: Int, initializingWith: (inout OutputSpan<Element>) throws(E) -> Void) throws(E)

  // Copying (Element: Copyable): result owns its storage; pair with take()
  @_lifetime(immortal) func clone() -> Self
  @_lifetime(immortal) func clone(capacity: Int) -> Self

  // Bulk access
  @_lifetime(borrow self) func nextSpan(after: inout Int, maximumCount: Int) -> Span<Element>
  @_lifetime(&self) mutating func nextMutableSpan(after: inout Int, maximumCount: Int) -> MutableSpan<Element>
  @_lifetime(borrow self) func previousSpan(before: inout Int, maximumCount: Int) -> Span<Element>
}

Conformances

TemporaryArray attempts to conform to the same protocols as UniqueArray, but
some of them do not yet support ~Escapable types.

extension TemporaryArray: Equatable where Element: Equatable & ~Copyable {
  static func == (_: borrowing Self, _: borrowing Self) -> Bool
}

extension TemporaryArray: Iterable_ where Element: ~Copyable {
  typealias IterableIterator_ = SpanIterator<Element>
  var underestimatedCount_: Int
  @_lifetime(borrow self) func makeIterableIterator_() -> IterableIterator_
}

extension TemporaryArray: Container              where Element: ~Copyable {}
extension TemporaryArray: BidirectionalContainer where Element: ~Copyable {}
extension TemporaryArray: RandomAccessContainer  where Element: ~Copyable {}

// Methods implemented; conformance pending protocol support for ~Copyable/~Escapable:
extension TemporaryArray /*: Hashable */ where Element: Hashable {
  func hash(into: inout Hasher)
}
extension TemporaryArray /*: CustomStringConvertible */ where Element: ~Copyable {
  var description: String
}
extension TemporaryArray /*: CustomDebugStringConvertible */ where Element: ~Copyable {
  var debugDescription: String
}

UniqueArray API omitted

// Value-populating initializers (Element: Copyable)
init(repeating: Element, count: Int)                        // use append(repeating:count:) instead
init(capacity: Int? = nil, copying: Span<Element>)          // use append(copying: Span<Element>) instead
init(capacity: Int? = nil, copying: some Sequence<Element>) // use append(copying: some Sequence<Element>) instead

// Non-proposal consumption helpers carried by the concrete UniqueArray type
mutating func consumeAll(consumingWith: (inout InputSpan<Element>) -> Void)
mutating func consumeLast(_: Int, consumingWith: (inout InputSpan<Element>) -> Void)
mutating func consume<R: RangeExpression<Index>>(_: R, consumingWith: (inout InputSpan<Element>) -> Void)

// Not available, blocked by ~Escapable: 
// DynamicContainer, RangeReplaceableContainer, MutableContainer

Checklist

• I've read the Contribution Guidelines
• My contributions are licensed under the Swift license.
• I've followed the coding style of the rest of the project.
• I've added tests covering all new code paths my change adds to the project (if appropriate).
• I've added benchmarks covering new functionality (if appropriate).
• I've verified that my change does not break any existing tests or introduce unexplained benchmark regressions.
• I've updated the documentation if necessary.

[MahdiBM]

As someone who had to hand-roll the same type as this in https://github.com/swift-dns/swift-idna, I'm already +1 on this.

I was also trying to extract the impl to https://github.com/swift-dns/swift-tiny-sequence but haven't yet had the time to properly do the whole work (The current impl which is specialized for swift-idna is here, the other one in swift-tiny-sequence is behind).

I'll later check to see if this branch's impl fits my usage (I see no reason why it shouldn't, the impls are similar, but just to be sure before the type is merged / tagged).

[dnadoba]

Great, implementation should be ready for experimentation. Let me know if you have any feedback.

[lorentey]

Looks good! Quick notes:

• As of Swift 6.4, we can have TemporaryArray conditionally conform to Equatable/Hashable. (RigidArray's conformances provide a usable pattern.) 6.4 is not out yet, but I think it's fair to assume it will have shipped with the generalized protocols by the time this gets into a tagged release.

  (FWIW, branch main already assumes a shipping 6.4, although if we need to make a release before we have a stable toolchain, we may end up having to adjust that.)

• We should probably replace take() with concrete initializers on {Rigid,Unique}Array that consume temporary arrays. It would not be wise to treat UniqueArray as The One True Array Type, at least not within this package. (RigidDeque and UniqueDeque can technically also provide the same initializers, but that may be a step too far.)

[dnadoba]

That makes sense.
Regarding the init on {Rigid,Unique}Array, I can't decide if I call the argument init(consuming: inout TemporaryArray) or moving. The elements are always moved but if the storage is currently on the heap it will also consume the storage. I'm leaning a bit more toward consuming which is what I will use for now.
Let me know what you think and I can change it to something else as well.