#include <cstdio>
#include <csetjmp>
#include <cstring>
#include <chrono>

using namespace std;

// Settings for stack and benchmark
const unsigned int NUM_RUNS = 100000;
const unsigned int STACK_SIZE = 512 * 8;
const unsigned char MAGIC_NUMBER = (unsigned char) 0xAB;

// Memory for custom stack and continuation semantics
unsigned char custom_stack[STACK_SIZE] = {0};
jmp_buf buffer;

// Example callback function and declaration of our assembly stack switching routine
extern "C" {
void custom_stack_callback(void *);

void fiber_call(void *);
void fiber_continue(void *);

void __attribute__ ((noinline)) callback() {
  static volatile int tmp;
  tmp = 0; // Force at least a single memory write
}
}

long __attribute__ ((noinline)) measure_function_call() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    callback();
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long __attribute__ ((noinline)) measure_stack_switch() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    custom_stack_callback(&custom_stack[STACK_SIZE - 16]);
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long __attribute__ ((noinline)) measure_continuation() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    if (setjmp(buffer) == 0) {
      custom_stack_callback(&custom_stack[STACK_SIZE - 16]);
    }
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long __attribute__ ((noinline)) measure_continuation_and_jump() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    if (setjmp(buffer) == 0) {
      custom_stack_callback(&custom_stack[STACK_SIZE - 16]);
      longjmp(buffer, 1);
    }
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

long __attribute__ ((noinline)) measure_fiber_call() {
  auto start_time = chrono::steady_clock::now();
  for (unsigned int i = 0; i < NUM_RUNS; i++) {
    void *stack_pointer = custom_stack;
    fiber_call(&stack_pointer);
  }
  auto end_time = chrono::steady_clock::now();
  return chrono::duration_cast<chrono::nanoseconds>(end_time - start_time).count();
}

int main() {
  memset(custom_stack, MAGIC_NUMBER, STACK_SIZE);

  auto time_cont_jump = measure_continuation_and_jump();
  auto time_cont = measure_continuation();
  auto time_stack = measure_stack_switch();
  auto time_func = measure_function_call();
  auto time_fiber = measure_fiber_call();

  for (unsigned int i = 0; i < STACK_SIZE; i++) {
    if (custom_stack[i] != MAGIC_NUMBER) {
      printf("Used stack size about %u bytes.\n\n\n", (STACK_SIZE - i));
      break;
    }
  }

  printf("Function Call    : %10ld, %5.5f\n", time_func, ((float) time_func / NUM_RUNS));
  printf("Stack Switching  : %10ld, %5.5f\n", time_stack, ((float) time_stack / NUM_RUNS));
  printf("Full Continuation: %10ld, %5.5f\n", time_cont, ((float) time_cont / NUM_RUNS));
  printf("Jump Continuation: %10ld, %5.5f\n", time_cont_jump, ((float) time_cont_jump / NUM_RUNS));
  printf("Fiber Call       : %10ld, %5.5f\n", time_fiber, ((float) time_fiber / NUM_RUNS));

  return 0;
}
