#include <socks/cpu.h>
#include <socks/percpu.h>
#include <socks/bitmap.h>

DECLARE_BITMAP(cpu_available, CPU_MAX);
DECLARE_BITMAP(cpu_online, CPU_MAX);

DEFINE_PERCPU_VAR(cpu_data_t, cpu_data);

cpu_data_t *get_this_cpu(void)
{
	return percpu_get(&cpu_data);
}

void put_cpu(cpu_data_t *cpu)
{
	percpu_put(cpu);
}

void cpu_set_available(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return;
	}

	bitmap_set(cpu_available, cpu_id);
}

void cpu_set_online(unsigned int cpu_id)
{
	if (cpu_id >= CPU_MAX) {
		return;
	}

	bitmap_set(cpu_online, cpu_id);
}

void preempt_disable(void)
{
	ml_cpu_block *ml_cpu = ml_this_cpu();
	cpu_data_t *cpu_data = ml_cpu_block_get_data(ml_cpu);
	if (!cpu_data) {
		return;
	}

	unsigned int preempt = READ_ONCE(cpu_data->c_preempt_count);
	preempt++;
	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
}

void preempt_enable(void)
{
	ml_cpu_block *ml_cpu = ml_this_cpu();
	cpu_data_t *cpu_data = ml_cpu_block_get_data(ml_cpu);
	if (!cpu_data) {
		return;
	}

	unsigned int preempt = READ_ONCE(cpu_data->c_preempt_count);

	if (preempt > 0) {
		preempt--;
	}

	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
}

unsigned int cpu_get_highest_available(void)
{
	return bitmap_highest_set(cpu_available, CPU_MAX);
}