withObservationTracking+AsyncSequence.swift

import ConcurrencyExtras

private struct ObservationTrackingState: Sendable {
	var isCancelled: Bool = false
	var updateTaskID: UUID?
}

public func withObservationTracking<T>(_ apply: @escaping () -> T) -> AsyncThrowingStream<T, any Error> {
	let (stream, continuation) = AsyncThrowingStream.makeStream(of: T.self)

	let lock = LockIsolated(ObservationTrackingState())
	continuation.onTermination = { _ in
		lock.withValue {
			$0.isCancelled = true
		}
	}

	@Sendable func update() {
		guard
			lock.withValue({ state -> Bool in
				guard !state.isCancelled else { return false }
				state.updateTaskID = nil
				return true
			})
		else { return }

		let value = withObservationTracking {
			apply()
		} onChange: {
			let taskID = UUID()
			guard
				lock.withValue({ state -> Bool in
					guard !state.isCancelled else { return false }
					state.updateTaskID = taskID
					return true
				})
			else { return }

			RunLoop.current.perform {
				guard lock.updateTaskID == taskID else { return }

				update()
			}
		}

		continuation.yield(value)
	}

	update()

	return stream
}