/**
* Copyright (C) Mellanox Technologies Ltd. 2020.  ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
*/

#ifndef UCS_PTR_MAP_INL_
#define UCS_PTR_MAP_INL_

#include "ptr_map.h"

#include <ucs/debug/log.h>

BEGIN_C_DECLS


/**
 * The flag is present in key if it is indirect key.
 */
#define UCS_PTR_MAP_KEY_INDIRECT_FLAG   UCS_BIT(0)


/**
 * Returns whether the key is indirect or not.
 *
 * @param [in]  key     Key to object pointer.
 *
 * @return 0 - direct, otherside - indirect.
 */
#define ucs_ptr_map_key_indirect(_key) ((_key) & UCS_PTR_MAP_KEY_INDIRECT_FLAG)


KHASH_IMPL(ucs_ptr_map_impl, ucs_ptr_map_key_t, void*, 1,
           kh_int64_hash_func, kh_int64_hash_equal);

/**
 * Initialize a pointer map.
 *
 * @param [in]  map     Map to initialize.
 * @return      UCS_OK on success otherwise error code as defined by
 *              @ref ucs_status_t.
 */
static inline ucs_status_t ucs_ptr_map_init(ucs_ptr_map_t *map)
{
    UCS_STATIC_ASSERT(!ucs_ptr_map_key_indirect(UCS_PTR_MAP_KEY_INVALID));
    map->next_id = 0;
    kh_init_inplace(ucs_ptr_map_impl, &map->hash);
    return UCS_OK;
}

/**
 * Destroy a pointer map.
 *
 * @param [in]  map     Map to destroy.
 */
static inline void ucs_ptr_map_destroy(ucs_ptr_map_t *map)
{
    size_t size = kh_size(&map->hash);

    if (size != 0) {
        ucs_warn("ptr map %p contains %zd elements on destroy", map, size);
    }

    kh_destroy_inplace(ucs_ptr_map_impl, &map->hash);
}

/**
 * Put a pointer into the map.
 *
 * @param [in]  map       Container.
 * @param [in]  ptr       Object pointer.
 * @param [in]  indirect  If nonzero, the pointer @a ptr is stored in the
 *                        internal hash table, associated with a unique indirect
 *                        key which is returned from this function. Otherwise,
 *                        the returned value is the integer representation of
 *                        the pointer @ptr.
 * @param [out] key       Key to object pointer @ptr if operation completed
 *                        successfully otherwise value is undefined.
 * @return UCS_OK on success, UCS_ERR_NO_PROGRESS if this key is direct and
 *         therefore no action was performed, otherwise error code as defined by
 *         @ref ucs_status_t.
 * @note @ptr must be aligned on @ref UCS_PTR_MAP_KEY_MIN_ALIGN.
 */
static UCS_F_ALWAYS_INLINE ucs_status_t
ucs_ptr_map_put(ucs_ptr_map_t *map, void *ptr, int indirect,
                ucs_ptr_map_key_t *key)
{
    khiter_t iter;
    int ret;

    if (ucs_likely(!indirect)) {
        *key = (uintptr_t)ptr;
        ucs_assert(!(*key & UCS_PTR_MAP_KEY_MIN_ALIGN));
        ucs_assert(*key != 0);
        return UCS_ERR_NO_PROGRESS;
    }

    *key = (map->next_id += UCS_PTR_MAP_KEY_MIN_ALIGN) |
           UCS_PTR_MAP_KEY_INDIRECT_FLAG;

    iter = kh_put(ucs_ptr_map_impl, &map->hash, *key, &ret);
    if (ucs_unlikely(ret == UCS_KH_PUT_FAILED)) {
        return UCS_ERR_NO_MEMORY;
    } else if (ucs_unlikely(ret == UCS_KH_PUT_KEY_PRESENT)) {
        return UCS_ERR_ALREADY_EXISTS;
    }

    kh_value(&map->hash, iter) = ptr;
    return UCS_OK;
}

/**
 * Get a pointer value from the map by its key.
 *
 * @param [in]  map     Container to get the pointer value from.
 * @param [in]  key     Key to look up in the container.
 * @param [in]  extract Whether to remove the key from the map.
 * @param [out] ptr_p   If successful, set to the pointer found in the map.
 *
 * @return UCS_OK if found, UCS_ERR_NO_PROGRESS if this key is direct and
 *         therefore no action was performed, UCS_ERR_NO_ELEM if not found.
 */
static UCS_F_ALWAYS_INLINE ucs_status_t
ucs_ptr_map_get(ucs_ptr_map_t *map, ucs_ptr_map_key_t key, int extract,
                void **ptr_p)
{
    khiter_t iter;

    if (ucs_likely(!ucs_ptr_map_key_indirect(key))) {
        *ptr_p = (void*)key;
        return UCS_ERR_NO_PROGRESS;
    }

    iter = kh_get(ucs_ptr_map_impl, &map->hash, key);
    if (ucs_unlikely(iter == kh_end(&map->hash))) {
        return UCS_ERR_NO_ELEM;
    }

    *ptr_p = kh_value(&map->hash, iter);
    if (extract) {
        kh_del(ucs_ptr_map_impl, &map->hash, iter);
    }
    return UCS_OK;
}

/**
 * Remove object pointer from the map by key.
 *
 * @param [in]  map     Container.
 * @param [in]  key     Key to object pointer.
 *
 * @return - UCS_OK on success
 *         - UCS_ERR_NO_PROGRESS if this key is direct and therefore no action
 *           was performed
 *         - UCS_ERR_NO_ELEM if the key is not found in the internal hash table.
 */
static UCS_F_ALWAYS_INLINE ucs_status_t
ucs_ptr_map_del(ucs_ptr_map_t *map, ucs_ptr_map_key_t key)
{
    void UCS_V_UNUSED *dummy;
    return ucs_ptr_map_get(map, key, 1, &dummy);
}

END_C_DECLS

#endif
