// SPDX-License-Identifier: GPL-2.0 OR Linux-OpenIB
/* Copyright (c) 2024 NVIDIA Corporation & Affiliates */

#include <mlx5_core.h>
#include "fs_pool.h"

int mlx5_fs_bulk_bitmap_alloc(struct mlx5_core_dev *dev,
                              struct mlx5_fs_bulk *fs_bulk)
{
        int i;

        fs_bulk->bitmask = kvcalloc(BITS_TO_LONGS(fs_bulk->bulk_len),
                                    sizeof(unsigned long), GFP_KERNEL);
        if (!fs_bulk->bitmask)
                return -ENOMEM;

        for (i = 0; i < fs_bulk->bulk_len; i++)
                set_bit(i, fs_bulk->bitmask);

        return 0;
}

void mlx5_fs_bulk_init(struct mlx5_fs_bulk *fs_bulk, int bulk_len)
{
        fs_bulk->bulk_len = bulk_len;
}

void mlx5_fs_bulk_cleanup(struct mlx5_fs_bulk *fs_bulk)
{
        kvfree(fs_bulk->bitmask);
}

int mlx5_fs_bulk_get_free_amount(struct mlx5_fs_bulk *bulk)
{
        return bitmap_weight(bulk->bitmask, bulk->bulk_len);
}

static int mlx5_fs_bulk_acquire_index(struct mlx5_fs_bulk *fs_bulk,
                                      struct mlx5_fs_pool_index *pool_index)
{
        int free_index = find_first_bit(fs_bulk->bitmask, fs_bulk->bulk_len);

        WARN_ON_ONCE(!pool_index || !fs_bulk);
        if (free_index >= fs_bulk->bulk_len)
                return -ENOSPC;

        clear_bit(free_index, fs_bulk->bitmask);
        pool_index->fs_bulk = fs_bulk;
        pool_index->index = free_index;
        return 0;
}

static int mlx5_fs_bulk_release_index(struct mlx5_fs_bulk *fs_bulk, int index)
{
        if (test_bit(index, fs_bulk->bitmask))
                return -EINVAL;

        set_bit(index, fs_bulk->bitmask);
        return 0;
}

void mlx5_fs_pool_init(struct mlx5_fs_pool *pool, struct mlx5_core_dev *dev,
                       const struct mlx5_fs_pool_ops *ops, void *pool_ctx)
{
        WARN_ON_ONCE(!ops || !ops->bulk_destroy || !ops->bulk_create ||
                     !ops->update_threshold);
        pool->dev = dev;
        pool->pool_ctx = pool_ctx;
        mutex_init(&pool->pool_lock);
        INIT_LIST_HEAD(&pool->fully_used);
        INIT_LIST_HEAD(&pool->partially_used);
        INIT_LIST_HEAD(&pool->unused);
        pool->available_units = 0;
        pool->used_units = 0;
        pool->threshold = 0;
        pool->ops = ops;
}

void mlx5_fs_pool_cleanup(struct mlx5_fs_pool *pool)
{
        struct mlx5_core_dev *dev = pool->dev;
        struct mlx5_fs_bulk *bulk;
        struct mlx5_fs_bulk *tmp;

        list_for_each_entry_safe(bulk, tmp, &pool->fully_used, pool_list)
                pool->ops->bulk_destroy(dev, bulk);
        list_for_each_entry_safe(bulk, tmp, &pool->partially_used, pool_list)
                pool->ops->bulk_destroy(dev, bulk);
        list_for_each_entry_safe(bulk, tmp, &pool->unused, pool_list)
                pool->ops->bulk_destroy(dev, bulk);
}

static struct mlx5_fs_bulk *
mlx5_fs_pool_alloc_new_bulk(struct mlx5_fs_pool *fs_pool)
{
        struct mlx5_core_dev *dev = fs_pool->dev;
        struct mlx5_fs_bulk *new_bulk;

        new_bulk = fs_pool->ops->bulk_create(dev, fs_pool->pool_ctx);
        if (new_bulk)
                fs_pool->available_units += new_bulk->bulk_len;
        fs_pool->ops->update_threshold(fs_pool);
        return new_bulk;
}

static void
mlx5_fs_pool_free_bulk(struct mlx5_fs_pool *fs_pool, struct mlx5_fs_bulk *bulk)
{
        struct mlx5_core_dev *dev = fs_pool->dev;

        fs_pool->available_units -= bulk->bulk_len;
        fs_pool->ops->bulk_destroy(dev, bulk);
        fs_pool->ops->update_threshold(fs_pool);
}

static int
mlx5_fs_pool_acquire_from_list(struct list_head *src_list,
                               struct list_head *next_list,
                               bool move_non_full_bulk,
                               struct mlx5_fs_pool_index *pool_index)
{
        struct mlx5_fs_bulk *fs_bulk;
        int err;

        if (list_empty(src_list))
                return -ENODATA;

        fs_bulk = list_first_entry(src_list, struct mlx5_fs_bulk, pool_list);
        err = mlx5_fs_bulk_acquire_index(fs_bulk, pool_index);
        if (move_non_full_bulk || mlx5_fs_bulk_get_free_amount(fs_bulk) == 0)
                list_move(&fs_bulk->pool_list, next_list);
        return err;
}

int mlx5_fs_pool_acquire_index(struct mlx5_fs_pool *fs_pool,
                               struct mlx5_fs_pool_index *pool_index)
{
        struct mlx5_fs_bulk *new_bulk;
        int err;

        mutex_lock(&fs_pool->pool_lock);

        err = mlx5_fs_pool_acquire_from_list(&fs_pool->partially_used,
                                             &fs_pool->fully_used, false,
                                             pool_index);
        if (err)
                err = mlx5_fs_pool_acquire_from_list(&fs_pool->unused,
                                                     &fs_pool->partially_used,
                                                     true, pool_index);
        if (err) {
                new_bulk = mlx5_fs_pool_alloc_new_bulk(fs_pool);
                if (!new_bulk) {
                        err = -ENOENT;
                        goto out;
                }
                err = mlx5_fs_bulk_acquire_index(new_bulk, pool_index);
                WARN_ON_ONCE(err);
                list_add(&new_bulk->pool_list, &fs_pool->partially_used);
        }
        fs_pool->available_units--;
        fs_pool->used_units++;

out:
        mutex_unlock(&fs_pool->pool_lock);
        return err;
}

int mlx5_fs_pool_release_index(struct mlx5_fs_pool *fs_pool,
                               struct mlx5_fs_pool_index *pool_index)
{
        struct mlx5_fs_bulk *bulk = pool_index->fs_bulk;
        int bulk_free_amount;
        int err;

        mutex_lock(&fs_pool->pool_lock);

        /* TBD would rather return void if there was no warn here in original code */
        err = mlx5_fs_bulk_release_index(bulk, pool_index->index);
        if (err)
                goto unlock;

        fs_pool->available_units++;
        fs_pool->used_units--;

        bulk_free_amount = mlx5_fs_bulk_get_free_amount(bulk);
        if (bulk_free_amount == 1)
                list_move_tail(&bulk->pool_list, &fs_pool->partially_used);
        if (bulk_free_amount == bulk->bulk_len) {
                list_del(&bulk->pool_list);
                if (fs_pool->available_units > fs_pool->threshold)
                        mlx5_fs_pool_free_bulk(fs_pool, bulk);
                else
                        list_add(&bulk->pool_list, &fs_pool->unused);
        }

unlock:
        mutex_unlock(&fs_pool->pool_lock);
        return err;
}