// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026 KylinSoft Corporation. */

#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include "bpf_misc.h"
#include "bpf_experimental.h"

#define NR_NODES 16

struct node_data {
        int data;
};

struct tree_node {
        struct bpf_rb_node node;
        u64 key;
        struct node_data __kptr * node_data;
};

struct tree_node_ref {
        struct bpf_refcount ref;
        struct bpf_rb_node node;
        u64 key;
        struct node_data __kptr * node_data;
};

#define private(name) SEC(".data." #name) __hidden __aligned(8)

private(A) struct bpf_rb_root root __contains(tree_node, node);
private(A) struct bpf_spin_lock lock;

private(B) struct bpf_rb_root root_r __contains(tree_node_ref, node);
private(B) struct bpf_spin_lock lock_r;

static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
{
        struct tree_node *node_a, *node_b;

        node_a = container_of(a, struct tree_node, node);
        node_b = container_of(b, struct tree_node, node);

        return node_a->key < node_b->key;
}

SEC("syscall")
__retval(0)
long rbtree_search_kptr(void *ctx)
{
        struct tree_node *tnode;
        struct bpf_rb_node *rb_n;
        struct node_data __kptr * node_data;
        int lookup_key  = NR_NODES / 2;
        int lookup_data = NR_NODES / 2;
        int i, data, ret = 0;

        for (i = 0; i < NR_NODES && can_loop; i++) {
                tnode = bpf_obj_new(typeof(*tnode));
                if (!tnode)
                        return __LINE__;

                node_data = bpf_obj_new(typeof(*node_data));
                if (!node_data) {
                        bpf_obj_drop(tnode);
                        return __LINE__;
                }

                tnode->key = i;
                node_data->data = i;

                node_data = bpf_kptr_xchg(&tnode->node_data, node_data);
                if (node_data)
                        bpf_obj_drop(node_data);

                bpf_spin_lock(&lock);
                bpf_rbtree_add(&root, &tnode->node, less);
                bpf_spin_unlock(&lock);
        }

        bpf_spin_lock(&lock);
        rb_n = bpf_rbtree_root(&root);
        while (rb_n && can_loop) {
                tnode = container_of(rb_n, struct tree_node, node);
                node_data = bpf_kptr_xchg(&tnode->node_data, NULL);
                if (!node_data) {
                        ret = __LINE__;
                        goto fail;
                }

                data = node_data->data;
                node_data = bpf_kptr_xchg(&tnode->node_data, node_data);
                if (node_data) {
                        bpf_spin_unlock(&lock);
                        bpf_obj_drop(node_data);
                        return __LINE__;
                }

                if (lookup_key == tnode->key) {
                        if (data == lookup_data)
                                break;

                        ret = __LINE__;
                        goto fail;
                }

                if (lookup_key < tnode->key)
                        rb_n = bpf_rbtree_left(&root, rb_n);
                else
                        rb_n = bpf_rbtree_right(&root, rb_n);
        }
        bpf_spin_unlock(&lock);

        while (can_loop) {
                bpf_spin_lock(&lock);
                rb_n = bpf_rbtree_first(&root);
                if (!rb_n) {
                        bpf_spin_unlock(&lock);
                        return 0;
                }

                rb_n = bpf_rbtree_remove(&root, rb_n);
                if (!rb_n) {
                        ret = __LINE__;
                        goto fail;
                }
                bpf_spin_unlock(&lock);

                tnode = container_of(rb_n, struct tree_node, node);

                node_data = bpf_kptr_xchg(&tnode->node_data, NULL);
                if (node_data)
                        bpf_obj_drop(node_data);

                bpf_obj_drop(tnode);
        }

        return 0;
fail:
        bpf_spin_unlock(&lock);
        return ret;
}

static bool less_r(struct bpf_rb_node *a, const struct bpf_rb_node *b)
{
        struct tree_node_ref *node_a, *node_b;

        node_a = container_of(a, struct tree_node_ref, node);
        node_b = container_of(b, struct tree_node_ref, node);

        return node_a->key < node_b->key;
}

