/*
 * Userspace read/write lock implementation, based on Linux rwsem-spinlock algorithm.
 * The linked-list part is a bit dirty but it's not the point here.
 *
 * gcc -o rwlock rwlock.c -g -Wall -W -pthread
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <pthread.h>
#include <assert.h>

int gcount;

int __xchg(int *v, int x)
{
    volatile int r;
    asm("movl %2, %%eax; xchg (%1), %%eax; movl %%eax, %0"
        : "=r"(r) : "r"(v), "r"(x) : "%eax");
    return r;
}

void down(int *s)
{
    while (__xchg(s, 0) == 0)
        usleep(10000);
}

void up(int *s)
{
    __xchg(s, 1);
}

struct waiter {
    int mutex;
    int wait_for_write;
    struct waiter *next, *prev;
};

struct rwlock {
    int activity;
    int mutex;
    struct waiter *waitlist_head;
    struct waiter *waitlist_tail;
    int wcount;
};

void rwlock_init(struct rwlock *lock)
{
    lock->activity = 0;
    lock->mutex = 1;
    lock->waitlist_head = NULL;
    lock->waitlist_tail = NULL;
    lock->wcount = 0;
}

void __add_waiter(struct rwlock *lock, struct waiter *w)
{
    w->prev = lock->waitlist_tail;
    w->next = NULL;
    if (w->prev)
        w->prev->next = w;
    lock->waitlist_tail = w;

    if (lock->waitlist_head == NULL)
        lock->waitlist_head = w;

    lock->wcount++;
}

struct waiter *__remove_waiter(struct rwlock *lock)
{
    struct waiter *w;

    w = lock->waitlist_head;
    if (!w)
        return NULL;

    if (w->next)
        w->next->prev = NULL;
    lock->waitlist_head = w->next;

    if (!w->next)
        lock->waitlist_tail = NULL;

    w->next = w->prev = NULL;

    lock->wcount--;

    return w;
}

/*
 * Wake up one or more waiters.
 * The rwlock is acquired and the waiter list is not empty
 *
 * - take the oldest waiter
 * - if it's waiting for a write lock, wake it up and return
 * - otherwise, wake it up as well as all next read lock waiters until we get
 *   on a write lock waiter or there is no more waiter.
 */
void do_wake(struct rwlock *lock)
{
    struct waiter *w;
    int woken = 0;

    w = lock->waitlist_head;
    if (w->wait_for_write) {
        lock->activity = -1;
        __remove_waiter(lock);
        up(&w->mutex);
        return;
    }

    while (w && !w->wait_for_write) {
        __remove_waiter(lock);
        up(&w->mutex);
        woken++;
        w = lock->waitlist_head;
    }

    lock->activity += woken;
}

/*
 * Acquire a read lock
 *
 * - if the lock is free or acquired by readers and no one is waiting for it,
 *   get it and return
 * - otherwise, wait for it
 */
void down_read(struct rwlock *lock)
{
    struct waiter w;

    down(&lock->mutex);
    if (lock->activity >= 0 && lock->waitlist_head == NULL) {
        lock->activity++;
        up(&lock->mutex);
        return;
    }

    w.mutex = 0;
    w.wait_for_write = 0;
    __add_waiter(lock, &w);

    up(&lock->mutex);
    down(&w.mutex);
}

/*
 * Acquire a write lock
 *
 * - if the lock is free and no one is waiting for it, get it and return
 * - otherwise, wait for it
 */
void down_write(struct rwlock *lock)
{
    struct waiter w;

    down(&lock->mutex);
    if (lock->activity == 0 && lock->waitlist_head == NULL) {
        lock->activity = -1;
        up(&lock->mutex);
        return;
    }

    w.mutex = 0;
    w.wait_for_write = 1;
    __add_waiter(lock, &w);

    up(&lock->mutex);
    down(&w.mutex);
}

/*
 * Release a read lock
 *
 * - if all readers have released the lock and the waiting list is not empty,
 *   wake up some waiters
 */
void up_read(struct rwlock *lock)
{
    down(&lock->mutex);
    if (--lock->activity == 0 && lock->waitlist_head != NULL)
        do_wake(lock);
    up(&lock->mutex);
}

/*
 * Release a write lock
 *
 * - if the waiting list is not empty, wake up some waiters
 */
void up_write(struct rwlock *lock)
{
    down(&lock->mutex);
    lock->activity = 0;
    if (lock->waitlist_head != NULL)
        do_wake(lock);
    up(&lock->mutex);
}

void *print(void *arg)
{
    struct rwlock *lock = arg;

    down_read(lock);
    printf("%d\n", gcount);
    up_read(lock);
    return NULL;
}

void *increment(void *arg)
{
    struct rwlock *lock = arg;

    /* yield to some readers */
    usleep(100);

    down_write(lock);
    gcount++;
    up_write(lock);
    return NULL;
}

int main()
{
    int i;
    struct rwlock *lock;
    pthread_t t[100], t2[2];

    gcount = 0;
    lock = malloc(sizeof(*lock));
    rwlock_init(lock);

    for (i = 0; i < 2; i++)
        pthread_create(&t2[i], NULL, increment, lock);

    for (i = 0; i < 100; i++)
        pthread_create(&t[i], NULL, print, lock);

    for (i = 0; i < 100; i++)
        pthread_join(t[i], NULL);

    for (i = 0; i < 2; i++)
        pthread_join(t2[i], NULL);

    printf("%d\n", gcount);

    free(lock);

    return 0;
}