Some more, and some fixes

9 files changed, 181 insertions, 51 deletions

diff --git a/src/ast/ast.h b/src/ast/ast.h
index 0da2c0d..e22b35c 100644
--- a/src/ast/ast.h
+++ b/src/ast/ast.h
@@ -242,8 +242,8 @@ struct ASTNode
         struct
         {
             char *var_name;
-            char *start;
-            char *end;
+            ASTNode *start;
+            ASTNode *end;
             char *step;
             ASTNode *body;
         } for_range;
diff --git a/src/codegen/codegen.c b/src/codegen/codegen.c
index 1592806..17b0816 100644
--- a/src/codegen/codegen.c
+++ b/src/codegen/codegen.c
@@ -11,14 +11,15 @@
 #include "zprep_plugin.h"
 
 // static function for internal use.
-static void codegen_match_internal(ParserContext *ctx, ASTNode *node, FILE *out)
+static char *g_current_func_ret_type = NULL;
+static void codegen_match_internal(ParserContext *ctx, ASTNode *node, FILE *out, int use_result)
 {
     int id = tmp_counter++;
     int is_self = (node->match_stmt.expr->type == NODE_EXPR_VAR &&
                    strcmp(node->match_stmt.expr->var_ref.name, "self") == 0);
 
     char *ret_type = infer_type(ctx, node);
-    int is_expr = (ret_type && strcmp(ret_type, "void") != 0);
+    int is_expr = (use_result && ret_type && strcmp(ret_type, "void") != 0);
 
     fprintf(out, "({ ");
     emit_auto_type(ctx, node->match_stmt.expr, node->token, out);
@@ -151,6 +152,11 @@ static void codegen_match_internal(ParserContext *ctx, ASTNode *node, FILE *out)
                 // Numeric pattern
                 fprintf(out, "_m_%d == %s", id, c->match_case.pattern);
             }
