From ad776c76a74c224ca2416437ce4cb34ef22e272b Mon Sep 17 00:00:00 2001
From: Christoffer Lerno <christoffer@aegik.com>
Date: Sat, 23 Sep 2023 00:20:15 +0200
Subject: [PATCH] Add benchfn and testfn allowing easy overwrite of test and
 benchmark runners. #990

---
 lib/std/core/runtime.c3          |  4 +-
 src/build/build_options.c        | 14 +++++++
 src/build/build_options.h        |  4 ++
 src/build/builder.c              |  4 +-
 src/build/project.c              | 12 ++++++
 src/compiler/compiler_internal.h |  2 -
 src/compiler/sema_decls.c        | 20 ---------
 src/compiler/semantic_analyser.c | 69 ++++++++++++++++++++++++++++++++
 src/compiler/symtab.c            |  4 --
 src/version.h                    |  2 +-
 10 files changed, 105 insertions(+), 30 deletions(-)

diff --git a/lib/std/core/runtime.c3 b/lib/std/core/runtime.c3
index 61486e0ce..b80250a84 100644
--- a/lib/std/core/runtime.c3
+++ b/lib/std/core/runtime.c3
@@ -136,7 +136,7 @@ fn bool run_benchmarks(BenchmarkUnit[] benchmarks)
 	return benchmark_count == benchmarks_passed;
 }
 
