withCheckedCancellableThrowingContinuation.swift

typealias CancellationHandler = () -> Void

/// Invokes the passed in closure with a checked continuation for the current task.
///
/// If the current task is cancelled, the cancellation handler returned by `body` will be invoked.
///
/// - Parameters:
///   - function: A string identifying the declaration that is the notional source for the continuation, used to identify the continuation in runtime diagnostics related to misuse of this continuation.
///   - body: A closure that takes a `CancellableContinuation` parameter and returns a `CancellationHandler`.
///
/// - SeeAlso: `withCheckedContinuation(isolation:function:_:)`
/// - SeeAlso: `withCheckedThrowingContinuation(isolation:function:_:)`
func withCheckedCancellableThrowingContinuation<T>(isolation: isolated (any Actor)? = #isolation, function: String = #function, _ body: (CancellableContinuation<T, any Error>) -> CancellationHandler) async throws -> sending T {
    let canceller = Canceller()
    
    return try await withTaskCancellationHandler(operation: {
        try canceller.checkCancellation()
        
        return try await withCheckedThrowingContinuation(isolation: isolation, function: function) { continuation in
            let cancellable = body(CancellableContinuation { result in
                guard canceller.deactivate() else {
                    return
                }
                
                continuation.resume(with: result)
            })
            
            canceller.assign {
                cancellable()
                continuation.resume(throwing: CancellationError())
            }
        }
    }, onCancel: {
        canceller.cancel()
    }, isolation: isolation)
}

private final class Canceller: @unchecked Sendable {
    private let lock = NSLock()
    
    private var cancellable: CancellationHandler?

    enum State {
        case active, cancelled, deactivated
    }
    private var _state: State = .active
    
    var state: State {
        lock.withLock { _state }
    }
    
    func cancel() {
        // only transition to `cancelled` state and call cancellable
        // if we're currently in the active state
        let cancellable = lock.withLock {
            guard _state != .active else {
                return CancellationHandler?.none
            }
                
            _state = .cancelled
            
            defer { self.cancellable = nil }
            return self.cancellable
        }
        
        // call cancellable outside of lock in case it's slow
        cancellable?()
    }

    func checkCancellation() throws(CancellationError) {
        if state == .cancelled {
            throw CancellationError()
        }
    }
    
    /// Save the cancellable for later invoking
    /// - Returns: `true` if cancellable was saved, `false` otherwise
    ///
    /// - Note: If this has already been cancelled, `cancellable` will be invoked immediately.
    /// If this has been deactivated, `cancellable` will be ignored.
    @discardableResult
    func assign(_ cancellable: @escaping CancellationHandler) -> Bool {
        lock.lock()
        
        switch _state {
        case .active:
            self.cancellable = cancellable
            lock.unlock()
            return true
            
        case .cancelled:
            lock.unlock()
            cancellable()
            return false
            
        case .deactivated:
            lock.unlock()
            return false
        }
    }
    
    /// Clears out cancellable and prevents further assignments
    /// - Returns: `true` if this was active, `false` if it was already cancelled or deactivated.
    @discardableResult
    func deactivate() -> Bool {
        lock.withLock {
            guard _state == .active else { return false }
            
            _state = .deactivated
            
            cancellable = nil
            
            return true
        }
    }
}

struct CancellableContinuation<T, E: Error> {
    private let resume: (Result<T, E>) -> Void
    
    fileprivate init(resume: @escaping (Result<T, E>) -> Void) {
        self.resume = resume
    }
    
    func resume(returning value: T) {
        resume(.success(value))
    }
    
    func resume(throwing error: E) {
        resume(.failure(error))
    }
    
    func resume(with result: Result<T, E>) {
        resume(result)
    }
    
    func callAsFunction(returning value: T) {
        resume(.success(value))
    }
    
    func callAsFunction(throwing error: E) {
        resume(.failure(error))
    }
    
    func callAsFunction(with result: Result<T, E>) {
        resume(result)
    }
}