/*
 * Copyright (c) 2017 Todd C. Miller <millert@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/time.h>

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <setjmp.h>
#include <time.h>
#include <unistd.h>
#include <err.h>

/*
 * Test program based on Bentley & McIlroy's "Engineering a Sort Function".
 * Uses mergesort(3) to check the results.
 *
 * Additional "killer" input taken from:
 *  David R. Musser's "Introspective Sorting and Selection Algorithms"
 *  http://calmerthanyouare.org/2014/06/11/algorithmic-complexity-attacks-and-libc-qsort.html
 *  M. D. McIlroy's "A Killer Adversary for Quicksort"
 */

#define timespecsub(tsp, usp, vsp)					\
	do {								\
		(vsp)->tv_sec = (tsp)->tv_sec - (usp)->tv_sec;		\
		(vsp)->tv_nsec = (tsp)->tv_nsec - (usp)->tv_nsec;	\
		if ((vsp)->tv_nsec < 0) {				\
			(vsp)->tv_sec--;				\
			(vsp)->tv_nsec += 1000000000L;			\
		}							\
	} while (0)

extern int mergesort(void *, size_t, size_t,
    int (*)(const void *, const void *));

/*
 * TODO:
 *	test with unaligned elements (char[]?)
 */
struct test_distribution {
	const char *name;
	void (*fill)(void *x, int n, int m);
	int (*test)(struct test_distribution *d, int n, void *x, void *y, void *z);
	int (*cmp)(const void *v1, const void *v2);
	int (*cmp_checked)(const void *v1, const void *v2);
};

#define MINIMUM(a, b)	(((a) < (b)) ? (a) : (b))
#define MAXIMUM(a, b)	(((a) > (b)) ? (a) : (b))

static size_t compares;
static size_t max_compares;
static size_t abrt_compares;
static sigjmp_buf cmpjmp;
static bool dump_table, timing, verbose;

extern int antiqsort(int n, int *a, int *ptr);

static int
cmp_i(const void *v1, const void *v2)
{
	const int a = *(const int *)v1;
	const int b = *(const int *)v2;

	return a > b ? 1 : a < b ? -1 : 0;
}

static int
cmp_checked_i(const void *v1, const void *v2)
{
	const int a = *(const int *)v1;
	const int b = *(const int *)v2;

	compares++;
	if (compares > abrt_compares)
		siglongjmp(cmpjmp, 1);
	return a > b ? 1 : a < b ? -1 : 0;
}

static int
cmp_ll(const void *v1, const void *v2)
{
	const long long a = *(const long long *)v1;
	const long long b = *(const long long *)v2;

	return a > b ? 1 : a < b ? -1 : 0;
}

static int
cmp_checked_ll(const void *v1, const void *v2)
{
	const long long a = *(const long long *)v1;
	const long long b = *(const long long *)v2;

	compares++;
	if (compares > abrt_compares)
		siglongjmp(cmpjmp, 1);
	return a > b ? 1 : a < b ? -1 : 0;
}

static int
cmp_d(const void *v1, const void *v2)
{
	const double a = *(const double *)v1;
	const double b = *(const double *)v2;

	return a > b ? 1 : a < b ? -1 : 0;
}

static int
cmp_checked_d(const void *v1, const void *v2)
{
	const double a = *(const double *)v1;
	const double b = *(const double *)v2;

	compares++;
	if (compares > abrt_compares)
		siglongjmp(cmpjmp, 1);
	return a > b ? 1 : a < b ? -1 : 0;
}

static int
check_result(char *sub, size_t es, char *got, char *expected, struct test_distribution *d,
    int m, int n)
{
	int i;

	if (verbose || compares > max_compares) {
		if (sub != NULL) {
			if (m != 0) {
				warnx("%s (%s): m: %d, n: %d, %zu compares, "
				    "max %zu(%zu)", d->name, sub, m, n,
				    compares, max_compares, abrt_compares);
			} else {
				warnx("%s (%s): n: %d, %zu compares, "
				    "max %zu(%zu)", d->name, sub, n,
				    compares, max_compares, abrt_compares);
			}
		} else {
			if (m != 0) {
				warnx("%s: m: %d, n: %d, %zu compares, "
				    "max %zu(%zu)", d->name, m, n,
				    compares, max_compares, abrt_compares);
			} else {
				warnx("%s: n: %d, %zu compares, "
				    "max %zu(%zu)", d->name, n,
				    compares, max_compares, abrt_compares);
			}
		}
	}

	for (i = 0; i < n; i++) {
		if (memcmp(got, expected, es) != 0)
			break;
		got += es;
		expected += es;
	}
	if (i == n)
		return 0;

	if (sub != NULL) {
		warnx("%s (%s): failure at %d, m: %d, n: %d",
		    d->name, sub, i, m, n);
	} else {
		warnx("%s: failure at %d, m: %d, n: %d",
		    d->name, i, m, n);
	}
	return 1;
}

