AsyncBufferedByteIterator.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift Async Algorithms open source project
//
// Copyright (c) 2022 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
//
//===----------------------------------------------------------------------===//

/// An `AsyncIterator` that provides a convenient and high-performance implementation
/// of a common architecture for `AsyncSequence` of `UInt8`, otherwise known as byte streams.
///
/// Bytes are read into an internal buffer of `capacity` bytes via the
/// `readFunction`. Invoking `next()` returns bytes from the internal buffer until it's
/// empty, and then suspends and awaits another invocation of `readFunction` to
/// refill. If `readFunction` returns 0 (indicating nothing was read), `next()` will
/// return `nil` from then on. Cancellation is checked before each invocation of
/// `readFunction`, which means that many calls to `next()` will not check for
/// cancellation.
///
/// A typical use of `AsyncBufferedByteIterator` looks something like this:
///
/// struct AsyncBytes: AsyncSequence {
/// public typealias Element = UInt8
/// var handle: ReadableThing
///
/// internal init(_ readable: ReadableThing) {
/// handle = readable
/// }
///
/// public func makeAsyncIterator() -> AsyncBufferedByteIterator {
/// return AsyncBufferedByteIterator(capacity: 16384) { buffer in
/// // This runs once every 16384 invocations of next()
/// return try await handle.read(into: buffer)
/// }
/// }
/// }
///
///
@available(AsyncAlgorithms 1.0,*)
publicstructAsyncBufferedByteIterator:AsyncIteratorProtocol{
publictypealiasElement=UInt8
@usableFromInlinevarbuffer:_AsyncBytesBuffer

 /// Creates an asynchronous buffered byte iterator with a specified capacity and read function.
 ///
 /// - Parameters:
 /// - capacity: The maximum number of bytes that a single invocation of `readFunction` may produce.
 /// This is the allocated capacity of the backing buffer for iteration; the value must be greater than 0.
 /// - readFunction: The function for refilling the buffer.
publicinit(
 capacity:Int,
 readFunction:@Sendable@escaping(UnsafeMutableRawBufferPointer)asyncthrows->Int
){
 buffer =_AsyncBytesBuffer(capacity: capacity, readFunction: readFunction)
}

 /// Reads a byte out of the buffer if available. When no bytes are available, this will trigger
 /// the read function to reload the buffer and then return the next byte from that buffer.
@inlinable@inline(__always)
publicmutatingfunc next()asyncthrows->UInt8?{
returntryawait buffer.next()
}
}

@available(*, unavailable)
extensionAsyncBufferedByteIterator:Sendable{}

@available(AsyncAlgorithms 1.0,*)
@frozen@usableFromInline
internalstruct_AsyncBytesBuffer{
@usableFromInline
finalclassStorage{
fileprivateletbuffer:UnsafeMutableRawBufferPointer

init(
 capacity:Int
){
precondition(capacity >0)
 buffer =UnsafeMutableRawBufferPointer.allocate(
 byteCount: capacity,
 alignment: MemoryLayout<AnyObject>.alignment
)
}

deinit{
 buffer.deallocate()
}
}

@usableFromInlineinternalletstorage:Storage
@usableFromInlineinternalvarnextPointer:UnsafeRawPointer
@usableFromInlineinternalvarendPointer:UnsafeRawPointer

internalletreadFunction:@Sendable(UnsafeMutableRawBufferPointer)asyncthrows->Int
internalvarfinished=false

@usableFromInlineinit(
 capacity:Int,
 readFunction:@Sendable@escaping(UnsafeMutableRawBufferPointer)asyncthrows->Int
){
lets=Storage(capacity: capacity)
self.readFunction = readFunction
 storage = s
 nextPointer =UnsafeRawPointer(s.buffer.baseAddress!)
 endPointer = nextPointer
}

@inline(never)@usableFromInline
internalmutatingfunc reloadBufferAndNext()asyncthrows->UInt8?{
if finished {
returnnil
}
tryTask.checkCancellation()
do{
letreadSize:Int=tryawaitreadFunction(storage.buffer)
if readSize ==0{
 finished =true
 nextPointer = endPointer
returnnil
}
 nextPointer =UnsafeRawPointer(storage.buffer.baseAddress!)
 endPointer = nextPointer + readSize
}catch{
 finished =true
 nextPointer = endPointer
throw error
}
returntryawaitnext()
}

@inlinable@inline(__always)
internalmutatingfunc next()asyncthrows->UInt8?{
if_fastPath(nextPointer != endPointer){
letbyte= nextPointer.load(fromByteOffset:0, as:UInt8.self)
 nextPointer = nextPointer +1
return byte
}
returntryawaitreloadBufferAndNext()
}
}