#include "test/jemalloc_test.h"

TEST_BEGIN(test_pinned_accessors) {
	edata_t edata;
	memset(&edata, 0, sizeof(edata));

	edata_arena_ind_set(&edata, 42);
	edata_slab_set(&edata, true);
	edata_committed_set(&edata, true);

	/* Default: not pinned. */
	edata_hook_flags_init(&edata, 0);
	expect_false(edata_pinned_get(&edata), "pinned should be false");
	expect_u_eq(0, edata_alloc_flags_get(&edata),
	    "alloc_flags should be 0");
	expect_u_eq(42, edata_arena_ind_get(&edata),
	    "arena_ind corrupted");
	expect_true(edata_slab_get(&edata), "slab corrupted");

	/* Set pinned. */
	edata_hook_flags_init(&edata, EXTENT_ALLOC_FLAG_PINNED);
	expect_true(edata_pinned_get(&edata), "pinned should be true");
	expect_u_eq(EXTENT_ALLOC_FLAG_PINNED, edata_alloc_flags_get(&edata),
	    "alloc_flags round-trip failed");
	expect_u_eq(42, edata_arena_ind_get(&edata),
	    "arena_ind corrupted");

	/* Split inheritance: trail gets lead's flags. */
	edata_t trail;
	memset(&trail, 0, sizeof(edata_t));
	edata_hook_flags_init(&trail, edata_alloc_flags_get(&edata));
	expect_true(edata_pinned_get(&trail),
	    "trail should inherit pinned from lead");
}
TEST_END

TEST_BEGIN(test_dirty_accounting) {
	unsigned arena_ind;
	size_t sz = sizeof(arena_ind);

	expect_d_eq(0, mallctl("arenas.create", &arena_ind, &sz, NULL, 0),
	    "arena creation failed");

	void *ptrs[16];
	for (unsigned i = 0; i < 16; i++) {
		ptrs[i] = mallocx(PAGE, MALLOCX_ARENA(arena_ind));
		expect_ptr_not_null(ptrs[i], "alloc %u failed", i);
	}
	for (unsigned i = 0; i < 16; i++) {
		dallocx(ptrs[i], MALLOCX_ARENA(arena_ind));
	}

	/* Default hooks: alloc_flags=0, so ecache_pinned must be empty. */
	tsd_t *tsd = tsd_fetch();
	tsdn_t *tsdn = tsd_tsdn(tsd);
	arena_t *arena = arena_get(tsdn, arena_ind, false);
	expect_ptr_not_null(arena, "arena_get failed");

	pac_t *pac = &arena->pa_shard.pac;
	expect_zu_eq(0, ecache_npages_get(&pac->ecache_pinned),
	    "ecache_pinned should be empty with default hooks");
}
TEST_END

/*
 * Custom alloc hook that sets EXTENT_ALLOC_FLAG_PINNED.
 * Passthrough to default hooks via ehooks_default_extent_hooks.
 */
static void *
pinned_extent_alloc(extent_hooks_t *extent_hooks, void *new_addr,
    size_t size, size_t alignment, bool *zero, bool *commit,
    unsigned arena_ind) {
	void *ret = ehooks_default_extent_hooks.alloc(
	    (extent_hooks_t *)&ehooks_default_extent_hooks,
	    new_addr, size, alignment, zero, commit, arena_ind);
	if (ret == NULL) {
		return NULL;
	}
	return (void *)((uintptr_t)ret | EXTENT_ALLOC_FLAG_PINNED);
}

static extent_hooks_t pinned_hooks = {
	pinned_extent_alloc,
	NULL, /* dalloc — force retain */
	NULL, /* destroy */
	NULL, /* commit */
	NULL, /* decommit */
	NULL, /* purge_lazy */
	NULL, /* purge_forced */
	NULL, /* split */
	NULL  /* merge */
};

TEST_BEGIN(test_pinned_stats) {
	unsigned arena_ind;
	size_t sz = sizeof(arena_ind);
	extent_hooks_t *hooks_ptr = &pinned_hooks;

	/* Create arena with pinned hooks. */
	expect_d_eq(0, mallctl("arenas.create", &arena_ind, &sz,
	    &hooks_ptr, sizeof(hooks_ptr)),
	    "arena creation failed");

	/* Allocate and free to populate ecache_pinned. */
	void *p = mallocx(PAGE * 4, MALLOCX_ARENA(arena_ind)
	    | MALLOCX_TCACHE_NONE);
	expect_ptr_not_null(p, "alloc failed");
	dallocx(p, MALLOCX_TCACHE_NONE);

	/* Refresh stats. */
	uint64_t epoch = 1;
	sz = sizeof(epoch);
	expect_d_eq(0, mallctl("epoch", &epoch, &sz, &epoch, sizeof(epoch)),
	    "epoch failed");

	/* Read total pinned stat. */
	char buf[128];
	size_t pinned_total;
	sz = sizeof(pinned_total);
	snprintf(buf, sizeof(buf), "stats.arenas.%u.pinned", arena_ind);
	expect_d_eq(0, mallctl(buf, &pinned_total, &sz, NULL, 0),
	    "stats.arenas.<i>.pinned read failed");
	expect_zu_gt(pinned_total, 0,
	    "pinned total should be > 0 after free to pinned arena");

	/* Destroy the arena. */
	snprintf(buf, sizeof(buf), "arena.%u.destroy", arena_ind);
	expect_d_eq(0, mallctl(buf, NULL, NULL, NULL, 0),
	    "arena destroy failed");
}
TEST_END

TEST_BEGIN(test_pinned_hook_arena_destroy) {
	unsigned arena_ind;
	size_t sz = sizeof(arena_ind);
	extent_hooks_t *hooks_ptr = &pinned_hooks;

	/* Create arena with pinned hooks. */
	expect_d_eq(0, mallctl("arenas.create", &arena_ind, &sz,
	    &hooks_ptr, sizeof(hooks_ptr)),
	    "arena creation failed");

	/* Allocate, shrink, and free through the pinned arena. */
	void *ptrs[8];
	for (unsigned i = 0; i < 8; i++) {
		ptrs[i] = mallocx(PAGE * 4, MALLOCX_ARENA(arena_ind)
		    | MALLOCX_TCACHE_NONE);
		expect_ptr_not_null(ptrs[i], "alloc %u failed", i);
		/* Shrink it to test pac_shrink_impl. */
		ptrs[i] = rallocx(ptrs[i], PAGE * 2, MALLOCX_ARENA(arena_ind)
		    | MALLOCX_TCACHE_NONE);
		expect_ptr_not_null(ptrs[i], "shrink %u failed", i);
	}
	for (unsigned i = 0; i < 8; i++) {
		dallocx(ptrs[i], MALLOCX_TCACHE_NONE);
	}

	/* Destroy the arena — must not crash or assert. */
	char buf[64];
	snprintf(buf, sizeof(buf), "arena.%u.destroy", arena_ind);
	expect_d_eq(0, mallctl(buf, NULL, NULL, NULL, 0),
	    "arena destroy failed");
}
TEST_END

int
main(void) {
	return test_no_reentrancy(
	    test_pinned_accessors,
	    test_dirty_accounting,
	    test_pinned_stats,
	    test_pinned_hook_arena_destroy);
}