[DNM] OutputSpan prototype (updated)

Runtimes/Core/core/CMakeLists.txt

@@ -154,6 +154,8 @@ add_library(swiftCore
   Span/RawSpan.swift
   Span/MutableSpan.swift
   Span/MutableRawSpan.swift
+  Span/OutputSpan.swift
+  Span/OutputRawSpan.swift
   StaticString.swift
   StaticPrint.swift
   Stride.swift

stdlib/public/core/Array.swift

@@ -1298,6 +1298,32 @@ extension Array: RangeReplaceableCollection {
     }
   }
 
+  /// Grows the array to ensure capacity for the specified number of elements,
+  /// then calls the closure with an OutputSpan covering the array's
+  /// uninitialized memory.
+  @_alwaysEmitIntoClient
+  @available(SwiftStdlib 6.2, *)
+  public mutating func append<E: Error>(
+    addingCapacity count: Int,
+    initializingWith initializer: (
+      _ span: inout OutputSpan<Element>
+    ) throws(E) -> Void
+  ) throws(E) {
+    _reserveCapacityImpl(
+      minimumCapacity: self.count + count, growForAppend: true
+    )
+    defer {
+      _endMutation()
+    }
+    let oldCount = _buffer.mutableCount
+    let newElements = unsafe _buffer.mutableFirstElementAddress + oldCount
+    let raw = unsafe UnsafeMutableBufferPointer(start: newElements, count: count)
+    var span = unsafe OutputSpan(buffer: raw, initializedCount: 0)
+    try initializer(&span)
+    let initialized = unsafe span.finalize(for: raw)
+    _buffer.mutableCount = oldCount + initialized
+  }
+
   @inlinable
   @_semantics("array.reserve_capacity_for_append")
   @_effects(notEscaping self.**)
@@ -1578,6 +1604,39 @@ extension Array {
       initializingWith: initializer)
   }
 
+  //FIXME: typed throws
+  /// Creates an array with the specified capacity, then calls the given
+  /// closure with an OutputSpan covering the array's uninitialized memory.
+  ///
+  /// - Note: While the resulting array may have a capacity larger than the
+  ///   requested amount, the `OutputSpan` passed to the closure will cover
+  ///   exactly the requested number of elements.
+  ///
+  /// - Parameters:
+  ///   - capacity: The number of elements to allocate
+  ///     space for in the new array.
+  ///   - initializer: A closure that initializes elements
+  ///     - Parameters:
+  ///       - span: An `OutputSpan` covering uninitialized memory with room
+  ///         for the specified number of elements.
+  @_alwaysEmitIntoClient
+  @available(SwiftStdlib 6.2, *)
+  public init(
+    capacity: Int,
+    initializingWith initializer: (
+      _ span: inout OutputSpan<Element>
+    ) throws -> Void
+  ) rethrows {
+    try unsafe self.init(
+      unsafeUninitializedCapacity: capacity,
+      initializingWith: { (buffer, count) in
+        var output = unsafe OutputSpan(buffer: buffer, initializedCount: 0)
+        try initializer(&output)
+        count = unsafe output.finalize(for: buffer)
+      }
+    )
+  }
+
   // Superseded by the typed-throws version of this function, but retained
   // for ABI reasons.
   @usableFromInline

stdlib/public/core/CMakeLists.txt

@@ -158,6 +158,8 @@ split_embedded_sources(
   EMBEDDED Sort.swift
   EMBEDDED Span/MutableRawSpan.swift
   EMBEDDED Span/MutableSpan.swift
+  EMBEDDED Span/OutputRawSpan.swift
+  EMBEDDED Span/OutputSpan.swift
   EMBEDDED Span/RawSpan.swift
   EMBEDDED Span/Span.swift
   EMBEDDED StaticString.swift

stdlib/public/core/ContiguousArray.swift

@@ -1106,6 +1106,23 @@ extension ContiguousArray {
       initializingWith: initializer))
   }
 
+  //FIXME: typed throws
+  @_alwaysEmitIntoClient
+  @available(SwiftStdlib 6.2, *)
+  public init(
+    capacity: Int,
+    initializingWith initializer: (inout OutputSpan<Element>) throws -> Void
+  ) rethrows {
+    try unsafe self.init(
+      unsafeUninitializedCapacity: capacity,
+      initializingWith: { (buffer, count) in
+        var output = unsafe OutputSpan(buffer: buffer, initializedCount: 0)
+        try initializer(&output)
+        count = unsafe output.finalize(for: buffer)
+      }
+    )
+  }
+
   // Superseded by the typed-throws version of this function, but retained
   // for ABI reasons.
   @usableFromInline

stdlib/public/core/GroupInfo.json

@@ -202,6 +202,8 @@
   "Span": [
     "MutableRawSpan.swift",
     "MutableSpan.swift",
+    "OutputRawSpan.swift",
+    "OutputSpan.swift",
     "RawSpan.swift",
     "Span.swift"
   ],

