/******************************************************************************/
/* SlunkCrypt, by LoRd_MuldeR <MuldeR2@GMX.de>                                */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

#ifndef SLUNKBUILD_NOTHREADS

#ifdef _WIN32
#  define _CRT_SECURE_NO_WARNINGS 1
#else
#  define _GNU_SOURCE 1
#endif

/* Internal */
#include "thread.h"
#include "compiler.h"

/* CRT */
#include <stdlib.h>
#include <errno.h>
#include <string.h>

/* PThread */
#if defined(_MSC_VER) && !defined(_DLL)
#  define PTW32_STATIC_LIB 1
#endif
#include <pthread.h>
#include <sched.h>

/* Platform */
#if defined(__linux__) || defined(__PTW32_VERSION) || defined(__CYGWIN__)
#  define HAVE_SCHED_GETAFFINITY 1
#elif defined(_WIN32)
#  define WIN32_LEAN_AND_MEAN 1
#  define HAVE_GETPROCESSAFFINITYMASK 1
#  include <Windows.h>
#else
#  include <unistd.h> /* fall back to sysconf() function */
#endif

/* States */
#define TSTATE_IDLE 0U
#define TSTATE_WORK 1U
#define TSTATE_EXIT 2U

// ==========================================================================
// Data types
// ==========================================================================

typedef struct
{
	thrdpl_worker_t worker;
	uint8_t *buffer;
	size_t thread_count, generation, remain, length;
	int stop_flag;
	pthread_mutex_t mutex;
	pthread_cond_t cond_0, cond_1;
}
thrdpl_shared_t;

typedef struct
{
	thrdpl_shared_t *shared;
	void *context;
	pthread_t thread;
}
thrdpl_thread_t;

struct thrdpl_data_t
{
	thrdpl_shared_t shared;
	thrdpl_thread_t thread_data[MAX_THREADS];
};

// ==========================================================================
// Utilities
// ==========================================================================

static INLINE size_t BOUND(const size_t min, const size_t value, const size_t max)
{
	return (value < min) ? min : ((value > max) ? max : value);
}

#define PTHRD_MUTEX_ENTER(X) do \
{ \
	int _retval; \
	while ((_retval = pthread_mutex_trylock((X))) != 0) \
	{ \
		if (_retval != EBUSY) { abort(); } \
	} \
} \
while(0)

#define PTHRD_MUTEX_LEAVE(X) do \
{ \
	if (pthread_mutex_unlock((X)) != 0) \
	{ \
		abort(); \
	} \
} \
while(0)

#define PTHRD_COND_SIGNAL(X) do \
{ \
	if (pthread_cond_signal((X)) != 0) \
	{ \
		abort(); \
	} \
} \
while(0)

#define PTHRD_COND_BRDCST(X) do \
{ \
	if (pthread_cond_broadcast((X)) != 0) \
	{ \
		abort(); \
	} \
} \
while(0)

#define PTHRD_COND_WAIT(X,Y) do \
{ \
	if (pthread_cond_wait((X), (Y)) != 0) \
	{ \
		abort(); \
	} \
} \
while(0)

#define CHECK_IF_STOPPED() do \
{ \
	if (shared->stop_flag) \
	{ \
		PTHRD_MUTEX_LEAVE(&shared->mutex); \
		return NULL; \
	} \
} \
while(0)

// ==========================================================================
// System info
// ==========================================================================

static size_t detect_available_cpu_count(void)
{
	long num_processors = 0L;
#if defined(HAVE_SCHED_GETAFFINITY)
	cpu_set_t cpu_mask;
	CPU_ZERO(&cpu_mask);
	if (sched_getaffinity(0, sizeof(cpu_set_t), &cpu_mask) == 0)
	{
		num_processors = CPU_COUNT(&cpu_mask);
	}
#elif defined(HAVE_GETPROCESSAFFINITYMASK)
	DWORD_PTR proc_mask, sys_mask;
	if (GetProcessAffinityMask(GetCurrentProcess(), &proc_mask, &sys_mask))
	{
		for (; proc_mask != 0U; proc_mask &= proc_mask - 1U)
		{
			++num_processors;
		}
	}
#else
	num_processors = sysconf(_SC_NPROCESSORS_ONLN);
#endif
	return (num_processors > 0) ? ((size_t)num_processors) : 1U;
}

// ==========================================================================
// Thread main
// ==========================================================================

