/*
 * SPDX-License-Identifier: GPL-2.0
 * Copyright (c) 2025 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2025 Tejun Heo <tj@kernel.org>
 * Copyright (c) 2025 Emil Tsalapatis <etsal@meta.com>
 */
#pragma once

#ifndef __BPF__
#define __arena
#endif /* __BPF__ */

struct scx_alloc_stats {
        __u64           chunk_allocs;
        __u64           data_allocs;
        __u64           alloc_ops;
        __u64           free_ops;
        __u64           active_allocs;
        __u64           arena_pages_used;
};

struct sdt_pool {
        void __arena    *slab;
        __u64           elem_size;
        __u64           max_elems;
        __u64           idx;
};

#ifndef div_round_up
#define div_round_up(a, b) (((a) + (b) - 1) / (b))
#endif

#ifndef round_up
#define round_up(a, b) (div_round_up((a), (b)) * (b))
#endif

typedef struct sdt_desc __arena sdt_desc_t;

enum sdt_consts {
        SDT_TASK_ENTS_PER_PAGE_SHIFT    = 9,
        SDT_TASK_LEVELS                 = 3,
        SDT_TASK_ENTS_PER_CHUNK         = 1 << SDT_TASK_ENTS_PER_PAGE_SHIFT,
        SDT_TASK_CHUNK_BITMAP_U64S      = div_round_up(SDT_TASK_ENTS_PER_CHUNK, 64),
        SDT_TASK_MIN_ELEM_PER_ALLOC     = 8,
};

union sdt_id {
        __s64                           val;
        struct {
                __s32                   idx;    /* index in the radix tree */
                __s32                   genn;   /* ++'d on recycle so that it forms unique'ish 64bit ID */
        };
};

struct sdt_chunk;

/*
 * Each index page is described by the following descriptor which carries the
 * bitmap. This way the actual index can host power-of-two numbers of entries
 * which makes indexing cheaper.
 */
struct sdt_desc {
        __u64                           allocated[SDT_TASK_CHUNK_BITMAP_U64S];
        __u64                           nr_free;
        struct sdt_chunk __arena        *chunk;
};

/*
 * Leaf node containing per-task data.
 */
struct sdt_data {
        union sdt_id                    tid;
        __u64                           payload[];
};

/*
 * Intermediate node pointing to another intermediate node or leaf node.
 */
struct sdt_chunk {
        union {
                sdt_desc_t * descs[SDT_TASK_ENTS_PER_CHUNK];
                struct sdt_data __arena *data[SDT_TASK_ENTS_PER_CHUNK];
        };
};

struct scx_allocator {
        struct sdt_pool pool;
        sdt_desc_t      *root;
};

struct scx_stats {
        int     seq;
        pid_t   pid;
        __u64   enqueue;
        __u64   exit;
        __u64   init;
        __u64   select_busy_cpu;
        __u64   select_idle_cpu;
};

#ifdef __BPF__

void __arena *scx_task_data(struct task_struct *p);
int scx_task_init(__u64 data_size);
void __arena *scx_task_alloc(struct task_struct *p);
void scx_task_free(struct task_struct *p);
void scx_arena_subprog_init(void);

int scx_alloc_init(struct scx_allocator *alloc, __u64 data_size);
u64 scx_alloc_internal(struct scx_allocator *alloc);
int scx_alloc_free_idx(struct scx_allocator *alloc, __u64 idx);

#endif /* __BPF__ */