ajeeb-style coroutines in C#

Program.cs

using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

int FRAME = 0;
var sched = new Schedule();

async Task DelayFrames(int frames)
{
    while (frames-- > 0)
    {
        await sched.Yield;
    }
}

async Task DelaySeconds(int seconds)
{
    var sw = System.Diagnostics.Stopwatch.StartNew();
    while (sw.Elapsed.Seconds < seconds)
    {
        await sched.Yield;
    }
}

async Task StaticDelayFrames(int frames)
{
    while (frames-- > 0)
    {
        await StaticSchedule.Yield;
    }
}

async Task StaticDelaySeconds(int seconds)
{
    var sw = System.Diagnostics.Stopwatch.StartNew();
    while (sw.Elapsed.Seconds < seconds)
    {
        await StaticSchedule.Yield;
    }
}

async Task Foo()
{
    var sw = System.Diagnostics.Stopwatch.StartNew();

    Console.WriteLine($"Foo 1 {sw.ElapsedMilliseconds} ({FRAME})");
    await DelaySeconds(3);

    Console.WriteLine($"Foo 2 {sw.ElapsedMilliseconds} ({FRAME})");
    await DelaySeconds(1);

    Console.WriteLine($"Foo 3 {sw.ElapsedMilliseconds} ({FRAME})");

    await Task.From(async () =>
    {
        Console.WriteLine($"Lambda 1 {sw.ElapsedMilliseconds} ({FRAME})");
        await DelaySeconds(4);
        Console.WriteLine($"Lambda 2 {sw.ElapsedMilliseconds} ({FRAME})");
    });

    Console.WriteLine($"Foo 4 {sw.ElapsedMilliseconds} ({FRAME})");
    await DelayFrames(20);
    Console.WriteLine($"Foo 5 {sw.ElapsedMilliseconds} ({FRAME})");
}

//// output
// Foo 1 0 (0)
// Foo 2 3014 (185)
// Foo 3 4016 (246)
// Lambda 1 4016 (246)
// Lambda 2 8030 (491)
// Foo 4 8030 (491)
// Foo 5 8359 (511)


async void Main()
{
    await Foo();
}

Main();

while (true)
{
    FRAME++;
    // StaticSchedule.Tick();
    sched.Tick();
    System.Threading.Thread.Sleep(1000 / 60);
}

//////////////// implementation ////////////////
class Schedule
{
    public Schedule()
    {
        Yield = new YieldAwaitable(this);
    }

    List<Action> ContinuationsThisFrame = new List<Action>();
    List<Action> ContinuationsNextFrame = new List<Action>();

    public void Tick()
    {
        // swap next frame and this frame
        var temp = ContinuationsThisFrame;
        ContinuationsThisFrame = ContinuationsNextFrame;
        ContinuationsNextFrame = temp;

        // clear next frame. this frame's continuations may enqueue new
        // continuations here.
        ContinuationsNextFrame.Clear();

        // invoke this frame's continuations
        foreach (var c in ContinuationsThisFrame)
        {
            c();
        }
        ContinuationsThisFrame.Clear();
    }

    public void EnqueueNextFrame(Action action)
    {
        ContinuationsNextFrame.Add(action);
    }

    public YieldAwaitable Yield { get; private set; }

    public struct YieldAwaitable
    {
        public YieldAwaitable(Schedule schedule)
        {
            CachedAwaiter = new YieldAwaiter(schedule);
        }

        YieldAwaiter CachedAwaiter;
        public YieldAwaiter GetAwaiter() => CachedAwaiter;

        public struct YieldAwaiter : ICriticalNotifyCompletion
        {
            Schedule schedule;

            public YieldAwaiter(Schedule schedule)
            {
                this.schedule = schedule;
            }

            public bool IsCompleted => false;

            public void GetResult() { }

            public void OnCompleted(Action continuation)
                => schedule.EnqueueNextFrame(continuation);