#define FILL_SAWTOOTH(x, n, m)	do {					\
	int i;								\
									\
	for (i = 0; i < n; i++)						\
		x[i] = i % m;						\
} while (0)

static void
fill_sawtooth_i(void *v, int n, int m)
{
	int *x = v;
	FILL_SAWTOOTH(x, n, m);
}

static void
fill_sawtooth_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_SAWTOOTH(x, n, m);
}

static void
fill_sawtooth_double(void *v, int n, int m)
{
	double *x = v;
	FILL_SAWTOOTH(x, n, m);
}

#define FILL_RANDOM(x, n, m)	do {					\
	int i;								\
									\
	for (i = 0; i < n; i++)						\
		x[i] = arc4random_uniform(m);				\
} while (0)


static void
fill_random_i(void *v, int n, int m)
{
	int *x = v;
	FILL_RANDOM(x, n, m);
}

static void
fill_random_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_RANDOM(x, n, m);
}

static void
fill_random_double(void *v, int n, int m)
{
	double *x = v;
	FILL_RANDOM(x, n, m);
}

#define FILL_STAGGER(x, n, m)	do {					\
	int i;								\
									\
	for (i = 0; i < n; i++)						\
		x[i] = (i * m + i) % n;					\
} while (0)


static void
fill_stagger_i(void *v, int n, int m)
{
	int *x = v;
	FILL_STAGGER(x, n, m);
}

static void
fill_stagger_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_STAGGER(x, n, m);
}

static void
fill_stagger_double(void *v, int n, int m)
{
	double *x = v;
	FILL_STAGGER(x, n, m);
}

#define FILL_PLATEAU(x, n, m)	do {					\
	int i;								\
									\
	for (i = 0; i < n; i++)						\
		x[i] = MINIMUM(i, m);					\
} while (0)

static void
fill_plateau_i(void *v, int n, int m)
{
	int *x = v;
	FILL_PLATEAU(x, n, m);
}

static void
fill_plateau_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_PLATEAU(x, n, m);
}

static void
fill_plateau_double(void *v, int n, int m)
{
	double *x = v;
	FILL_PLATEAU(x, n, m);
}

#define FILL_SHUFFLE(x, n, m)	do {					\
	int i, j, k;							\
									\
	for (i = 0, j = 0, k = 1; i < n; i++)				\
		x[i] = arc4random_uniform(m) ? (j += 2) : (k += 2);	\
} while (0)

static void
fill_shuffle_i(void *v, int n, int m)
{
	int *x = v;
	FILL_SHUFFLE(x, n, m);
}

static void
fill_shuffle_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_SHUFFLE(x, n, m);
}

static void
fill_shuffle_double(void *v, int n, int m)
{
	double *x = v;
	FILL_SHUFFLE(x, n, m);
}

#define FILL_BSD_KILLER(x, n, m)	do {				\
	int i, k;							\
									\
	/* 4.4BSD insertion sort optimization killer, Mats Linander */	\
	k = n / 2;							\
	for (i = 0; i < n; i++) {					\
		if (i < k)						\
			x[i] = k - i;					\
		else if (i > k)						\
			x[i] = n + k + 1 - i;				\
		else							\
			x[i] = k + 1;					\
	}								\
} while (0)

static void
fill_bsd_killer_i(void *v, int n, int m)
{
	int *x = v;
	FILL_BSD_KILLER(x, n, m);

}

static void
fill_bsd_killer_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_BSD_KILLER(x, n, m);

}

static void
fill_bsd_killer_double(void *v, int n, int m)
{
	double *x = v;
	FILL_BSD_KILLER(x, n, m);

}

#define FILL_MED3_KILLER(x, n, m)	do {				\
	int i, k;							\
									\
	/*								\
	 * Median of 3 killer, as seen in David R Musser's		\
	 * "Introspective Sorting and Selection Algorithms"		\
	 */								\
	if (n & 1) {							\
		/* odd size, store last at end and make even. */	\
		x[n - 1] = n;						\
		n--;							\
	}								\
	k = n / 2;							\
	for (i = 0; i < n; i++) {					\
		if (i & 1) {						\
			x[i] = k + x[i - 1];				\
		} else {						\
			x[i] = i + 1;					\
		}							\
	}								\
} while (0)

