#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <sys/time.h>
#include <string.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#define MB(a)		((a) * 1024UL * 1024UL)

static struct timeval start_tv, stop_tv;

// Calculate the difference between two time values.
void tvsub(struct timeval *tdiff, struct timeval *t1, struct timeval *t0)
{
	tdiff->tv_sec = t1->tv_sec - t0->tv_sec;
	tdiff->tv_usec = t1->tv_usec - t0->tv_usec;
	if (tdiff->tv_usec < 0)
		tdiff->tv_sec--, tdiff->tv_usec += 1000000;
}

// Start timing now.
void start()
{
	(void) gettimeofday(&start_tv, (struct timezone *) 0);
}

// Stop timing and return real time in microseconds.
unsigned long long stop()
{
	struct timeval tdiff;

	(void) gettimeofday(&stop_tv, (struct timezone *) 0);
	tvsub(&tdiff, &stop_tv, &start_tv);
	return (tdiff.tv_sec * 1000000 + tdiff.tv_usec);
}

void test_write(unsigned long *p, size_t size)
{
	int i;
	unsigned long *wp, tmp;
	unsigned long long timeval;

	start();
	for (i=0, wp=p; i<(size/sizeof(wp)); i++)
		*wp++ = 1;
	timeval = stop();
	printf("Write: %10llu usec\n", timeval);
}

void test_read(unsigned long *p, size_t size)
{
	int i;
	unsigned long *wp, tmp;
	unsigned long long timeval;

	start();
	for (i=0, wp=p; i<(size/sizeof(wp)); i++)
		tmp = *wp++;
	timeval = stop();
	printf("Read : %10llu usec\n", timeval);
}

int main(int argc, char **argv)
{
	int fd, i, opt, ret;
	int oflags, mprot, mflags = 0;
	int is_read_only = 0, is_mlock = 0, is_mlockall = 0;
	int mlock_skip = 0, read_test = 0, write_test = 0;
	void *mptr = NULL;
	unsigned long *p;
	struct stat stat;
	size_t size, cpy_size;
	const char *file_name = NULL;

	while ((opt = getopt(argc, argv, "LRMSApsrw")) != -1) {
		switch (opt) {
		case 'L':
			file_name = "/mnt/pmem0/4Gfile";
			break;
		case 'R':
			printf("> mmap: read-only\n");
			is_read_only = 1;
			break;
		case 'M':
			printf("> mlock\n");
			is_mlock = 1;
			break;
		case 'S':
			printf("> mlock - skip first ite\n");
			mlock_skip = 1;
			break;
		case 'A':
			printf("> mlockall\n");
			is_mlockall = 1;
			break;
		case 'p':
			printf("> MAP_POPULATE\n");
			mflags |= MAP_POPULATE;
			break;
		case 's':
			printf("> MAP_SHARED\n");
			mflags |= MAP_SHARED;
			break;
		case 'r':
			printf("> read-test\n");
			read_test = 1;
			break;
		case 'w':
			printf("> write-test\n");
			write_test = 1;
			break;
		}
	}

	if (!file_name) {
		file_name = "/mnt/pmem1/32Kfile";
	}

	if (!(mflags & MAP_SHARED)) {
		printf("> MAP_PRIVATE\n");
		mflags |= MAP_PRIVATE;
	}

	if (is_read_only) {
		oflags = O_RDONLY;
		mprot = PROT_READ;
	} else {
		oflags = O_RDWR;
		mprot = PROT_READ|PROT_WRITE;
	}

	fd = open(file_name, oflags);
	if (fd == -1) {
		perror("open failed");
		exit(1);
	}

	ret = fstat(fd, &stat);
	if (ret < 0) {
		perror("fstat failed");
		exit(1);
	}
	size = stat.st_size;

	printf("> open %s size 0x%x flags 0x%x\n", file_name, size, oflags);

	ret = posix_memalign(&mptr, MB(2), size);
	if (ret ==0)
		free(mptr);

	printf("> mmap mprot 0x%x flags 0x%x\n", mprot, mflags);
	p = mmap(mptr, size, mprot, mflags, fd, 0x0);
	if (!p) {
		perror("mmap failed");
		exit(1);
	}
	if ((long unsigned)p & (MB(2)-1))
		printf("> mmap: NOT 2MB aligned: 0x%p\n", p);
	else
		printf("> mmap: 2MB aligned: 0x%p\n", p);

#if 0	/* SIZE LIMIT */
	if (size >= MB(2))
		cpy_size = MB(32);
	else
#endif
		cpy_size = size;

	for (i=0; i<3; i++) {

		if (is_mlock && !mlock_skip) {
			printf("> mlock 0x%p\n", p);
			ret = mlock(p, size);
			if (ret < 0) {
				perror("mlock failed");
				exit(1);
			}
		} else if (is_mlockall) {
			printf("> mlockall\n");
			ret = mlockall(MCL_CURRENT|MCL_FUTURE);
			if (ret < 0) {
				perror("mlockall failed");
				exit(1);
			}
		}

		printf("===== %d =====\n", i+1);
		if (write_test)
			test_write(p, cpy_size);
		if (read_test)
			test_read(p, cpy_size);

		if (is_mlock && !mlock_skip) {
			printf("> munlock 0x%p\n", p);
			ret = munlock(p, size);
			if (ret < 0) {
				perror("munlock failed");
				exit(1);
			}
		} else if (is_mlockall) {
			printf("> munlockall\n");
			ret = munlockall();
			if (ret < 0) {
				perror("munlockall failed");
				exit(1);
			}
		}

		/* skip, if requested, only the first iteration */
		mlock_skip = 0;
	}

	printf("> munmap 0x%p\n", p);
	munmap(p, size);
}