Syntriax.Engine/Engine.Core/Systems/NestedCoroutineManager.cs

using System.Collections;
using System.Collections.Generic;

namespace Engine.Core;

public class NestedCoroutineManager : Behaviour, IUpdate, ICoroutineManager
{
    private readonly List<CoroutineStack> stacks = [];
    private readonly Pool<CoroutineStack> pool = new(() => new());

    public IEnumerator StartCoroutine(IEnumerator enumerator)
    {
        CoroutineStack stack = pool.Get();
        stack.EntryPoint = enumerator;
        stack.Stack.Push(enumerator);

        stacks.Add(stack);

        return enumerator;
    }

    public void StopCoroutine(IEnumerator enumerator)
    {
        for (int i = 0; i < stacks.Count; i++)
            if (stacks[i].EntryPoint == enumerator)
            {
                RemoveCoroutineAt(i);
                return;
            }
    }

    private void RemoveCoroutineAt(int i)
    {
        stacks[i].Reset();
        stacks.RemoveAt(i);
    }

    void IUpdate.Update()
    {
        for (int i = stacks.Count - 1; i >= 0; i--)
        {
            Stack<IEnumerator> stack = stacks[i].Stack;

            if (stack.Count == 0)
            {
                RemoveCoroutineAt(i);
                continue;
            }

            IEnumerator top = stack.Peek();

            if (top.Current is ICoroutineYield coroutineYield && coroutineYield.Yield())
                continue;

            if (top.Current is IEnumerator nested)
            {
                stack.Push(nested);
                continue;
            }

            if (!top.MoveNext())
            {
                stack.Pop();
                if (stack.Count != 0)
                    stack.Peek().MoveNext();
                continue;
            }
        }
    }

    private class CoroutineStack
    {
        public IEnumerator EntryPoint = null!;
        public Stack<IEnumerator> Stack = new();

        public void Reset() { EntryPoint = null!; Stack.Clear(); }
    }

    public NestedCoroutineManager() => Priority = int.MinValue;
}