[Swift] withUnsafeContinuationとexecutorに関するバグ

# TL;DR

「最初の行からasync呼び出しがあるisolated-async function」にかかる、executorへのhopを省略する最適化が`@_unsafeInheritExecutor`を想定できていない問題

apple/swiftに報告済み

https://github.com/apple/swift/issues/69659

# 検証環境

- Swift compiler version info: swift-driver version: 1.87.1 Apple Swift version 5.9 (swiftlang-5.9.0.128.108 clang-1500.0.40.1)
- Xcode version info: Xcode 15.0.1 Build version 15A507
- Deployment target: arm64-apple-macosx14.0

# 実験コード

```swift
@MainActor
func doSomething() async {
    await withUnsafeContinuation { continuation in
        if #available(macOS 14.0, *) {
            MainActor.shared.assertIsolated()
        }

        continuation.resume(returning: ())
    }
}

await Task.detached {
    await doSomething()
}.value
```

`MainActor.shared.assertIsolated()`が失敗する。

`withUnsafeContinuation`は`@_unsafeInheritExecutor`によって呼び出し元のexecutorを引き継ぐ仕様のはずが、引き継いでいないことがわかります。

期待する仕様としては以下のはずが

```swift
@MainActor
func doSomething() async {
    // ここはMainActor
    await withUnsafeContinuation { continuation in
        // ここもMainActor（`@_unsafeInheritExecutor`）
        if #available(macOS 14.0, *) {
            MainActor.shared.assertIsolated()
        }

        continuation.resume(returning: ())
    }
}

await Task.detached {
    // ここはどこのActor executorでもない
    await doSomething()
}.value
```

実際には以下になっていると言うことです。

```swift
@MainActor
func doSomething() async {
    await withUnsafeContinuation { continuation in
        // ここがMainActorになっていない
        if #available(macOS 14.0, *) {
            MainActor.shared.assertIsolated()
        }

        continuation.resume(returning: ())
    }
}
```

# なぜ？

SILを覗くとわかります。

<details><summary>フル出力</summary>


https://gist.github.com/kntkymt/3f8af33b98f238128a0eba638b74b81d

</details>

`doSomething()`定義部分に答えはありました。

raw出力（最適化なし）

```swift
// swiftc -emit-silgen main.swift
// doSomething()
sil hidden [ossa] @$s4main11doSomethingyyYaF : $@convention(thin) @async () -> () {
bb0:
  %0 = metatype $@thick MainActor.Type            // user: %2
  // function_ref static MainActor.shared.getter
  %1 = function_ref @$sScM6sharedScMvgZ : $@convention(method) (@thick MainActor.Type) -> @owned MainActor // user: %2
  %2 = apply %1(%0) : $@convention(method) (@thick MainActor.Type) -> @owned MainActor // users: %14, %3
  %3 = begin_borrow %2 : $MainActor               // users: %13, %10, %4

  // 1. MainActorのexecutorにhop
  hop_to_executor %3 : $MainActor                 // id: %4
  %5 = alloc_stack $()                            // users: %12, %11, %9

  // 2. withUnsafeContinuationを呼ぶ
  // function_ref closure #1 in doSomething()
  %6 = function_ref @$s4main11doSomethingyyYaFySccyyts5NeverOGXEfU_ : $@convention(thin) @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> // user: %7
  %7 = thin_to_thick_function %6 : $@convention(thin) @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> to $@noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> // user: %9
  // function_ref withUnsafeContinuation<A>(_:)
  %8 = function_ref @$ss22withUnsafeContinuationyxySccyxs5NeverOGXEYalF : $@convention(thin) @async <τ_0_0> (@guaranteed @noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <τ_0_0>) -> @out τ_0_0 // user: %9
  %9 = apply %8<()>(%5, %7) : $@convention(thin) @async <τ_0_0> (@guaranteed @noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <τ_0_0>) -> @out τ_0_0

  // 3. MainActorのexecutorに再度hop（戻ってくる）
  hop_to_executor %3 : $MainActor                 // id: %10
  %11 = load [trivial] %5 : $*()
  dealloc_stack %5 : $*()                         // id: %12
  end_borrow %3 : $MainActor                      // id: %13
  destroy_value %2 : $MainActor                   // id: %14
  %15 = tuple ()                                  // user: %16
  return %15 : $()                                // id: %16
} // end sil function '$s4main11doSomethingyyYaF'
```

canonical出力（最適化あり）

