#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdint.h>
#include <stdbool.h>
#include <sys/mman.h>

static size_t pagesize;
static int pagemap_fd;

static void clear_softdirty(void)
{
	int fd = open("/proc/self/clear_refs", O_WRONLY);
	const char *ctrl = "4";
	int ret;

	if (fd < 0) {
		fprintf(stderr, "open() failed\n");
		exit(1);
	}
	ret = write(fd, ctrl, strlen(ctrl));
	close(fd);
	if (ret != strlen(ctrl)) {
		fprintf(stderr, "write() failed\n");
		exit(1);
	}
}

static uint64_t pagemap_get_entry(int fd, char *start)
{
	const unsigned long pfn = (unsigned long)start / pagesize;
	uint64_t entry;
	int ret;

	ret = pread(fd, &entry, sizeof(entry), pfn * sizeof(entry));
	if (ret != sizeof(entry)) {
		fprintf(stderr, "pread() failed\n");
		exit(1);
	}

	return entry;
}

static bool pagemap_is_softdirty(int fd, char *start)
{
	uint64_t entry = pagemap_get_entry(fd, start);

	return entry & 0x0080000000000000ull;
}

void main(void)
{
	char *mem, *mem2;

	pagesize = getpagesize();
	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
	if (pagemap_fd < 0) {
		fprintf(stderr, "open() failed\n");
		exit(1);
	}

	/* Map 2 pages. */
	mem = mmap(0, 2 * pagesize, PROT_READ|PROT_WRITE,
		   MAP_PRIVATE|MAP_ANON, -1, 0);
	if (mem == MAP_FAILED) {
		fprintf(stderr, "mmap() failed\n");
		exit(1);
	}

	/* Populate both pages. */
	memset(mem, 1, 2 * pagesize);

	if (!pagemap_is_softdirty(pagemap_fd, mem))
		fprintf(stderr, "Page #1 should be softdirty\n");
	if (!pagemap_is_softdirty(pagemap_fd, mem + pagesize))
		fprintf(stderr, "Page #2 should be softdirty\n");

	/*
	 * Start softdirty tracking. Clear VM_SOFTDIRTY and clear the softdirty
	 * PTE bit.
	 */
	clear_softdirty();

	if (pagemap_is_softdirty(pagemap_fd, mem))
		fprintf(stderr, "Page #1 should not be softdirty\n");
	if (pagemap_is_softdirty(pagemap_fd, mem + pagesize))
		fprintf(stderr, "Page #2 should not be softdirty\n");

	/*
	 * Remap the second page. The VMA gets VM_SOFTDIRTY set. Both VMAs
	 * get merged such that the resulting VMA has VM_SOFTDIRTY set.
	 */
	mem2 = mmap(mem + pagesize, pagesize, PROT_READ|PROT_WRITE,
		    MAP_PRIVATE|MAP_ANON|MAP_FIXED, -1, 0);
	if (mem2 == MAP_FAILED) {
		fprintf(stderr, "mmap() failed\n");
		exit(1);
	}

	/* Protect + unprotect. */
	mprotect(mem, 2 * pagesize, PROT_READ);
	mprotect(mem, 2 * pagesize, PROT_READ|PROT_WRITE);

	/* Modify both pages. */
	memset(mem, 2, 2 * pagesize);

	if (!pagemap_is_softdirty(pagemap_fd, mem))
		fprintf(stderr, "Page #1 should be softdirty\n");
	if (!pagemap_is_softdirty(pagemap_fd, mem + pagesize))
		fprintf(stderr, "Page #2 should be softdirty\n");
}