[C#]UniRxのAsyncMethodBuilderについてのメモ

UniRxのソースを呼んでいるとAsyncMethodBuilderの実装が以下のようになっていました。

// UniRx AsyncUniTaskMethodBuilder.cs より引用
// 5. AwaitOnCompleted
[DebuggerHidden]
public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
    where TAwaiter : INotifyCompletion
    where TStateMachine : IAsyncStateMachine
{
    if (moveNext == null)
    {
        if (promise == null)
        {
            promise = new UniTaskCompletionSource(); // built future.
        }

        var runner = new MoveNextRunner<TStateMachine>();
        moveNext = runner.Run;
        runner.StateMachine = stateMachine; // set after create delegate.
    }

    awaiter.OnCompleted(moveNext);
}

set after create delegateと書かれている部分でなぜそうでなければいけないか疑問だったのですが、
ようやく理解できたので備忘録としてメモしておきます。

---

メモリ初期状態。

# スタック領域
StateMachine
  - AsyncMethodBuilder
  - その他ローカル変数

# ヒープ領域
なし

Awaiterに後続処理を依頼するためにはActionを渡す必要がある。
Actionはクラスオブジェクトなのでヒープ領域に作成しなければならない。
よって最終的には以下のようなメモリレイアウトになる。

# スタック領域
StateMachine
  - AsyncMethodBuilder
  - その他ローカル変数

# ヒープ領域
Action
  - StateMachine Pointer
  - Method Address

StateMachine(スタック領域からのコピー, Actionから参照される)
  - AsyncMethodBuilder
  - その他ローカル変数

これを踏まえて再度UniRxのコードを見る。

// UniRx AsyncUniTaskMethodBuilder.cs より引用
var runner = new MoveNextRunner<TStateMachine>();
moveNext = runner.Run;
runner.StateMachine = stateMachine; // set after create delegate.

もしrunner.StateMachineへの代入がdelegateを作るより前に行われた場合、

1. StateMachineをスタック領域からヒープ領域にコピー
2. スタック領域のAsyncMethodBuilderのmoveNextを書き換える

の順番で処理が行われる。

# スタック領域
StateMachine
  - AsyncMethodBuilder <- ここのmoveNextが書き換わる
  - その他ローカル変数

# ヒープ領域
StateMachine
  - AsyncMethodBuilder <- こっちは書き換える前にコピーしてしまったのでそのまま
  - その他ローカル変数

ヒープ領域にコピーした側のmoveNextを書き換えなければ明らかにまずい。
なのでスタック領域のmoveNextを先に書き換えてその後ヒープ領域にコピーする必要がある。

StateMachineのMoveNextを直接AsyncMethodBuilderのmoveNextに入れることはできない。
delegateを作った時点でAsyncMethodBuilderがヒープ領域にコピーされてしまうからである。
必要条件としてdelegateを作った後にヒープ領域にコピーできる必要がある。
UniRxではそれを実現するためにMoveNextRunnerというラッパーを挟んでいる。

なおDebugビルドの場合はStateMachineがclassとして作られるので順番が違っても動いてしまう。