// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#pragma once

#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/mman.h>

#include <bpf/libbpf.h>
#include <bpf/bpf.h>

static inline int libarena_run_prog(int prog_fd)
{
        LIBBPF_OPTS(bpf_test_run_opts, opts);
        int ret;

        ret = bpf_prog_test_run_opts(prog_fd, &opts);
        if (ret)
                return ret;

        return opts.retval;
}

static inline bool libarena_is_test_prog(const char *name)
{
        return strstr(name, "test_") == name;
}

static inline bool libarena_is_asan_test_prog(const char *name)
{
        return strstr(name, "asan_test") == name;
}

static inline bool libarena_is_parallel_test_prog(const char *name)
{
        return strstr(name, "parallel_test") == name;
}


static inline int libarena_run_prog_args(int prog_fd, void *args, size_t argsize)
{
        LIBBPF_OPTS(bpf_test_run_opts, opts);
        int ret;

        opts.ctx_in = args;
        opts.ctx_size_in = argsize;

        ret = bpf_prog_test_run_opts(prog_fd, &opts);

        return ret ?: opts.retval;
}

static inline int libarena_get_arena_base(int arena_get_info_fd,
                                          void **arena_base)
{
        LIBBPF_OPTS(bpf_test_run_opts, opts);
        struct arena_get_info_args args = { .arena_base = NULL };
        int ret;

        opts.ctx_in = &args;
        opts.ctx_size_in = sizeof(args);

        ret = bpf_prog_test_run_opts(arena_get_info_fd, &opts);
        if (ret)
                return ret;
        if (opts.retval)
                return opts.retval;

        *arena_base = args.arena_base;
        return 0;
}

static inline int libarena_get_globals_pages(int arena_get_globals_fd,
                                             size_t arena_all_pages,
                                             u64 *globals_pages)
{
        size_t pgsize = sysconf(_SC_PAGESIZE);
        void *arena_base;
        ssize_t i;
        u8 *vec;
        int ret;

        ret = libarena_get_arena_base(arena_get_globals_fd, &arena_base);
        if (ret)
                return ret;

        if (!arena_base)
                return -EINVAL;

        vec = calloc(arena_all_pages, sizeof(*vec));
        if (!vec)
                return -ENOMEM;

        if (mincore(arena_base, arena_all_pages * pgsize, vec) < 0) {
                ret = -errno;
                free(vec);
                return ret;
        }

        *globals_pages = 0;
        for (i = arena_all_pages - 1; i >= 0; i--) {
                if (!(vec[i] & 0x1))
                        break;
                *globals_pages += 1;
        }

        free(vec);
        return 0;
}

static inline int libarena_asan_init(int arena_asan_init_fd,
                                     int asan_init_fd,
                                     size_t arena_all_pages)
{
        LIBBPF_OPTS(bpf_test_run_opts, opts);
        struct asan_init_args args;
        u64 globals_pages;
        int ret;

        ret = libarena_get_globals_pages(arena_asan_init_fd,
                                         arena_all_pages, &globals_pages);
        if (ret)
                return ret;

        args = (struct asan_init_args){
                .arena_all_pages = arena_all_pages,
                .arena_globals_pages = globals_pages,
        };

        opts.ctx_in = &args;
        opts.ctx_size_in = sizeof(args);

        ret = bpf_prog_test_run_opts(asan_init_fd, &opts);
        if (ret)
                return ret;
        return opts.retval;
}