/* SPDX-License-Identifier: GPL-2.0 */
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */

#ifndef __BENCH_BPF_TIMING_BPF_H__
#define __BENCH_BPF_TIMING_BPF_H__

#include <stdbool.h>
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include <bpf_may_goto.h>

#ifndef BENCH_NR_SAMPLES
#define BENCH_NR_SAMPLES        4096
#endif
#ifndef BENCH_NR_CPUS
#define BENCH_NR_CPUS           256
#endif
#define BENCH_CPU_MASK          (BENCH_NR_CPUS - 1)

__u64 timing_samples[BENCH_NR_CPUS][BENCH_NR_SAMPLES];
__u32 timing_idx[BENCH_NR_CPUS];

volatile __u32 batch_iters;
volatile __u32 timing_enabled;

static __always_inline void bench_record_sample(__u64 elapsed_ns)
{
        __u32 cpu, idx;

        if (!timing_enabled)
                return;

        cpu = bpf_get_smp_processor_id() & BENCH_CPU_MASK;
        idx = timing_idx[cpu];

        if (idx >= BENCH_NR_SAMPLES)
                return;

        timing_samples[cpu][idx] = elapsed_ns;
        timing_idx[cpu] = idx + 1;
}

/*
 * @body:  expression to time; return value (int) stored in __bench_result.
 * @reset: undo body's side-effects so each iteration starts identically.
 *         May reference __bench_result.  Use ({}) for empty reset.
 *
 * Runs batch_iters timed iterations, then one untimed iteration whose
 * return value the macro evaluates to (for validation).
 */
#define BENCH_BPF_LOOP(body, reset) ({                                  \
        __u64 __bench_start = bpf_ktime_get_ns();                       \
        __u32 __bench_i;                                                \
        int __bench_result;                                             \
                                                                        \
        for (__bench_i = 0;                                             \
             __bench_i < batch_iters && can_loop;                       \
             __bench_i++) {                                             \
                __bench_result = (body);                                \
                reset;                                                  \
        }                                                               \
                                                                        \
        bench_record_sample(bpf_ktime_get_ns() - __bench_start);        \
                                                                        \
        __bench_result = (body);                                        \
        __bench_result;                                                 \
})

#endif /* __BENCH_BPF_TIMING_BPF_H__ */