/*
 * DF-4522 semantic demo kld.
 *
 * Proves, on the RUNNING DragonFly kernel, the shim interaction behind
 * DF-4522:
 *
 *   (1) sys/dev/drm/include/linux/rbtree.h's rb_insert_color_cached()
 *       (rbtree.h:100-101) ignores its `leftmost' argument and never writes
 *       root->rb_leftmost; rb_erase_cached() (rbtree.h:102-103) neither.
 *       Therefore amdgpu_vm.c's DragonFly fini walk
 *
 *            while (vm->va.rb_leftmost) { ... }        (amdgpu_vm.c:3325)
 *
 *       executes ZERO iterations even when the tree is populated, and every
 *       remaining struct amdgpu_bo_va_mapping leaks (kfree at :3333 never
 *       runs).
 *
 *   (2) The proposed fix loop (iterate on rb_first(), which is a real RB_MIN
 *       walk on this port, rbtree.h:97/104) drains the identical tree.
 *
 * Everything below the "VERBATIM" markers is copied byte-for-byte from
 * sys/dev/drm/amd/amdgpu/amdgpu_vm.c (amdgpu-side interval-tree shim and the
 * amdgpu_vm_fini loop body).  The only adaptations, each marked: the mapping
 * struct is reduced to its tree-relevant fields, allocations use a private
 * MALLOC_DEFINE, and list_del(&mapping->list) is omitted because the reduced
 * struct has no list member (irrelevant to the tree semantics under test).
 *
 * Build:  see Makefile / build.sh in this directory.
 * Run:    kldload ./df4522demo.ko ; dmesg | tail ; kldunload df4522demo
 * Safe:   touches no hardware, no drm state; allocates 128 small nodes and
 *         frees all of them before returning.
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/malloc.h>

/* the REAL DragonFly drm shim header under test: */
#include <linux/rbtree.h>

/* Same generation line the drm core uses (sys/dev/drm/linux_compat.c:40).
 * drm.ko is not loaded on the test guest (no GPU), so no duplicate-symbol
 * interaction. */
RB_GENERATE(linux_root, rb_node, __entry, panic_cmp);

/* verbatim body: sys/dev/drm/linux_compat.c:33-37 */
int
panic_cmp(struct rb_node *one, struct rb_node *two)
{
	panic("no cmp");
}

/*
 * Minimal stand-in for struct amdgpu_bo_va_mapping: only the fields the
 * interval-tree shim touches (see amdgpu/amdgpu_vm.h:124-134).
 */
struct amdgpu_bo_va_mapping {
	struct	rb_node rb;		/* -> __subtree_last, rb entry */
	uint64_t start;			/* first VPAGE in the mapping */
	uint64_t last;			/* last VPAGE in the mapping */
};

/* verbatim: amdgpu_vm.c:62-63 */
#define START(node) ((node)->start)
#define LAST(node) ((node)->last)

/* ---- VERBATIM amdgpu_vm.c:100-105 ---- */
static void
amdgpu_vm_it_remove(struct amdgpu_bo_va_mapping *node,
    struct rb_root_cached *root)
{
	rb_erase_cached(&node->rb, root);
}

/* ---- VERBATIM amdgpu_vm.c:107-127 (DragonFly interval-tree insert) ---- */
static void
amdgpu_vm_it_insert(struct amdgpu_bo_va_mapping *node,
    struct rb_root_cached *root)
{
	struct rb_node **iter = &root->rb_root.rb_node;
	struct rb_node *parent = NULL;
	struct amdgpu_bo_va_mapping *iter_node;

	while (*iter) {
		parent = *iter;
		iter_node = rb_entry(*iter, struct amdgpu_bo_va_mapping, rb);

		if (node->start < iter_node->start)
			iter = &(*iter)->rb_left;
		else
			iter = &(*iter)->rb_right;
	}

	rb_link_node(&node->rb, parent, iter);
	rb_insert_color_cached(&node->rb, root, false);
}

MALLOC_DEFINE(M_DF4522, "df4522demo", "DF-4522 rb_leftmost demo");

