//===--- OutputSpan.swift -------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2024 - 2026 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#if SPAN_COMPATIBILITY_STUB
import Swift
#endif

/// `OutputSpan` is a reference to a contiguous region of memory that starts
/// with some number of initialized `Element` instances followed by
/// uninitialized memory. It provides operations to access the items it stores,
/// as well as to add new elements and to remove existing ones.
@safe
@frozen
@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
public struct OutputSpan<Element: ~Copyable>: ~Copyable, ~Escapable {
  @usableFromInline
  internal let _pointer: UnsafeMutableRawPointer?

  /// The total number of elements that this output span can contain.
  public let capacity: Int

  @usableFromInline
  internal var _count: Int

  @export(implementation)
  deinit {
    if _count > 0 {
      unsafe _start().withMemoryRebound(
        to: Element.self, capacity: _count
      ) {
        [ workaround = _count ] in
        _ = unsafe $0.deinitialize(count: workaround)
      }
    }
  }

  /// Create an OutputSpan with zero capacity.
  @export(implementation)
  @_lifetime(immortal)
  public init() {
    unsafe _pointer = nil
    capacity = 0
    _count = 0
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan: @unchecked Sendable where Element: Sendable & ~Copyable {}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  @export(implementation)
  @_transparent
  @unsafe
  internal func _start() -> UnsafeMutableRawPointer {
    unsafe _pointer._unsafelyUnwrappedUnchecked
  }

  @export(implementation)
  @_transparent
  @unsafe
  internal func _tail() -> UnsafeMutableRawPointer {
    // NOTE: `_pointer` must be known to be not-nil.
    unsafe _start().advanced(by: _count &* MemoryLayout<Element>.stride)
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// The number of initialized elements in this span.
  @export(implementation)
  @_semantics("fixed_storage.get_count")
  public var count: Int { _assumeNonNegative(_count) }

  /// The number of additional elements that can be added to this span.
  @export(implementation)
  @_transparent
  public var freeCapacity: Int { capacity &- _count }

  /// A Boolean value indicating whether the span is empty.
  @export(implementation)
  @_transparent
  public var isEmpty: Bool { _count == 0 }

  /// A Boolean value indicating whether the span is full.
  @export(implementation)
  @_transparent
  public var isFull: Bool { _count == capacity }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable  {

  @unsafe
  @export(implementation)
  @_lifetime(borrow buffer)
  internal init(
    _uncheckedBuffer buffer: UnsafeMutableBufferPointer<Element>,
    initializedCount: Int
  ) {
    unsafe _pointer = .init(buffer.baseAddress)
    capacity = buffer.count
    _count = initializedCount
  }

  /// Unsafely create an OutputSpan over partly-initialized memory.
  ///
  /// The memory in `buffer` must remain valid throughout the lifetime
  /// of the newly-created `OutputSpan`. Its prefix must contain
  /// `initializedCount` initialized instances, followed by uninitialized
  /// memory. The default value of `initializedCount` is 0, representing
  /// the common case of a completely uninitialized `buffer`.
  ///
  /// - Parameters:
  ///   - buffer: an `UnsafeMutableBufferPointer` to be initialized
  ///   - initializedCount: the number of initialized elements
  ///                       at the beginning of `buffer`.
  @unsafe
  @export(implementation)
  @_lifetime(borrow buffer)
  public init(
    buffer: UnsafeMutableBufferPointer<Element>,
    initializedCount: Int
  ) {
    _precondition(buffer._isWellAligned(), "Misaligned OutputSpan")
    if let baseAddress = buffer.baseAddress {
      _precondition(
        unsafe baseAddress.advanced(by: buffer.count) >= baseAddress,
        "Buffer must not wrap around the address space"
      )
    }
    _precondition(
      0 <= initializedCount && initializedCount <= buffer.count,
      "OutputSpan count is not within capacity"
    )
    unsafe self.init(
      _uncheckedBuffer: buffer, initializedCount: initializedCount
    )
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan {

  /// Unsafely create an OutputSpan over partly-initialized memory.
  ///
  /// The memory in `buffer` must remain valid throughout the lifetime
  /// of the newly-created `OutputSpan`. Its prefix must contain
  /// `initializedCount` initialized instances, followed by uninitialized
  /// memory. The default value of `initializedCount` is 0, representing
  /// the common case of a completely uninitialized `buffer`.
  ///
  /// - Parameters:
  ///   - buffer: a slice of an `UnsafeMutableBufferPointer` to be initialized
  ///   - initializedCount: the number of initialized elements
  ///                       at the beginning of `buffer`.
  @unsafe
  @export(implementation)
  @_lifetime(borrow buffer)
  public init(
    buffer: borrowing Slice<UnsafeMutableBufferPointer<Element>>,
    initializedCount: Int
  ) {
    let rebased = unsafe UnsafeMutableBufferPointer(rebasing: buffer)
    let os = unsafe OutputSpan(
      buffer: rebased, initializedCount: initializedCount
    )
    self = unsafe _overrideLifetime(os, borrowing: buffer)
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// The type that represents an initialized index in an `OutputSpan`.
  public typealias Index = Int

  /// The range of initialized indices for this `OutputSpan`.
  @export(implementation)
  public var indices: Range<Index> {
    unsafe Range(_uncheckedBounds: (0, count))
  }

  // SILOptimizer looks for fixed_storage.check_index semantics for bounds check optimizations.
  @_semantics("fixed_storage.check_index")
  @inline(__always)
  @export(implementation)
  internal func _checkIndex(_ index: Index) {
    _precondition(indices.contains(index), "index out of bounds")
  }

  /// Accesses the element at the specified index.
  ///
  /// - Parameter index: A valid index into this span.
  ///
  /// - Complexity: O(1)
  @export(implementation)
  public subscript(_ index: Index) -> Element {
    @_transparent
    borrow {
      _checkIndex(index)
      return unsafe self[unchecked: index]
    }
    @_transparent
    @_lifetime(self: copy self)
    mutate {
      _checkIndex(index)
      return unsafe &(self[unchecked: index])
    }
  }

  /// Accesses the element at the specified index.
  ///
  /// This subscript does not validate `index`; this is an unsafe operation.
  ///
  /// - Parameter index: A valid index into this span.
  ///
  /// - Complexity: O(1)
  @unsafe
  @export(implementation)
  public subscript(unchecked index: Index) -> Element {
    @_unsafeSelfDependentResult
    borrow {
      unsafe UnsafePointer<Element>(
        _unsafeAddressOfElement(unchecked: index)).pointee
    }
    @_unsafeSelfDependentResult
    @_lifetime(self: copy self)
    mutate {
      let address = unsafe _unsafeAddressOfElement(unchecked: index)
      return unsafe &(UnsafeMutablePointer(address).pointee)
    }
  }

  @unsafe
  @export(implementation)
  internal func _unsafeAddressOfElement(
    unchecked index: Index
  ) -> Builtin.RawPointer {
    let elementOffset = index &* MemoryLayout<Element>.stride
    return unsafe _start().advanced(by: elementOffset)._rawValue
  }

  /// Exchange the elements at the two given indices.
  ///
  /// - Parameter i: A valid index into this span.
  /// - Parameter j: A valid index into this span.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func swapAt(_ i: Index, _ j: Index) {
    _precondition(indices.contains(i))
    _precondition(indices.contains(j))
    unsafe swapAt(unchecked: i, unchecked: j)
  }

  /// Exchange the elements at the two given indices.
  ///
  /// This function does not validate `i` or `j`; this is an unsafe operation.
  ///
  /// - Parameter i: A valid index into this span.
  /// - Parameter j: A valid index into this span.
  @unsafe
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func swapAt(unchecked i: Index, unchecked j: Index) {
    let ri = unsafe _unsafeAddressOfElement(unchecked: i)
    let pi = unsafe UnsafeMutablePointer<Element>(ri)
    let rj = unsafe _unsafeAddressOfElement(unchecked: j)
    let pj = unsafe UnsafeMutablePointer<Element>(rj)
    let temporary = unsafe pi.move()
    unsafe pi.initialize(to: pj.move())
    unsafe pj.initialize(to: consume temporary)
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// Append a single element to this span.
  ///
  /// - Parameter value: The element to append.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func append(_ value: consuming Element) {
    _precondition(_count < capacity, "OutputSpan capacity overflow")
    unsafe _tail().initializeMemory(as: Element.self, to: value)
    _count &+= 1
  }

  /// Remove the last initialized element from this span.
  ///
  /// Returns the last element. The `OutputSpan` must not be empty.
  ///
  /// - Returns: The removed element.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func removeLast() -> Element {
    _precondition(!isEmpty, "OutputSpan underflow")
    _count &-= 1
    return unsafe _tail().withMemoryRebound(to: Element.self, capacity: 1) {
      unsafe $0.move()
    }
  }

  /// Remove the last n elements of this span, returning the memory
  /// they occupy to the uninitialized state.
  ///
  /// `n` must not be greater than `count`.
  ///
  /// - Parameter n: The number of elements to remove.
  ///     `n` must not be negative or greater than `count`.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func removeLast(_ n: Int) {
    _precondition(n >= 0, "Can't remove a negative number of elements")
    _precondition(n <= _count, "OutputSpan underflow")
    _count &-= n
    unsafe _tail().withMemoryRebound(to: Element.self, capacity: n) {
      _ = unsafe $0.deinitialize(count: n)
    }
  }

  /// Remove all this span's elements and return its memory
  /// to the uninitialized state.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func removeAll() {
    guard count > 0 else { return }
    _ = unsafe _start().withMemoryRebound(to: Element.self, capacity: _count) {
      unsafe $0.deinitialize(count: _count)
    }
    _count = 0
  }
}

//MARK: bulk-append functions
@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan {

  /// Repeatedly append an element to this span.
  ///
  /// - Parameters:
  ///   - repeatedValue: The element to append repeatedly.
  ///   - count: The number of times to append `repeatedValue`.
  ///       `count` must not exceed `freeCapacity`.
  @export(implementation)
  @_lifetime(self: copy self)
  public mutating func append(repeating repeatedValue: Element, count: Int) {
    _precondition(count <= freeCapacity, "OutputSpan capacity overflow")
    unsafe _tail().initializeMemory(
      as: Element.self, repeating: repeatedValue, count: count
    )
    _count &+= count
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// Borrow the underlying initialized memory for read-only access.
  @export(implementation)
  @_transparent
  public var span: Span<Element> {
    @_lifetime(borrow self)
    borrowing get {
      let pointer = unsafe _pointer?.assumingMemoryBound(to: Element.self)
      let span = unsafe Span(
        _unchecked: pointer._unsafelyUnwrappedUnchecked,
        count: _count
      )
      return unsafe _overrideLifetime(span, borrowing: self)
    }
  }

  /// Exclusively borrow the underlying initialized memory for mutation.
  @export(implementation)
  @_transparent
  public var mutableSpan: MutableSpan<Element> {
    @_lifetime(&self)
    mutating get {
      let pointer = unsafe _pointer?.assumingMemoryBound(to: Element.self)
      let span = unsafe MutableSpan(
        _unchecked: pointer._unsafelyUnwrappedUnchecked,
        count: _count
      )
      return unsafe _overrideLifetime(span, mutating: &self)
    }
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// Call the given closure with the unsafe buffer pointer addressed by this
  /// OutputSpan and a mutable reference to its count of initialized elements.
  ///
  /// This method provides a way to process or populate an `OutputSpan` using
  /// unsafe operations, such as dispatching to code written in legacy
  /// (memory-unsafe) languages.
  ///
  /// The supplied closure may process the buffer in any way it wants; however,
  /// when it finishes (whether by returning or throwing), it must leave the
  /// buffer in a state that satisfies the invariants of the output span:
  ///
  /// 1. The inout integer passed in as the second argument must be the exact
  ///     number of initialized items in the buffer passed in as the first
  ///     argument.
  /// 2. These initialized elements must be located in a single contiguous
  ///     region starting at the beginning of the buffer. The rest of the buffer
  ///     must hold uninitialized memory.
  ///
  /// This function cannot verify these two invariants, and therefore
  /// this is an unsafe operation. Violating the invariants of `OutputSpan`
  /// may result in undefined behavior.
  ///
  /// - Parameter body: A closure that can read from and write to the
  ///     buffer and update the initialized count.
  /// - Returns: The return value of the `body` closure.
  @export(implementation)
  @_transparent
  @_lifetime(self: copy self)
  @unsafe
  public mutating func withUnsafeMutableBufferPointer<E: Error, R: ~Copyable>(
    _ body: (
      UnsafeMutableBufferPointer<Element>,
      _ initializedCount: inout Int
    ) throws(E) -> R
  ) throws(E) -> R {
    let bytes = unsafe UnsafeMutableRawBufferPointer(
      start: _pointer, count: capacity &* MemoryLayout<Element>.stride
    )
    var initializedCount = _count
    defer {
      _precondition(
        0 <= initializedCount && initializedCount <= capacity,
        "OutputSpan capacity overflow"
      )
      self._count = initializedCount
    }
    return try unsafe bytes.withMemoryRebound(to: Element.self) {
      buffer throws(E) -> R in
      try unsafe body(buffer, &initializedCount)
    }
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  /// Consume the output span and return the number of initialized elements.
  ///
  /// This method should be invoked in the scope where the `OutputSpan` was
  /// created, when it is time to commit the contents of the updated buffer
  /// back into the construct being initialized.
  ///
  /// The context that created the output span is expected to remember what
  /// memory region the span is addressing. This consuming method expects to
  /// receive a copy of the same buffer pointer as a (loose) proof of ownership.
  ///
  /// - Parameter buffer: The buffer we expect the `OutputSpan` to reference.
  ///      This must be the same region of memory passed to
  ///      the `OutputSpan` initializer.
  /// - Returns: The number of initialized elements in the same buffer, as
  ///      tracked by the consumed `OutputSpan` instance.
  @unsafe
  @export(implementation)
  public consuming func finalize(
    for buffer: UnsafeMutableBufferPointer<Element>
  ) -> Int {
    _precondition(
      unsafe UnsafeMutableRawPointer(buffer.baseAddress) == self._pointer
      && buffer.count == self.capacity,
      "OutputSpan identity mismatch"
    )
    let count = self._count
    discard self
    return count
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan {
  /// Consume the output span and return the number of initialized elements.
  ///
  /// This method should be invoked in the scope where the `OutputSpan` was
  /// created, when it is time to commit the contents of the updated buffer
  /// back into the construct being initialized.
  ///
  /// The context that created the output span is expected to remember what
  /// memory region the span is addressing. This consuming method expects to
  /// receive a copy of the same buffer pointer as a (loose) proof of ownership.
  ///
  /// - Parameter buffer: The buffer we expect the `OutputSpan` to reference.
  ///      This must be the same region of memory passed to
  ///      the `OutputSpan` initializer.
  /// - Returns: The number of initialized elements in the same buffer, as
  ///      tracked by the consumed `OutputSpan` instance.
  @unsafe
  @export(implementation)
  public consuming func finalize(
    for buffer: Slice<UnsafeMutableBufferPointer<Element>>
  ) -> Int {
    unsafe finalize(for: UnsafeMutableBufferPointer(rebasing: buffer))
  }
}

#if !SPAN_COMPATIBILITY_STUB
@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: ~Copyable {
  @export(implementation)
  @_lifetime(self: copy self)
  internal mutating func _append(
    moving source: UnsafeMutableBufferPointer<Element>
  ) {
    guard source.count > 0 else { return }
    unsafe self.withUnsafeMutableBufferPointer { dst, dstCount in
      let dstEnd = dstCount + source.count
      precondition(dstEnd <= dst.count, "OutputSpan capacity overflow")
      unsafe dst
        .extracting(unchecked: Range(uncheckedBounds: (dstCount, dstEnd)))
        ._moveInitializeAll(fromContentsOf: source)
      dstCount &+= source.count
    }
  }

  @export(implementation)
  @_lifetime(self: copy self)
  @_lifetime(source: copy source)
  internal mutating func _append(moving source: inout OutputSpan<Element>) {
    unsafe source.withUnsafeMutableBufferPointer { src, srcCount in
      let range = unsafe Range(uncheckedBounds: (0, srcCount))
      let items = unsafe src.extracting(unchecked: range)
      unsafe self._append(moving: items)
      srcCount = 0
    }
  }
}

@available(SwiftCompatibilitySpan 5.0, *)
@_originallyDefinedIn(module: "Swift;CompatibilitySpan", SwiftCompatibilitySpan 6.2)
extension OutputSpan where Element: Copyable {
  @export(implementation)
  @_lifetime(self: copy self)
  internal mutating func _append(copying source: UnsafeBufferPointer<Element>) {
    unsafe self.withUnsafeMutableBufferPointer { dst, dstCount in
      let dstEnd = dstCount + source.count
      precondition(dstEnd <= dst.count, "OutputSpan capacity overflow")
      unsafe dst
        .extracting(unchecked: Range(uncheckedBounds: (dstCount, dstEnd)))
        ._initializeAll(fromContentsOf: source)
      dstCount &+= source.count
    }
  }

  @export(implementation)
  @_lifetime(self: copy self)
  internal mutating func _append(copying source: Span<Element>) {
    source.withUnsafeBufferPointer { buffer in
      unsafe self._append(copying: buffer)
    }
  }
}

@available(SwiftStdlib 6.4, *)
extension OutputSpan: Iterable where Element: ~Copyable {
  @available(SwiftStdlib 6.4, *)
  public typealias Failure = Never

  @export(implementation)
  public var underestimatedCount: Int {
    self.count
  }

  @available(SwiftStdlib 6.4, *)
  @export(implementation)
  @_lifetime(borrow self)
  public func makeBorrowingIterator() -> Span<Element>.BorrowingIterator {
    .init(self.span)
  }
}
#endif
