#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <ucontext.h>

typedef struct
{
  int index;
  void (* handler) (void);
  ucontext_t context;
  uint32_t already_running : 1;
  volatile uint32_t should_return_to_task : 1;
} task_t;

#define MAX_TASKS 5
static task_t scheduler_task = { 0 };
static task_t tasks[MAX_TASKS] = { { 0 } };
static task_t *running_task;
static int number_of_tasks = 0;

static void
init_scheduler (void)
{
  getcontext (&scheduler_task.context);
  scheduler_task.context.uc_link = NULL;
}

static void
yield (void)
{
  running_task->should_return_to_task = false;
  running_task->already_running = true;
  getcontext (&running_task->context);

  /* If false we arrived here from the task calling yield(), not
   * from jumping to the above getcontext() call
   */
  if (running_task->should_return_to_task)
    return;

  running_task = &tasks[(running_task->index + 1) % number_of_tasks];

  setcontext (&scheduler_task.context);
  abort ();
}

static void
schedule (void (* task_handler) (void))
{

  tasks[number_of_tasks].index = number_of_tasks;
  tasks[number_of_tasks].handler = task_handler;
  tasks[number_of_tasks].already_running = false;

  getcontext (&tasks[number_of_tasks].context);

  tasks[number_of_tasks].context.uc_link = &scheduler_task.context;
  tasks[number_of_tasks].context.uc_stack.ss_sp = malloc (SIGSTKSZ);
  tasks[number_of_tasks].context.uc_stack.ss_size = SIGSTKSZ;

  number_of_tasks++;
}

static void
dispatch (void)
{
  if (!running_task->already_running)
    makecontext (&running_task->context, running_task->handler, 0);
  else
    running_task->should_return_to_task = true;

  setcontext (&running_task->context);

  abort ();
}

void
run_tasks (void)
{
  running_task = &tasks[0];
  while (true)
    {
      getcontext (&scheduler_task.context);
      dispatch ();
    }
}

void
task0 (void)
{
  char i;
  while (true)
    {
      for (i = 0; i <= 9; i++)
        printf ("%d\n", i);
      yield ();
    }
}

void
task1 (void)
{
  char c;
  while (true)
    {
      for (c = 'a'; c <= 'z'; c++)
        printf ("%c\n", c);
      yield ();
    }
}

void
task2 (void)
{
  char c;
  while (true)
    {
      for (c = 'A'; c <= 'Z'; c++)
        printf ("%c\n", c);
      yield ();
    }
}

int
main (int    argc,
      char **argv)
{
  init_scheduler ();
  schedule (task0);
  schedule (task1);
  schedule (task2);
  run_tasks ();
  return 0;
}