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

#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include "bpf_kfuncs.h"
#include "../test_kmods/bpf_testmod_kfunc.h"

#if (defined(__TARGET_ARCH_x86) || defined(__TARGET_ARCH_arm64)) && \
	defined(__BPF_FEATURE_STACK_ARGUMENT)

const volatile bool has_stack_arg = true;

struct bpf_iter_testmod_seq {
	u64 :64;
	u64 :64;
};

extern int bpf_iter_testmod_seq_new(struct bpf_iter_testmod_seq *it, s64 value, int cnt) __ksym;
extern void bpf_iter_testmod_seq_destroy(struct bpf_iter_testmod_seq *it) __ksym;

struct timer_map_value {
	struct bpf_timer timer;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, int);
	__type(value, struct timer_map_value);
} kfunc_timer_map SEC(".maps");

SEC("tc")
int test_stack_arg_scalar(struct __sk_buff *skb)
{
	return bpf_kfunc_call_stack_arg(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
}

SEC("tc")
int test_stack_arg_ptr(struct __sk_buff *skb)
{
	struct prog_test_pass1 p = { .x0 = 10, .x1 = 20 };

	return bpf_kfunc_call_stack_arg_ptr(1, 2, 3, 4, 5, 6, 7, 8, 9, &p);
}

SEC("tc")
int test_stack_arg_mix(struct __sk_buff *skb)
{
	struct prog_test_pass1 p = { .x0 = 10 };
	struct prog_test_pass1 q = { .x1 = 20 };

	return bpf_kfunc_call_stack_arg_mix(1, 2, 3, 4, 5, 6, 7, &p, 8, &q);
}

/* 1+2+3+4+5+6+7+8+9+sizeof(pkt_v4) = 45+54 = 99 */
SEC("tc")
int test_stack_arg_dynptr(struct __sk_buff *skb)
{
	struct bpf_dynptr ptr;

	bpf_dynptr_from_skb(skb, 0, &ptr);
	return bpf_kfunc_call_stack_arg_dynptr(1, 2, 3, 4, 5, 6, 7, 8, 9, &ptr);
}

/* 1 + 2 + 3 + 4 + 5 + (1 + 2 + ... + 16) = 15 + 136 = 151 */
SEC("tc")
int test_stack_arg_mem(struct __sk_buff *skb)
{
	char buf[16] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};

	return bpf_kfunc_call_stack_arg_mem(1, 2, 3, 4, 5, buf, sizeof(buf));
}

/* 1+2+3+4+5+6+7+8+9+100 = 145 */
SEC("tc")
int test_stack_arg_iter(struct __sk_buff *skb)
{
	struct bpf_iter_testmod_seq it;
	u64 ret;

	bpf_iter_testmod_seq_new(&it, 100, 10);
	ret = bpf_kfunc_call_stack_arg_iter(1, 2, 3, 4, 5, 6, 7, 8, 9, &it);
	bpf_iter_testmod_seq_destroy(&it);
	return ret;
}

const char cstr[] = "hello";

/* 1+2+3+4+5+6+7+8+9 = 45 */
SEC("tc")
int test_stack_arg_const_str(struct __sk_buff *skb)
{
	return bpf_kfunc_call_stack_arg_const_str(1, 2, 3, 4, 5, 6, 7, 8, 9,
						  cstr);
}

/* 1+2+3+4+5+6+7+8+9 = 45 */
SEC("tc")
int test_stack_arg_timer(struct __sk_buff *skb)
{
	struct timer_map_value *val;
	int key = 0;

	val = bpf_map_lookup_elem(&kfunc_timer_map, &key);
	if (!val)
		return 0;
	return bpf_kfunc_call_stack_arg_timer(1, 2, 3, 4, 5, 6, 7, 8, 9,
					      &val->timer);
}

#else

const volatile bool has_stack_arg = false;

SEC("tc")
int test_stack_arg_scalar(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_ptr(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_mix(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_dynptr(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_mem(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_iter(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_const_str(struct __sk_buff *skb)
{
	return 0;
}

SEC("tc")
int test_stack_arg_timer(struct __sk_buff *skb)
{
	return 0;
}

#endif

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