rte_member_heap.h

/* SPDX-License-Identifier: BSD-3-Clause
 * Copyright(c) 2020 Intel Corporation
 * Copyright(c) 2020, Alan Liu <zaoxingliu@gmail.com>
 */

#ifndef RTE_MEMBER_HEAP_H
#define RTE_MEMBER_HEAP_H

#include <rte_ring_elem.h>
#include "rte_member.h"

#define LCHILD(x) (2 * x + 1)
#define RCHILD(x) (2 * x + 2)
#define PARENT(x) ((x - 1) / 2)

#define HASH_BKT_SIZE 16
#define HASH_HP_MULTI 4
#define HASH_RESIZE_MULTI 2

struct hash_bkt {
    uint16_t sig[HASH_BKT_SIZE];
    uint16_t idx[HASH_BKT_SIZE];
};

struct hash {
    uint16_t bkt_cnt;
    uint16_t num_item;
    uint32_t seed;
    struct hash_bkt buckets[0];
};

struct node {
    void *key;
    uint64_t count;
};

struct minheap {
    uint32_t key_len;
    uint32_t size;
    uint32_t socket;
    struct hash *hashtable;
    struct node *elem;
};

static int
hash_table_insert(const void *key, int value, int key_len, struct hash *table)
{
    uint32_t hash = MEMBER_HASH_FUNC(key, key_len, table->seed);
    uint16_t idx = hash % table->bkt_cnt;
    uint16_t sig = hash >> 16;
    int i;

    for (i = 0; i < HASH_BKT_SIZE; i++) {
        if (table->buckets[idx].idx[i] == 0) {
            table->buckets[idx].idx[i] = value;
            table->buckets[idx].sig[i] = sig;
            table->num_item++;
            return 0;
        }
    }

    return -ENOMEM;
}

static int
hash_table_update(const void *key, int old_value, int value, int key_len, struct hash *table)
{
    uint32_t hash = MEMBER_HASH_FUNC(key, key_len, table->seed);
    uint16_t idx = hash % table->bkt_cnt;
    uint16_t sig = hash >> 16;
    int i;

    for (i = 0; i < HASH_BKT_SIZE; i++) {
        if (table->buckets[idx].sig[i] == sig && table->buckets[idx].idx[i] == old_value) {
            table->buckets[idx].idx[i] = value;
            return 0;
        }
    }

    return -1;
}

static int
hash_table_del(const void *key, uint16_t value, int key_len, struct hash *table)
{
    uint32_t hash = MEMBER_HASH_FUNC(key, key_len, table->seed);
    uint16_t idx = hash % table->bkt_cnt;
    uint16_t sig = hash >> 16;
    int i;

    for (i = 0; i < HASH_BKT_SIZE; i++) {
        if (table->buckets[idx].sig[i] == sig && table->buckets[idx].idx[i] == value) {
            table->buckets[idx].idx[i] = 0;
            table->num_item--;
            return 0;
        }
    }

    return -1;
}

static int
hash_table_lookup(const void *key, int key_len, struct minheap *hp)
{
    struct hash *table = hp->hashtable;
    uint32_t hash = MEMBER_HASH_FUNC(key, key_len, table->seed);
    uint16_t idx = hash % table->bkt_cnt;
    uint16_t sig = hash >> 16;
    int i;

    for (i = 0; i < HASH_BKT_SIZE; i++) {
        if (table->buckets[idx].sig[i] == sig && table->buckets[idx].idx[i] != 0) {
            uint32_t hp_idx = table->buckets[idx].idx[i] - 1;

            if (memcmp(hp->elem[hp_idx].key, key, hp->key_len) == 0)
                return hp_idx;
        }
    }

    return -ENOENT; /* key doesn't exist */
}

static int
resize_hash_table(struct minheap *hp)
{
    uint32_t i;
    uint32_t new_bkt_cnt;

    while (1) {
        new_bkt_cnt = hp->hashtable->bkt_cnt * HASH_RESIZE_MULTI;

        RTE_MEMBER_LOG(ERR, "Sketch Minheap HT load factor is [%f]\n",
            hp->hashtable->num_item / ((float)hp->hashtable->bkt_cnt * HASH_BKT_SIZE));
        RTE_MEMBER_LOG(ERR, "Sketch Minheap HT resize happen!\n");
        rte_free(hp->hashtable);
        hp->hashtable = rte_zmalloc_socket(NULL, sizeof(struct hash) +
                        new_bkt_cnt * sizeof(struct hash_bkt),
                        RTE_CACHE_LINE_SIZE, hp->socket);

        if (hp->hashtable == NULL) {
            RTE_MEMBER_LOG(ERR, "Sketch Minheap HT allocation failed\n");
            return -ENOMEM;
        }

        hp->hashtable->bkt_cnt = new_bkt_cnt;

        for (i = 0; i < hp->size; ++i) {
            if (hash_table_insert(hp->elem[i].key,
                i + 1, hp->key_len, hp->hashtable) < 0) {
                RTE_MEMBER_LOG(ERR,
                    "Sketch Minheap HT resize insert fail!\n");
                break;
            }
        }
        if (i == hp->size)
            break;
    }

    return 0;
}