#define NMAP	64

static void
populate(struct rb_root_cached *va)
{
	struct amdgpu_bo_va_mapping *mapping;
	int i;

	for (i = 0; i < NMAP; i++) {
		/* adapted: amdgpu_vm.c:2526 uses kmalloc(M_DRM, GFP_KERNEL) */
		mapping = kmalloc(sizeof(*mapping), M_DF4522, M_WAITOK | M_ZERO);
		mapping->start = i * 16;
		mapping->last = i * 16 + 15;
		/* amdgpu_vm_bo_insert_map -> amdgpu_vm_it_insert, amdgpu_vm.c:2462 */
		amdgpu_vm_it_insert(mapping, va);
	}
}

static void
df4522_demo_run(void)
{
	struct amdgpu_bo_va_mapping *mapping;
	struct rb_node *rn;
	struct rb_root_cached va = LINUX_RB_ROOT_CACHED; /* amdgpu_vm.c:3055 */
	struct rb_root_cached va2 = LINUX_RB_ROOT_CACHED;
	int bugged_iters = 0, bugged_freed = 0;
	int fixed_iters = 0, fixed_freed = 0;

	populate(&va);
	populate(&va2);

	kprintf("DF4522: inserted %d mappings per tree; tree_nonempty=%d "
	    "rb_leftmost=%p rb_first=%p\n",
	    NMAP, !RB_EMPTY_ROOT(&va.rb_root), va.rb_leftmost,
	    rb_first(&va.rb_root));

	/* ---- VERBATIM amdgpu_vm.c:3325-3334 (amdgpu_vm_fini walk) ---- */
	while (va.rb_leftmost) {
		mapping = container_of((void *)va.rb_leftmost,
		    struct amdgpu_bo_va_mapping, rb);
		/* list_del(&mapping->list); -- omitted, see file comment */
		amdgpu_vm_it_remove(mapping, &va);
		kfree(mapping, M_DF4522);	/* adapted M_DRM -> M_DF4522 */
		bugged_freed++;
		bugged_iters++;
	}

	kprintf("DF4522: BUGGED  fini loop: iterations=%d freed=%d LEAKED=%d "
	    "(tree still nonempty=%d)\n",
	    bugged_iters, bugged_freed, NMAP - bugged_freed,
	    !RB_EMPTY_ROOT(&va.rb_root));

	/* ---- the fixed loop from fix.diff (rb_first = real RB_MIN walk) ---- */
	while (!RB_EMPTY_ROOT(&va2.rb_root)) {
		rn = rb_first(&va2.rb_root);
		mapping = rb_entry(rn, struct amdgpu_bo_va_mapping, rb);
		amdgpu_vm_it_remove(mapping, &va2);
		kfree(mapping, M_DF4522);
		fixed_freed++;
		fixed_iters++;
	}

	kprintf("DF4522: FIXED  fini loop: iterations=%d freed=%d LEAKED=%d "
	    "(tree empty=%d)\n",
	    fixed_iters, fixed_freed, NMAP - fixed_freed,
	    RB_EMPTY_ROOT(&va2.rb_root));

	/* leave zero garbage behind: drain the leaked tree with the fixed loop */
	while (!RB_EMPTY_ROOT(&va.rb_root)) {
		rn = rb_first(&va.rb_root);
		mapping = rb_entry(rn, struct amdgpu_bo_va_mapping, rb);
		amdgpu_vm_it_remove(mapping, &va);
		kfree(mapping, M_DF4522);
	}

	kprintf("DF4522: VERDICT buggy_loop_leaked=%d fixed_loop_leaked=%d\n",
	    NMAP - bugged_freed, NMAP - fixed_freed);
}

static int
df4522_modevent(module_t mod, int type, void *data)
{
	switch (type) {
	case MOD_LOAD:
		df4522_demo_run();
		return 0;
	case MOD_UNLOAD:
		return 0;
	default:
		return 0;
	}
}

static moduledata_t df4522_mod = {
	"df4522demo",
	df4522_modevent,
	NULL
};

DECLARE_MODULE(df4522demo, df4522_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);