/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */

/*

 * See COPYRIGHT in top-level directory.

 */

/*

 * The original version of this code was contributed by Milind Chabbi

 * based on the work when he was at Rice University. It relies on the

 * HMCS lock description in [1] and the fast path described in [2].

 *

 * [1] Chabbi, Milind, Michael Fagan, and John Mellor-Crummey. "High

 * performance locks for multi-level NUMA systems." In Proceedings of

 * the ACM SIGPLAN Symposium on Principles and Practice of Parallel

 * Programming (PPoPP'15), ACM, 2015.

 *

 * [2] Chabbi, Milind, and John Mellor-Crummey. "Contention-conscious,

 * locality-preserving locks." In Proceedings of the 21st ACM SIGPLAN

 * Symposium on Principles and Practice of Parallel Programming (PPoPP'16,

 * p. 22. ACM, 2016.

 */

#include "lock/zm_lock_types.h"

#include <hwloc.h>

#ifndef DEFAULT_THRESHOLD

#define DEFAULT_THRESHOLD 256

#endif

#ifndef HMCS_DEFAULT_MAX_LEVELS

#define HMCS_DEFAULT_MAX_LEVELS 3

#endif

#define WAIT (0xffffffff)

#define COHORT_START (0x1)

#define ACQUIRE_PARENT (0xcffffffc)

#ifndef TRUE

#define TRUE 1

#else

#error "TRUE already defined"

#endif

#ifndef FALSE

#define FALSE 0

#else

#error "TRUE already defined"

#endif

/* Atomic operation shorthands. The memory ordering defaults to:

 * 1- Acquire ordering for loads

 * 2- Release ordering for stores

 * 3- Acquire+Release ordering for read-modify-write operations

 * */

#define LOAD(addr)                  zm_atomic_load(addr, zm_memord_acquire)

#define STORE(addr, val)            zm_atomic_store(addr, val, zm_memord_release)

#define SWAP(addr, desire)          zm_atomic_exchange_ptr(addr, desire, zm_memord_acq_rel)

#define CAS(addr, expect, desire)   zm_atomic_compare_exchange_strong(addr,\

                                                                      expect,\

                                                                      desire,\

                                                                      zm_memord_acq_rel,\

                                                                      zm_memord_acquire)

struct hnode{

    unsigned threshold __attribute__((aligned(ZM_CACHELINE_SIZE)));

    struct hnode * parent __attribute__((aligned(ZM_CACHELINE_SIZE)));

    zm_atomic_ptr_t lock __attribute__((aligned(ZM_CACHELINE_SIZE)));

    zm_mcs_qnode_t node __attribute__((aligned(ZM_CACHELINE_SIZE)));

}__attribute__((aligned(ZM_CACHELINE_SIZE)));

struct leaf{

    struct hnode * cur_node;

    struct hnode * root_node;

    zm_mcs_qnode_t I;

    int curDepth;

    int took_fast_path;

};

struct lock{

    // Assumes tids range from [0.. max_threads)

    // Assumes that tid 0 is close to tid and so on.

    struct leaf ** leaf_nodes __attribute__((aligned(ZM_CACHELINE_SIZE)));

    hwloc_topology_t topo;

    int levels;

};

static zm_thread_local int tid = -1;

/* TODO: automate hardware topology detection

 * instead of the below hard-coded method */

/* Check the actual affinity mask assigned to the thread */

static void check_affinity(hwloc_topology_t topo) {

    hwloc_cpuset_t cpuset = hwloc_bitmap_alloc();

    int set_length;

    hwloc_get_cpubind(topo, cpuset, HWLOC_CPUBIND_THREAD);

    set_length = hwloc_get_nbobjs_inside_cpuset_by_type(topo, cpuset, HWLOC_OBJ_PU);

    hwloc_bitmap_free(cpuset);

    if(set_length != 1) {

        printf("IZEM:HMCS:ERROR: thread bound to more than one HW thread!\n");

        exit(EXIT_FAILURE);

    }

}

static inline void reuse_qnode(zm_mcs_qnode_t *I){

    STORE(&I->status, WAIT);

    STORE(&I->next, ZM_NULL);

}

