// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/*
 * Copyright (c) 2025-2026 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2025-2026 Emil Tsalapatis <etsal@meta.com>
 */

#include <bpf_atomic.h>

#include <libarena/common.h>

#include <libarena/asan.h>
#include <libarena/spmc.h>

static inline
u64 spmc_arr_size(volatile struct spmc_arr __arena *spmc_arr)
{
	return SPMC_ARR_BASESZ << spmc_arr->order;
}

static inline
u64 spmc_arr_get(volatile struct spmc_arr __arena *spmc_arr, u64 ind)
{
	u64 ret = READ_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)]);

	return ret;
}

static inline
void spmc_arr_put(volatile struct spmc_arr __arena *spmc_arr, u64 ind, u64 value)
{
	WRITE_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)], value);
}

static inline
void spmc_arr_copy(volatile struct spmc_arr __arena *dst,
		   volatile struct spmc_arr __arena *src, u64 b, u64 t)
{
	u64 i;

	for (i = t; i < b && can_loop; i++)
		spmc_arr_put(dst, i, spmc_arr_get(src, i));
}

static inline
int spmc_order_init(struct spmc __arena *spmc, int order)
{
	volatile struct spmc_arr __arena *arr = &spmc->arr[order];

	if (unlikely(!spmc))
		return -EINVAL;

	if (order >= SPMC_ARR_ORDERS)
		return -E2BIG;

	/* Already allocated? */
	if (arr->data)
		return 0;

	arr->data = arena_malloc((SPMC_ARR_BASESZ << order) * sizeof(*arr->data));
	if (!arr->data)
		return -ENOMEM;

	return 0;
}

__weak
int spmc_owned_add(struct spmc __arena *spmc, u64 val)
{
	volatile struct spmc_arr __arena *newarr;
	volatile struct spmc_arr __arena *arr;
	ssize_t sz;
	u64 b, t;
	int ret;

	if (unlikely(!spmc))
		return -EINVAL;

	/* 
	 * Bottom must always be read first, also
	 * see spmc_steal().
	 */
	b = smp_load_acquire(&spmc->bottom);
	t = READ_ONCE(spmc->top);
	arr = READ_ONCE(spmc->cur);

	sz = b - t;
	if (sz >= spmc_arr_size(arr) - 1) {
		ret = spmc_order_init(spmc, arr->order + 1);
		if (ret)
			return ret;

		newarr = &spmc->arr[arr->order + 1];

		spmc_arr_copy(newarr, arr, b, t);
		smp_store_release(&spmc->cur, newarr);
		arr = newarr;
	}

	spmc_arr_put(arr, b, val);
	smp_store_release(&spmc->bottom, b + 1);

	return 0;
}


__weak
int spmc_owned_remove(struct spmc __arena *spmc, u64 *val)
{
	volatile struct spmc_arr __arena *arr;
	int ret = 0;
	ssize_t sz;
	u64 value;
	u64 b, t;

	if (unlikely(!spmc || !val))
		return -EINVAL;

	b = READ_ONCE(spmc->bottom) - 1;
	WRITE_ONCE(spmc->bottom, b);
	smp_mb();

	t = READ_ONCE(spmc->top);
	arr = READ_ONCE(spmc->cur);

	sz = b - t;
	if (sz < 0) {
		WRITE_ONCE(spmc->bottom, t);
		return -ENOENT;
	}

	value = spmc_arr_get(arr, b);
	if (sz > 0) {
		*val = value;
		return 0;
	}

	if (cmpxchg(&spmc->top, t, t + 1) != t)
		ret = -EAGAIN;

	WRITE_ONCE(spmc->bottom, t + 1);

	if (ret)
		return ret;

	*val = value;

	return 0;
}

__weak
int spmc_steal(struct spmc __arena *spmc, u64 *val)
{
	volatile struct spmc_arr __arena *arr;
	ssize_t sz;
	u64 value;
	u64 b, t;

	if (unlikely(!spmc || !val))
		return -EINVAL;

	/*
	 * It is important that t is read before b for
	 * stealers to avoid racing with the owner.
	 * Races between stealers are dealt with using
	 * CAS to increment the top value below.
	 */
	t = smp_load_acquire(&spmc->top);
	b = smp_load_acquire(&spmc->bottom);

	sz = b - t;
	if (sz <= 0)
		return -ENOENT;

	arr = smp_load_acquire(&spmc->cur);
	value = spmc_arr_get(arr, t);

	if (cmpxchg(&spmc->top, t, t + 1) != t)
		return -EAGAIN;

	*val = value;

	return 0;
}


__weak
struct spmc __arena *spmc_create(void)
{
	/*
	 * Marked as volatile because otherwise the array
	 * reference in the internal loop gets demoted to
	 * scalar and the program fails verification.
	 */
	struct spmc __arena *volatile spmc;
	int ret, i;

	spmc = arena_malloc(sizeof(*spmc));
	if (!spmc)
		return NULL;

	spmc->bottom = 0;
	spmc->top = 0;

	for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++) {
		spmc->arr[i].data = NULL;
		spmc->arr[i].order = i;
	}

	ret = spmc_order_init((struct spmc __arena *)spmc, 0);
	if (ret) {
		arena_free(spmc);
		return NULL;
	}

	spmc->cur = &spmc->arr[0];

	return (struct spmc __arena *)spmc;
}

__weak
int spmc_destroy(struct spmc __arena *spmc)
{
	int i;

	if (unlikely(!spmc))
		return -EINVAL;

	for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++)
		arena_free(spmc->arr[i].data);

	arena_free(spmc);

	return 0;
}