SEC("syscall")
__retval(0)
long rbtree_search_kptr_ref(void *ctx)
{
        struct tree_node_ref *tnode_r, *tnode_m;
        struct bpf_rb_node *rb_n;
        struct node_data __kptr * node_data;
        int lookup_key  = NR_NODES / 2;
        int lookup_data = NR_NODES / 2;
        int i, data, ret = 0;

        for (i = 0; i < NR_NODES && can_loop; i++) {
                tnode_r = bpf_obj_new(typeof(*tnode_r));
                if (!tnode_r)
                        return __LINE__;

                node_data = bpf_obj_new(typeof(*node_data));
                if (!node_data) {
                        bpf_obj_drop(tnode_r);
                        return __LINE__;
                }

                tnode_r->key = i;
                node_data->data = i;

                node_data = bpf_kptr_xchg(&tnode_r->node_data, node_data);
                if (node_data)
                        bpf_obj_drop(node_data);

                /* Unused reference */
                tnode_m = bpf_refcount_acquire(tnode_r);
                if (!tnode_m)
                        return __LINE__;

                bpf_spin_lock(&lock_r);
                bpf_rbtree_add(&root_r, &tnode_r->node, less_r);
                bpf_spin_unlock(&lock_r);

                bpf_obj_drop(tnode_m);
        }

        bpf_spin_lock(&lock_r);
        rb_n = bpf_rbtree_root(&root_r);
        while (rb_n && can_loop) {
                tnode_r = container_of(rb_n, struct tree_node_ref, node);
                node_data = bpf_kptr_xchg(&tnode_r->node_data, NULL);
                if (!node_data) {
                        ret = __LINE__;
                        goto fail;
                }

                data = node_data->data;
                node_data = bpf_kptr_xchg(&tnode_r->node_data, node_data);
                if (node_data) {
                        bpf_spin_unlock(&lock_r);
                        bpf_obj_drop(node_data);
                        return __LINE__;
                }

                if (lookup_key == tnode_r->key) {
                        if (data == lookup_data)
                                break;

                        ret = __LINE__;
                        goto fail;
                }

                if (lookup_key < tnode_r->key)
                        rb_n = bpf_rbtree_left(&root_r, rb_n);
                else
                        rb_n = bpf_rbtree_right(&root_r, rb_n);
        }
        bpf_spin_unlock(&lock_r);

        while (can_loop) {
                bpf_spin_lock(&lock_r);
                rb_n = bpf_rbtree_first(&root_r);
                if (!rb_n) {
                        bpf_spin_unlock(&lock_r);
                        return 0;
                }

                rb_n = bpf_rbtree_remove(&root_r, rb_n);
                if (!rb_n) {
                        ret = __LINE__;
                        goto fail;
                }
                bpf_spin_unlock(&lock_r);

                tnode_r = container_of(rb_n, struct tree_node_ref, node);

                node_data = bpf_kptr_xchg(&tnode_r->node_data, NULL);
                if (node_data)
                        bpf_obj_drop(node_data);

                bpf_obj_drop(tnode_r);
        }

        return 0;
fail:
        bpf_spin_unlock(&lock_r);
        return ret;
}

SEC("syscall")
__failure __msg("R1 type=scalar expected=map_value, ptr_, ptr_")
long non_own_ref_kptr_xchg_no_lock(void *ctx)
{
        struct tree_node *tnode;
        struct bpf_rb_node *rb_n;
        struct node_data __kptr * node_data;
        int data;

        bpf_spin_lock(&lock);
        rb_n = bpf_rbtree_first(&root);
        if (!rb_n) {
                bpf_spin_unlock(&lock);
                return __LINE__;
        }
        bpf_spin_unlock(&lock);

        tnode = container_of(rb_n, struct tree_node, node);
        node_data = bpf_kptr_xchg(&tnode->node_data, NULL);
        if (!node_data)
                return __LINE__;

        data = node_data->data;
        if (data < 0)
                return __LINE__;

        node_data = bpf_kptr_xchg(&tnode->node_data, node_data);
        if (node_data)
                return __LINE__;

        return 0;
}

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