static void* new_hnode() {

    int err;

    void *storage;

    err = posix_memalign(&storage, ZM_CACHELINE_SIZE, sizeof(struct hnode));

    if (err != 0) {

        printf("posix_memalign failed in HMCS : new_hnode \n");

        exit(EXIT_FAILURE);

    }

    return storage;

}

/* TODO: Macro or Template this for fast comprison */

static inline unsigned get_threshold(struct hnode *L) {

    return L->threshold;

}

static inline void normal_mcs_release_with_value(struct hnode * L, zm_mcs_qnode_t *I, unsigned val){

    zm_mcs_qnode_t *succ = (zm_mcs_qnode_t *)LOAD(&I->next);

    if(succ) {

        STORE(&succ->status, val);

        return;

    }

    zm_mcs_qnode_t *tmp = I;

    if (CAS(&(L->lock), (zm_ptr_t*)&tmp,ZM_NULL))

        return;

    while(succ == NULL)

        succ = (zm_mcs_qnode_t *)LOAD(&I->next); /* SPIN */

    STORE(&succ->status, val);

    return;

}

static inline void acquire_root(struct hnode * L, zm_mcs_qnode_t *I) {

    // Prepare the node for use.

    reuse_qnode(I);

    zm_mcs_qnode_t *pred = (zm_mcs_qnode_t*) SWAP(&(L->lock), (zm_ptr_t)I);

    if(!pred) {

        // I am the first one at this level

        return;

    }

    STORE(&pred->next, I);

    while(LOAD(&I->status) == WAIT)

        ; /* SPIN */

    return;

}

static inline void tryacq_root(struct hnode * L, zm_mcs_qnode_t *I, int *success) {

    zm_ptr_t expected = ZM_NULL;

    // Prepare the node for use.

    reuse_qnode(I);

    *success = CAS(&(L->lock), &expected, (zm_ptr_t)I);

    return;

}

static inline void release_root(struct hnode * L, zm_mcs_qnode_t *I) {

    // Top level release is usual MCS

    // At the top level MCS we always writr COHORT_START since

    // 1. It will release the lock

    // 2. Will never grow large

    // 3. Avoids a read from I->status

    normal_mcs_release_with_value(L, I, COHORT_START);

}

static inline int nowaiters_root(struct hnode * L, zm_mcs_qnode_t *I) {

    return (LOAD(&I->next) == ZM_NULL);

}

static inline void acquire_helper(int level, struct hnode * L, zm_mcs_qnode_t *I) {

    // Trivial case = root level

    if (level == 1)

        acquire_root(L, I);

    else {

        // Prepare the node for use.

        reuse_qnode(I);

        zm_mcs_qnode_t* pred = (zm_mcs_qnode_t*)SWAP(&(L->lock), (zm_ptr_t)I);

        if(!pred) {

            // I am the first one at this level

            // begining of cohort

            STORE(&I->status, COHORT_START);

            // acquire at next level if not at the top level

            acquire_helper(level - 1, L->parent, &(L->node));

            return;

        } else {

            STORE(&pred->next, I);

            for(;;){

                unsigned myStatus = LOAD(&I->status);

                if(myStatus < ACQUIRE_PARENT) {

                    return;

                }

                if(myStatus == ACQUIRE_PARENT) {

                    // beginning of cohort

                    STORE(&I->status, COHORT_START);

                    // This means this level is acquired and we can start the next level

                    acquire_helper(level - 1, L->parent, &(L->node));

                    return;

                }

                // spin back; (I->status == WAIT)

            }

        }

    }

}

static inline void release_helper(int level, struct hnode * L, zm_mcs_qnode_t *I) {

    // Trivial case = root level

    if (level == 1) {

        release_root(L, I);

    } else {

        unsigned cur_count = LOAD(&(I->status)) ;

        zm_mcs_qnode_t * succ;

        // Lower level releases

        if(cur_count == get_threshold(L)) {

            // NO KNOWN SUCCESSORS / DESCENDENTS

            // reached threshold and have next level

            // release to next level

            release_helper(level - 1, L->parent, &(L->node));

            //COMMIT_ALL_WRITES();

            // Tap successor at this level and ask to spin acquire next level lock

            normal_mcs_release_with_value(L, I, ACQUIRE_PARENT);

            return;

        }

        succ = (zm_mcs_qnode_t*)LOAD(&(I->next));

        // Not reached threshold

        if(succ) {

            STORE(&succ->status, cur_count + 1);

            return; // released

        }

        // No known successor, so release

        release_helper(level - 1, L->parent, &(L->node));

        // Tap successor at this level and ask to spin acquire next level lock

        normal_mcs_release_with_value(L, I, ACQUIRE_PARENT);

    }

}

