5 files changed, 64 insertions, 3 deletions

diff --git a/src/cli.c b/src/cli.c
index 7fe2119..56f6098 100644
--- a/src/cli.c
+++ b/src/cli.c
@@ -18,6 +18,8 @@ CliArgs* parse_cli_args(int argc, char** argv) {
     bool opt_MMD = false;
     StrArray include_dirs;
     strings_init(&include_dirs);
+    StrArray defines;
+    strings_init(&defines);
 
     for (int i = 1; i < argc; ++i) {
         if (argv[i][0] != '-') {
@@ -46,6 +48,19 @@ CliArgs* parse_cli_args(int argc, char** argv) {
                 fatal_error("-I requires directory");
             }
             strings_push(&include_dirs, dir);
+        } else if (c == 'D') {
+            const char* def = NULL;
+            if (argv[i][2] != '\0') {
+                // -DFOO or -DFOO=value format
+                def = argv[i] + 2;
+            } else if (argc > i + 1) {
+                // -D FOO or -D FOO=value format
+                def = argv[i + 1];
+                ++i;
+            } else {
+                fatal_error("-D requires macro definition");
+            }
+            strings_push(&defines, def);
         } else if (c == 'o') {
             if (argc <= i + 1) {
                 fatal_error("-o requires filename");
@@ -84,6 +99,7 @@ CliArgs* parse_cli_args(int argc, char** argv) {
     a->gcc_command = NULL;
     a->generate_deps = opt_MMD;
     a->include_dirs = include_dirs;
+    a->defines = defines;
 
     if (!a->only_compile && str_ends_with(a->input_filename, ".o")) {
         a->totally_deligate_to_gcc = true;
diff --git a/src/cli.h b/src/cli.h
index d97874c..6758323 100644
--- a/src/cli.h
+++ b/src/cli.h
@@ -14,6 +14,7 @@ typedef struct {
     bool wasm;
     const char* gcc_command;
     StrArray include_dirs;
+    StrArray defines;
 } CliArgs;
 
 CliArgs* parse_cli_args(int argc, char** argv);
diff --git a/src/main.c b/src/main.c
index 44a317c..54fe87a 100644
--- a/src/main.c
+++ b/src/main.c
@@ -21,7 +21,7 @@ int main(int argc, char** argv) {
     StrArray included_files;
     strings_init(&included_files);
 
-    TokenArray* pp_tokens = preprocess(source, &included_files, &cli_args->include_dirs);
+    TokenArray* pp_tokens = preprocess(source, &included_files, &cli_args->include_dirs, &cli_args->defines);
 
     if (cli_args->preprocess_only) {
         FILE* output_file = cli_args->output_filename ? fopen(cli_args->output_filename, "w") : stdout;
diff --git a/src/preprocess.c b/src/preprocess.c
index b06de6d..1ae8eac 100644
--- a/src/preprocess.c
+++ b/src/preprocess.c
@@ -104,6 +104,49 @@ static void add_predefined_macros(MacroArray* macros) {
     define_macro_to_1(macros, "__LP64__");
 }
 
+// Accept "FOO" or "FOO=value"
+static void define_macro_from_string(MacroArray* macros, const char* def) {
+    Macro* m = macros_push_new(macros);
+    m->kind = MacroKind_obj;
+
+    const char* eq = strchr(def, '=');
+    if (eq) {
+        // FOO=value format
+        size_t name_len = eq - def;
+        char* name = calloc(name_len + 1, sizeof(char));
+        memcpy(name, def, name_len);
+        m->name = name;
+
+        const char* value = eq + 1;
+        tokens_init(&m->replacements, 1);
+        Token* tok = tokens_push_new(&m->replacements);
+
+        // Try to parse as integer
+        char* num_end;
+        long int_val = strtol(value, &num_end, 10);
+        if (value[0] != '\0' && *num_end == '\0') {
+            tok->kind = TokenKind_literal_int;
+            tok->value.integer = int_val;
+        } else {
+            tok->kind = TokenKind_ident;
+            tok->value.string = value;
+        }
+    } else {
+        // FOO format (equivalent to FOO=1)
+        m->name = def;
+        tokens_init(&m->replacements, 1);
+        Token* tok = tokens_push_new(&m->replacements);
+        tok->kind = TokenKind_literal_int;
+        tok->value.integer = 1;
+    }
+}
+
+static void add_user_defines(MacroArray* macros, StrArray* user_defines) {
+    for (size_t i = 0; i < user_defines->len; ++i) {
+        define_macro_from_string(macros, user_defines->data[i]);
+    }
+}
+
 typedef struct {
     TokenArray tokens;
 } MacroArg;
@@ -1170,9 +1213,10 @@ static TokenArray* do_preprocess(InFile* src, int depth, MacroArray* macros, Str
     return pp->pp_tokens;
 }
 
-TokenArray* preprocess(InFile* src, StrArray* included_files, StrArray* user_include_dirs) {
+TokenArray* preprocess(InFile* src, StrArray* included_files, StrArray* user_include_dirs, StrArray* user_defines) {
     MacroArray* macros = macros_new();
     add_predefined_macros(macros);
+    add_user_defines(macros, user_defines);
     strings_push(included_files, src->loc.filename);
     return do_preprocess(src, 0, macros, included_files, user_include_dirs);
 }
diff --git a/src/preprocess.h b/src/preprocess.h
index 4bf9834..0365a57 100644
--- a/src/preprocess.h
+++ b/src/preprocess.h
@@ -5,7 +5,7 @@
 #include "io.h"
 #include "token.h"
 
-TokenArray* preprocess(InFile* src, StrArray* included_files, StrArray* user_include_dirs);
+TokenArray* preprocess(InFile* src, StrArray* included_files, StrArray* user_include_dirs, StrArray* user_defines);
 void concat_adjacent_string_literals(TokenArray* pp_tokens);
 void print_token_to_file(FILE* output_file, TokenArray* pp_tokens);