AsyncStreamを`for await`する`task`をラップしようとしてハマった話

# はじめに

`RxSwift.PublishRelay` や `Combine.PassthroughSubject` に代わるものを作ろうとすると`AsyncStream`を使用することになります。

例えば以下のコードを見てみましょう。

```swift
private var event: ((any Sendable) -> Void)?

var stream: AsyncStream<any Sendable> {
    AsyncStream { cont in
        cont.onTermination = { print($0) }
        event = { cont.yield($0) }
    }
}

let msgs = ["foo", "bar", "baz"]

struct MyView1: View {
    @State private var count = 0
    @Binding var navigationPath: NavigationPath

    var body: some View {
        VStack(spacing: 16) {
            Button("Hello World!") {
                defer { count += 1 }
                event?(msgs[count % msgs.count])
            }
            Text("count: \(count)")
            Button("Submit") {
                navigationPath.append("submit")
            }
            Button("Reset") {
                count = 0
            }
        }
        .padding(20)
        .task {
            for await event in stream {
                print(event) // 流れて来たイベント
            }
        }
        .navigationDestination(for: String.self) {
            switch $0 {
            case "submit":
                MyView2(
                    count: count,
                    navigationPath: $navigationPath
                )
            default: EmptyView()
            }
        }
    }
}

struct MyView2: View {
    var count: Int
    @Binding var navigationPath: NavigationPath

    var body: some View {
        Text("count: \(count)")
    }
}

struct RootView: View {
    @State private var navigationPath = NavigationPath()

    var body: some View {
        NavigationStack(path: $navigationPath) {
            MyView1(navigationPath: $navigationPath)
        }
    }
}
```

このコードをPlaygroundで動かした結果がこちら

![task.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/135304/dde044ec-2324-a5df-7883-016baf99b619.gif)

`stream`に流した値が`task`の中で処理されていることが分かります。
また、画面遷移によって`task`がキャンセルされるため、`stream`内の `cont.onTermination` が発火していることも確認できます。

## 問題のコード

`PassthroughSubject`を使っていた時は `onReceive` を使って

```swift
.onReceive(subject) {
    // イベントを処理する
}
```

と書いていました。
このview側のコードをなるべく変えないよう以下のextensionを切ってみました。

```swift
extension View {
    func onReceive<E: Sendable>(
        _ stream: AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
            for await event in stream {
                action(event)
            }
        }
    }
}
```

そしてMyView1のbodyを以下のように書き換えます

```diff_swift
VStack(spacing: 16) {
    Button("Hello World!") {
        defer { count += 1 }
        event?(msgs[count % msgs.count])
    }
    Text("count: \(count)")
    Button("Submit") {
        navigationPath.append("submit")
    }
    Button("Reset") {
        count = 0
    }
}
.padding(20)
- .task {
-     for await event in stream {
-         print(event) // 流れて来たイベント
-     }
- }
+ .onReceive(stream) { print($0) }
.navigationDestination(for: String.self) {
    switch $0 {
    case "submit":
        MyView2(
            count: count,
            navigationPath: $navigationPath
        )
    default: EmptyView()
    }
}
```

一見問題無いように見えますが、実際に動かすと想定しない動きをします。

![on_receive.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/135304/8478918a-945d-021a-5dfa-415875a6113d.gif)

# 原因

`task`の中はclosureなので、画面表示されてから初めて実行されそのタイミングで`stream`が発行されます。
また、`task`は画面表示後の1度しか実行されない(`@State`の変更に左右されない)ため想定通りの動きになります。
一方、今回作成した `onReceive` は初期表示時の他`@State`が変更される度に呼ばれてしまうため、その度に`stream`が発行されてしまいますが `task` の中は実行されないためそのまま捨てられてしまいます。

# 解決策
この現象を回避する手段として、`onReceive`に渡す`stream`を遅延評価させる方法があります。
具体的には値をそのまま渡すのではなくClosureに閉じ込めて渡します。

```diff_swift
extension View {
    func onReceive<E: Sendable>(
-       _ stream: AsyncStream<E>,
+       _ stream: @escaping () -> AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
-           for await event in stream {
+           for await event in stream() {
                action(event)
            }
        }
    }
}
```

呼び出し方は

```diff_swift
- .onReceive(stream) { print($0) }
+ .onReceive({stream}) { print($0) }
```

または

```diff_swift
- .onReceive(stream) { print($0) }
+ .onReceive { stream } action: { print($0) }
```

のようになります。

...ちょっと見栄えが悪いですね😅

## シン・解決策
Swiftには `@autoclosure` という独特の機能があります。
この機能を使うことで、呼び出し元に手を加えることなく遅延評価を実現することができます。

```diff_swift
extension View {
    func onReceive<E: Sendable>(
-       _ stream: AsyncStream<E>,
+       _ stream: @autoclosure @escaping () -> AsyncStream<E>,
        action: @escaping (E) -> Void
    ) -> some View {
        task {
-           for await event in stream {
+           for await event in stream() {
                action(event)
            }
        }
    }
}
```

そして最終的にできたコードで動かした結果がこちら
![autoclosure.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/135304/15bd8d21-0661-d44a-35cc-8c31b155f78a.gif)

想定通り動くようになりました:tada::tada::tada:

# まとめ
`task(_:)`をラップして関数を作る場合、遅延評価をしないと意図しない動きをする場合があることに注意しましょう。

# 参考
https://docs.swift.org/swift-book/documentation/the-swift-programming-language/closures/#Autoclosures