/*
 * Copyright (c) 2019 The Fuchsia Authors
 */
#include <malloc.h>
#include <common.h>
#include <command.h>

#include <u-boot/sha256.h>

static void run_sha256(void)
{
	const size_t buffer_size = (10 * 1024 * 1024);
	uint8_t *buffer = malloc(buffer_size);
	if (!buffer) {
		printf("Error: out of memory!\n");
		return;
	}

	uint8_t digest[SHA256_SUM_LEN];
	int repeat = 10;
	sha256_context ctx;

	unsigned long start = timer_get_us();
	sha256_starts(&ctx);
	for (int i = 0; i < repeat; i++)
		sha256_update(&ctx, buffer, buffer_size);
	sha256_finish(&ctx, digest);

	printf("sha256 - %lluMB/s\n",
	       ((buffer_size/1024/1024) * repeat * 1000000LL) / (timer_get_us() - start));

	free(buffer);
}

struct benchmark {
	char *name;
	void (*run)(void);
};

struct benchmark benchmarks[] = {
	{ "sha256", run_sha256 },
};

static struct benchmark *find_benchmark(const char *name)
{
	for (size_t i = 0; i < ARRAY_SIZE(benchmarks); i++)
		if (!strcmp(name, benchmarks[i].name))
			return &benchmarks[i];
	return NULL;
}

static void list_all_benchmarks(void)
{
	printf("List of Benchmarks:\n");
	for (size_t i = 0; i < ARRAY_SIZE(benchmarks); i++)
		printf("- %s\n", benchmarks[i].name);
	printf("\n(run all with `benchmark all`)\n");
	printf("(run selected with `benchmark name1 name2 ...`)\n");
}

static void run_benchmark_by_name(const char *name)
{
	struct benchmark *bm = find_benchmark(name);
	if (!bm) {
		printf("%s - not found and ignored.\n", name);
	} else {
		bm->run();
	}
}

static void run_all_benchmarks(void)
{
	for (size_t i = 0; i < ARRAY_SIZE(benchmarks); i++)
		run_benchmark_by_name(benchmarks[i].name);
}

static int do_benchmark(cmd_tbl_t *cmdtp, int flag, int argc,
			char *const argv[])
{
	if (argc == 1) {
		list_all_benchmarks();
		return CMD_RET_SUCCESS;
	}

	if (!strcmp(argv[1], "all")) {
		run_all_benchmarks();
		return CMD_RET_SUCCESS;
	}

	for (int i = 1; i < argc; i++) {
		run_benchmark_by_name(argv[i]);
	}

	return CMD_RET_SUCCESS;
}

U_BOOT_CMD(benchmark, /* maxargs = */ 100, /* repeatable = */ 1, do_benchmark,
	   "List/run benchmarks.", "");