#include <stdio.h>
#include <assert.h>
#include "mcs-lock.h"
#include "mpi.h"

static int MCS_LOCKRANK = MPI_KEYVAL_INVALID;
static int MCS_COMM = MPI_KEYVAL_INVALID;
enum { nextRank = 0, lockTail = 1 };

void MCSLockInit(MPI_Comm comm, MPI_Win * win)
{
    int *lmem, rank;
    MPI_Comm newcomm;
    MPI_Aint winsize;

    if (MCS_LOCKRANK == MPI_KEYVAL_INVALID)
        MPI_Win_create_keyval(MPI_WIN_NULL_COPY_FN, MPI_WIN_NULL_DELETE_FN, &MCS_LOCKRANK,
                              (void *) 0);

    if (MCS_COMM == MPI_KEYVAL_INVALID)
        MPI_Win_create_keyval(MPI_WIN_NULL_COPY_FN, MPI_WIN_NULL_DELETE_FN, &MCS_COMM,
                              (void *) 0);

    MPI_Comm_dup(comm, &newcomm);
    MPI_Comm_rank(newcomm, &rank);

    winsize = 2 * sizeof(int);
    if (rank == 0)
        winsize += sizeof(int);
    MPI_Win_allocate(winsize, sizeof(int), MPI_INFO_NULL, newcomm, &lmem, win);
    lmem[nextRank] = -1;
    if (rank == 0)
        lmem[lockTail] = -1;

    MPI_Win_set_attr(*win, MCS_LOCKRANK, (void *) (MPI_Aint) rank);
    MPI_Win_set_attr(*win, MCS_COMM, (void *) (MPI_Aint) newcomm);
    MPI_Barrier(newcomm);
}

void MCSLockAcquire(MPI_Win win)
{
    int flag, myrank, predecessor, *lmem;
    void *attrval;
    MPI_Comm comm;
    int fetch_blocked, dummy;

    MPI_Win_get_attr(win, MCS_LOCKRANK, &attrval, &flag);
    myrank = (int) (MPI_Aint) attrval;
    MPI_Win_get_attr(win, MPI_WIN_BASE, &lmem, &flag);
    MPI_Win_get_attr(win, MCS_COMM, &attrval, &flag);
    comm = (MPI_Comm) (MPI_Aint) attrval;

    MPI_Win_lock_all(0, win);
    MPI_Fetch_and_op(&myrank, &predecessor, MPI_INT, 0, lockTail, MPI_REPLACE, win);
    MPI_Win_flush(0, win);

    if (predecessor != -1) {
        /* We didn't get the lock.  Add us to the tail of the list */
        MPI_Accumulate(&myrank, 1, MPI_INT, predecessor, nextRank, 1, MPI_INT, MPI_REPLACE, win);
        MPI_Recv(NULL, 0, MPI_INT, predecessor, 0, comm, MPI_STATUS_IGNORE);
    }

    MPI_Win_unlock_all(win);
}

void MCSLockRelease(MPI_Win win)
{
    int nullrank = -1, zero = 0, myrank, curtail, flag, *lmem;
    void *attrval;
    MPI_Comm comm;
    int predecessor;
    int fetch_nextrank, dummy;

    MPI_Win_get_attr(win, MCS_LOCKRANK, &attrval, &flag);
    myrank = (int) (MPI_Aint) attrval;
    MPI_Win_get_attr(win, MPI_WIN_BASE, &lmem, &flag);
    MPI_Win_get_attr(win, MCS_COMM, &attrval, &flag);
    comm = (MPI_Comm) (MPI_Aint) attrval;

    MPI_Win_lock_all(0, win);

    /* check to see who is next */
    MPI_Fetch_and_op(&dummy, &fetch_nextrank, MPI_INT, myrank, nextRank, MPI_NO_OP, win);
    MPI_Win_flush(myrank, win);

    if (fetch_nextrank == -1) {
        /* no next rank waiting for lock: there are two possibilities.
         * One is that no one else wants a lock.  Second is that the
         * next rank hasn't updated our nextRank location yet. */

        /* Possibility 1: check if no one else is waiting for a
         * lock */
        MPI_Compare_and_swap(&nullrank, &myrank, &curtail, MPI_INT,
                             0, lockTail, win);

        if (curtail == myrank) {
            MPI_Win_unlock_all(win);
            return;
        }

        /* Possibility 2: the next rank hasn't updated our nextRank
         * location yet.  wait for it. */
        do {
            MPI_Fetch_and_op(&dummy, &fetch_nextrank, MPI_INT, myrank, nextRank, MPI_NO_OP, win);
            MPI_Win_flush(myrank, win);
        } while (fetch_nextrank == -1);
    }

    lmem[nextRank] = -1;
    MPI_Win_sync(win);

    MPI_Send(NULL, 0, MPI_INT, fetch_nextrank, 0, comm);

    MPI_Win_unlock_all(win);
}

void MCSLockFinalize(MPI_Win win)
{
    void *attrval;
    int flag;
    MPI_Comm comm;

    MPI_Win_get_attr(win, MCS_COMM, &attrval, &flag);
    comm = (MPI_Comm) (MPI_Aint) attrval;

    MPI_Comm_free(&comm);
}