#include "test/jemalloc_test.h"

#include "jemalloc/internal/rtree.h"

#define INVALID_ARENA_IND ((1U << MALLOCX_ARENA_BITS) - 1)

rtree_node_alloc_t *rtree_node_alloc_orig;
rtree_node_dalloc_t *rtree_node_dalloc_orig;
rtree_leaf_alloc_t *rtree_leaf_alloc_orig;
rtree_leaf_dalloc_t *rtree_leaf_dalloc_orig;

/* Potentially too large to safely place on the stack. */
rtree_t test_rtree;

static rtree_node_elm_t *
rtree_node_alloc_intercept(tsdn_t *tsdn, rtree_t *rtree, size_t nelms) {
	rtree_node_elm_t *node;

	if (rtree != &test_rtree) {
		return rtree_node_alloc_orig(tsdn, rtree, nelms);
	}

	malloc_mutex_unlock(tsdn, &rtree->init_lock);
	node = (rtree_node_elm_t *)calloc(nelms, sizeof(rtree_node_elm_t));
	assert_ptr_not_null(node, "Unexpected calloc() failure");
	malloc_mutex_lock(tsdn, &rtree->init_lock);

	return node;
}

static void
rtree_node_dalloc_intercept(tsdn_t *tsdn, rtree_t *rtree,
    rtree_node_elm_t *node) {
	if (rtree != &test_rtree) {
		rtree_node_dalloc_orig(tsdn, rtree, node);
		return;
	}

	free(node);
}

static rtree_leaf_elm_t *
rtree_leaf_alloc_intercept(tsdn_t *tsdn, rtree_t *rtree, size_t nelms) {
	rtree_leaf_elm_t *leaf;

	if (rtree != &test_rtree) {
		return rtree_leaf_alloc_orig(tsdn, rtree, nelms);
	}

	malloc_mutex_unlock(tsdn, &rtree->init_lock);
	leaf = (rtree_leaf_elm_t *)calloc(nelms, sizeof(rtree_leaf_elm_t));
	assert_ptr_not_null(leaf, "Unexpected calloc() failure");
	malloc_mutex_lock(tsdn, &rtree->init_lock);

	return leaf;
}

static void
rtree_leaf_dalloc_intercept(tsdn_t *tsdn, rtree_t *rtree,
    rtree_leaf_elm_t *leaf) {
	if (rtree != &test_rtree) {
		rtree_leaf_dalloc_orig(tsdn, rtree, leaf);
		return;
	}

	free(leaf);
}

TEST_BEGIN(test_rtree_read_empty) {
	tsdn_t *tsdn;

	tsdn = tsdn_fetch();

	rtree_t *rtree = &test_rtree;
	rtree_ctx_t rtree_ctx;
	rtree_ctx_data_init(&rtree_ctx);
	assert_false(rtree_new(rtree, false), "Unexpected rtree_new() failure");
	assert_ptr_null(rtree_edata_read(tsdn, rtree, &rtree_ctx, PAGE,
	    false), "rtree_edata_read() should return NULL for empty tree");
	rtree_delete(tsdn, rtree);
}
TEST_END

#undef NTHREADS
#undef NITERS
#undef SEED

TEST_BEGIN(test_rtree_extrema) {
	edata_t edata_a, edata_b;
	edata_init(&edata_a, INVALID_ARENA_IND, NULL, SC_LARGE_MINCLASS,
	    false, sz_size2index(SC_LARGE_MINCLASS), 0,
	    extent_state_active, false, false, true, EXTENT_NOT_HEAD);
	edata_init(&edata_b, INVALID_ARENA_IND, NULL, 0, false, SC_NSIZES, 0,
	    extent_state_active, false, false, true, EXTENT_NOT_HEAD);

	tsdn_t *tsdn = tsdn_fetch();

	rtree_t *rtree = &test_rtree;
	rtree_ctx_t rtree_ctx;
	rtree_ctx_data_init(&rtree_ctx);
	assert_false(rtree_new(rtree, false), "Unexpected rtree_new() failure");

	assert_false(rtree_write(tsdn, rtree, &rtree_ctx, PAGE, &edata_a,
	    edata_szind_get(&edata_a), edata_slab_get(&edata_a)),
	    "Unexpected rtree_write() failure");
	rtree_szind_slab_update(tsdn, rtree, &rtree_ctx, PAGE,
	    edata_szind_get(&edata_a), edata_slab_get(&edata_a));
	assert_ptr_eq(rtree_edata_read(tsdn, rtree, &rtree_ctx, PAGE, true),
	    &edata_a,
	    "rtree_edata_read() should return previously set value");

	assert_false(rtree_write(tsdn, rtree, &rtree_ctx, ~((uintptr_t)0),
	    &edata_b, edata_szind_get_maybe_invalid(&edata_b),
	    edata_slab_get(&edata_b)), "Unexpected rtree_write() failure");
	assert_ptr_eq(rtree_edata_read(tsdn, rtree, &rtree_ctx,
	    ~((uintptr_t)0), true), &edata_b,
	    "rtree_edata_read() should return previously set value");

	rtree_delete(tsdn, rtree);
}
TEST_END