static inline int nowaiters_helper(int level, struct hnode * L, zm_mcs_qnode_t *I) {

    if (level == 1 ) {

        return nowaiters_root(L,I);

    } else {

        if(LOAD(&I->next) != ZM_NULL)

            return FALSE;

        else

            return nowaiters_helper(level - 1, L->parent, &(L->node));

    }

}

static void* new_leaf(struct hnode *h, int depth) {

    int err;

    struct leaf *leaf;

    err = posix_memalign((void **) &leaf, ZM_CACHELINE_SIZE, sizeof(struct leaf));

    if (err != 0) {

        printf("posix_memalign failed in HMCS : new_leaf \n");

        exit(EXIT_FAILURE);

    }

    leaf->cur_node = h;

    leaf->curDepth = depth;

    leaf->took_fast_path = FALSE;

    struct hnode *tmp, *root_node;

    for(tmp = leaf->cur_node; tmp->parent != NULL; tmp = tmp->parent);

    root_node = tmp;

    leaf->root_node = root_node;

    return leaf;

}

static inline void acquire_from_leaf(int level, struct leaf *L){

    if((zm_ptr_t)L->cur_node->lock == ZM_NULL

    && (zm_ptr_t)L->root_node->lock == ZM_NULL) {

        // go FP

        L->took_fast_path = TRUE;

        acquire_root(L->root_node, &L->I);

        return;

    }

    acquire_helper(level, L->cur_node, &L->I);

    return;

}

static inline void tryacq_from_leaf(int level, struct leaf *L, int *success){

    *success = 0;

    if((zm_ptr_t)L->cur_node->lock == ZM_NULL

    && (zm_ptr_t)L->root_node->lock == ZM_NULL) {

        tryacq_root(L->root_node, &L->I, success);

        if (*success)

            L->took_fast_path = TRUE;

    }

    return;

}

static inline void release_from_leaf(int level, struct leaf *L){

    //myrelease(cur_node, I);

    if(L->took_fast_path) {

        release_root(L->root_node, &L->I);

        L->took_fast_path = FALSE;

        return;

    }

    release_helper(level, L->cur_node, &L->I);

    return;

}

static inline int nowaiters_from_leaf(int level, struct leaf *L){

    // Shouldnt this be nowaiters(root_node, I)?

    if(L->took_fast_path) {

        return nowaiters_root(L->cur_node, &L->I);

    }

    return nowaiters_helper(level, L->cur_node, &L->I);

}

static int get_hwthread_id(hwloc_topology_t topo){

    hwloc_cpuset_t cpuset = hwloc_bitmap_alloc();

    hwloc_obj_t obj;

    hwloc_get_cpubind(topo, cpuset, HWLOC_CPUBIND_THREAD);

    obj = hwloc_get_obj_inside_cpuset_by_type(topo, cpuset, HWLOC_OBJ_PU, 0);

    hwloc_bitmap_free(cpuset);

    return obj->logical_index;

}

