/*
 * Copyright (c) 2013      Mellanox Technologies, Inc.
 *                         All rights reserved.
 * $COPYRIGHT$
 *
 * Additional copyrights may follow
 *
 * $HEADER$
 */
#include "oshmem_config.h"

#include "oshmem/util/oshmem_util.h"
#include "oshmem/proc/proc.h"
#include "oshmem/mca/memheap/memheap.h"
#include "oshmem/mca/memheap/base/base.h"

#include <stdio.h>

static int _dereg_segment(map_segment_t *s);
static int _reg_segment(map_segment_t *s, int *num_btl);

int mca_memheap_base_reg(mca_memheap_map_t *memheap_map)
{
    int ret = OSHMEM_SUCCESS;
    int i;

    for (i = 0; i < memheap_map->n_segments; i++) {
        map_segment_t *s = &memheap_map->mem_segs[i];

        MEMHEAP_VERBOSE(5,
                        "register seg#%02d: 0x%p - 0x%p %llu bytes type=0x%X id=0x%X",
                        i,
                        s->start,
                        s->end,
                        (long long)((uintptr_t)s->end - (uintptr_t)s->start),
                        s->type,
                        s->shmid);
        ret = _reg_segment(s, &memheap_map->num_transports);
    }

    return ret;
}

int mca_memheap_base_dereg(mca_memheap_map_t *memheap_map)
{
    int ret = OSHMEM_SUCCESS;
    int i;

    for (i = 0; i < memheap_map->n_segments; i++) {
        map_segment_t *s = &memheap_map->mem_segs[i];

        if (!s->is_active)
            continue;

        MEMHEAP_VERBOSE(5,
                        "deregistering segment#%d: %p - %p %llu bytes",
                        i,
                        s->start,
                        s->end,
                        (long long)((uintptr_t)s->end - (uintptr_t)s->start));
        ret = _dereg_segment(s);
    }

    return ret;
}

static int _dereg_segment(map_segment_t *s)
{
    int rc = OSHMEM_SUCCESS;
    int j;
    int nprocs, my_pe;

    nprocs = oshmem_num_procs();
    my_pe = oshmem_my_proc_id();

    MCA_SPML_CALL(deregister(s->mkeys));

    if (s->mkeys_cache) {
        for (j = 0; j < nprocs; j++) {
            if (j == my_pe)
                continue;
            if (s->mkeys_cache[j]) {
                if (s->mkeys_cache[j]->len) {
                    free(s->mkeys_cache[j]->u.data);
                    s->mkeys_cache[j]->len = 0;
                }
                free(s->mkeys_cache[j]);
                s->mkeys_cache[j] = NULL;
            }
        }
        free(s->mkeys_cache);
        s->mkeys_cache = NULL;
    }

    s->is_active = 0;

    return rc;
}

static int _reg_segment(map_segment_t *s, int *num_btl)
{
    int rc = OSHMEM_SUCCESS;
    int my_pe;
    int nprocs;

    nprocs = oshmem_num_procs();
    my_pe = oshmem_my_proc_id();

    s->mkeys_cache = (mca_spml_mkey_t **) calloc(nprocs,
                                                 sizeof(mca_spml_mkey_t *));
    if (NULL == s->mkeys_cache) {
        MEMHEAP_ERROR("Failed to allocate memory for remote segments");
        rc = OSHMEM_ERROR;
    }

    if (!rc) {
        s->mkeys = MCA_SPML_CALL(register((void *)(unsigned long)s->start,
                        (uintptr_t)s->end - (uintptr_t)s->start,
                        MEMHEAP_SHM_CODE(s->type, s->shmid),
                        num_btl));
        if (NULL == s->mkeys) {
            free(s->mkeys_cache);
            s->mkeys_cache = NULL;

            MEMHEAP_ERROR("Failed to register segment");
            rc = OSHMEM_ERROR;
        }
    }

    if (OSHMEM_SUCCESS == rc) {
        s->mkeys_cache[my_pe] = s->mkeys;
        s->is_active = 1;
    }

    return rc;
}