stdlib/public/core/InlineArray.swift

@@ -233,6 +233,25 @@ extension InlineArray where Element: ~Copyable {
     }
 #else
     fatalError()
+#endif
+  }
+
+  @available(SwiftStdlib 6.2, *)
+  @_alwaysEmitIntoClient
+  public init<E: Error>(
+    initializingWith initializer: (inout OutputSpan<Element>) throws(E) -> Void
+  ) throws(E) {
+#if $BuiltinEmplaceTypedThrows
+    self = try Builtin.emplace { (rawPtr) throws(E) -> () in
+      let buffer = unsafe Self._initializationBuffer(start: rawPtr)
+      _internalInvariant(Self.count == buffer.count)
+      var output = unsafe OutputSpan(buffer: buffer, initializedCount: 0)
+      try initializer(&output)
+      let initialized = unsafe output.finalize(for: buffer)
+      _precondition(count == initialized, "InlineArray initialization underflow")
+    }
+#else
+    fatalError()
 #endif
   }
 }
@@ -483,13 +502,33 @@ extension InlineArray where Element: ~Copyable {
   }
 }
 
+@available(SwiftStdlib 6.2, *)
+extension InlineArray where Element: BitwiseCopyable {
+
+  @available(SwiftStdlib 6.2, *)
+  @_alwaysEmitIntoClient
+  public consuming func _consume<E: Error, Result: ~Copyable>(
+    _ body: (inout OutputSpan<Element>) throws(E) -> Result
+  ) throws(E) -> Result {
+    var span = unsafe OutputSpan(
+      _uncheckedBuffer: _mutableBuffer, initializedCount: count
+    )
+    defer {
+      _precondition(span.count == 0, "must consume all elements")
+    }
+    //FIXME: discard self
+    return try body(&span)
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: - Unsafe APIs
 //===----------------------------------------------------------------------===//
 
 @available(SwiftStdlib 6.2, *)
 extension InlineArray where Element: ~Copyable {
-  // FIXME: @available(*, deprecated, renamed: "span.withUnsafeBufferPointer(_:)")
+
+  @available(*, deprecated, renamed: "span.withUnsafeBufferPointer(_:)")
   @available(SwiftStdlib 6.2, *)
   @_alwaysEmitIntoClient
   @_transparent
@@ -499,7 +538,7 @@ extension InlineArray where Element: ~Copyable {
     try unsafe body(_buffer)
   }
 
-  // FIXME: @available(*, deprecated, renamed: "mutableSpan.withUnsafeMutableBufferPointer(_:)")
+  @available(*, deprecated, renamed: "mutableSpan.withUnsafeMutableBufferPointer(_:)")
   @available(SwiftStdlib 6.2, *)
   @_alwaysEmitIntoClient
   @_transparent

stdlib/public/core/Span/MutableSpan.swift

@@ -71,11 +71,7 @@ extension MutableSpan where Element: ~Copyable {
   public init(
     _unsafeElements buffer: UnsafeMutableBufferPointer<Element>
   ) {
-    _precondition(
-      ((Int(bitPattern: buffer.baseAddress) &
-        (MemoryLayout<Element>.alignment &- 1)) == 0),
-      "baseAddress must be properly aligned to access Element"
-    )
+    _precondition(buffer._isWellAligned(), "Misaligned MutableSpan")
     let ms = unsafe MutableSpan<Element>(_unchecked: buffer)
     self = unsafe _overrideLifetime(ms, borrowing: buffer)
   }
@@ -120,17 +116,16 @@ extension MutableSpan where Element: BitwiseCopyable {
   public init(
     _unsafeBytes buffer: UnsafeMutableRawBufferPointer
   ) {
+    let baseAddress = unsafe buffer.baseAddress?
+                               .assumingMemoryBound(to: Element.self)
     _precondition(
-      ((Int(bitPattern: buffer.baseAddress) &
-        (MemoryLayout<Element>.alignment &- 1)) == 0),
-      "baseAddress must be properly aligned to access Element"
+      unsafe baseAddress?._isWellAligned() ?? true, "Misaligned MutableSpan"
     )
     let (byteCount, stride) = (buffer.count, MemoryLayout<Element>.stride)
     let (count, remainder) = byteCount.quotientAndRemainder(dividingBy: stride)
     _precondition(remainder == 0, "Span must contain a whole number of elements")
     let elements = unsafe UnsafeMutableBufferPointer<Element>(
-      start: buffer.baseAddress?.assumingMemoryBound(to: Element.self),
-      count: count
+      start: baseAddress, count: count
     )
     let ms = unsafe MutableSpan(_unsafeElements: elements)
     self = unsafe _overrideLifetime(ms, borrowing: buffer)