// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Cloudflare */

#include <vmlinux.h>
#include <errno.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_core_read.h>
#include "bpf_misc.h"
#include "bpf_atomic.h"
#include "progs/lpm_trie.h"

#define BPF_OBJ_NAME_LEN 16U
#define MAX_ENTRIES 100000000
#define NR_LOOPS 10000

char _license[] SEC("license") = "GPL";

/* Filled by userspace. See fill_map() in bench_lpm_trie_map.c */
struct {
	__uint(type, BPF_MAP_TYPE_LPM_TRIE);
	__type(key, struct trie_key);
	__type(value, __u32);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__uint(max_entries, MAX_ENTRIES);
} trie_map SEC(".maps");

long hits;
long duration_ns;

/* Configured from userspace */
__u32 nr_entries;
__u32 prefixlen;
bool random;
__u8 op;

static __u64 latency_free_start;

SEC("fentry/bpf_map_free_deferred")
int BPF_PROG(trie_free_entry, struct work_struct *work)
{
	struct bpf_map *map = container_of(work, struct bpf_map, work);
	char name[BPF_OBJ_NAME_LEN];
	u32 map_type;

	map_type = BPF_CORE_READ(map, map_type);
	if (map_type != BPF_MAP_TYPE_LPM_TRIE)
		return 0;

	/*
	 * Ideally we'd have access to the map ID but that's already
	 * freed before we enter trie_free().
	 */
	BPF_CORE_READ_STR_INTO(&name, map, name);
	if (bpf_strncmp(name, BPF_OBJ_NAME_LEN, "trie_free_map"))
		return 0;

	latency_free_start = bpf_ktime_get_ns();

	return 0;
}

SEC("fexit/bpf_map_free_deferred")
int BPF_PROG(trie_free_exit, struct work_struct *work)
{
	__u64 val;

	if (!latency_free_start)
		return 0;

	val = bpf_ktime_get_ns() - latency_free_start;
	latency_free_start = 0;

	__sync_add_and_fetch(&duration_ns, val);
	__sync_add_and_fetch(&hits, 1);

	return 0;
}

static __u32 cur_key;

static __always_inline void generate_key(struct trie_key *key)
{
	key->prefixlen = prefixlen;

	if (random)
		key->data = bpf_get_prandom_u32() % nr_entries;
	else
		key->data = cur_key++ % nr_entries;
}

static int noop(__u32 index, __u32 *unused)
{
	return 0;
}

static int baseline(__u32 index, __u32 *unused)
{
	struct trie_key key;
	__u32 blackbox = 0;

	generate_key(&key);
	/* Avoid compiler optimizing out the modulo */
	barrier_var(blackbox);
	blackbox = READ_ONCE(key.data);

	return 0;
}

static int lookup(__u32 index, int *retval)
{
	struct trie_key key;

	generate_key(&key);
	if (!bpf_map_lookup_elem(&trie_map, &key)) {
		*retval = -ENOENT;
		return 1;
	}

	return 0;
}

static int insert(__u32 index, int *retval)
{
	struct trie_key key;
	u32 val = 1;
	int err;

	generate_key(&key);
	err = bpf_map_update_elem(&trie_map, &key, &val, BPF_NOEXIST);
	if (err) {
		*retval = err;
		return 1;
	}

	/* Is this the last entry? */
	if (key.data == nr_entries - 1) {
		/* For atomicity concerns, see the comment in delete() */
		*retval = LPM_BENCH_REINIT_MAP;
		return 1;
	}

	return 0;
}

static int update(__u32 index, int *retval)
{
	struct trie_key key;
	u32 val = 1;
	int err;

	generate_key(&key);
	err = bpf_map_update_elem(&trie_map, &key, &val, BPF_EXIST);
	if (err) {
		*retval = err;
		return 1;
	}

	return 0;
}

static int delete(__u32 index, int *retval)
{
	struct trie_key key;
	int err;

	generate_key(&key);
	err = bpf_map_delete_elem(&trie_map, &key);
	if (err) {
		*retval = err;
		return 1;
	}

	/* Do we need to refill the map? */
	if (key.data == nr_entries - 1) {
		/*
		 * Atomicity isn't required because DELETE only supports
		 * one producer running concurrently. What we need is a
		 * way to track how many entries have been deleted from
		 * the trie between consecutive invocations of the BPF
		 * prog because a single bpf_loop() call might not
		 * delete all entries, e.g. when NR_LOOPS < nr_entries.
		 */
		*retval = LPM_BENCH_REINIT_MAP;
		return 1;
	}

	return 0;
}

SEC("xdp")
int BPF_PROG(run_bench)
{
	int err = LPM_BENCH_SUCCESS;
	u64 start, delta;
	int loops;

	start = bpf_ktime_get_ns();

	switch (op) {
	case LPM_OP_NOOP:
		loops = bpf_loop(NR_LOOPS, noop, NULL, 0);
		break;
	case LPM_OP_BASELINE:
		loops = bpf_loop(NR_LOOPS, baseline, NULL, 0);
		break;
	case LPM_OP_LOOKUP:
		loops = bpf_loop(NR_LOOPS, lookup, &err, 0);
		break;
	case LPM_OP_INSERT:
		loops = bpf_loop(NR_LOOPS, insert, &err, 0);
		break;
	case LPM_OP_UPDATE:
		loops = bpf_loop(NR_LOOPS, update, &err, 0);
		break;
	case LPM_OP_DELETE:
		loops = bpf_loop(NR_LOOPS, delete, &err, 0);
		break;
	default:
		bpf_printk("invalid benchmark operation\n");
		return -1;
	}

	delta = bpf_ktime_get_ns() - start;

	__sync_add_and_fetch(&duration_ns, delta);
	__sync_add_and_fetch(&hits, loops);

	return err;
}