static void *worker_thread_main(void *const arg)
{
	thrdpl_thread_t *const data = (thrdpl_thread_t*) arg;
	thrdpl_shared_t *const shared = (thrdpl_shared_t*) data->shared;

	size_t previous = 0U;

	PTHRD_MUTEX_ENTER(&shared->mutex);
	CHECK_IF_STOPPED();

	for (;;)
	{
		while (shared->generation == previous)
		{
			PTHRD_COND_WAIT(&shared->cond_0, &shared->mutex);
			CHECK_IF_STOPPED();
		}

		previous = shared->generation;
		PTHRD_MUTEX_LEAVE(&shared->mutex);

		shared->worker(shared->thread_count, data->context, shared->buffer, shared->length);

		PTHRD_MUTEX_ENTER(&shared->mutex);
		CHECK_IF_STOPPED();

		if (!(--shared->remain))
		{
			PTHRD_COND_SIGNAL(&shared->cond_1);
		}
	}
}

// ==========================================================================
// Thread pool API
// ==========================================================================

thrdpl_t *slunkcrypt_thrdpl_create(const size_t count, const thrdpl_worker_t worker)
{
	size_t i;
	thrdpl_t *thrdpl = NULL;

	const size_t cpu_count = BOUND(1U, (count > 0U) ? count : detect_available_cpu_count(), MAX_THREADS);
	if (cpu_count < 2U)
	{
		return NULL;
	}

	if (!(thrdpl = (thrdpl_t*)malloc(sizeof(thrdpl_t))))
	{
		return NULL;
	}

	memset(thrdpl, 0, sizeof(thrdpl_t));
	thrdpl->shared.worker = worker;

	if (pthread_mutex_init(&thrdpl->shared.mutex, NULL) != 0)
	{
		goto failure;
	}

	if (pthread_cond_init(&thrdpl->shared.cond_0, NULL) != 0)
	{
		pthread_mutex_destroy(&thrdpl->shared.mutex);
		goto failure;
	}

	if (pthread_cond_init(&thrdpl->shared.cond_1, NULL) != 0)
	{
		pthread_cond_destroy(&thrdpl->shared.cond_0);
		pthread_mutex_destroy(&thrdpl->shared.mutex);
		goto failure;
	}

	for (i = 0U; i < cpu_count; ++i)
	{
		thrdpl->thread_data[i].shared = &thrdpl->shared;
		if (pthread_create(&thrdpl->thread_data[i].thread, NULL, worker_thread_main, &thrdpl->thread_data[i]) != 0)
		{
			slunkcrypt_thrdpl_destroy(thrdpl);
			return NULL;
		}
		++thrdpl->shared.thread_count;
	}

	return thrdpl;

failure:
	free(thrdpl);
	return NULL;
}

size_t slunkcrypt_thrdpl_count(const thrdpl_t *const thrdpl)
{
	return thrdpl->shared.thread_count;
}

void slunkcrypt_thrdpl_init(thrdpl_t *const thrdpl, const size_t index, void *const context)
{
	thrdpl->thread_data[index].context = context;
}

void slunkcrypt_thrdpl_exec(thrdpl_t *const thrdpl, uint8_t *const buffer, const size_t length)
{
	PTHRD_MUTEX_ENTER(&thrdpl->shared.mutex);

	if (thrdpl->shared.stop_flag || (thrdpl->shared.remain != 0U))
	{
		abort(); /*this is not supposed to happen!*/
	}

	thrdpl->shared.buffer = buffer;
	thrdpl->shared.length = length;
	thrdpl->shared.remain = thrdpl->shared.thread_count;

	++thrdpl->shared.generation;
	PTHRD_COND_BRDCST(&thrdpl->shared.cond_0);

	while (thrdpl->shared.remain)
	{
		PTHRD_COND_WAIT(&thrdpl->shared.cond_1, &thrdpl->shared.mutex);
	}

	PTHRD_MUTEX_LEAVE(&thrdpl->shared.mutex);
}

void slunkcrypt_thrdpl_destroy(thrdpl_t *const thrdpl)
{
	size_t i;
	PTHRD_MUTEX_ENTER(&thrdpl->shared.mutex);

	if (!thrdpl->shared.stop_flag)
	{
		thrdpl->shared.stop_flag = 1;
		PTHRD_COND_BRDCST(&thrdpl->shared.cond_0);
	}

	PTHRD_MUTEX_LEAVE(&thrdpl->shared.mutex);

	for (i = 0U; i < thrdpl->shared.thread_count; ++i)
	{
		pthread_join(thrdpl->thread_data[i].thread, NULL);
	}

	pthread_cond_destroy(&thrdpl->shared.cond_0);
	pthread_cond_destroy(&thrdpl->shared.cond_1);
	pthread_mutex_destroy(&thrdpl->shared.mutex);

	free(thrdpl);
}

#endif /*SLUNKBUILD_NOTHREADS*/
