#define DEBUG 1
#include <thread.h>
#include <sys/synch.h>
#include <sys/types.h>
#include <sys/debug.h>
#include <stdio.h>
#include <stdlib.h>

typedef struct zrlock {
	mutex_t zr_mtx;
	volatile int32_t zr_refcount;
	cond_t zr_cv;
	uint16_t zr_pad;
} zrlock_t;

void zrl_init(zrlock_t *);
void zrl_destroy(zrlock_t *);
#define	zrl_add(_z)	zrl_add_impl((_z), __func__)
void zrl_add_impl(zrlock_t *, const char *);
void zrl_remove(zrlock_t *);
int zrl_tryenter(zrlock_t *);
void zrl_exit(zrlock_t *);
int zrl_is_zero(zrlock_t *);
int zrl_is_locked(zrlock_t *);

#define	ZRL_LOCKED	-1
#define	ZRL_DESTROYED	-2

void
zrl_init(zrlock_t *zrl)
{
	mutex_init(&zrl->zr_mtx, USYNC_THREAD, NULL);
	zrl->zr_refcount = 0;
	cond_init(&zrl->zr_cv, USYNC_THREAD, NULL);
}

void
zrl_destroy(zrlock_t *zrl)
{
	ASSERT0(zrl->zr_refcount);

	mutex_destroy(&zrl->zr_mtx);
	zrl->zr_refcount = ZRL_DESTROYED;
	cond_destroy(&zrl->zr_cv);
}

void
zrl_add_impl_old(zrlock_t *zrl, const char *zc)
{
	uint32_t n = (uint32_t)zrl->zr_refcount;

	while (n != ZRL_LOCKED) {
		uint32_t cas = atomic_cas_32(
		    (uint32_t *)&zrl->zr_refcount, n, n + 1);
		if (cas == n) {
			ASSERT3S((int32_t)n, >=, 0);
			return;
		}
		n = cas;
	}

	mutex_lock(&zrl->zr_mtx);
	while (zrl->zr_refcount == ZRL_LOCKED) {
		cond_wait(&zrl->zr_cv, &zrl->zr_mtx);
	}
	ASSERT3S(zrl->zr_refcount, >=, 0);
	zrl->zr_refcount++;
	mutex_unlock(&zrl->zr_mtx);
}

void
zrl_add_impl(zrlock_t *zrl, const char *zc)
{
	uint32_t n = (uint32_t)zrl->zr_refcount;

	while (1) {
			while (n != ZRL_LOCKED) {
					uint32_t cas = atomic_cas_32(
						(uint32_t *)&zrl->zr_refcount, n, n + 1);
					if (cas == n) {
							ASSERT3S((int32_t)n, >=, 0);
							return;
					}
					n = cas;
			}
			mutex_lock(&zrl->zr_mtx);
			while (zrl->zr_refcount == ZRL_LOCKED) {
					cond_wait(&zrl->zr_cv, &zrl->zr_mtx);
			}
			n = (uint32_t)zrl->zr_refcount;
			mutex_unlock(&zrl->zr_mtx);
	}
}

void
zrl_remove(zrlock_t *zrl)
{
	uint32_t n;

	n = atomic_dec_32_nv((uint32_t *)&zrl->zr_refcount);
	ASSERT3S((int32_t)n, >=, 0);
}

int
zrl_tryenter(zrlock_t *zrl)
{
	uint32_t n = (uint32_t)zrl->zr_refcount;

	if (n == 0) {
		uint32_t cas = atomic_cas_32(
		    (uint32_t *)&zrl->zr_refcount, 0, ZRL_LOCKED);
		if (cas == 0) {
			return (1);
		}
	}

	ASSERT3S((int32_t)n, >, ZRL_DESTROYED);

	return (0);
}

void
zrl_exit(zrlock_t *zrl)
{
	ASSERT3S(zrl->zr_refcount, ==, ZRL_LOCKED);

	mutex_lock(&zrl->zr_mtx);
	zrl->zr_refcount = 0;
	cond_broadcast(&zrl->zr_cv);
	mutex_unlock(&zrl->zr_mtx);
}

int
zrl_refcount(zrlock_t *zrl)
{
	ASSERT3S(zrl->zr_refcount, >, ZRL_DESTROYED);

	int n = (int)zrl->zr_refcount;
	return (n <= 0 ? 0 : n);
}

