Move arg parsing to inline logic in main() function

4 files changed, 31 insertions, 36 deletions

diff --git a/builtins/tomo.h b/builtins/tomo.h
index f0455ad1..bc8523df 100644
--- a/builtins/tomo.h
+++ b/builtins/tomo.h
@@ -29,4 +29,6 @@
 #include "text.h"
 #include "types.h"
 
+extern bool USE_COLOR;
+
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/compile.c b/compile.c
index de553edc..000bcaf8 100644
--- a/compile.c
+++ b/compile.c
@@ -1917,10 +1917,11 @@ CORD compile_type_info(env_t *env, type_t *t)
     }
 }
 
-static CORD compile_main_arg_parser(env_t *env, const char *module_name, type_t *main_fn_type)
+CORD compile_cli_arg_call(env_t *env, CORD fn_name, type_t *fn_type)
 {
-    CORD code = CORD_all("void ", module_name, "$main$run(int argc, char *argv[]) {\n");
-    auto fn_info = Match(main_fn_type, FunctionType);
+    auto fn_info = Match(fn_type, FunctionType);
+    if (!fn_info->args)
+        return "if (argc > 1) errx(1, \"This program doesn't take any arguments.\");\n";
     env_t *main_env = fresh_scope(env);
 
     CORD usage = CORD_EMPTY;
@@ -1942,9 +1943,9 @@ static CORD compile_main_arg_parser(env_t *env, const char *module_name, type_t
                 usage = CORD_all(usage, "<", flag, ">");
         }
     }
-    code = CORD_all(code, "CORD usage = CORD_all(\"Usage: \", argv[0], ", usage ? Text$quoted(usage, false) : "CORD_EMPTY", ");\n",
-                    "#define USAGE_ERR(...) errx(1, CORD_to_const_char_star(CORD_all(__VA_ARGS__)))\n"
-                    "#define IS_FLAG(str, flag) (strncmp(str, flag, strlen(flag) == 0 && (str[strlen(flag)] == 0 || str[strlen(flag)] == '=')) == 0)\n");
+    CORD code = CORD_all("CORD usage = CORD_all(\"Usage: \", argv[0], ", usage ? Text$quoted(usage, false) : "CORD_EMPTY", ");\n",
+                         "#define USAGE_ERR(...) errx(1, CORD_to_const_char_star(CORD_all(__VA_ARGS__)))\n"
+                         "#define IS_FLAG(str, flag) (strncmp(str, flag, strlen(flag) == 0 && (str[strlen(flag)] == 0 || str[strlen(flag)] == '=')) == 0)\n");
 
     // Declare args:
     for (arg_t *arg = fn_info->args; arg; arg = arg->next) {
@@ -2068,12 +2069,12 @@ static CORD compile_main_arg_parser(env_t *env, const char *module_name, type_t
     code = CORD_all(code, "for (; i < argc; i++) {\n"
                     "if (argv[i])\nUSAGE_ERR(\"Unexpected argument: \", Text$quoted(argv[i], false), \"\\n\", usage);\n}\n");
 
-    code = CORD_all(code, module_name, "$main(");
+    code = CORD_all(code, fn_name, "(");
     for (arg_t *arg = fn_info->args; arg; arg = arg->next) {
         code = CORD_all(code, "$", arg->name);
         if (arg->next) code = CORD_all(code, ", ");
     }
-    code = CORD_all(code, ");\n}\n");
+    code = CORD_all(code, ");\n");
     return code;
 }
 
@@ -2125,22 +2126,9 @@ module_code_t compile_file(ast_t *ast)
         }
     }
 
