Intrusive, noncopyable version of the Swift ManagedBuffer type

IntrusiveManagedBuffer.swift

/// A type that has some number of elements of a given element type that follow
/// it directly in memory.
public protocol TrailingElements: ~Copyable {
    /// The element type of the data that follows the value in memory.
    ///
    /// We require this data to be BitwiseCopyable because it is not managed
    /// in any way.
    associatedtype Element: BitwiseCopyable

    /// The number of elements following the value.
    var count: Int { get }
}

/// A value that manages a contiguous block of memory of a given type, starting
/// with a value of the given Header type. The memory is of type Header.Element,
/// and the number of elements is derivable from the header value, so it is
/// not stored separately.
public struct IntrusiveManagedBuffer<Header: TrailingElements>: ~Copyable
    where Header: ~Copyable
{
    /// The underlying storage for the Header instance, followed by the array.
    public let pointer: UnsafeMutablePointer<Header>

    /// The element type stored within the buffer.
    public typealias Element = Header.Element

    /// Allocate an intrusive managed buffer with the given header.
    public init(header: consuming Header) {
        // Allocate enough memory for the header + the number of elements.
        let totalSize = Self.allocationSize(header: header)
        let pointer = UnsafeMutableRawPointer.allocate(
            byteCount: totalSize,
            alignment: MemoryLayout<Header>.alignment
        )

        self.pointer = pointer.assumingMemoryBound(to: Header.self)
        self.pointer.initialize(to: header)
    }

    /// Take ownership over a pointer to memory containing the header followed
    /// by the trailing elements.
    ///
    /// Once this managed buffer instance is no longer used, this memory will
    /// be freed.
    public init(consuming pointer: UnsafeMutablePointer<Header>) {
        self.pointer = pointer
    }

    /// Take ownership over the value, returning its pointer. The underlying
    /// storage will not be freed.
    public consuming func takePointer() -> UnsafeMutablePointer<Header> {
        let pointer = self.pointer
        discard self
        return pointer
    }

    deinit {
        // Deinitialize and deallocate the pointer itself.
        pointer.deinitialize(count: 1).deallocate()
    }

    /// Access the header portion of the value.
    public var header: Header {
        unsafeAddress {
            UnsafePointer(pointer)
        }

        unsafeMutableAddress {
            pointer
        }
    }

    /// The number of elements in the value.
    var count: Int { header.count }
    var startIndex: Int { 0 }
    var endIndex: Int { count }
    var indices: Range<Int> { 0..<count }

    subscript(index: Int) -> Element {
        get {
            precondition(index >= 0 && index < count)
            return rawElements[index]
        }

        set {
            precondition(index >= 0 && index < count)
            rawElements[index] = newValue
        }
    }

    /// Access the flexible array elements.
    var rawElements: UnsafeMutableBufferPointer<Element> {
        UnsafeMutableBufferPointer(
            start: UnsafeMutableRawPointer(pointer.advanced(by: 1))
                     .assumingMemoryBound(to: Element.self),
            count: header.count
        )
    }

    @available(macOS 14.4, *)
    public var span: Span<Element> {
        @_lifetime(self)
        get {
            _overrideLifetime(rawElements.span, borrowing: self)
        }
    }

    @available(macOS 14.4, *)
    public var mutableSpan: MutableSpan<Element> {
        @_lifetime(self)
        mutating get {
            let elements = self.rawElements.mutableSpan
            return _overrideLifetime(elements, mutating: &self)
        }
    }

    /// Determine the allocation size needed for the given header value.
    public static func allocationSize(header: borrowing Header) -> Int {
        MemoryLayout<Header>.size + MemoryLayout<Element>.stride * header.count
    }
}

extension IntrusiveManagedBuffer where Header: Copyable {
    /// Create a temporary intrusive managed buffer for the given header,
    /// providing that instance to the given `body` to operate on for the
    /// duration. The temporary is allocated on the stack (unless it is very
    /// large).
    public static func withTemporary<R, E>(
        header: Header,
        body: (inout IntrusiveManagedBuffer<Header>) throws(E) -> R
    ) throws(E) -> R {
        // Allocate temporary storage large enough for the value we need.
        let result: Result<R, E> = withUnsafeTemporaryAllocation(
            byteCount: allocationSize(header: header),
            alignment: MemoryLayout<Header>.alignment
        ) { buffer in
            /// Create a managed buffer over that temporary storage.
            let pointer = buffer.baseAddress!.assumingMemoryBound(to: Header.self)
            pointer.initialize(to: header)
            var managedBuffer = IntrusiveManagedBuffer(consuming: pointer)

            do throws(E) {
                let result = try body(&managedBuffer)

                // Tell the managed buffer not to free the storage.
                let finalPointer = managedBuffer.takePointer()
                precondition(finalPointer == pointer)

                return .success(result)
            } catch {
                // Tell the managed buffer not to free the storage.
                let finalPointer = managedBuffer.takePointer()
                precondition(finalPointer == pointer)

                return .failure(error)
            }
        }

        return try result.get()
    }
}