int
zrl_is_zero(zrlock_t *zrl)
{
	ASSERT3S(zrl->zr_refcount, >, ZRL_DESTROYED);

	return (zrl->zr_refcount <= 0);
}

int
zrl_is_locked(zrlock_t *zrl)
{
	ASSERT3S(zrl->zr_refcount, >, ZRL_DESTROYED);

	return (zrl->zr_refcount == ZRL_LOCKED);
}

zrlock_t zrlock;

void *
func_zrl_add_remove_old(void *arg)
{
	int *p = (int *) arg;
	int i;
	for(i = 0; i < *p; i++)
	{
		zrl_add_impl_old(&zrlock, NULL);
		zrl_remove(&zrlock);
	} 
}

void *
func_zrl_add_remove_new(void *arg)
{
	int *p = (int *) arg;
	int i;
	for(i = 0; i < *p; i++)
	{
		zrl_add_impl(&zrlock, NULL);
		zrl_remove(&zrlock);
	} 
}

void *
func_zrl_tryenter_exit(void *arg)
{
	int *p = (int *) arg;
	int i;
	for(i = 0; i < *p; i++)
	{
		if(!zrl_tryenter(&zrlock))
			continue;
		zrl_exit(&zrlock);
	}
}

int main(int argc, char **argv)
{
	pthread_t *tids_add_remove = NULL;
	pthread_t *tids_enter_exit = NULL;
	int nthreads_add_remove = 10;
	int nthread_enter_exit = 1;
	int loops = 1000;
	int i;
	int c;
	int ret = 0;
	boolean_t new_zrl_add_impl = B_FALSE;

	while ((c = getopt(argc, argv, ":na:e:l:")) != -1) {
		switch(c) {
		case 'n':
			new_zrl_add_impl = B_TRUE;
			break;
		case 'a':
			nthreads_add_remove = atoi(optarg);
			break;
		case 'e':
			nthread_enter_exit = atoi(optarg);
			break;
		case 'l':
			loops = atoi(optarg);
			break;
		case ':':
			printf("-%c without arg\n", optopt);
			exit(1);
			break;
		default:
			exit(2);
		}
	}

	printf("-a %d: nthreads_add_remove = %d\n", nthreads_add_remove,
	    nthreads_add_remove);
	printf("-e %d: nthread_enter_exit = %d\n", nthread_enter_exit,
	    nthread_enter_exit);
	printf("-l %d: loops = %d\n", loops, loops);
	printf("Run new zrl_add_impl: %s\n", new_zrl_add_impl ? "YES" : "NO");

	zrl_init(&zrlock);
	tids_add_remove = calloc(nthreads_add_remove, sizeof(pthread_t));
	tids_enter_exit = calloc(nthread_enter_exit, sizeof(pthread_t));

	for(i = 0; i < nthreads_add_remove; i++) {
		if(pthread_create(&tids_add_remove[i], NULL, 
			new_zrl_add_impl ? func_zrl_add_remove_new : func_zrl_add_remove_old, &loops) != 0) {
			fprintf(stderr, "ERROR: failed to create thread\n");
			ret = 3;
			goto done;
		}
	}

	for(i = 0; i < nthread_enter_exit; i++) {
		if(pthread_create(&tids_enter_exit[i], NULL, func_zrl_tryenter_exit, &loops) != 0) {
			fprintf(stderr, "ERROR: failed to create thread\n");
			ret = 3;
			goto done;
		}
	}

	for(i = 0; i < nthreads_add_remove; i++) {
		(void) pthread_join(tids_add_remove[i], NULL);
	}

	for(i = 0; i < nthread_enter_exit; i++) {
		(void) pthread_join(tids_enter_exit[i], NULL);
	}
	
	ASSERT3S(zrlock.zr_refcount, ==, 0);
	
done:
	zrl_destroy(&zrlock);
	if(tids_add_remove) free(tids_add_remove);
	if(tids_enter_exit) free(tids_enter_exit);

	printf("DONE.\n");
	return(ret);
}

/*
build: gcc -m64 -o zrlock-test zrlock-test.c

./zrlock-test 
assertion failed for thread 0xfffffd7ffef03240, thread-id 9: (int32_t)n >= 0 (0xffffffffffffffff >= 0x0), file zrlock-test.c, line 101

./zrlock-test -n
Run new zrl_add_impl: YES
*/