```swift
// swiftc -emit-sil main.swift
// doSomething()
sil hidden @$s4main11doSomethingyyYaF : $@convention(thin) @async () -> () {
bb0:
  %0 = metatype $@thick MainActor.Type            // user: %2
  // function_ref static MainActor.shared.getter
  %1 = function_ref @$sScM6sharedScMvgZ : $@convention(method) (@thick MainActor.Type) -> @owned MainActor // user: %2
  %2 = apply %1(%0) : $@convention(method) (@thick MainActor.Type) -> @owned MainActor // users: %10, %8
  %3 = alloc_stack $()                            // users: %9, %7

  // MainActor executorへのhopがない！！

  // 1. withUnsafeContinuationを呼ぶ
  // function_ref closure #1 in doSomething()
  %4 = function_ref @$s4main11doSomethingyyYaFySccyyts5NeverOGXEfU_ : $@convention(thin) @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> // user: %5
  %5 = thin_to_thick_function %4 : $@convention(thin) @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> to $@noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <()> // user: %7
  // function_ref withUnsafeContinuation<A>(_:)
  %6 = function_ref @$ss22withUnsafeContinuationyxySccyxs5NeverOGXEYalF : $@convention(thin) @async <τ_0_0> (@guaranteed @noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <τ_0_0>) -> @out τ_0_0 // user: %7
  %7 = apply %6<()>(%3, %5) : $@convention(thin) @async <τ_0_0> (@guaranteed @noescape @callee_guaranteed @substituted <τ_0_0> (UnsafeContinuation<τ_0_0, Never>) -> () for <τ_0_0>) -> @out τ_0_0
  hop_to_executor %2 : $MainActor                 // id: %8
  dealloc_stack %3 : $*()                         // id: %9
  strong_release %2 : $MainActor                  // id: %10
  %11 = tuple ()                                  // user: %12
  return %11 : $()                                // id: %12
} // end sil function '$s4main11doSomethingyyYaF'
```

最適化なしのSILでは、`withUnsafeContinuation`を呼ぶ前にMainActor executorへのhopがあるのに対して、最適化ありのSILではそれがありません。

このことから、最適化によってこのバグは引き起こされていることが推測できます。

# なんのための最適化？

SILの最適化コードを読んでいないのでここからは推測ですが

基本的には、最初の行から（別のActor-isolated若しくはnon isolatedな）async関数への呼び出しがあるisolated async functionは、呼び出し前にisolatedされたActorのexecutorへのhopが必要ありません。
なぜなら、最初の行の呼び出しで別のexecutorへhopされるので、呼び出し前にisolatedなActorのexecutorにhopしても無駄になるからです。

```swift
@MainActor
func A() async {
    // ここでMainActorのexecutorにhopしても
    await B()　// ここの先で別のexecutorにhopされるので無駄になる
}

func B() async {
}
```

そのため、この無駄なhopを減らしてパフォーマンスを上げる最適化が存在すると推測しています。

一方で、`@_unsafeInheritExecutor`が絡んでくると話は変わります。
`@_unsafeInheritExecutor`は呼び出し元のexecutorを引き継ぐので、上記の最適化を当てはめてしまうと、問題になります。
今回のバグは、この最適化が`@_unsafeInheritExecutor`を想定できていないのが原因だと考えています。

```swift
@MainActor
func A() async {
    // ここでMainActorのexecutorにhopしないと
    await B()　// ここの先がMainActorのexecutorで実行できない
}

@_unsafeInheritExecutor
func B() async {
}
```

つまり
「最初の行からasync呼び出しがあるisolated-async function」にかかる、executorへのhopを省略する最適化が`@_unsafeInheritExecutor`を想定できていない問題
と言うわけです。

## 補足

証拠の一つに、実は`doSomething`の先頭に何か別の処理を一行以上書くと期待通り動作します。

```swift
@MainActor
func doSomething() async {
    print("Hello") // なんか一行以上書く
    await withUnsafeContinuation { continuation in
        if #available(macOS 14.0, *) {
            MainActor.shared.assertIsolated() // OK
        }

        continuation.resume(returning: ())
    }
}

await Task.detached {
    await doSomething()
}.value
```

また、class/structで記述した際に、同じ型内の別のメソッドをどこがで使用すれば動きます。
（これはよくわかりませんが、最適化の条件に「型内の他の要素を触っているか」が含まれていそうです。）

```swift
class Hoge {
    func hoge() {}

    @MainActor
    func doSomething() async {
        await withUnsafeContinuation { continuation in
            hoge()
            continuation.resume(returning: ())
        }
    }
}

await Task.detached {
    let hoge = Hoge()
    await hoge.doSomething()
}.value
```

以下の記事のシチュエーションでは、この現象によってバグが回避できています。

https://qiita.com/kntkymt/items/60602d6326ed92c2f8a0

# どうすればいい？

`withCheckedContinuation`を用いれば大丈夫だと思います。
`withCheckedContinuation`は、名前の通りContinuationの使用をCheckするためのコードが呼び出し前に挿入されるため、「最初の行からasync呼び出しがあるisolated-async function」に該当せず、今回の最適化がかかりません。

# 併せて読みたい

https://qiita.com/kntkymt/items/a588bf316e7478134e75

# 余談

推測ではなく最適化のコードを読んで追記したい