TEST_BEGIN(test_rtree_bits) {
	tsdn_t *tsdn = tsdn_fetch();

	uintptr_t keys[] = {PAGE, PAGE + 1,
	    PAGE + (((uintptr_t)1) << LG_PAGE) - 1};

	edata_t edata;
	edata_init(&edata, INVALID_ARENA_IND, NULL, 0, false, SC_NSIZES, 0,
	    extent_state_active, false, false, true, EXTENT_NOT_HEAD);

	rtree_t *rtree = &test_rtree;
	rtree_ctx_t rtree_ctx;
	rtree_ctx_data_init(&rtree_ctx);
	assert_false(rtree_new(rtree, false), "Unexpected rtree_new() failure");

	for (unsigned i = 0; i < sizeof(keys)/sizeof(uintptr_t); i++) {
		assert_false(rtree_write(tsdn, rtree, &rtree_ctx, keys[i],
		    &edata, SC_NSIZES, false),
		    "Unexpected rtree_write() failure");
		for (unsigned j = 0; j < sizeof(keys)/sizeof(uintptr_t); j++) {
			assert_ptr_eq(rtree_edata_read(tsdn, rtree, &rtree_ctx,
			    keys[j], true), &edata,
			    "rtree_edata_read() should return previously set "
			    "value and ignore insignificant key bits; i=%u, "
			    "j=%u, set key=%#"FMTxPTR", get key=%#"FMTxPTR, i,
			    j, keys[i], keys[j]);
		}
		assert_ptr_null(rtree_edata_read(tsdn, rtree, &rtree_ctx,
		    (((uintptr_t)2) << LG_PAGE), false),
		    "Only leftmost rtree leaf should be set; i=%u", i);
		rtree_clear(tsdn, rtree, &rtree_ctx, keys[i]);
	}

	rtree_delete(tsdn, rtree);
}
TEST_END

TEST_BEGIN(test_rtree_random) {
#define NSET 16
#define SEED 42
	sfmt_t *sfmt = init_gen_rand(SEED);
	tsdn_t *tsdn = tsdn_fetch();
	uintptr_t keys[NSET];
	rtree_t *rtree = &test_rtree;
	rtree_ctx_t rtree_ctx;
	rtree_ctx_data_init(&rtree_ctx);

	edata_t edata;
	edata_init(&edata, INVALID_ARENA_IND, NULL, 0, false, SC_NSIZES, 0,
	    extent_state_active, false, false, true, EXTENT_NOT_HEAD);

	assert_false(rtree_new(rtree, false), "Unexpected rtree_new() failure");

	for (unsigned i = 0; i < NSET; i++) {
		keys[i] = (uintptr_t)gen_rand64(sfmt);
		rtree_leaf_elm_t *elm = rtree_leaf_elm_lookup(tsdn, rtree,
		    &rtree_ctx, keys[i], false, true);
		assert_ptr_not_null(elm,
		    "Unexpected rtree_leaf_elm_lookup() failure");
		rtree_leaf_elm_write(tsdn, rtree, elm, &edata, SC_NSIZES,
		    false);
		assert_ptr_eq(rtree_edata_read(tsdn, rtree, &rtree_ctx,
		    keys[i], true), &edata,
		    "rtree_edata_read() should return previously set value");
	}
	for (unsigned i = 0; i < NSET; i++) {
		assert_ptr_eq(rtree_edata_read(tsdn, rtree, &rtree_ctx,
		    keys[i], true), &edata,
		    "rtree_edata_read() should return previously set value, "
		    "i=%u", i);
	}

	for (unsigned i = 0; i < NSET; i++) {
		rtree_clear(tsdn, rtree, &rtree_ctx, keys[i]);
		assert_ptr_null(rtree_edata_read(tsdn, rtree, &rtree_ctx,
		    keys[i], true),
		   "rtree_edata_read() should return previously set value");
	}
	for (unsigned i = 0; i < NSET; i++) {
		assert_ptr_null(rtree_edata_read(tsdn, rtree, &rtree_ctx,
		    keys[i], true),
		    "rtree_edata_read() should return previously set value");
	}

	rtree_delete(tsdn, rtree);
	fini_gen_rand(sfmt);
#undef NSET
#undef SEED
}
TEST_END

int
main(void) {
	rtree_node_alloc_orig = rtree_node_alloc;
	rtree_node_alloc = rtree_node_alloc_intercept;
	rtree_node_dalloc_orig = rtree_node_dalloc;
	rtree_node_dalloc = rtree_node_dalloc_intercept;
	rtree_leaf_alloc_orig = rtree_leaf_alloc;
	rtree_leaf_alloc = rtree_leaf_alloc_intercept;
	rtree_leaf_dalloc_orig = rtree_leaf_dalloc;
	rtree_leaf_dalloc = rtree_leaf_dalloc_intercept;

	return test(
	    test_rtree_read_empty,
	    test_rtree_extrema,
	    test_rtree_bits,
	    test_rtree_random);
}