UnityのCoroutineを自作する

基本の実装

UnityにはStartCoroutineという関数があります。
これはIEnumeratorを受け取って、MoveNext()を毎フレーム呼び出すようになっています。
これと同じような仕組みを自作します。

public class CoroutineEx {
    IEnumerator routine;

    public bool IsDone { get; private set; }

    // コンストラクタ
    public CoroutineEx(IEnumerator routine) {
        this.routine = routine;
    }

    // 更新処理
    public void Update() {
        IsDone = !routine.MoveNext();
    }
}

基本の実装は上記で完了します。
引数として実行したいIEnumeratorを受け取り、毎フレームUpdate関数が呼び出されるものとします。

実行が完了したかどうかをIsDoneに設定して確認できるようにしています。

実装の拡張

もう少し便利に利用するため以下の仕様を追加します。

• MoveNext処理でエラーが発生した場合、エラーを取得できるようにする。
• yield returnで別のIEnumeratorが返されたとき、それを待機するようにする。
• AsyncOperationが返されたとき、isDoneがtrueになるまで待機する。
• その他のオブジェクトが返されたとき、そのオブジェクトを取得できるようにする。

public class CoroutineEx {
    public Stack<object> stack = new Stack<object>();

    public object Current { get; private set; }

    public Exception Error { get; private set; }

    public bool IsDone { get; private set; }

    // コンストラクタ
    public CoroutineEx(IEnumerator routine) {
        stack.Push(routine);
    }
}

• 発生したエラーを取得できるようにするため、Errorプロパティを定義。
• 複数のIEnumeratorを待機することを可能にするため、routineはStackで持つように変更。
• MoveNextで返された最後のオブジェクトを取得するため、Currentプロパティを定義。

更新処理の拡張

更新処理では、routineの型によって処理を分岐します。

public void Update() {
    if (stack.Count == 0) {
        Done();
        return;
    }

    var peek = stack.Peek();
    if (peek == null) {
        stack.Pop();
        Current = null;
    }
    else if (peek is IEnumerator) {
        var e = (IEnumerator)peek;
        try {
            if (e.MoveNext()) {
                stack.Push(e.Current);
                Update();
            } else {
                stack.Pop();
                Update();
            }
        }
        catch (Exception error) {
            stack.Clear();
            Error = error;
            Done();
        }
    }
    else if (peek is AsyncOperation) {
        if (!((AsyncOperation)peek).isDone) {
            Current = null;
        } else {
            stack.Pop();
            Update();
        }
    }
    else {
        stack.Pop();
        Current = peek;
    }
}

void Done() {
    Current = null;
    IsDone = true;
}

stackが空になっている場合は処理が完了しているので、IsDoneをtrueにします。

peek == null

IEnumeratorをコルーチンとして利用するとき、yield return nullによって
頻繁にnullを受け取るので、nullの場合の分岐を先頭に入れています。

peek is IEnumerator

IEnumeratorの場合はエラーハンドリングをしつつMoveNextを呼び出します。

yield returnで返された値がCurrentから取得できるので、stackにPushしてUpdateを再帰呼び出しします。

IEnumeratorが完了していた場合はstackから取り除き、再帰呼び出しします。
この際stackが空になっていればDoneが呼び出されて処理が完了します。

peek is AsyncOperation

AsyncOperationを受け取った場合はstackの頭にAsyncOperationを残しつつ、
Currentにはnullを設定して処理を継続させます。

完了していればstackから取り除き、Updateを再帰します。

else

その他のオブジェクトの場合、Currentに設定して終わりです。

IEnumeratorインターフェースの実装

CoroutineExをUnityのStartCoroutineに渡しても動作するように、IEnumeratorインターフェースを実装します。

public class CoroutineEx : IEnumerator {
    bool IEnumerator.MoveNext() {
        Update();
        return !IsDone;
    }

    void IEnumerator.Reset() { }
}

IEnumeratorインターフェースのうち、Currentは実装済みのため省略します。

MoveNextは毎フレーム呼び出される処理なので、Updateを呼び出し、
処理を継続するかどうかのboolとして!IsDoneを返します。

Resetはコルーチンとしては利用しない関数のため、実装は空です。

実装のまとめ

ここまでの実装により完成したコードを貼ります。

public class CoroutineEx : IEnumerator {
    public Stack<object> stack = new Stack<object>();

    public object Current { get; private set; }

    public Exception Error { get; private set; }

    public bool IsDone { get; private set; }

    // コンストラクタ
    public CoroutineEx(IEnumerator routine) {
        stack.Push(routine);
    }

    // 更新処理
    public void Update() {
        if (stack.Count == 0) {
            Done();
            return;
        }

        var peek = stack.Peek();
        if (peek == null) {
            stack.Pop();
            Current = null;
        }
        else if (peek is IEnumerator) {
            var e = (IEnumerator)peek;
            try {
                if (e.MoveNext()) {
                    stack.Push(e.Current);
                    Update();
                } else {
                    stack.Pop();
                    Update();
                }
            }
            catch (Exception error) {
                stack.Clear();
                Error = error;
                Done();
            }
        }
        else if (peek is AsyncOperation) {
            if (!((AsyncOperation)peek).isDone) {
                Current = null;
            } else {
                stack.Pop();
                Update();
            }
        }
        else {
            stack.Pop();
            Current = peek;
        }
    }

    void Done() {
        Current = null;
        IsDone = true;
    }

    bool IEnumerator.MoveNext() {
        Update();
        return !IsDone;
    }

    void IEnumerator.Reset() { }
}

上記実装に加えて

• 処理を中断する関数を実装したり
• 完了やエラーのコールバックを実装したり
• 完了後に継続して行う処理を追加できるようにしたり

など、Unity標準のCoroutineではできなかった処理が
拡張次第で書けるようになります。

そもそも標準のCoroutineが何の変数も関数もないオブジェクトで
持ってても意味無いので、こういう拡張クラスがあると便利です。

割と勢いで書いたので編集リクエスト歓迎です。