// SPDX-License-Identifier: GPL-2.0
#include <fcntl.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "vm_util.h"
#include "thp_settings.h"

#define THP_SYSFS "/sys/kernel/mm/transparent_hugepage/"
#define MAX_SETTINGS_DEPTH 4
static struct thp_settings settings_stack[MAX_SETTINGS_DEPTH];
static int settings_index;
static struct thp_settings saved_settings;
static char dev_queue_read_ahead_path[PATH_MAX];

static const char * const thp_enabled_strings[] = {
	"never",
	"always",
	"inherit",
	"madvise",
	NULL
};

static const char * const thp_defrag_strings[] = {
	"always",
	"defer",
	"defer+madvise",
	"madvise",
	"never",
	NULL
};

static const char * const shmem_enabled_strings[] = {
	"never",
	"always",
	"within_size",
	"advise",
	"inherit",
	"deny",
	"force",
	NULL
};

int read_file(const char *path, char *buf, size_t buflen)
{
	int fd;
	ssize_t numread;

	fd = open(path, O_RDONLY);
	if (fd == -1)
		return 0;

	numread = read(fd, buf, buflen - 1);
	if (numread < 1) {
		close(fd);
		return 0;
	}

	buf[numread] = '\0';
	close(fd);

	return (unsigned int) numread;
}

unsigned long read_num(const char *path)
{
	char buf[21];

	if (read_file(path, buf, sizeof(buf)) < 0) {
		perror("read_file()");
		exit(EXIT_FAILURE);
	}

	return strtoul(buf, NULL, 10);
}

void write_num(const char *path, unsigned long num)
{
	char buf[21];

	sprintf(buf, "%ld", num);
	write_file(path, buf, strlen(buf) + 1);
}

int thp_read_string(const char *name, const char * const strings[])
{
	char path[PATH_MAX];
	char buf[256];
	char *c;
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}

	if (!read_file(path, buf, sizeof(buf))) {
		perror(path);
		exit(EXIT_FAILURE);
	}

	c = strchr(buf, '[');
	if (!c) {
		printf("%s: Parse failure\n", __func__);
		exit(EXIT_FAILURE);
	}

	c++;
	memmove(buf, c, sizeof(buf) - (c - buf));

	c = strchr(buf, ']');
	if (!c) {
		printf("%s: Parse failure\n", __func__);
		exit(EXIT_FAILURE);
	}
	*c = '\0';

	ret = 0;
	while (strings[ret]) {
		if (!strcmp(strings[ret], buf))
			return ret;
		ret++;
	}

	printf("Failed to parse %s\n", name);
	exit(EXIT_FAILURE);
}

void thp_write_string(const char *name, const char *val)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	write_file(path, val, strlen(val) + 1);
}

unsigned long thp_read_num(const char *name)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	return read_num(path);
}

void thp_write_num(const char *name, unsigned long num)
{
	char path[PATH_MAX];
	int ret;

	ret = snprintf(path, PATH_MAX, THP_SYSFS "%s", name);
	if (ret >= PATH_MAX) {
		printf("%s: Pathname is too long\n", __func__);
		exit(EXIT_FAILURE);
	}
	write_num(path, num);
}

void thp_read_settings(struct thp_settings *settings)
{
	unsigned long orders = thp_supported_orders();
	unsigned long shmem_orders = thp_shmem_supported_orders();
	char path[PATH_MAX];
	int i;

	*settings = (struct thp_settings) {
		.thp_enabled = thp_read_string("enabled", thp_enabled_strings),
		.thp_defrag = thp_read_string("defrag", thp_defrag_strings),
		.shmem_enabled =
			thp_read_string("shmem_enabled", shmem_enabled_strings),
		.use_zero_page = thp_read_num("use_zero_page"),
	};
	settings->khugepaged = (struct khugepaged_settings) {
		.defrag = thp_read_num("khugepaged/defrag"),
		.alloc_sleep_millisecs =
			thp_read_num("khugepaged/alloc_sleep_millisecs"),
		.scan_sleep_millisecs =
			thp_read_num("khugepaged/scan_sleep_millisecs"),
		.max_ptes_none = thp_read_num("khugepaged/max_ptes_none"),
		.max_ptes_swap = thp_read_num("khugepaged/max_ptes_swap"),
		.max_ptes_shared = thp_read_num("khugepaged/max_ptes_shared"),
		.pages_to_scan = thp_read_num("khugepaged/pages_to_scan"),
	};
	if (dev_queue_read_ahead_path[0])
		settings->read_ahead_kb = read_num(dev_queue_read_ahead_path);

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & orders)) {
			settings->hugepages[i].enabled = THP_NEVER;
			continue;
		}
		snprintf(path, PATH_MAX, "hugepages-%ukB/enabled",
			(getpagesize() >> 10) << i);
		settings->hugepages[i].enabled =
			thp_read_string(path, thp_enabled_strings);
	}

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & shmem_orders)) {
			settings->shmem_hugepages[i].enabled = SHMEM_NEVER;
			continue;
		}
		snprintf(path, PATH_MAX, "hugepages-%ukB/shmem_enabled",
			(getpagesize() >> 10) << i);
		settings->shmem_hugepages[i].enabled =
			thp_read_string(path, shmem_enabled_strings);
	}
}