static void set_hierarchy(struct lock *L, int *max_threads, int** particip_per_level) {

    int max_depth, levels = 0, max_levels = HMCS_DEFAULT_MAX_LEVELS, explicit_levels = 0;

    char tmp[20];

    char *s = getenv("ZM_HMCS_MAX_LEVELS");

    if (s != NULL)

        max_levels = atoi(s);

    int depths[max_levels];

    int idx = 0;

    /* advice to users: run `hwloc-ls -s --no-io --no-icaches` and choose

     * depths of interest in ascending order. The first depth must be `0' */

    s = getenv("ZM_HMCS_EXPLICIT_LEVELS");

    if (s != NULL) {

        strcpy(tmp, s);

        explicit_levels = 1;

        char* token;

        token = strtok(tmp,",");

        while(token != NULL) {

            depths[idx] = atoi(token);

            if (idx == 0)

                assert(depths[idx] == 0 && "the first depth must be machine level (i.e., depth 0), run `hwloc-ls -s --no-io --no-icaches` and choose appropriate depth values");

            idx++;

            token = strtok(NULL,",");

        }

        assert(idx == max_levels);

    }

    hwloc_topology_init(&L->topo);

    hwloc_topology_load(L->topo);

    *max_threads = hwloc_get_nbobjs_by_type(L->topo, HWLOC_OBJ_PU);

    max_depth = hwloc_topology_get_depth(L->topo);

    assert(max_depth >= 2); /* At least Machine and Core levels exist */

    *particip_per_level = (int*) malloc(max_levels * sizeof(int));

    int prev_nobjs = -1;

    if(!explicit_levels) {

        for (int d = max_depth - 2; d > 1; d--) {

            int cur_nobjs = hwloc_get_nbobjs_by_depth(L->topo, d);

            /* Check if this level has a hierarchical impact */

            if(cur_nobjs != prev_nobjs) {

                prev_nobjs = cur_nobjs;

                (*particip_per_level)[levels] = (*max_threads)/cur_nobjs;

                levels++;

                if(levels >= max_levels - 1)

                    break;

            }

        }

        (*particip_per_level)[levels] = *max_threads;

        levels++;

    } else {

        for(int i = max_levels - 1; i >= 0; i--){

            int d = depths[i];

            int cur_nobjs = hwloc_get_nbobjs_by_depth(L->topo, d);

            /* Check if this level has a hierarchical impact */

            if(cur_nobjs != prev_nobjs) {

                prev_nobjs = cur_nobjs;

                (*particip_per_level)[levels] = (*max_threads)/cur_nobjs;

                levels++;

            } else {

                assert(0 && "plz choose levels that have a hierarchical impact");

            }

        }

    }

    L->levels = levels;

}

static void free_hierarchy(int* particip_per_level){

    free(particip_per_level);

}

static void* new_lock(){

    struct lock *L;

    posix_memalign((void **) &L, ZM_CACHELINE_SIZE, sizeof(struct lock));

    int max_threads;

    int *participants_at_level;

    set_hierarchy(L, &max_threads, &participants_at_level);

    // Total locks needed = participantsPerLevel[1] + participantsPerLevel[2] + .. participantsPerLevel[levels-1] + 1

    // Save affinity

    hwloc_cpuset_t cpuset = hwloc_bitmap_alloc();

    hwloc_get_cpubind(L->topo, cpuset, HWLOC_CPUBIND_THREAD);

    int total_locks_needed = 0;

    int levels = L->levels;

    for (int i=0; i < levels; i++) {

        total_locks_needed += max_threads / participants_at_level[i] ;

    }

    struct hnode ** lock_locations;

    posix_memalign((void **) &lock_locations, ZM_CACHELINE_SIZE, sizeof(struct hnode*) * total_locks_needed);

    struct leaf ** leaf_nodes;

    posix_memalign((void **) &leaf_nodes, ZM_CACHELINE_SIZE, sizeof(struct leaf*) * max_threads);

    int threshold = DEFAULT_THRESHOLD;

    char *s = getenv("ZM_HMCS_THRESHOLD");

    if (s != NULL)

        threshold = atoi(s);

    hwloc_obj_t obj;

    for(int tid = 0 ; tid < max_threads; tid ++){

        obj = hwloc_get_obj_by_type (L->topo, HWLOC_OBJ_PU, tid);

        hwloc_set_cpubind(L->topo, obj->cpuset, HWLOC_CPUBIND_THREAD);

        // Pin me to hw-thread-id = tid

        int last_lock_location_end = 0;

        for(int cur_level = 0 ; cur_level < levels; cur_level++){

            if (tid%participants_at_level[cur_level] == 0) {

                // master, initialize the lock

                int lock_location = last_lock_location_end + tid/participants_at_level[cur_level];

                last_lock_location_end += max_threads/participants_at_level[cur_level];

                struct hnode * cur_hnode = new_hnode();

                cur_hnode->threshold = threshold;

                cur_hnode->parent = NULL;

                cur_hnode->lock = ZM_NULL;

                lock_locations[lock_location] = cur_hnode;

            }

        }

    }

    // setup parents

    for(int tid = 0 ; tid < max_threads; tid ++){

        obj = hwloc_get_obj_by_type (L->topo, HWLOC_OBJ_PU, tid);

        hwloc_set_cpubind(L->topo, obj->cpuset, HWLOC_CPUBIND_THREAD);

        int last_lock_location_end = 0;

        for(int cur_level = 0 ; cur_level < levels - 1; cur_level++){

            if (tid%participants_at_level[cur_level] == 0) {

                int lock_location = last_lock_location_end + tid/participants_at_level[cur_level];

                last_lock_location_end += max_threads/participants_at_level[cur_level];

                int parentLockLocation = last_lock_location_end + tid/participants_at_level[cur_level+1];

                lock_locations[lock_location]->parent = lock_locations[parentLockLocation];

            }

        }

        leaf_nodes[tid] = (struct leaf*)new_leaf(lock_locations[tid/participants_at_level[0]], levels);

    }

    free(lock_locations);

    free_hierarchy(participants_at_level);

    // Restore affinity

    hwloc_set_cpubind(L->topo, cpuset, HWLOC_CPUBIND_THREAD);

    L->leaf_nodes = leaf_nodes;

    hwloc_bitmap_free(cpuset);

    return L;

}

