#include "test/jemalloc_test.h"

static int  data_cleanup_count;

void
data_cleanup(int *data) {
	if (data_cleanup_count == 0) {
		expect_x_eq(*data, MALLOC_TSD_TEST_DATA_INIT,
		    "Argument passed into cleanup function should match tsd "
		    "value");
	}
	++data_cleanup_count;

	/*
	 * Allocate during cleanup for two rounds, in order to assure that
	 * jemalloc's internal tsd reinitialization happens.
	 */
	bool reincarnate = false;
	switch (*data) {
	case MALLOC_TSD_TEST_DATA_INIT:
		*data = 1;
		reincarnate = true;
		break;
	case 1:
		*data = 2;
		reincarnate = true;
		break;
	case 2:
		return;
	default:
		not_reached();
	}

	if (reincarnate) {
		void *p = mallocx(1, 0);
		expect_ptr_not_null(p, "Unexpeced mallocx() failure");
		dallocx(p, 0);
	}
}

static void *
thd_start(void *arg) {
	int   d = (int)(uintptr_t)arg;
	void *p;

	/*
	 * Test free before tsd init -- the free fast path (which does not
	 * explicitly check for NULL) has to tolerate this case, and fall back
	 * to free_default.
	 */
	free(NULL);

	tsd_t *tsd = tsd_fetch();
	expect_x_eq(tsd_test_data_get(tsd), MALLOC_TSD_TEST_DATA_INIT,
	    "Initial tsd get should return initialization value");

	p = malloc(1);
	expect_ptr_not_null(p, "Unexpected malloc() failure");

	tsd_test_data_set(tsd, d);
	expect_x_eq(tsd_test_data_get(tsd), d,
	    "After tsd set, tsd get should return value that was set");

	d = 0;
	expect_x_eq(tsd_test_data_get(tsd), (int)(uintptr_t)arg,
	    "Resetting local data should have no effect on tsd");

	tsd_test_callback_set(tsd, &data_cleanup);

	free(p);
	return NULL;
}

TEST_BEGIN(test_tsd_main_thread) {
	thd_start((void *)(uintptr_t)0xa5f3e329);
}
TEST_END

TEST_BEGIN(test_tsd_sub_thread) {
	thd_t thd;

	data_cleanup_count = 0;
	thd_create(&thd, thd_start, (void *)MALLOC_TSD_TEST_DATA_INIT);
	thd_join(thd, NULL);
	/*
	 * We reincarnate twice in the data cleanup, so it should execute at
	 * least 3 times.
	 */
	expect_x_ge(data_cleanup_count, 3,
	    "Cleanup function should have executed multiple times.");
}
TEST_END

static void *
thd_start_reincarnated(void *arg) {
	tsd_t *tsd = tsd_fetch();
	assert(tsd);

	void *p = malloc(1);
	expect_ptr_not_null(p, "Unexpected malloc() failure");

	/* Manually trigger reincarnation. */
	expect_ptr_not_null(tsd_arena_get(tsd), "Should have tsd arena set.");
	tsd_cleanup((void *)tsd);
	expect_ptr_null(
	    *tsd_arenap_get_unsafe(tsd), "TSD arena should have been cleared.");
	expect_u_eq(tsd_state_get(tsd), tsd_state_purgatory,
	    "TSD state should be purgatory\n");

	free(p);
	expect_u_eq(tsd_state_get(tsd), tsd_state_reincarnated,
	    "TSD state should be reincarnated\n");
	p = mallocx(1, MALLOCX_TCACHE_NONE);
	expect_ptr_not_null(p, "Unexpected malloc() failure");
	expect_ptr_null(*tsd_arenap_get_unsafe(tsd),
	    "Should not have tsd arena set after reincarnation.");

	free(p);
	tsd_cleanup((void *)tsd);
	expect_ptr_null(*tsd_arenap_get_unsafe(tsd),
	    "TSD arena should have been cleared after 2nd cleanup.");

	return NULL;
}

TEST_BEGIN(test_tsd_reincarnation) {
	thd_t thd;
	thd_create(&thd, thd_start_reincarnated, NULL);
	thd_join(thd, NULL);
}
TEST_END

static void *
thd_start_dalloc_only(void *arg) {
	void **ptrs = (void **)arg;

	tsd_t *tsd = tsd_fetch_min();
	if (tsd_state_get(tsd) != tsd_state_minimal_initialized) {
		/* Allocation happened implicitly. */
		expect_u_eq(tsd_state_get(tsd), tsd_state_nominal,
		    "TSD state should be nominal");
		return NULL;
	}

	void *ptr;
	for (size_t i = 0; (ptr = ptrs[i]) != NULL; i++) {
		/* Offset by 1 because of the manual tsd_fetch_min above. */
		if (i + 1 < TSD_MIN_INIT_STATE_MAX_FETCHED) {
			expect_u_eq(tsd_state_get(tsd),
			    tsd_state_minimal_initialized,
			    "TSD should be minimal initialized");
		} else {
			/* State may be nominal or nominal_slow. */
			expect_true(tsd_nominal(tsd), "TSD should be nominal");
		}
		free(ptr);
	}

	return NULL;
}

static void
test_sub_thread_n_dalloc(size_t nptrs) {
	void **ptrs = (void **)malloc(sizeof(void *) * (nptrs + 1));
	for (size_t i = 0; i < nptrs; i++) {
		ptrs[i] = malloc(8);
	}
	ptrs[nptrs] = NULL;

	thd_t thd;
	thd_create(&thd, thd_start_dalloc_only, (void *)ptrs);
	thd_join(thd, NULL);
	free(ptrs);
}

TEST_BEGIN(test_tsd_sub_thread_dalloc_only) {
	test_sub_thread_n_dalloc(1);
	test_sub_thread_n_dalloc(16);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED - 2);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED - 1);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED + 1);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED + 2);
	test_sub_thread_n_dalloc(TSD_MIN_INIT_STATE_MAX_FETCHED * 2);
}
TEST_END

int
main(void) {
	/* Ensure tsd bootstrapped. */
	if (nallocx(1, 0) == 0) {
		malloc_printf("Initialization error");
		return test_status_fail;
	}

	return test_no_reentrancy(test_tsd_main_thread, test_tsd_sub_thread,
	    test_tsd_sub_thread_dalloc_only, test_tsd_reincarnation);
}