/* find the item in the given minheap */
static int
rte_member_minheap_find(struct minheap *hp, const void *key)
{
    int idx = hash_table_lookup(key, hp->key_len, hp);
    return idx;
}

static int
rte_member_minheap_init(struct minheap *heap, int size,
            uint32_t socket, uint32_t seed)
{
    heap->elem = rte_zmalloc_socket(NULL, sizeof(struct node) * size,
                RTE_CACHE_LINE_SIZE, socket);
    if (heap->elem == NULL) {
        RTE_MEMBER_LOG(ERR, "Sketch Minheap elem allocation failed\n");
        return -ENOMEM;
    }

    uint32_t hash_bkt_cnt = rte_align32pow2(size * HASH_HP_MULTI) / HASH_BKT_SIZE;

    if (hash_bkt_cnt == 0)
        hash_bkt_cnt = 1;

    heap->hashtable = rte_zmalloc_socket(NULL, sizeof(struct hash) +
                    hash_bkt_cnt * sizeof(struct hash_bkt),
                    RTE_CACHE_LINE_SIZE, socket);

    if (heap->hashtable == NULL) {
        RTE_MEMBER_LOG(ERR, "Sketch Minheap HT allocation failed\n");
        rte_free(heap->elem);
        return -ENOMEM;
    }

    heap->hashtable->seed = seed;
    heap->hashtable->bkt_cnt = hash_bkt_cnt;
    heap->socket = socket;

    return 0;
}

/* swap the minheap nodes */
static __rte_always_inline void
rte_member_heap_swap(struct node *n1, struct node *n2)
{
    struct node temp = *n1;
    *n1 = *n2;
    *n2 = temp;
}

/* heapify function */
static void
rte_member_heapify(struct minheap *hp, uint32_t idx, bool update_hash)
{
    uint32_t smallest;

    if (LCHILD(idx) < hp->size &&
            hp->elem[LCHILD(idx)].count < hp->elem[idx].count)
        smallest = LCHILD(idx);
    else
        smallest = idx;

    if (RCHILD(idx) < hp->size &&
            hp->elem[RCHILD(idx)].count < hp->elem[smallest].count)
        smallest = RCHILD(idx);

    if (smallest != idx) {
        rte_member_heap_swap(&(hp->elem[idx]), &(hp->elem[smallest]));

        if (update_hash) {
            if (hash_table_update(hp->elem[smallest].key, idx + 1, smallest + 1,
                    hp->key_len, hp->hashtable) < 0) {
                RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
                return;
            }

            if (hash_table_update(hp->elem[idx].key, smallest + 1, idx + 1,
                    hp->key_len, hp->hashtable) < 0) {
                RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
                return;
            }
        }
        rte_member_heapify(hp, smallest, update_hash);
    }
}

/* insert a node into the minheap */
static int
rte_member_minheap_insert_node(struct minheap *hp, const void *key,
                   int counter, void *key_slot,
                   struct rte_ring *free_key_slot)
{
    struct node nd;
    uint32_t slot_id;

    if (rte_ring_sc_dequeue_elem(free_key_slot, &slot_id, sizeof(uint32_t)) != 0) {
        RTE_MEMBER_LOG(ERR, "Minheap get empty keyslot failed\n");
        return -1;
    }

    nd.count = counter;
    nd.key = RTE_PTR_ADD(key_slot, slot_id * hp->key_len);

    memcpy(nd.key, key, hp->key_len);

    uint32_t i = (hp->size)++;

    while (i && nd.count < hp->elem[PARENT(i)].count) {
        hp->elem[i] = hp->elem[PARENT(i)];
        if (hash_table_update(hp->elem[i].key, PARENT(i) + 1, i + 1,
                hp->key_len, hp->hashtable) < 0) {
            RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
            return -1;
        }
        i = PARENT(i);
    }
    hp->elem[i] = nd;

