Awaiting.swift

import Foundation

/// @Awaiting wraps optional properties and projects an async function that waits
/// when `wrappedValue` is nil, until some value is set.
///
/// Throws `CancellationError` when task is cancelled.
///
/// ```
///  @Awaiting var name: String?
///  ...
///  let name = try await $name()
///  ```
@propertyWrapper
final class Awaiting<Element> {

  init(wrappedValue: Element?) {
    self.wrappedValue = wrappedValue
  }

  var wrappedValue: Element? {
    didSet {
      guard let wrappedValue = wrappedValue else { return }
      lock.lock()
      let waiting = self.waiting
      lock.unlock()
      for waiter in waiting {
        waiter.resume(with: wrappedValue)
      }
    }
  }

  var projectedValue: () async throws -> Element {
    getValue
  }

  private var waiting = Set<Waiter>()
  private let lock = NSLock()

  private func getValue() async throws -> Element {
    lock.lock()
    if let value = wrappedValue {
      lock.unlock()
      return value
    }

    let waiter = Waiter()
    waiting.insert(waiter)
    lock.unlock()

    defer {
      lock.lock()
      waiting.remove(waiter)
      lock.unlock()
    }

    return try await withTaskCancellationHandler(
      operation: waiter.getValue,
      onCancel: waiter.cancel
    )
  }

  private final class Waiter: Hashable {

    private var continuation: CheckedContinuation<Element, Error>?
    private var result: Result<Element, Error>?

    @Sendable
    func getValue() async throws -> Element {
      try await withCheckedThrowingContinuation {
        if let result = result {
          $0.resume(with: result)
        } else {
          self.continuation = $0
        }
      }
    }

    @Sendable
    func resume(with value: Element) {
      resume(with: .success(value))
    }

    @Sendable
    func cancel() {
      resume(with: .failure(CancellationError()))
    }

    private func resume(with result: Result<Element, Error>) {
      assert(self.result == nil, "Cannot resume twice")
      if let continuation = continuation {
        continuation.resume(with: result)
      } else {
        self.result = result
      }
    }

    func hash(into hasher: inout Hasher) {
      ObjectIdentifier(self).hash(into: &hasher)
    }

    static func == (lhs: Awaiting<Element>.Waiter, rhs: Awaiting<Element>.Waiter) -> Bool {
      lhs === rhs
    }
  }
}