LockIsolatedValueStream.swift

import ConcurrencyExtras
import struct Foundation.UUID

@dynamicMemberLookup
final class LockIsolatedValueStream<Value>: Sendable where Value: Sendable {
  init(_ value: Value) {
    self._lock = LockIsolated(value)
  }

  var value: Value { _lock.value }

  subscript<Subject: Sendable>(dynamicMember keyPath: KeyPath<Value, Subject>) -> Subject {
    self._lock.value[keyPath: keyPath]
  }

  @discardableResult
  func withValue<T: Sendable>(_ operation: @Sendable (inout Value) throws -> T) rethrows -> T {
    try _lock.withValue { value in
      var copy = value
      defer {
        value = copy
        yield(value)
      }
      return try operation(&copy)
    }
  }

  func setValue(_ newValue: @autoclosure @Sendable () throws -> Value) rethrows {
    try _lock.withValue { value in
      value = try newValue()
      yield(value)
    }
  }

  var stream: any AsyncSequence<Value, Never> {
    AsyncStream { continuation in
      let id = addContinuation(continuation)
      continuation.onTermination = { [weak self] _ in
        self?.removeContinuation(id)
      }
      continuation.yield(value)
    }
  }

  private let _lock: LockIsolated<Value>
  private let _continuations = LockIsolated<[UUID: AsyncStream<Value>.Continuation]>([:])

  private func addContinuation(_ continuation: AsyncStream<Value>.Continuation) -> UUID {
    _continuations.withValue {
      let id = UUID()
      $0[id] = continuation
      return id
    }
  }

  private func removeContinuation(_ id: UUID) {
    _continuations.withValue { $0[id] = nil }
  }

  private func yield(_ value: Value) {
    _continuations.withValue { [value] continuations in
      for continuation in continuations.values {
        continuation.yield(value)
      }
    }
  }
}