+            else if (c->match_case.pattern[0] == '\'')
+            {
+                // Char literal pattern
+                fprintf(out, "_m_%d == %s", id, c->match_case.pattern);
+            }
             else
             {
                 fprintf(out, "1");
@@ -290,7 +296,7 @@ void codegen_expression(ParserContext *ctx, ASTNode *node, FILE *out)
     switch (node->type)
     {
     case NODE_MATCH:
-        codegen_match_internal(ctx, node, out);
+        codegen_match_internal(ctx, node, out, 1);
         break;
     case NODE_EXPR_BINARY:
         if (strncmp(node->binary.op, "??", 2) == 0 && strlen(node->binary.op) == 2)
@@ -848,13 +854,13 @@ void codegen_expression(ParserContext *ctx, ASTNode *node, FILE *out)
     case NODE_TRY:
     {
         char *type_name = "Result";
-        if (node->try_stmt.expr->type_info && node->try_stmt.expr->type_info->name)
+        if (g_current_func_ret_type)
         {
-            type_name = node->try_stmt.expr->type_info->name;
+            type_name = g_current_func_ret_type;
         }
-        else if (node->try_stmt.expr->resolved_type)
+        else if (node->try_stmt.expr->type_info && node->try_stmt.expr->type_info->name)
         {
-            type_name = node->try_stmt.expr->resolved_type;
+            type_name = node->try_stmt.expr->type_info->name;
         }
 
         if (strcmp(type_name, "__auto_type") == 0 || strcmp(type_name, "unknown") == 0)
@@ -862,14 +868,57 @@ void codegen_expression(ParserContext *ctx, ASTNode *node, FILE *out)
             type_name = "Result";
         }
 
+        char *search_name = type_name;
+        if (strncmp(search_name, "struct ", 7) == 0)
+        {
+            search_name += 7;
+        }
+
+        int is_enum = 0;
+        StructRef *er = ctx->parsed_enums_list;
+        while (er)
+        {
+            if (er->node && er->node->type == NODE_ENUM &&
+                strcmp(er->node->enm.name, search_name) == 0)
+            {
+                is_enum = 1;
+                break;
+            }
+            er = er->next;
+        }
+        if (!is_enum)
+        {
+            ASTNode *ins = ctx->instantiated_structs;
+            while (ins)
+            {
+                if (ins->type == NODE_ENUM && strcmp(ins->enm.name, search_name) == 0)
+                {
+                    is_enum = 1;
+                    break;
+                }
+                ins = ins->next;
+            }
+        }
+
         fprintf(out, "({ ");
         emit_auto_type(ctx, node->try_stmt.expr, node->token, out);
         fprintf(out, " _try = ");
         codegen_expression(ctx, node->try_stmt.expr, out);
-        fprintf(out,
-                "; if (_try.tag == %s_Err_Tag) return (%s_Err(_try.data.Err)); "
-                "_try.data.Ok; })",
-                type_name, type_name);
+
+        if (is_enum)
+        {
+            fprintf(out,
+                    "; if (_try.tag == %s_Err_Tag) return (%s_Err(_try.data.Err)); "
+                    "_try.data.Ok; })",
+                    search_name, search_name);
+        }
+        else
+        {
+            fprintf(out,
+                    "; if (!_try.is_ok) return %s_Err(_try.err); "
+                    "_try.val; })",
+                    search_name);
+        }
         break;
     }
     case NODE_RAW_STMT:
@@ -1149,6 +1198,10 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
     }
     switch (node->type)
     {
+    case NODE_MATCH:
+        codegen_match_internal(ctx, node, out, 0); // 0 = statement context
+        fprintf(out, ";\n");
+        break;
     case NODE_FUNCTION:
         if (!node->func.body)
         {
@@ -1320,11 +1373,14 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         }
         fprintf(out, "%s %s(%s)\n", node->func.ret_type, node->func.name, node->func.args);
         fprintf(out, "{\n");
+        char *prev_ret = g_current_func_ret_type;
+        g_current_func_ret_type = node->func.ret_type;
         codegen_walker(ctx, node->func.body, out);
         for (int i = defer_count - 1; i >= 0; i--)
         {
             codegen_node_single(ctx, defer_stack[i], out);
         }
+        g_current_func_ret_type = prev_ret;
         fprintf(out, "}\n");
         break;
 
@@ -1500,8 +1556,6 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
             if (node->var_decl.init_expr)
             {
                 inferred = infer_type(ctx, node->var_decl.init_expr);
-                fprintf(stderr, "DEBUG: var '%s' inferred = '%s'\n", node->var_decl.name,
-                        inferred ? inferred : "(null)");
             }
 
             if (inferred && strcmp(inferred, "__auto_type") != 0)
@@ -1682,14 +1736,18 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
         fprintf(out, "for (");
         if (strstr(g_config.cc, "tcc"))
         {
-            fprintf(out, "__typeof__((%s)) %s = ", node->for_range.start, node->for_range.var_name);
+            fprintf(out, "__typeof__((");
+            codegen_expression(ctx, node->for_range.start, out);
+            fprintf(out, ")) %s = ", node->for_range.var_name);
         }
         else
         {
             fprintf(out, "__auto_type %s = ", node->for_range.var_name);
         }
-        fprintf(out, "%s; %s < %s; %s", node->for_range.start, node->for_range.var_name,
-                node->for_range.end, node->for_range.var_name);
+        codegen_expression(ctx, node->for_range.start, out);
+        fprintf(out, "; %s < ", node->for_range.var_name);
+        codegen_expression(ctx, node->for_range.end, out);
+        fprintf(out, "; %s", node->for_range.var_name);
         if (node->for_range.step)
         {
             fprintf(out, " += %s) ", node->for_range.step);
diff --git a/src/codegen/codegen_utils.c b/src/codegen/codegen_utils.c
index 5dcbf19..8969276 100644
--- a/src/codegen/codegen_utils.c
+++ b/src/codegen/codegen_utils.c
@@ -194,26 +194,38 @@ char *infer_type(ParserContext *ctx, ASTNode *node)
             }
         }
 
