//  Shared mutex can be used in shared memory for interprocess
//  synchronization.
//

#ifndef KIT_SHARED_MUTEX_H
#define KIT_SHARED_MUTEX_H

#include "atomic.h"
#include "threads.h"

#include <assert.h>
#include <string.h>

#ifdef __cplusplus
extern "C" {
#endif

#if defined(__GNUC__) || defined(__clang__)
#  pragma GCC diagnostic push
#  pragma GCC diagnostic ignored "-Wunused-function"
#  pragma GCC diagnostic ignored "-Wunknown-pragmas"
#  pragma GCC            push_options
#  pragma GCC            optimize("O3")
#endif

enum {
  KIT_SHARED_MUTEX_READY = 1,
  KIT_SHARED_MUTEX_LOCKED,
  KIT_SHARED_MUTEX_WRITER
};

typedef union {
  struct {
    i8 _Atomic state;
    i32        readers;
  };
  struct {
    u8 _pad[16];
  };
} kit_shared_mutex_t;

static void kit_shared_mutex_init(kit_shared_mutex_t *m) {
  assert(m != NULL);
  memset(m, 0, sizeof *m);

  atomic_store_explicit(&m->state, KIT_SHARED_MUTEX_READY,
                        memory_order_relaxed);
}

static b8 kit_shared_try_lock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  for (;;) {
    i8 prev_state = KIT_SHARED_MUTEX_READY;
    if (atomic_compare_exchange_strong_explicit(
            &m->state, &prev_state, KIT_SHARED_MUTEX_LOCKED,
            memory_order_seq_cst, memory_order_seq_cst))
      break;
    if (prev_state == KIT_SHARED_MUTEX_WRITER)
      return 0;
    //  FIXME
    //  Check performance
    thrd_yield();
  }

  assert(m->readers >= 0);
  ++m->readers;

  i8 prev_state = KIT_SHARED_MUTEX_LOCKED;
  if (!atomic_compare_exchange_strong_explicit(
          &m->state, &prev_state, KIT_SHARED_MUTEX_READY,
          memory_order_seq_cst, memory_order_seq_cst))
    assert(0);

  return 1;
}

static void kit_shared_lock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  while (!kit_shared_try_lock(m))
    //  FIXME
    //  Check performance
    thrd_yield();
}

static void kit_shared_unlock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  for (;;) {
    i8 prev_state = KIT_SHARED_MUTEX_READY;
    if (atomic_compare_exchange_strong_explicit(
            &m->state, &prev_state, KIT_SHARED_MUTEX_LOCKED,
            memory_order_seq_cst, memory_order_seq_cst))
      break;
    //  FIXME
    //  Check performance
    thrd_yield();
  }

  assert(m->readers > 0);
  m->readers--;

  i8 prev_state = KIT_SHARED_MUTEX_LOCKED;
  if (!atomic_compare_exchange_strong_explicit(
          &m->state, &prev_state, KIT_SHARED_MUTEX_READY,
          memory_order_seq_cst, memory_order_seq_cst))
    assert(0);
}

static b8 kit_unique_try_lock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  i8 prev_state = KIT_SHARED_MUTEX_READY;
  if (!atomic_compare_exchange_strong_explicit(
          &m->state, &prev_state, KIT_SHARED_MUTEX_WRITER,
          memory_order_seq_cst, memory_order_seq_cst))
    return 0;

  i8 is_locked = (m->readers == 0);

  prev_state = KIT_SHARED_MUTEX_WRITER;
  if (!is_locked &&
      !atomic_compare_exchange_strong_explicit(
          &m->state, &prev_state, KIT_SHARED_MUTEX_READY,
          memory_order_seq_cst, memory_order_seq_cst))
    assert(0);

  return is_locked;
}

static void kit_unique_lock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  while (!kit_unique_try_lock(m))
    //  FIXME
    //  Check performance
    thrd_yield();
}

static void kit_unique_unlock(kit_shared_mutex_t *m) {
  assert(m != NULL);

  i8 prev_state = KIT_SHARED_MUTEX_WRITER;
  if (!atomic_compare_exchange_strong_explicit(
          &m->state, &prev_state, KIT_SHARED_MUTEX_READY,
          memory_order_seq_cst, memory_order_seq_cst))
    assert(0);
}

#if defined(__GNUC__) || defined(__clang__)
#  pragma GCC            pop_options
#  pragma GCC diagnostic pop
#endif

#ifdef __cplusplus
}
#endif

#define shared_mutex_t kit_shared_mutex_t
#define shared_mutex_init kit_shared_mutex_init
#define shared_try_lock kit_shared_try_lock
#define shared_lock kit_shared_lock
#define shared_unlock kit_shared_unlock
#define unique_try_lock kit_unique_try_lock
#define unique_lock kit_unique_lock
#define unique_unlock kit_unique_unlock
#define shared_mutex_init kit_shared_mutex_init
#define shared_try_lock kit_shared_try_lock
#define shared_lock kit_shared_lock
#define shared_unlock kit_shared_unlock
#define unique_try_lock kit_unique_try_lock
#define unique_lock kit_unique_lock
#define unique_unlock kit_unique_unlock

#endif