-fn bool __run_default_benchmark_runner()
+fn bool default_benchmark_runner()
 {
 	@pool()
 	{
@@ -245,7 +245,7 @@ fn bool run_tests(TestUnit[] tests)
 	return test_count == tests_passed;
 }
 
-fn bool __run_default_test_runner()
+fn bool default_test_runner()
 {
 	@pool()
 	{
diff --git a/src/build/build_options.c b/src/build/build_options.c
index 3cb42ad56..040a940c8 100644
--- a/src/build/build_options.c
+++ b/src/build/build_options.c
@@ -137,6 +137,8 @@ static void usage(void)
 	OUTPUT("  --link-libc=<yes|no>      - Link libc other default libraries (default: yes).");
 	OUTPUT("  --emit-stdlib=<yes|no>    - Output files for the standard library. (default: yes)");
 	OUTPUT("  --panicfn <name>          - Override the panic function name.");
+	OUTPUT("  --testfn <name>           - Override the test runner function name.");
+	OUTPUT("  --benchfn <name>          - Override the benchmark runner function name.");
 	OUTPUT("");
 	OUTPUT("  --reloc=<option>          - Relocation model: none, pic, PIC, pie, PIE.");
 	OUTPUT("  --x86cpu=<option>         - Set general level of x64 cpu: baseline, ssse3, sse4, avx1, avx2-v1, avx2-v2 (Skylake/Zen1+), avx512 (Icelake/Zen4+), native.");
@@ -851,6 +853,18 @@ static void parse_option(BuildOptions *options)
 				options->panicfn = next_arg();
 				return;
 			}
+			if (match_longopt("testfn"))
+			{
+				if (at_end() || next_is_opt()) error_exit("error: --testfn needs a function name.");
+				options->testfn = next_arg();
+				return;
+			}
+			if (match_longopt("benchfn"))
+			{
+				if (at_end() || next_is_opt()) error_exit("error: --benchfn needs a function name.");
+				options->benchfn = next_arg();
+				return;
+			}
 			if (match_longopt("macossdk"))
 			{
 				if (at_end() || next_is_opt()) error_exit("error: --macossdk needs a directory.");
diff --git a/src/build/build_options.h b/src/build/build_options.h
index d5add48be..7e9b741b5 100644
--- a/src/build/build_options.h
+++ b/src/build/build_options.h
@@ -422,6 +422,8 @@ typedef struct BuildOptions_
 	bool read_stdin;
 	bool print_output;
 	const char *panicfn;
+	const char *benchfn;
+	const char *testfn;
 	const char *cc;
 	const char *build_dir;
 	const char *llvm_out;
@@ -541,6 +543,8 @@ typedef struct
 	uint32_t symtab_size;
 	uint32_t switchrange_max_size;
 	const char *panicfn;
+	const char *benchfn;
+	const char *testfn;
 	const char *cc;
 	const char *cflags;
 	const char **csource_dirs;
diff --git a/src/build/builder.c b/src/build/builder.c
index e527d2a2a..363f18f85 100644
--- a/src/build/builder.c
+++ b/src/build/builder.c
@@ -278,7 +278,9 @@ static void update_build_target_from_options(BuildTarget *target, BuildOptions *
 	target->emit_llvm = options->emit_llvm;
 	target->build_threads = options->build_threads;
 	target->emit_asm = options->emit_asm;
-	target->panicfn = options->panicfn;
+	if (options->panicfn) target->panicfn = options->panicfn;
+	if (options->testfn) target->testfn = options->testfn;
+	if (options->benchfn) target->benchfn = options->benchfn;
 	target->benchmarking = options->benchmarking;
 	target->testing = options->testing;
 	if (options->macos.sdk) target->macos.sdk = options->macos.sdk;
diff --git a/src/build/project.c b/src/build/project.c
index 497767821..7d90f9d5c 100644
--- a/src/build/project.c
+++ b/src/build/project.c
@@ -9,6 +9,7 @@
 
 const char *project_default_keys[] = {
 		"authors",
+		"benchfn",
 		"c-sources",
 		"cc",
 		"cflags",
@@ -41,6 +42,7 @@ const char *project_default_keys[] = {
 		"system-linker",
 		"target",
 		"targets",
+		"testfn",
 		"trap-on-wrap",
 		"use-stdlib",
 		"version",
@@ -55,6 +57,7 @@ const char *project_default_keys[] = {
 const int project_default_keys_count = sizeof(project_default_keys) / sizeof(char*);
 
 const char* project_target_keys[] = {
+		"benchfn",
 		"c-sources-add",
 		"c-sources-override",
 		"cc",
@@ -94,6 +97,7 @@ const char* project_target_keys[] = {
 		"symtab",
 		"system-linker",
 		"target",
+		"testfn",
 		"trap-on-wrap",
 		"type",
 		"use-stdlib",
@@ -450,6 +454,14 @@ static void load_into_build_target(JSONObject *json, const char *type, BuildTarg
 	const char *panicfn = get_valid_string(json, "panicfn", type, false);
 	target->panicfn = panicfn;
 
+	// testfn
+	const char *testfn = get_valid_string(json, "testfn", type, false);
+	target->testfn = testfn;
+
+	// testfn
+	const char *benchfn = get_valid_string(json, "benchfn", type, false);
+	target->benchfn = benchfn;
+
 	// link-libc
 	target->link_libc = (LinkLibc)get_valid_bool(json, "link-libc", type, target->link_libc);
 
diff --git a/src/compiler/compiler_internal.h b/src/compiler/compiler_internal.h
index 2df4a3b3f..1746b3fbd 100644
--- a/src/compiler/compiler_internal.h
+++ b/src/compiler/compiler_internal.h
@@ -1887,8 +1887,6 @@ extern const char *type_property_list[NUMBER_OF_TYPE_PROPERTIES];
 extern const char *kw_std__core;
 extern const char *kw_std__core__types;
 extern const char *kw_std__io;
-extern const char *kw___run_default_benchmark_runner;
-extern const char *kw___run_default_test_runner;
 extern const char *kw_typekind;
 extern const char *kw_FILE_NOT_FOUND;
 extern const char *kw_IoError;
diff --git a/src/compiler/sema_decls.c b/src/compiler/sema_decls.c
index 3584f8ca5..167032ead 100644
--- a/src/compiler/sema_decls.c
+++ b/src/compiler/sema_decls.c
@@ -2528,26 +2528,6 @@ static inline bool sema_analyse_func(SemaContext *context, Decl *decl, bool *era
 
 	if (*erase_decl) return true;
 
-	if (decl->name == kw___run_default_benchmark_runner)
-	{
-		if (global_context.benchmark_func)
-		{
-			SEMA_ERROR(decl, "Multiple benchmark runners defined.");
-			return false;
-		}
-		global_context.benchmark_func = decl;
-		if (active_target.benchmarking) decl->no_strip = true;
-	}
-	if (decl->name == kw___run_default_test_runner)
-	{
-		if (global_context.test_func)
-		{
-			SEMA_ERROR(decl, "Multiple test runners defined.");
-			return false;
-		}
-		global_context.test_func = decl;
-		if (active_target.testing) decl->no_strip = true;
-	}
 	bool is_test = decl->func_decl.attr_test;
 	bool is_benchmark = decl->func_decl.attr_benchmark;
 	Signature *sig = &decl->func_decl.signature;
diff --git a/src/compiler/semantic_analyser.c b/src/compiler/semantic_analyser.c
index 3ef8096f6..defdb1821 100644
--- a/src/compiler/semantic_analyser.c
+++ b/src/compiler/semantic_analyser.c
@@ -245,6 +245,7 @@ static void assign_panicfn(void)
 	{
 		global_context.panic_var = NULL;
 		global_context.panicf = NULL;
+		return;
 	}
 
 	const char *panicfn = active_target.panicfn ? active_target.panicfn : "std::core::builtin::panic";
@@ -301,6 +302,72 @@ static void assign_panicfn(void)
 	global_context.panicf = panicf_decl;
 }
 
+static void assign_testfn(void)
+{
+	if (!active_target.testing) return;
+	if (!active_target.testfn && no_stdlib())
+	{
+		global_context.test_func = NULL;
+		return;
+	}
+	const char *testfn = active_target.testfn ? active_target.testfn : "std::core::runtime::default_test_runner";
+	Path *path;
+	const char *ident;
+	TokenType type;
+	if (sema_splitpathref(testfn, strlen(testfn), &path, &ident) != TOKEN_IDENT || path == NULL || !ident)
+	{
+		error_exit("'%s' is not a valid test function.", testfn);
+	}
+	Decl *decl = sema_find_decl_in_modules(global_context.module_list, path, ident);
+	if (!decl)
+	{
+		error_exit("Test function '%s::%s' could not be found.", path->module, ident);
+	}
+	if (decl->decl_kind != DECL_FUNC)
+	{
+		error_exit("'%s::%s' is not a function.", path->module, ident);
+	}
+	if (!type_func_match(type_get_ptr(decl->type->canonical), type_bool, 0))
+	{
+		error_exit("Expected test runner to have the signature fn void().");
+	}
+	global_context.test_func = decl;
+	decl->no_strip = true;
+}
+
+static void assign_benchfn(void)
+{
+	if (!active_target.benchmarking) return;
+	if (!active_target.benchfn && no_stdlib())
+	{
+		global_context.benchmark_func = NULL;
+		return;
+	}
+	const char *testfn = active_target.benchfn ? active_target.benchfn : "std::core::runtime::default_benchmark_runner";
+	Path *path;
+	const char *ident;
+	TokenType type;
+	if (sema_splitpathref(testfn, strlen(testfn), &path, &ident) != TOKEN_IDENT || path == NULL || !ident)
+	{
+		error_exit("'%s' is not a valid benchmark function.", testfn);
+	}
+	Decl *decl = sema_find_decl_in_modules(global_context.module_list, path, ident);
+	if (!decl)
+	{
+		error_exit("Benchmark function '%s::%s' could not be found.", path->module, ident);
+	}
+	if (decl->decl_kind != DECL_FUNC)
+	{
+		error_exit("'%s::%s' is not a function.", path->module, ident);
+	}
+	if (!type_func_match(type_get_ptr(decl->type->canonical), type_bool, 0))
+	{
+		error_exit("Expected benchmark function to have the signature fn void().");
+	}
+	global_context.benchmark_func = decl;
+	decl->no_strip = true;
+}
+
 /**
  * Perform the entire semantic analysis.
  */
@@ -358,6 +425,8 @@ RESOLVE_LAMBDA:;
 	halt_on_error();
 
 	assign_panicfn();
+	assign_testfn();
+	assign_benchfn();
 
 	if (strip_unused())
 	{
diff --git a/src/compiler/symtab.c b/src/compiler/symtab.c
index 7c478df95..9bee1f7a3 100644
--- a/src/compiler/symtab.c
+++ b/src/compiler/symtab.c
@@ -73,8 +73,6 @@ const char *kw_std;
 const char *kw_std__core;
 const char *kw_std__core__types;
 const char *kw_std__io;
-const char *kw___run_default_benchmark_runner;
-const char *kw___run_default_test_runner;
 const char *kw_type;
 const char *kw_typekind;
 const char *kw_winmain;
@@ -167,8 +165,6 @@ void symtab_init(uint32_t capacity)
 	kw_std__core = KW_DEF("std::core");
 	kw_std__core__types = KW_DEF("std::core::types");
 	kw_std__io = KW_DEF("std::io");
-	kw___run_default_benchmark_runner = KW_DEF("__run_default_benchmark_runner");
-	kw___run_default_test_runner = KW_DEF("__run_default_test_runner");
 	kw_type = KW_DEF("type");
 	kw_winmain = KW_DEF("wWinMain");
 	kw_wmain = KW_DEF("wmain");
diff --git a/src/version.h b/src/version.h
index e896da68d..45705f8a8 100644
--- a/src/version.h
+++ b/src/version.h
@@ -1 +1 @@
-#define COMPILER_VERSION "0.4.657"
\ No newline at end of file
+#define COMPILER_VERSION "0.4.658"
\ No newline at end of file