-    env->code->fndefs = CORD_all(env->code->fndefs, "void ", name, "$main$run(int argc, char *argv[]);\n");
-
-    binding_t *main_fn = get_binding(env, "main");
-    if (!main_fn) {
-        env->code->funcs = CORD_all(env->code->funcs, "public void ", name, "$main$run(int argc, char *argv[]) {\n"
-                                     "(void)argc;\n"
-                                     "(void)argv;\n"
-                                     "}\n");
-    } else if (main_fn->type->tag != FunctionType) {
-        compiler_err(NULL, NULL, NULL, "The name 'main' is bound to something that isn't a function, it's %T", main_fn->type);
-    } else {
-        env->code->funcs = CORD_all(env->code->funcs, compile_main_arg_parser(env, name, main_fn->type));
-    }
-
     return (module_code_t){
         .module_name=name,
+        .env=env,
         .object_files=env->code->object_files,
         .header=CORD_all(
             // CORD_asprintf("#line 1 %r\n", Text$quoted(ast->file->filename, false)),
diff --git a/compile.h b/compile.h
index 6be4a997..205d5517 100644
--- a/compile.h
+++ b/compile.h
@@ -11,6 +11,7 @@
 
 typedef struct {
     const char *module_name;
+    env_t *env;
     CORD header, c_file, object_files;
 } module_code_t;
 
@@ -23,5 +24,6 @@ CORD compile(env_t *env, ast_t *ast);
 void compile_namespace(env_t *env, const char *ns_name, ast_t *block);
 CORD compile_statement(env_t *env, ast_t *ast);
 CORD compile_type_info(env_t *env, type_t *t);
+CORD compile_cli_arg_call(env_t *env, CORD fn_name, type_t *fn_type);
 
 // vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/tomo.c b/tomo.c
index a2acdc29..74b6b09e 100644
--- a/tomo.c
+++ b/tomo.c
@@ -29,9 +29,9 @@ static const char *ldflags = "-Wl,-rpath '-Wl,$ORIGIN' -L/usr/local/lib";
 static const char *cc;
 
 static array_t get_file_dependencies(const char *filename);
-static int transpile(const char *filename, bool force_retranspile);
+static int transpile(const char *filename, bool force_retranspile, module_code_t *module_code);
 static int compile_object_file(const char *filename, bool force_recompile);
-static int compile_executable(const char *filename, const char *object_files);
+static int compile_executable(const char *filename, const char *object_files, module_code_t *module_code);
 
 int main(int argc, char *argv[])
 {
@@ -104,7 +104,8 @@ int main(int argc, char *argv[])
     cc = getenv("CC");
     if (!cc) cc = "tcc";
 
-    int transpile_status = transpile(filename, true);
+    module_code_t module_code;
+    int transpile_status = transpile(filename, true, &module_code);
     if (mode == MODE_TRANSPILE || transpile_status != 0)
         return transpile_status;
 
@@ -122,7 +123,7 @@ int main(int argc, char *argv[])
 
     const char *object_files = CORD_to_const_char_star(object_files_cord);
     assert(object_files);
-    int executable_status = compile_executable(filename, object_files);
+    int executable_status = compile_executable(filename, object_files, &module_code);
     if (mode == MODE_COMPILE_EXE || executable_status != 0)
         return executable_status;
 
@@ -156,7 +157,7 @@ static void build_file_dependency_graph(const char *filename, table_t *dependenc
     Array$insert(deps, &base_filename, 0, $ArrayInfo(&$Text));
     Table$str_set(dependencies, base_filename, deps);
 
-    transpile(base_filename, false);
+    transpile(base_filename, false, NULL);
 
     const char *to_scan[] = {
         heap_strf("%s.h", base_filename),
@@ -215,7 +216,7 @@ static bool stale(const char *filename, const char *relative_to)
     return target_stat.st_mtime < relative_to_stat.st_mtime;
 }
 
-int transpile(const char *filename, bool force_retranspile)
+int transpile(const char *filename, bool force_retranspile, module_code_t *module_code)
 {
     const char *tm_file = filename;
     const char *c_filename = heap_strf("%s.c", tm_file);
@@ -245,13 +246,13 @@ int transpile(const char *filename, bool force_retranspile)
         pclose(out);
     }
 
-    module_code_t module = compile_file(ast);
+    *module_code = compile_file(ast);
 
     FILE *h_file = fopen(h_filename, "w");
     if (!h_file)
         errx(1, "Couldn't open file: %s", h_filename);
     CORD_put("#pragma once\n", h_file);
-    CORD_put(module.header, h_file);
+    CORD_put(module_code->header, h_file);
     if (fclose(h_file))
         errx(1, "Failed to close file: %s", h_filename);
     if (verbose)
@@ -265,7 +266,7 @@ int transpile(const char *filename, bool force_retranspile)
     FILE *c_file = fopen(c_filename, "w");
     if (!c_file)
         errx(1, "Couldn't open file: %s", c_filename);
-    CORD_put(CORD_all("#include \"", module.module_name, ".tm.h\"\n\n", module.c_file), c_file);
+    CORD_put(CORD_all("#include \"", module_code->module_name, ".tm.h\"\n\n", module_code->c_file), c_file);
     if (fclose(c_file))
         errx(1, "Failed to close file: %s", c_filename);
     if (verbose)
@@ -303,7 +304,7 @@ int compile_object_file(const char *filename, bool force_recompile)
     return WIFEXITED(status) ? WEXITSTATUS(status) : EXIT_FAILURE;
 }
 
-int compile_executable(const char *filename, const char *object_files)
+int compile_executable(const char *filename, const char *object_files, module_code_t *module_code)
 {
     const char *bin_name = file_base_name(filename);
     const char *run = heap_strf("%s | %s %s %s %s %s -x c - -o %s", autofmt, cc, cflags, ldflags, ldlibs, object_files, bin_name);
@@ -311,7 +312,7 @@ int compile_executable(const char *filename, const char *object_files)
         printf("%s\n", run);
     FILE *runner = popen(run, "w");
 
-    const char *module_name = file_base_name(filename);
+    binding_t *main_binding = get_binding(module_code->env, "main");
     CORD program = CORD_all(
         "#include <tomo/tomo.h>\n"
         "#include \"", filename, ".h\"\n"
@@ -322,9 +323,11 @@ int compile_executable(const char *filename, const char *object_files)
         "GC_INIT();\n"
         "USE_COLOR = getenv(\"COLOR\") ? strcmp(getenv(\"COLOR\"), \"1\") == 0 : isatty(STDOUT_FILENO);\n"
         "srand(arc4random_uniform(UINT32_MAX));\n"
-        "srand48(arc4random_uniform(UINT32_MAX));\n",
-        module_name, "$main$run(argc, argv);\n",
-        "return 0;\n"
+        "srand48(arc4random_uniform(UINT32_MAX));\n"
+        "\n",
+        main_binding && main_binding->type->tag == FunctionType ?
+            CORD_all(compile_cli_arg_call(module_code->env, main_binding->code, main_binding->type), "return 0;\n")
+            : "errx(1, \"No main function is defined!\");\n",
         "}\n"
     );