static void
fill_med3_killer_i(void *v, int n, int m)
{
	int *x = v;
	FILL_MED3_KILLER(x, n, m);
}

static void
fill_med3_killer_ll(void *v, int n, int m)
{
	long long *x = v;
	FILL_MED3_KILLER(x, n, m);
}

static void
fill_med3_killer_double(void *v, int n, int m)
{
	double *x = v;
	FILL_MED3_KILLER(x, n, m);
}

static void
print_timing(struct test_distribution *d, char *sub, int m, int n, double elapsed)
{
	if (sub != NULL) {
		if (m != 0) {
			warnx("%s (%s): m: %d, n: %d, %f seconds",
			    d->name, sub, m, n, elapsed);
		} else {
			warnx("%s (%s): n: %d, %f seconds",
			    d->name, sub, n, elapsed);
		}
	} else {
		if (m != 0) {
			warnx("%s: m: %d, n: %d, %f seconds",
			    d->name, m, n, elapsed);
		} else {
			warnx("%s: n: %d, %f seconds",
			    d->name, n, elapsed);
		}
	}
}

static int
do_test(struct test_distribution *d, char *sub, int m, int n, size_t es, void *y, void *z)
{
	int ret = 0;
	struct timespec before, after;

	compares = 0;
	if (sigsetjmp(cmpjmp, 1) != 0) {
		if (sub != NULL) {
			warnx("%s (%s): qsort aborted after %zu compares, m: %d, n: %d",
			    d->name, sub, compares, m, n);
		} else {
			warnx("%s: qsort aborted after %zu compares, m: %d, n: %d",
			    d->name, compares, m, n);
		}
		ret = 1;
	} else {
		if (timing)
			(void) clock_gettime(CLOCK_MONOTONIC, &before);
		qsort(y, n, es, d->cmp_checked);
		if (timing) {
			double elapsed;
			(void) clock_gettime(CLOCK_MONOTONIC, &after);
			timespecsub(&after, &before, &after);
			elapsed = after.tv_sec + after.tv_nsec / 1000000000.0;
			print_timing(d, sub, m, n, elapsed);
		}
		if (check_result(sub, es, y, z, d, m, n) != 0)
			ret = 1;
	}
	return ret;
}

