/*
 * unalign_check - check the CPU behaviour on different alignments
 * Copyright (C) 2021 Matteo Croce <mcroce@linux.microsoft.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <time.h>
#include <sys/mman.h>

#define ACT_READ	0
#define ACT_WRITE	1
#define ACT_XOR		2
#define ACT_COPY	3

#define READ(SIZE)							\
	case SIZE / 8: {						\
		volatile uint##SIZE##_t *buf2 = (uint##SIZE##_t *)buf;	\
		int i;							\
		for (i = 0; i < count; i++)				\
			(void)buf2[i];					\
		break;							\
	}

#define WRITE(SIZE)							\
	case SIZE / 8: {						\
		volatile uint##SIZE##_t *buf2 = (uint##SIZE##_t *)buf;	\
		int i;							\
		for (i = 0; i < count; i++)				\
			buf2[i] = (uint##SIZE##_t)0xaabbccdd11223344;	\
		break;							\
	}

#define XOR(SIZE)							\
	case SIZE / 8: {						\
		volatile uint##SIZE##_t *buf2 = (uint##SIZE##_t *)buf;	\
		int i;							\
		for (i = 0; i < count; i++)				\
			buf2[i] = ~buf2[i];				\
		break;							\
	}

#define COPY(SIZE)							\
	case SIZE / 8: {						\
		volatile uint##SIZE##_t *buf2 = (uint##SIZE##_t *)buf;	\
		int i;							\
		for (i = 0; i < count / 2; i++)				\
			buf2[i] = buf2[i + count / 2];			\
		for (i = count / 2; i < count; i++)			\
			buf2[i] = buf2[i - count / 2];			\
		break;							\
	}

static void do_read(void *buf, size_t count, int size)
{
	switch (size) {
		READ(8);
		READ(16);
		READ(32);
		READ(64);
	}
}

static void do_write(void *buf, size_t count, int size)
{
	switch (size) {
		WRITE(8);
		WRITE(16);
		WRITE(32);
		WRITE(64);
	}
}

static void do_xor(void *buf, size_t count, int size)
{
	switch (size) {
		XOR(8);
		XOR(16);
		XOR(32);
		XOR(64);
	}
}

static void do_copy(void *buf, size_t count, int size)
{
	switch (size) {
		COPY(8);
		COPY(16);
		COPY(32);
		COPY(64);
	}
}

static uint64_t time_sub(struct timespec *since, struct timespec *to)
{
	if (to->tv_sec == since->tv_sec)
		return to->tv_nsec - since->tv_nsec;

	return (to->tv_sec - since->tv_sec) * 1000000000 + to->tv_nsec - since->tv_nsec;
}

static void __attribute__ ((noreturn)) usage(char *argv0, int ret)
{
	fprintf(ret ? stderr : stdout,
		"usage: %s [-rwxc1234h] [-l length]  [-u unalignment]\n"
		"\n"
		"Options:\n"
		"	-r	read memory (default)\n"
		"	-w	write memory\n"
		"	-x	xor memory\n"
		"	-c	copy memory\n"
		"	-l SIZE	use SIZE Mb for the test (default 100)\n"
		"	-u BYTE	unalign buffer by BYTE bytes (default 0)\n"
		"	-1	read 1 byte at time\n"
		"	-2	read 2 bytes at time\n"
		"	-4	read 4 bytes at time (default)\n"
		"	-8	read 8 bytes at time\n"
		"	-h	this help\n",
		argv0);
	exit(ret);
}

static const char *actions[] = {
	"read",
	"write",
	"xor",
	"copy",
};

int main(int argc, char *argv[])
{
	struct timespec before, after;
	uint64_t elapsed;
	int action = ACT_READ;
	size_t len = 100 * 1024 * 1024;
	int shift = 0;
	int size = sizeof(long);
	char *buf;
	int c;

	while((c = getopt(argc, argv, "hrwxc1248l:u:")) != -1) {
		switch (c) {
		case 'r':
			action = ACT_READ;
			break;
		case 'w':
			action = ACT_WRITE;
			break;
		case 'x':
			action = ACT_XOR;
			break;
		case 'c':
			action = ACT_COPY;
			break;
		case 'l':
			len = atol(optarg) * 1024 * 1024;
			if (len <= 0) {
				fprintf(stderr, "Invalid size %s\n", optarg);
				return 1;
			}
			break;
		case 'u':
			shift = atoi(optarg);
			break;
		case '1':
		case '2':
		case '4':
		case '8':
			size = c - '0';
			break;
		case 'h':
		default:
			usage(argv[0], c != 'h');
		}
	}

	shift %= size;

	if (optind != argc)
		usage(argv[0], 1);

	buf = malloc(len);
	if (!buf) {
		perror("malloc");
		return 1;
	}

	if (mlock(buf, len)) {
		perror("mlock");
		return 1;
	}

	clock_gettime(CLOCK_MONOTONIC, &before);
	switch (action) {
	case ACT_READ:
		do_read(buf + shift, (len - shift) / size, size);
		break;
	case ACT_WRITE:
		do_write(buf + shift, (len - shift) / size, size);
		break;
	case ACT_XOR:
		do_xor(buf + shift, (len - shift) / size, size);
		break;
	case ACT_COPY:
		do_copy(buf + shift, (len - shift) / size, size);
		break;
	}
	clock_gettime(CLOCK_MONOTONIC, &after);

	elapsed = time_sub(&before, &after);

	printf(	"size:		%lu Mb\n"
		"%s size:	%d bit\n"
		"unalignment:	%d byte\n"
		"elapsed time:	%.2f sec\n"
		"throughput:	%.2f Mb/s\n",
		len / 1024 / 1024,
		actions[action], size * 8,
		shift,
		elapsed / 1E9,
		(len / 1024 / 1024) / (elapsed / 1E9));

	return 0;
}