    if (hash_table_insert(key, i + 1, hp->key_len, hp->hashtable) < 0) {
        if (resize_hash_table(hp) < 0) {
            RTE_MEMBER_LOG(ERR, "Minheap Hash Table resize failed\n");
            return -1;
        }
    }

    return 0;
}

/* delete a key from the minheap */
static int
rte_member_minheap_delete_node(struct minheap *hp, const void *key,
                   void *key_slot, struct rte_ring *free_key_slot)
{
    int idx = rte_member_minheap_find(hp, key);
    uint32_t offset = RTE_PTR_DIFF(hp->elem[idx].key, key_slot) / hp->key_len;

    if (hash_table_del(key, idx + 1, hp->key_len, hp->hashtable) < 0) {
        RTE_MEMBER_LOG(ERR, "Minheap Hash Table delete failed\n");
        return -1;
    }

    rte_ring_sp_enqueue_elem(free_key_slot, &offset, sizeof(uint32_t));

    if (idx == (int)(hp->size - 1)) {
        hp->size--;
        return 0;
    }

    hp->elem[idx] = hp->elem[hp->size - 1];

    if (hash_table_update(hp->elem[idx].key, hp->size, idx + 1,
                hp->key_len, hp->hashtable) < 0) {
        RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
        return -1;
    }
    hp->size--;
    rte_member_heapify(hp, idx, true);

    return 0;
}

/* replace a min node with a new key. */
static int
rte_member_minheap_replace_node(struct minheap *hp,
                const void *new_key,
                int new_counter)
{
    struct node nd;
    void *recycle_key = NULL;

    recycle_key = hp->elem[0].key;

    if (hash_table_del(recycle_key, 1, hp->key_len, hp->hashtable) < 0) {
        RTE_MEMBER_LOG(ERR, "Minheap Hash Table delete failed\n");
        return -1;
    }

    hp->elem[0] = hp->elem[hp->size - 1];

    if (hash_table_update(hp->elem[0].key, hp->size, 1,
                hp->key_len, hp->hashtable) < 0) {
        RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
        return -1;
    }
    hp->size--;

    rte_member_heapify(hp, 0, true);

    nd.count = new_counter;
    nd.key = recycle_key;

    memcpy(nd.key, new_key, hp->key_len);

    uint32_t i = (hp->size)++;

    while (i && nd.count < hp->elem[PARENT(i)].count) {
        hp->elem[i] = hp->elem[PARENT(i)];
        if (hash_table_update(hp->elem[i].key, PARENT(i) + 1, i + 1,
                hp->key_len, hp->hashtable) < 0) {
            RTE_MEMBER_LOG(ERR, "Minheap Hash Table update failed\n");
            return -1;
        }
        i = PARENT(i);
    }

    hp->elem[i] = nd;

    if (hash_table_insert(new_key, i + 1, hp->key_len, hp->hashtable) < 0) {
        RTE_MEMBER_LOG(ERR, "Minheap Hash Table replace insert failed\n");
        if (resize_hash_table(hp) < 0) {
            RTE_MEMBER_LOG(ERR, "Minheap Hash Table replace resize failed\n");
            return -1;
        }
    }

    return 0;
}

/* sort the heap into a descending array */
static void
rte_member_heapsort(struct minheap *hp, struct node *result_array)
{
    struct minheap new_hp;

    /* build a new heap for using the given array */
    new_hp.size = hp->size;
    new_hp.key_len = hp->key_len;
    new_hp.elem = result_array;
    memcpy(result_array, hp->elem, hp->size * sizeof(struct node));

    /* sort the new heap */
    while (new_hp.size > 1) {
        rte_member_heap_swap(&(new_hp.elem[0]), &(new_hp.elem[new_hp.size - 1]));
        new_hp.size--;
        rte_member_heapify(&new_hp, 0, false);
    }
}

static void
rte_member_minheap_free(struct minheap *hp)
{
    if (hp == NULL)
        return;

    rte_free(hp->elem);
    rte_free(hp->hashtable);
}

static void
rte_member_minheap_reset(struct minheap *hp)
{
    if (hp == NULL)
        return;

    memset(hp->elem, 0, sizeof(struct node) * hp->size);
    hp->size = 0;

    memset((char *)hp->hashtable + sizeof(struct hash), 0,
            hp->hashtable->bkt_cnt * sizeof(struct hash_bkt));
    hp->hashtable->num_item = 0;
}

#endif /* RTE_MEMBER_HEAP_H */