#define TEST_PERTURBED(d, n, x, y, z)	do {				\
	int i, j, m;							\
	const int es = sizeof(x[0]);					\
									\
	for (m = 1; m < 2 * n; m *= 2) {				\
		/* Fill in x[n] modified by m */			\
		d->fill(x, n, m);					\
									\
		/* Test on copy of x[] */				\
		for (i = 0; i < n; i++)					\
			y[i] = z[i] = x[i];				\
		if (mergesort(z, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		if (do_test(d, "copy", m, n, es, y, z) != 0)		\
			ret = 1;					\
									\
		/* Test on reversed copy of x[] */			\
		for (i = 0, j = n - 1; i < n; i++, j--)			\
			y[i] = z[i] = x[j];				\
		if (mergesort(z, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		if (do_test(d, "reversed", m, n, es, y, z) != 0)	\
			ret = 1;					\
									\
		/* Test with front half of x[] reversed */		\
		for (i = 0, j = (n / 2) - 1; i < n / 2; i++, j--)	\
			y[i] = z[i] = x[j];				\
		for (; i < n; i++)					\
			y[i] = z[i] = x[i];				\
		if (mergesort(z, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		if (do_test(d, "front reversed", m, n, es, y, z) != 0)	\
			ret = 1;					\
									\
		/* Test with back half of x[] reversed */		\
		for (i = 0; i < n / 2; i++)				\
			y[i] = z[i] = x[i];				\
		for (j = n - 1; i < n; i++, j--)			\
			y[i] = z[i] = x[j];				\
		if (mergesort(z, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		if (do_test(d, "back reversed", m, n, es, y, z) != 0)	\
			ret = 1;					\
									\
		/* Test on sorted copy of x[] */			\
		if (mergesort(x, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		for (i = 0; i < n; i++)					\
			y[i] = x[i];					\
		if (do_test(d, "sorted", m, n, es, y, x) != 0)		\
			ret = 1;					\
									\
		/* Test with i%5 added to x[i] (dither) */		\
		for (i = 0, j = n - 1; i < n; i++, j--)			\
			y[i] = z[i] = x[j] + (i % 5);			\
		if (mergesort(z, n, es, d->cmp) != 0)			\
			err(1, NULL);					\
		if (do_test(d, "dithered", m, n, es, y, z) != 0)	\
			ret = 1;					\
	}								\
} while (0)

static int
test_perturbed_i(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	int *x = vx;
	int *y = vx;
	int *z = vz;
	int ret = 0;

	TEST_PERTURBED(d, n, x, y, z);
	return ret;
}

static int
test_perturbed_ll(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	long long *x = vx;
	long long *y = vx;
	long long *z = vz;
	int ret = 0;

	TEST_PERTURBED(d, n, x, y, z);
	return ret;
}

static int
test_perturbed_double(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	double *x = vx;
	double *y = vx;
	double *z = vz;
	int ret = 0;

	TEST_PERTURBED(d, n, x, y, z);
	return ret;
}

/*
 * Like TEST_PERTURBED but because the input is tailored to cause
 * quicksort to go quadratic we don't perturb the input.
 */
#define TEST_SIMPLE(d, n, x, y, z)	do {				\
	int i;								\
									\
	/* Fill in x[n] */						\
	d->fill(x, n, 0);						\
									\
	/* Make a copy of x[] */					\
	for (i = 0; i < n; i++)						\
		y[i] = z[i] = x[i];					\
	if (mergesort(z, n, sizeof(z[0]), d->cmp) != 0)			\
		err(1, NULL);						\
	if (do_test(d, NULL, 0, n, sizeof(x[0]), y, z) != 0)		\
		ret = 1;						\
} while (0)

static int
test_simple_i(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	int *x = vx;
	int *y = vx;
	int *z = vz;
	int ret = 0;

	TEST_SIMPLE(d, n, x, y, z);
	return ret;
}

static int
test_simple_ll(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	long long *x = vx;
	long long *y = vx;
	long long *z = vz;
	int ret = 0;

	TEST_SIMPLE(d, n, x, y, z);
	return ret;
}

static int
test_simple_double(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	double *x = vx;
	double *y = vx;
	double *z = vz;
	int ret = 0;

	TEST_SIMPLE(d, n, x, y, z);
	return ret;
}

/*
 * Use D. McIlroy's "Killer Adversary for Quicksort"
 * We don't compare the results in this case.
 */
static int
test_antiqsort(struct test_distribution *d, int n, void *vx, void *vy, void *vz)
{
	struct timespec before, after;
	int *x = vx;
	int *y = vx;
	int i, ret = 0;

	/*
	 * We expect antiqsort to generate > 1.5 * nlgn compares.
	 * If introspection is not used, it will be > 10 * nlgn compares.
	 */
	if (timing)
		(void) clock_gettime(CLOCK_MONOTONIC, &before);
	i = antiqsort(n, x, y);
	if (timing)
		(void) clock_gettime(CLOCK_MONOTONIC, &after);
	if (i > abrt_compares)
		ret = 1;
	if (dump_table) {
		printf("/* %d items, %d compares */\n", n, i);
		printf("static const int table_%d[] = {", n);
		for (i = 0; i < n - 1; i++) {
			if ((i % 12) == 0)
				printf("\n\t");
			printf("%4d, ", x[i]);
		}
		printf("%4d\n};\n", x[i]);
	} else {
		if (timing) {
			double elapsed;
			timespecsub(&after, &before, &after);
			elapsed = after.tv_sec + after.tv_nsec / 1000000000.0;
			print_timing(d, NULL, 0, n, elapsed);
		}
		if (verbose || ret != 0) {
			warnx("%s: n: %d, %d compares, max %zu(%zu)",
			    d->name, n, i, max_compares, abrt_compares);
		}
	}

	return ret;
}

static struct test_distribution distributions[] = {
	{ "sawtooth_i", fill_sawtooth_i, test_perturbed_i, cmp_i, cmp_checked_i },
	{ "sawtooth_ll", fill_sawtooth_ll, test_perturbed_ll, cmp_ll, cmp_checked_ll },
	{ "sawtooth_d", fill_sawtooth_double, test_perturbed_double, cmp_d, cmp_checked_d },
	{ "random_i", fill_random_i, test_perturbed_i, cmp_i, cmp_checked_i },
	{ "random_ll", fill_random_ll, test_perturbed_ll, cmp_ll, cmp_checked_ll },
	{ "random_d", fill_random_double, test_perturbed_double, cmp_d, cmp_checked_d },
	{ "stagger_i", fill_stagger_i, test_perturbed_i, cmp_i, cmp_checked_i },
	{ "stagger_ll", fill_stagger_ll, test_perturbed_ll, cmp_ll, cmp_checked_ll },
	{ "stagger_d", fill_stagger_double, test_perturbed_double, cmp_d, cmp_checked_d },
	{ "plateau_i", fill_plateau_i, test_perturbed_i, cmp_i, cmp_checked_i },
	{ "plateau_ll", fill_plateau_ll, test_perturbed_ll, cmp_ll, cmp_checked_ll },
	{ "plateau_d", fill_plateau_double, test_perturbed_double, cmp_d, cmp_checked_d },
	{ "shuffle_i", fill_shuffle_i, test_perturbed_i, cmp_i, cmp_checked_i },
	{ "shuffle_ll", fill_shuffle_ll, test_perturbed_ll, cmp_ll, cmp_checked_ll },
	{ "shuffle_d", fill_shuffle_double, test_perturbed_double, cmp_d, cmp_checked_d },
	{ "bsd_killer_i", fill_bsd_killer_i, test_simple_i, cmp_i, cmp_checked_i },
	{ "bsd_killer_ll", fill_bsd_killer_ll, test_simple_ll, cmp_ll, cmp_checked_ll },
	{ "bsd_killer_d", fill_bsd_killer_double, test_simple_double, cmp_d, cmp_checked_d },
	{ "med3_killer_i", fill_med3_killer_i, test_simple_i, cmp_i, cmp_checked_i },
	{ "med3_killer_ll", fill_med3_killer_ll, test_simple_ll, cmp_ll, cmp_checked_ll },
	{ "med3_killer_d", fill_med3_killer_double, test_simple_double, cmp_d, cmp_checked_d },
	{ "antiqsort", NULL, test_antiqsort, cmp_i, cmp_checked_i },
	{ NULL }
};

static int
run_tests(int n, const char *name)
{
	void *x, *y, *z;
	int i, nlgn = 0;
	int ret = 0;
	size_t es;
	struct test_distribution *d;

	/*
	 * We expect A*n*lg(n) compares where A is between 1 and 2.
	 * For A > 1.5, print a warning.
	 * For A > 10 abort the test since qsort has probably gone quadratic.
	 */
	for (i = n - 1; i > 0; i >>= 1)
	    nlgn++;
	nlgn *= n;
	max_compares = nlgn * 1.5;
	abrt_compares = nlgn * 10;

	/* Allocate enough space to store ints or doubles. */
	es = MAXIMUM(sizeof(int), sizeof(double));
	x = reallocarray(NULL, n, es);
	y = reallocarray(NULL, n, es);
	z = reallocarray(NULL, n, es);
	if (x == NULL || y == NULL || z == NULL)
		err(1, NULL);

	for (d = distributions; d->name != NULL; d++) {
		if (name != NULL && strncmp(name, d->name, strlen(name)) != 0)
			continue;
		if (d->test(d, n, x, y, z) != 0)
			ret = 1;
	}

	free(x);
	free(y);
	free(z);
	return ret;
}

void
usage(void)
{
        fprintf(stderr, "usage: qsort_test [-dvt] [-n test_name] [num ...]\n");
	exit(1);
}

int
main(int argc, char *argv[])
{
	char *nums[] = { "100", "1023", "1024", "1025", "4095", "4096", "4097" };
	struct test_distribution *d;
	int ch, n, ret = 0;
	char *name = NULL;

        while ((ch = getopt(argc, argv, "dn:tv")) != -1) {
                switch (ch) {
                case 'd':
                        dump_table = true;
                        break;
                case 'n':
			for (d = distributions; d->name != NULL; d++) {
				if (strncmp(optarg, d->name, strlen(optarg)) == 0)
					break;
			}
			if (d->name == NULL)
				errx(1, "unknown test %s", optarg);
                        name = optarg;
                        break;
                case 't':
                        timing = true;
                        break;
                case 'v':
                        verbose = true;
                        break;
                default:
                        usage();
                        break;
                }
        }
        argc -= optind;
        argv += optind;

	if (argc == 0) {
		argv = nums;
		argc = sizeof(nums) / sizeof(nums[0]);
	}

	while (argc > 0) {
		n = atoi(*argv);
		if (run_tests(n, name) != 0)
			ret = 1;
		argc--;
		argv++;
	}

	return ret;
}