-        // Check if it's a function pointer or built-in registered as a symbol.
         if (node->call.callee->type == NODE_EXPR_VAR)
         {
             Symbol *sym = find_symbol_entry(ctx, node->call.callee->var_ref.name);
-            if (sym && sym->type_info)
+            if (sym && sym->type_info && sym->type_info->kind == TYPE_FUNCTION &&
+                sym->type_info->inner)
             {
-                // If it's a function pointer type, return return type.
-                if (sym->type_info->kind == TYPE_FUNCTION && sym->type_info->inner)
-                {
-                    return type_to_string(sym->type_info->inner);
-                }
-                if (sym->type_info->kind == TYPE_POINTER && sym->type_info->inner &&
-                    sym->type_info->inner->kind == TYPE_FUNCTION && sym->type_info->inner->inner)
+                return type_to_string(sym->type_info->inner);
+            }
+        }
+    }
+
+    if (node->type == NODE_TRY)
+    {
+        char *inner_type = infer_type(ctx, node->try_stmt.expr);
+        if (inner_type)
+        {
+            // Extract T from Result<T> or Option<T>
+            char *start = strchr(inner_type, '<');
+            if (start)
+            {
+                start++; // Skip <
+                char *end = strrchr(inner_type, '>');
+                if (end && end > start)
                 {
-                    return type_to_string(sym->type_info->inner->inner);
+                    int len = end - start;
+                    char *extracted = xmalloc(len + 1);
+                    strncpy(extracted, start, len);
+                    extracted[len] = 0;
+                    return extracted;
                 }
             }
         }
-
-        return NULL;
     }
 
     if (node->type == NODE_EXPR_MEMBER)
diff --git a/src/lexer/token.c b/src/lexer/token.c
index ebed001..6ada798 100644
--- a/src/lexer/token.c
+++ b/src/lexer/token.c
@@ -79,7 +79,7 @@ Token lexer_next(Lexer *l)
         }
         l->pos += len;
         l->col += len;
-        return (Token){TOK_COMMENT, s, len, start_line, start_col};
+        return lexer_next(l);
     }
 
     // Identifiers.
@@ -150,6 +150,14 @@ Token lexer_next(Lexer *l)
         {
             return (Token){TOK_AWAIT, s, 5, start_line, start_col};
         }
+        if (len == 3 && strncmp(s, "and", 3) == 0)
+        {
+            return (Token){TOK_AND, s, 3, start_line, start_col};
+        }
+        if (len == 2 && strncmp(s, "or", 2) == 0)
+        {
+            return (Token){TOK_OR, s, 2, start_line, start_col};
+        }
 
         // F-Strings
         if (len == 1 && s[0] == 'f' && s[1] == '"')
diff --git a/src/main.c b/src/main.c
index 602f06b..d23cd0b 100644
--- a/src/main.c
+++ b/src/main.c
@@ -192,6 +192,13 @@ int main(int argc, char **argv)
     // Initialize Plugin Manager
     zptr_plugin_mgr_init();
 
+    // Parse context init
+    ParserContext ctx;
+    memset(&ctx, 0, sizeof(ctx));
+
+    // Scan for build directives (e.g. //> link: -lm)
+    scan_build_directives(&ctx, src);
+
     // Register built-in plugins
     extern ZPlugin brainfuck_plugin;
     extern ZPlugin befunge_plugin;
@@ -210,9 +217,6 @@ int main(int argc, char **argv)
     Lexer l;
     lexer_init(&l, src);
 
