Improved support for CLI arg parsing

2 files changed, 295 insertions, 0 deletions

diff --git a/stdlib/stdlib.c b/stdlib/stdlib.c
index db60ca98..801283a7 100644
--- a/stdlib/stdlib.c
+++ b/stdlib/stdlib.c
@@ -10,11 +10,15 @@
 #include <sys/random.h>
 #include <time.h>
 
+#include "bools.h"
 #include "files.h"
 #include "integers.h"
+#include "optionals.h"
 #include "metamethods.h"
 #include "patterns.h"
+#include "paths.h"
 #include "siphash.h"
+#include "stdlib.h"
 #include "tables.h"
 #include "text.h"
 #include "util.h"
@@ -36,6 +40,289 @@ public void tomo_init(void)
        errx(1, "Couldn't set printf specifier");
 }
 
+static bool parse_single_arg(const TypeInfo *info, char *arg, void *dest)
+{
+    while (info->tag == OptionalInfo)
+        info = info->OptionalInfo.type;
+
+    if (info == &Int$info) {
+        OptionalInt_t parsed = Int$from_str(arg);
+        if (parsed.small != 0)
+            *(OptionalInt_t*)dest = parsed;
+        return parsed.small != 0;
+    } else if (info == &Int64$info) {
+        OptionalInt64_t parsed = Int64$from_text(Text$from_str(arg));
+        if (!parsed.is_null)
+            *(OptionalInt64_t*)dest = parsed;
+        return !parsed.is_null;
+    } else if (info == &Int32$info) {
+        OptionalInt32_t parsed = Int32$from_text(Text$from_str(arg));
+        if (!parsed.is_null)
+            *(OptionalInt32_t*)dest = parsed;
+        return !parsed.is_null;
+    } else if (info == &Int16$info) {
+        OptionalInt16_t parsed = Int16$from_text(Text$from_str(arg));
+        if (!parsed.is_null)
+            *(OptionalInt16_t*)dest = parsed;
+        return !parsed.is_null;
+    } else if (info == &Int8$info) {
+        OptionalInt8_t parsed = Int8$from_text(Text$from_str(arg));
+        if (!parsed.is_null)
+            *(OptionalInt8_t*)dest = parsed;
+        return !parsed.is_null;
+    } else if (info == &Bool$info) {
+        OptionalBool_t parsed = Bool$from_text(Text$from_str(arg));
+        if (parsed != NULL_BOOL)
+            *(OptionalBool_t*)dest = parsed;
+        return parsed != NULL_BOOL;
+    } else if (info == &Num$info) {
+        OptionalNum_t parsed = Num$from_text(Text$from_str(arg));
+        if (!isnan(parsed))
+            *(OptionalNum_t*)dest = parsed;
+        return !isnan(parsed);
+    } else if (info == &Num32$info) {
+        OptionalNum32_t parsed = Num32$from_text(Text$from_str(arg));
+        if (!isnan(parsed))
+            *(OptionalNum32_t*)dest = parsed;
+        return !isnan(parsed);
+    } else if (info == &Path$info) {
+        Path_t path = Text$from_str(arg);
+        if (Text$equal_values(path, Path("~"))) {
+            path = Path("~/");
+        } else if (Text$equal_values(path, Path("."))) {
+            path = Path("./");
+        } else if (Text$equal_values(path, Path(".."))) {
+            path = Path("../");
+        } else if (!Text$starts_with(path, Text("./"))
+            && !Text$starts_with(path, Text("../"))
+            && !Text$starts_with(path, Text("/"))
+            && !Text$starts_with(path, Text("~/"))) {
+            path = Text$concat(Text("./"), path);
+        }
+        *(OptionalText_t*)dest = path;
+        return true;
+    } else if (info->tag == TextInfo) {
+        *(OptionalText_t*)dest = Text$from_str(arg);
+        return true;
+    } else {
+        Text_t t = generic_as_text(NULL, false, info);
+        errx(1, "Unsupported type for argument parsing: %k", &t);
+    }
+}
+
+static Array_t parse_array(const TypeInfo *item_info, int n, char *args[])
+{
+    int64_t padded_size = item_info->size;
+    if ((padded_size % item_info->align) > 0)
+        padded_size = padded_size + item_info->align - (padded_size % item_info->align);
+
+    Array_t items = {
+        .stride=padded_size,
+        .length=n,
+        .data=GC_MALLOC((size_t)(padded_size*n)),
+    };
+    for (int i = 0; i < n; i++) {
+        bool success = parse_single_arg(item_info, args[i], items.data + items.stride*i);
+        if (!success)
+            errx(1, "Couldn't parse argument: %s", args[i]);
+    }
+    return items;
+}
+
+#pragma GCC diagnostic ignored "-Wstack-protector"
+public void tomo_parse_args(Text_t usage, Text_t help, int spec_len, cli_arg_t spec[spec_len], int argc, char *argv[])
+{
+    bool populated_args[spec_len] = {};
+    bool used_args[argc] = {};
+    for (int i = 1; i < argc; ) {
+        if (argv[i][0] == '-' && argv[i][1] == '-') {
+            if (argv[i][2] == '\0') { // "--" signals the rest of the arguments are literal
+                used_args[i] = true;
+                i += 1;
+                break;
+            }
+
+            for (int s = 0; s < spec_len; s++) {
+                const TypeInfo *non_opt_type = spec[s].type;
+                while (non_opt_type->tag == OptionalInfo)
+                    non_opt_type = non_opt_type->OptionalInfo.type;
+
+                if (non_opt_type == &Bool$info
+                    && strncmp(argv[i], "--no-", strlen("--no-")) == 0
+                    && strcmp(argv[i] + strlen("--no-"), spec[s].name) == 0) {
+                    *(OptionalBool_t*)spec[s].dest = false;
+                    populated_args[s] = true;
+                    used_args[i] = true;
+                    goto next_arg;
+                }
+
+                if (strncmp(spec[s].name, argv[i] + 2, strlen(spec[s].name)) != 0)
+                    continue;
+
+                char after_name = argv[i][2+strlen(spec[s].name)];
+                if (after_name == '\0') { // --foo val
+                    used_args[i] = true;
+                    if (non_opt_type->tag == ArrayInfo || non_opt_type->tag == TableInfo) {
+                        int num_args = 0;
+                        while (i + 1 + num_args < argc) {
+                            if (argv[i+1+num_args][0] == '-')
+                                break;
+                            used_args[i+1+num_args] = true;
+                            num_args += 1;
+                        }
+                        populated_args[s] = true;
+                        const TypeInfo *item_type = non_opt_type->tag == ArrayInfo ? non_opt_type->ArrayInfo.item : non_opt_type->TableInfo.key;
+                        Array_t items = parse_array(item_type, num_args, &argv[i+1]);
+                        if (non_opt_type->tag == ArrayInfo) {
+                            *(OptionalArray_t*)spec[s].dest = items;
+                        } else {
+                            *(OptionalTable_t*)spec[s].dest = Table$from_entries(items, non_opt_type);
+                        }
+                    } else if (non_opt_type == &Bool$info) { // --flag
+                        populated_args[s] = true;
+                        *(OptionalBool_t*)spec[s].dest = true;
+                    } else {
+                        used_args[i+1] = true;
+                        populated_args[s] = parse_single_arg(spec[s].type, argv[i+1], spec[s].dest);
+                        if (!populated_args[s])
+                            errx(1, "Couldn't parse argument: %s %s\n%k", argv[i], argv[i+1], &usage);
+                    }
+                    goto next_arg;
+                } else if (after_name == '=') { // --foo=val
+                    used_args[i] = true;
+                    populated_args[s] = parse_single_arg(spec[s].type, 2 + argv[i] + strlen(spec[s].name) + 1, spec[s].dest);
+                    if (!populated_args[s])
+                        errx(1, "Couldn't parse argument: %s\n%k", argv[i], &usage);
+                    goto next_arg;
+                } else {
+                    continue;
+                }
+            }
+
+            if (streq(argv[i], "--help")) {
+                say(help, true);
+                exit(0);
+            }
+            errx(1, "Unrecognized argument: %s\n%k", argv[i], &usage);
+        } else if (argv[i][0] == '-' && argv[i][1] && argv[i][1] != '-') { // Single flag args
+            used_args[i] = true;
+            for (char *f = argv[i] + 1; *f; f++) {
+                for (int s = 0; s < spec_len; s++) {
+                    if (spec[s].name[0] != *f || strlen(spec[s].name) > 1)
+                        continue;
+
+                    const TypeInfo *non_opt_type = spec[s].type;
+                    while (non_opt_type->tag == OptionalInfo)
+                        non_opt_type = non_opt_type->OptionalInfo.type;
+
+                    if (non_opt_type->tag == ArrayInfo || non_opt_type->tag == TableInfo) {
+                        if (f[1]) errx(1, "No value provided for -%c\n%k", *f, &usage);
+                        int num_args = 0;
+                        while (i + 1 + num_args < argc) {
+                            if (argv[i+1+num_args][0] == '-')
+                                break;
+                            used_args[i+1+num_args] = true;
+                            num_args += 1;
+                        }
+                        populated_args[s] = true;
+                        const TypeInfo *item_type = non_opt_type->tag == ArrayInfo ? non_opt_type->ArrayInfo.item : non_opt_type->TableInfo.key;
+                        Array_t items = parse_array(item_type, num_args, &argv[i+1]);
+                        if (non_opt_type->tag == ArrayInfo) {
+                            *(OptionalArray_t*)spec[s].dest = items;
+                        } else {
+                            *(OptionalTable_t*)spec[s].dest = Table$from_entries(items, non_opt_type);
+                        }
+                    } else if (non_opt_type == &Bool$info) { // -f
+                        populated_args[s] = true;
+                        *(OptionalBool_t*)spec[s].dest = true;
+                    } else {
+                        if (f[1] || i+1 >= argc) errx(1, "No value provided for -%c\n%k", *f, &usage);
+                        used_args[i+1] = true;
+                        populated_args[s] = parse_single_arg(spec[s].type, argv[i+1], spec[s].dest);
+                        if (!populated_args[s])
+                            errx(1, "Couldn't parse argument: %s %s\n%k", argv[i], argv[i+1], &usage);
+                    }
+                    goto next_flag;
+                }
+
+                if (*f == 'h') {
+                    say(help, true);
+                    exit(0);
+                }
+                errx(1, "Unrecognized flag: -%c\n%k", *f, &usage);
+              next_flag:;
+            }
+        } else {
+            // Handle positional args later
+            i += 1;
+            continue;
+        }
+
+      next_arg:
+        while (used_args[i] && i < argc)
+            i += 1;
+    }
+
+    // Get remaining positional arguments
+    bool ignore_dashes = false;
+    for (int i = 1, s = 0; i < argc; i++) {
+        if (!ignore_dashes && streq(argv[i], "--")) {
+            ignore_dashes = true;
+            continue;
+        }
+        if (used_args[i]) continue;
+
+        while (populated_args[s]) {
+          next_non_bool_flag:
+            ++s;
+            if (s >= spec_len)
+                errx(1, "Extra argument: %s\n%k", argv[i], &usage);
+        }
+
+        const TypeInfo *non_opt_type = spec[s].type;
+        while (non_opt_type->tag == OptionalInfo)
+            non_opt_type = non_opt_type->OptionalInfo.type;
+
+        // You can't specify boolean flags positionally
+        if (non_opt_type == &Bool$info)
+            goto next_non_bool_flag;
+
+        if (non_opt_type->tag == ArrayInfo || non_opt_type->tag == TableInfo) {
+            int num_args = 0;
+            while (i + num_args < argc) {
+                if (!ignore_dashes && argv[i+num_args][0] == '-')
+                    break;
+                used_args[i+num_args] = true;
+                num_args += 1;
+            }
+            populated_args[s] = true;
+            const TypeInfo *item_type = non_opt_type->tag == ArrayInfo ? non_opt_type->ArrayInfo.item : non_opt_type->TableInfo.key;
+            Array_t items = parse_array(item_type, num_args, &argv[i]);
+            if (non_opt_type->tag == ArrayInfo) {
+                *(OptionalArray_t*)spec[s].dest = items;
+            } else {
+                *(OptionalTable_t*)spec[s].dest = Table$from_entries(items, non_opt_type);
+            }
+        } else {
+            populated_args[s] = parse_single_arg(spec[s].type, argv[i], spec[s].dest);
+        }
+
+        if (!populated_args[s])
+            errx(1, "Invalid value for %s: %s\n%k", spec[s].name, argv[i], &usage);
+    }
+
+    for (int s = 0; s < spec_len; s++) {
+        if (!populated_args[s] && spec[s].required) {
+            if (spec[s].type->tag == ArrayInfo)
+                *(OptionalArray_t*)spec[s].dest = (Array_t){};
+            else if (spec[s].type->tag == TableInfo)
+                *(OptionalTable_t*)spec[s].dest = (Table_t){};
+            else
+                errx(1, "The required argument '%s' was not provided\n%k", spec[s].name, &usage);
+        }
+    }
+}
+
 void print_stack_trace(FILE *out, int start, int stop)
 {
     // Print stack trace:
diff --git a/stdlib/stdlib.h b/stdlib/stdlib.h
index 2037fae4..b8e78e9b 100644
--- a/stdlib/stdlib.h
+++ b/stdlib/stdlib.h
@@ -12,7 +12,15 @@
 
 extern bool USE_COLOR;
 
+typedef struct {
+    const char *name;
+    bool required;
+    const TypeInfo *type;
+    void *dest;
+} cli_arg_t;
+
 void tomo_init(void);
+void tomo_parse_args(Text_t usage, Text_t help, int spec_len, cli_arg_t spec[spec_len], int argc, char *argv[]);
 __attribute__((format(printf, 1, 2)))
 _Noreturn void fail(const char *fmt, ...);
 __attribute__((format(printf, 4, 5)))