static void search_nodes_rec(struct hnode *node, struct hnode **nodes_to_free, int *num_ptrs, int max_threads) {

    int i;

    if(node != NULL) {

        for(i = 0; i < *num_ptrs; i++) {

            if(node == nodes_to_free[i])

                break; /* already marked to be free'd */

        }

        if(i == *num_ptrs) { /* newly encountered pointer */

            nodes_to_free[*num_ptrs] = node;

            (*num_ptrs)++;

            assert(*num_ptrs < 2*max_threads);

        }

        search_nodes_rec(node->parent, nodes_to_free, num_ptrs, max_threads);

    }

}

static void free_lock(struct lock* L) {

    int max_threads = hwloc_get_nbobjs_by_type(L->topo, HWLOC_OBJ_PU);

    int num_ptrs = 0;

    struct hnode **nodes_to_free = (struct hnode**) malloc(2*max_threads*sizeof(struct hnode*));

    for (int tid = 0; tid < max_threads; tid++) {

        search_nodes_rec(L->leaf_nodes[tid]->cur_node, nodes_to_free, &num_ptrs, max_threads);

        free(L->leaf_nodes[tid]);

    }

    free(L->leaf_nodes);

    for(int i = 0; i < num_ptrs; i++)

        free(nodes_to_free[i]);

    free(nodes_to_free);

    hwloc_topology_destroy(L->topo);

    free(L);

}

static inline void hmcs_acquire(struct lock *L){

    if (zm_unlikely(tid == -1)) {

        check_affinity(L->topo);

        tid = get_hwthread_id(L->topo);

    }

    acquire_from_leaf(L->levels, L->leaf_nodes[tid]);

}

static inline void hmcs_tryacq(struct lock *L, int *success){

    if (zm_unlikely(tid == -1)) {

        check_affinity(L->topo);

        tid = get_hwthread_id(L->topo);

    }

    tryacq_from_leaf(L->levels, L->leaf_nodes[tid], success);

}

static inline void hmcs_release(struct lock *L){

    release_from_leaf(L->levels, L->leaf_nodes[tid]);

}

static inline int hmcs_nowaiters(struct lock *L){

    return nowaiters_from_leaf(L->levels, L->leaf_nodes[tid]);

}

int zm_hmcs_init(zm_hmcs_t * handle) {

    void *p = new_lock();

    *handle  = (zm_hmcs_t) p;

    return 0;

}

int zm_hmcs_destroy(zm_hmcs_t *L) {

    free_lock((struct lock*)(*L));

    return 0;

}

int zm_hmcs_acquire(zm_hmcs_t L){

    hmcs_acquire((struct lock*)L);

    return 0;

}

int zm_hmcs_tryacq(zm_hmcs_t L, int *success){

    hmcs_tryacq((struct lock*)L, success);

    return 0;

}

int zm_hmcs_release(zm_hmcs_t L){

    hmcs_release((struct lock*)L);

    return 0;

}

int zm_hmcs_nowaiters(zm_hmcs_t L){

    return hmcs_nowaiters((struct lock*)L);

}