-    // Parse
-    ParserContext ctx;
-    memset(&ctx, 0, sizeof(ctx));
     ctx.hoist_out = tmpfile(); // Temp file for plugin hoisting
     if (!ctx.hoist_out)
     {
@@ -262,8 +266,9 @@ int main(int argc, char **argv)
     // TCC-specific adjustments?
     // Already handled by user passing --cc tcc
 
-    snprintf(cmd, sizeof(cmd), "%s %s %s -o %s out.c -lm -lpthread -I./src", g_config.cc,
-             g_config.gcc_flags, g_config.is_freestanding ? "-ffreestanding" : "", outfile);
+    snprintf(cmd, sizeof(cmd), "%s %s %s %s %s -o %s out.c -lm -lpthread -I./src %s", g_config.cc,
+             g_config.gcc_flags, g_cflags, g_config.is_freestanding ? "-ffreestanding" : "", "",
+             outfile, g_link_flags);
 
     if (g_config.verbose)
     {
diff --git a/src/parser/parser_expr.c b/src/parser/parser_expr.c
index 469d623..b451a26 100644
--- a/src/parser/parser_expr.c
+++ b/src/parser/parser_expr.c
@@ -245,6 +245,16 @@ Precedence get_token_precedence(Token t)
         return PREC_OR;
     }
 
+    if (t.type == TOK_OR)
+    {
+        return PREC_OR;
+    }
+
+    if (t.type == TOK_AND)
+    {
+        return PREC_AND;
+    }
+
     if (t.type == TOK_QQ_EQ)
     {
         return PREC_ASSIGNMENT;
@@ -1247,6 +1257,7 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                             free(acc);
                             acc = tmp;
                         }
+
                         else
                         {
                             char *tmp = xmalloc(strlen(mod->base_name) + suffix.len + 2);
@@ -3096,6 +3107,14 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
         {
             bin->binary.op = xstrdup(">");
         }
+        else if (op.type == TOK_AND)
+        {
+            bin->binary.op = xstrdup("&&");
+        }
+        else if (op.type == TOK_OR)
+        {
+            bin->binary.op = xstrdup("||");
+        }
         else
         {
             bin->binary.op = token_strdup(op);
@@ -3465,6 +3484,27 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 char *t2 = type_to_string(rhs->type_info);
                 // Skip type check if either operand is void* (escape hatch type)
                 int skip_check = (strcmp(t1, "void*") == 0 || strcmp(t2, "void*") == 0);
+
+                // Allow comparing pointers/strings with integer literal 0 (NULL)
+                if (!skip_check)
+                {
+                    int lhs_is_ptr =
+                        (lhs->type_info->kind == TYPE_POINTER ||
+                         lhs->type_info->kind == TYPE_STRING || (t1 && strstr(t1, "*")));
+                    int rhs_is_ptr =
+                        (rhs->type_info->kind == TYPE_POINTER ||
+                         rhs->type_info->kind == TYPE_STRING || (t2 && strstr(t2, "*")));
+
+                    if (lhs_is_ptr && rhs->type == NODE_EXPR_LITERAL && rhs->literal.int_val == 0)
+                    {
+                        skip_check = 1;
+                    }
+                    if (rhs_is_ptr && lhs->type == NODE_EXPR_LITERAL && lhs->literal.int_val == 0)
+                    {
+                        skip_check = 1;
+                    }
+                }
+
                 if (!skip_check && !type_eq(lhs->type_info, rhs->type_info))
                 {
                     char msg[256];
@@ -3493,8 +3533,10 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                     if (strcmp(bin->binary.op, "+") == 0 || strcmp(bin->binary.op, "-") == 0)
                     {
                         int lhs_is_ptr = (lhs->type_info->kind == TYPE_POINTER ||
+                                          lhs->type_info->kind == TYPE_STRING ||
                                           (t1 && strstr(t1, "*") != NULL));
                         int rhs_is_ptr = (rhs->type_info->kind == TYPE_POINTER ||
+                                          rhs->type_info->kind == TYPE_STRING ||
                                           (t2 && strstr(t2, "*") != NULL));
                         int lhs_is_int =
                             (lhs->type_info->kind == TYPE_INT || lhs->type_info->kind == TYPE_I32 ||
diff --git a/src/parser/parser_stmt.c b/src/parser/parser_stmt.c
index 30c9e90..1b8637e 100644
--- a/src/parser/parser_stmt.c
+++ b/src/parser/parser_stmt.c
@@ -1168,10 +1168,6 @@ ASTNode *parse_var_decl(ParserContext *ctx, Lexer *l)
                 type_obj = type_new(TYPE_STRUCT);
                 type_obj->name = xstrdup(type);
             }
-
-            // fprintf(stderr, "DEBUG PVarDecl: Var '%s' inferred type '%s'
-            // (init->type_info present: %d)\n", name, type, init && init->type_info ?
-            // 1 : 0);
         }
     }
 
@@ -1553,21 +1549,18 @@ ASTNode *parse_for(ParserContext *ctx, Lexer *l)
 
         if (in_tok.type == TOK_IDENT && strncmp(in_tok.start, "in", 2) == 0)
         {
-            Token start_tok = lexer_next(l);
-            if (lexer_next(l).type == TOK_DOTDOT)
+            ASTNode *start_expr = parse_expression(ctx, l);
+            if (lexer_peek(l).type == TOK_DOTDOT)
             {
-                Token end_tok = lexer_next(l);
+                lexer_next(l); // consume ..
+                ASTNode *end_expr = parse_expression(ctx, l);
 
                 ASTNode *n = ast_create(NODE_FOR_RANGE);
                 n->for_range.var_name = xmalloc(var.len + 1);
                 strncpy(n->for_range.var_name, var.start, var.len);
                 n->for_range.var_name[var.len] = 0;
-                n->for_range.start = xmalloc(start_tok.len + 1);
-                strncpy(n->for_range.start, start_tok.start, start_tok.len);
-                n->for_range.start[start_tok.len] = 0;
-                n->for_range.end = xmalloc(end_tok.len + 1);
-                strncpy(n->for_range.end, end_tok.start, end_tok.len);
-                n->for_range.end[end_tok.len] = 0;
+                n->for_range.start = start_expr;
+                n->for_range.end = end_expr;
 
                 if (lexer_peek(l).type == TOK_IDENT && strncmp(lexer_peek(l).start, "step", 4) == 0)
                 {
@@ -2118,6 +2111,10 @@ ASTNode *parse_statement(ParserContext *ctx, Lexer *l)
             else if (next_type == TOK_DOTDOT)
             {
                 lexer_next(l); // consume ..
+                if (lexer_peek(l).type == TOK_SEMICOLON)
+                {
+                    lexer_next(l); // consume optional ;
+                }
             }
 
             ASTNode *n = ast_create(NODE_RAW_STMT);
diff --git a/src/parser/parser_utils.c b/src/parser/parser_utils.c
index fcaab7f..ef66104 100644
--- a/src/parser/parser_utils.c
+++ b/src/parser/parser_utils.c
@@ -2148,9 +2148,10 @@ char *rewrite_expr_methods(ParserContext *ctx, char *raw)
 
             dest -= strlen(acc);
 
-            if (*src == '(')
+            Module *mod = find_module(ctx, acc);
+            if (mod && mod->is_c_header)
             {
-                dest += sprintf(dest, "%s_%s", acc, field);
+                dest += sprintf(dest, "%s", field);
             }
             else
             {
diff --git a/src/zprep.h b/src/zprep.h
index 299ec13..17e7e02 100644
--- a/src/zprep.h
+++ b/src/zprep.h
@@ -64,6 +64,8 @@ typedef enum
     TOK_DCOLON,
     TOK_TRAIT,
     TOK_IMPL,
+    TOK_AND,
+    TOK_OR,
     TOK_FOR,
     TOK_COMPTIME,
     TOK_ELLIPSIS,
@@ -181,5 +183,10 @@ typedef struct
 } CompilerConfig;
 
 extern CompilerConfig g_config;
+extern char g_link_flags[];
+extern char g_cflags[];
+
+struct ParserContext;
+void scan_build_directives(struct ParserContext *ctx, const char *src);
 
 #endif