void thp_write_settings(struct thp_settings *settings)
{
	struct khugepaged_settings *khugepaged = &settings->khugepaged;
	unsigned long orders = thp_supported_orders();
	unsigned long shmem_orders = thp_shmem_supported_orders();
	char path[PATH_MAX];
	int enabled;
	int i;

	thp_write_string("enabled", thp_enabled_strings[settings->thp_enabled]);
	thp_write_string("defrag", thp_defrag_strings[settings->thp_defrag]);
	thp_write_string("shmem_enabled",
			shmem_enabled_strings[settings->shmem_enabled]);
	thp_write_num("use_zero_page", settings->use_zero_page);

	thp_write_num("khugepaged/defrag", khugepaged->defrag);
	thp_write_num("khugepaged/alloc_sleep_millisecs",
			khugepaged->alloc_sleep_millisecs);
	thp_write_num("khugepaged/scan_sleep_millisecs",
			khugepaged->scan_sleep_millisecs);
	thp_write_num("khugepaged/max_ptes_none", khugepaged->max_ptes_none);
	thp_write_num("khugepaged/max_ptes_swap", khugepaged->max_ptes_swap);
	thp_write_num("khugepaged/max_ptes_shared", khugepaged->max_ptes_shared);
	thp_write_num("khugepaged/pages_to_scan", khugepaged->pages_to_scan);

	if (dev_queue_read_ahead_path[0])
		write_num(dev_queue_read_ahead_path, settings->read_ahead_kb);

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & orders))
			continue;
		snprintf(path, PATH_MAX, "hugepages-%ukB/enabled",
			(getpagesize() >> 10) << i);
		enabled = settings->hugepages[i].enabled;
		thp_write_string(path, thp_enabled_strings[enabled]);
	}

	for (i = 0; i < NR_ORDERS; i++) {
		if (!((1 << i) & shmem_orders))
			continue;
		snprintf(path, PATH_MAX, "hugepages-%ukB/shmem_enabled",
			(getpagesize() >> 10) << i);
		enabled = settings->shmem_hugepages[i].enabled;
		thp_write_string(path, shmem_enabled_strings[enabled]);
	}
}

struct thp_settings *thp_current_settings(void)
{
	if (!settings_index) {
		printf("Fail: No settings set");
		exit(EXIT_FAILURE);
	}
	return settings_stack + settings_index - 1;
}

void thp_push_settings(struct thp_settings *settings)
{
	if (settings_index >= MAX_SETTINGS_DEPTH) {
		printf("Fail: Settings stack exceeded");
		exit(EXIT_FAILURE);
	}
	settings_stack[settings_index++] = *settings;
	thp_write_settings(thp_current_settings());
}

void thp_pop_settings(void)
{
	if (settings_index <= 0) {
		printf("Fail: Settings stack empty");
		exit(EXIT_FAILURE);
	}
	--settings_index;
	thp_write_settings(thp_current_settings());
}

void thp_restore_settings(void)
{
	thp_write_settings(&saved_settings);
}

void thp_save_settings(void)
{
	thp_read_settings(&saved_settings);
}

void thp_set_read_ahead_path(char *path)
{
	if (!path) {
		dev_queue_read_ahead_path[0] = '\0';
		return;
	}

	strncpy(dev_queue_read_ahead_path, path,
		sizeof(dev_queue_read_ahead_path));
	dev_queue_read_ahead_path[sizeof(dev_queue_read_ahead_path) - 1] = '\0';
}

static unsigned long __thp_supported_orders(bool is_shmem)
{
	unsigned long orders = 0;
	char path[PATH_MAX];
	char buf[256];
	int ret, i;
	char anon_dir[] = "enabled";
	char shmem_dir[] = "shmem_enabled";

	for (i = 0; i < NR_ORDERS; i++) {
		ret = snprintf(path, PATH_MAX, THP_SYSFS "hugepages-%ukB/%s",
			       (getpagesize() >> 10) << i, is_shmem ? shmem_dir : anon_dir);
		if (ret >= PATH_MAX) {
			printf("%s: Pathname is too long\n", __func__);
			exit(EXIT_FAILURE);
		}

		ret = read_file(path, buf, sizeof(buf));
		if (ret)
			orders |= 1UL << i;
	}

	return orders;
}

unsigned long thp_supported_orders(void)
{
	return __thp_supported_orders(false);
}

unsigned long thp_shmem_supported_orders(void)
{
	return __thp_supported_orders(true);
}

bool thp_available(void)
{
	if (access(THP_SYSFS, F_OK) != 0)
		return false;
	return true;
}

bool thp_is_enabled(void)
{
	if (!thp_available())
		return false;

	int mode = thp_read_string("enabled", thp_enabled_strings);

	/* THP is considered enabled if it's either "always" or "madvise" */
	return mode == 1 || mode == 3;
}