            public void UnsafeOnCompleted(Action continuation)
                => schedule.EnqueueNextFrame(continuation);
        }
    }

}

static class StaticSchedule
{
    static List<Action> ContinuationsThisFrame = new List<Action>();
    static List<Action> ContinuationsNextFrame = new List<Action>();
    public static void Tick()
    {
        // swap next frame and this frame
        var temp = ContinuationsThisFrame;
        ContinuationsThisFrame = ContinuationsNextFrame;
        ContinuationsNextFrame = temp;

        // clear next frame. this frame's continuations may enqueue new
        // continuations here.
        ContinuationsNextFrame.Clear();

        // invoke this frame's continuations
        foreach (var c in ContinuationsThisFrame)
        {
            c();
        }
        ContinuationsThisFrame.Clear();
    }

    public static void EnqueueNextFrame(Action action)
    {
        ContinuationsNextFrame.Add(action);
    }

    public static StaticYield Yield = default(StaticYield);

    public struct StaticYield
    {
        public static StaticYield Yield = default(StaticYield);
        static YieldAwaiter CachedAwaiter = default(YieldAwaiter);
        public YieldAwaiter GetAwaiter() => CachedAwaiter;

        public struct YieldAwaiter : ICriticalNotifyCompletion
        {
            public bool IsCompleted => false;

            public void GetResult() { }

            public void OnCompleted(Action continuation)
                => StaticSchedule.EnqueueNextFrame(continuation);

            public void UnsafeOnCompleted(Action continuation)
                => StaticSchedule.EnqueueNextFrame(continuation);
        }
    }
}

[AsyncMethodBuilder(typeof(TaskAsyncMethodBuilder))]
public class Task
{
    IAsyncStateMachine stateMachine;

    public Task(IAsyncStateMachine stateMachine)
    {
        this.stateMachine = stateMachine;
    }

    public static Task From(Func<Task> func)
    {
        return func();
    }

    public void MoveNext() => stateMachine.MoveNext();

    public bool Completed { get; set; }
    public Action Continuation { get; set; }

    public TaskAwaiter GetAwaiter()
    {
        return new TaskAwaiter(this);
    }

    public Exception Exception { get; private set; }
    public void SetException(Exception e) => Exception = e;
    public void SetResult() => Completed = true;
}

public struct TaskAwaiter : ICriticalNotifyCompletion
{
    Task task;

    public TaskAwaiter(Task task)
    {
        this.task = task;
    }

    public bool IsCompleted => task.Completed;

    public void OnCompleted(Action continuation)
    {
        task.Continuation = continuation;
    }

    public void GetResult() { /* task.Result*/ }

    public void UnsafeOnCompleted(Action continuation)
    {
        task.Continuation = continuation;
    }
}

public struct TaskAsyncMethodBuilder
{
    IAsyncStateMachine stateMachine;
    Task task;
    Action stateMachineMoveNext;

    public static TaskAsyncMethodBuilder Create() => new TaskAsyncMethodBuilder();

    public void SetResult()
    {
        task.SetResult();
        task.Continuation(); // should this enqueue?
    }

    public void SetException(Exception e)
    {
        task.SetException(e);
    }

    public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : INotifyCompletion where TStateMachine : IAsyncStateMachine
    {
        awaiter.OnCompleted(stateMachineMoveNext);
    }

    public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
        where TAwaiter : ICriticalNotifyCompletion where TStateMachine : IAsyncStateMachine
    {
        awaiter.UnsafeOnCompleted(stateMachineMoveNext);
    }

    public void SetStateMachine(IAsyncStateMachine stateMachine)
    {
        this.stateMachine = stateMachine;
    }

    public void Start<TStateMachine>(ref TStateMachine stateMachine)
        where TStateMachine : IAsyncStateMachine
    {
        this.stateMachine = stateMachine;
        task = new Task(stateMachine);
        stateMachineMoveNext = stateMachine.MoveNext;
        stateMachine.MoveNext(); // should this enqueue?
    }

    public Task Task => task;
}