Merge branch 'main' into pr-109

1 files changed, 318 insertions, 47 deletions

diff --git a/src/parser/parser_expr.c b/src/parser/parser_expr.c
index 98b1c5d..da5379e 100644
--- a/src/parser/parser_expr.c
+++ b/src/parser/parser_expr.c
@@ -203,7 +203,7 @@ static void check_format_string(ASTNode *call, Token t)
         return;
     }
 
-    if (fmt_arg->type != NODE_EXPR_LITERAL || fmt_arg->literal.type_kind != TOK_STRING)
+    if (fmt_arg->type != NODE_EXPR_LITERAL || fmt_arg->literal.type_kind != 2)
     {
         return;
     }
@@ -605,11 +605,17 @@ static void find_declared_vars(ASTNode *node, char ***decls, int *count)
         (*count)++;
     }
 
-    if (node->type == NODE_MATCH_CASE && node->match_case.binding_name)
+    if (node->type == NODE_MATCH_CASE)
     {
-        *decls = xrealloc(*decls, sizeof(char *) * (*count + 1));
-        (*decls)[*count] = xstrdup(node->match_case.binding_name);
-        (*count)++;
+        if (node->match_case.binding_names)
+        {
+            for (int i = 0; i < node->match_case.binding_count; i++)
+            {
+                *decls = xrealloc(*decls, sizeof(char *) * (*count + 1));
+                (*decls)[*count] = xstrdup(node->match_case.binding_names[i]);
+                (*count)++;
+            }
+        }
     }
 
     switch (node->type)
@@ -783,6 +789,8 @@ ASTNode *parse_lambda(ParserContext *ctx, Lexer *l)
 
     lexer_next(l);
 
+    Type *t = type_new(TYPE_FUNCTION);
+    t->args = xmalloc(sizeof(Type *) * 16);
     char **param_names = xmalloc(sizeof(char *) * 16);
     char **param_types = xmalloc(sizeof(char *) * 16);
     int num_params = 0;
@@ -814,9 +822,11 @@ ASTNode *parse_lambda(ParserContext *ctx, Lexer *l)
 
         lexer_next(l);
 
-        char *param_type_str = parse_type(ctx, l);
-        param_types[num_params] = param_type_str;
+        Type *typef = parse_type_formal(ctx, l);
+        t->args[t->arg_count] = typef;
+        param_types[num_params] = type_to_string(typef);
         num_params++;
+        t->arg_count = num_params;
     }
     lexer_next(l);
 
@@ -824,7 +834,16 @@ ASTNode *parse_lambda(ParserContext *ctx, Lexer *l)
     if (lexer_peek(l).type == TOK_ARROW)
     {
         lexer_next(l);
-        return_type = parse_type(ctx, l);
+
+        t->inner = parse_type_formal(ctx, l);
+        return_type = type_to_string(t->inner);
+    }
+
+    enter_scope(ctx);
+
+    for (int i = 0; i < num_params; i++)
+    {
+        add_symbol(ctx, param_names[i], param_types[i], t->args[i]);
     }
 
     ASTNode *body = NULL;
@@ -845,9 +864,13 @@ ASTNode *parse_lambda(ParserContext *ctx, Lexer *l)
     lambda->lambda.num_params = num_params;
     lambda->lambda.lambda_id = ctx->lambda_counter++;
     lambda->lambda.is_expression = 0;
+    lambda->type_info = t;
+    lambda->resolved_type = type_to_string(t);
     register_lambda(ctx, lambda);
     analyze_lambda_captures(ctx, lambda);
 
+    exit_scope(ctx);
+
     return lambda;
 }
 
@@ -991,7 +1014,7 @@ static ASTNode *create_fstring_block(ParserContext *ctx, const char *content)
             char *fmt_str = xmalloc(strlen(fmt) + 3);
             sprintf(fmt_str, "%%%s", fmt);
             arg_fmt = ast_create(NODE_EXPR_LITERAL);
-            arg_fmt->literal.type_kind = 2;
+            arg_fmt->literal.type_kind = LITERAL_STRING;
             arg_fmt->literal.string_val = fmt_str;
             arg_fmt->type_info = type_new(TYPE_STRING);
         }
@@ -1045,7 +1068,7 @@ static ASTNode *create_fstring_block(ParserContext *ctx, const char *content)
 static ASTNode *parse_int_literal(Token t)
 {
     ASTNode *node = ast_create(NODE_EXPR_LITERAL);
-    node->literal.type_kind = 0;
+    node->literal.type_kind = LITERAL_INT;
     node->type_info = type_new(TYPE_INT);
     char *s = token_strdup(t);
     unsigned long long val;
@@ -1066,7 +1089,7 @@ static ASTNode *parse_int_literal(Token t)
 static ASTNode *parse_float_literal(Token t)
 {
     ASTNode *node = ast_create(NODE_EXPR_LITERAL);
-    node->literal.type_kind = 1;
+    node->literal.type_kind = LITERAL_FLOAT;
     node->literal.float_val = atof(t.start);
     node->type_info = type_new(TYPE_F64);
     return node;
@@ -1076,7 +1099,7 @@ static ASTNode *parse_float_literal(Token t)
 static ASTNode *parse_string_literal(Token t)
 {
     ASTNode *node = ast_create(NODE_EXPR_LITERAL);
-    node->literal.type_kind = TOK_STRING;
+    node->literal.type_kind = LITERAL_STRING;
     node->literal.string_val = xmalloc(t.len);
     strncpy(node->literal.string_val, t.start + 1, t.len - 2);
     node->literal.string_val[t.len - 2] = 0;
@@ -1099,7 +1122,7 @@ static ASTNode *parse_fstring_literal(ParserContext *ctx, Token t)
 static ASTNode *parse_char_literal(Token t)
 {
     ASTNode *node = ast_create(NODE_EXPR_LITERAL);
-    node->literal.type_kind = TOK_CHAR;
+    node->literal.type_kind = LITERAL_CHAR;
     node->literal.string_val = token_strdup(t);
     node->type_info = type_new(TYPE_I8);
     return node;
@@ -1310,18 +1333,47 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
             char *pattern = token_strdup(p);
             int is_default = (strcmp(pattern, "_") == 0);
 
-            char *binding = NULL;
-            int is_destructure = 0;
-            skip_comments(l);
+            // Handle Destructuring: Ok(v) or Rect(w, h)
+            char **bindings = NULL;
+            int *binding_refs = NULL;
+            int binding_count = 0;
+            int is_destructure = 0; // Initialize here
+
+            // Assuming pattern_count is 1 for now, or needs to be determined
+            // For single identifier patterns, pattern_count would be 1.
+            // This logic needs to be adjusted if `pattern_count` is not available or needs to be
+            // calculated. For now, assuming `pattern_count == 1` is implicitly true for single
+            // token patterns.
             if (!is_default && lexer_peek(l).type == TOK_LPAREN)
             {
-                lexer_next(l);
-                Token b = lexer_next(l);
-                if (b.type != TOK_IDENT)
+                lexer_next(l);                           // eat (
+                bindings = xmalloc(sizeof(char *) * 8);  // Initial capacity
+                binding_refs = xmalloc(sizeof(int) * 8); // unused but consistent
+
+                while (1)
                 {
-                    zpanic_at(b, "Expected binding name");
+                    int is_r = 0;
+                    if (lexer_peek(l).type == TOK_IDENT && lexer_peek(l).len == 3 &&
+                        strncmp(lexer_peek(l).start, "ref", 3) == 0)
+                    {
+                        lexer_next(l); // eat ref
+                        is_r = 1;
+                    }
+                    Token b = lexer_next(l);
+                    if (b.type != TOK_IDENT)
+                    {
+                        zpanic_at(b, "Expected binding");
+                    }
+                    bindings[binding_count] = token_strdup(b);
+                    binding_refs[binding_count] = is_r;
+                    binding_count++;
+                    if (lexer_peek(l).type == TOK_COMMA)
+                    {
+                        lexer_next(l);
+                        continue;
+                    }
+                    break;
                 }
-                binding = token_strdup(b);
                 if (lexer_next(l).type != TOK_RPAREN)
                 {
                     zpanic_at(lexer_peek(l), "Expected )");
@@ -1343,6 +1395,17 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                 zpanic_at(lexer_peek(l), "Expected '=>'");
             }
 
+            // Create scope for the case to hold the binding
+            enter_scope(ctx);
+            if (binding_count > 0)
+            {
+                for (int i = 0; i < binding_count; i++)
+                {
+                    add_symbol(ctx, bindings[i], NULL,
+                               NULL); // Let inference handle it or default to void*?
+                }
+            }
+
             ASTNode *body;
             skip_comments(l);
             Token pk = lexer_peek(l);
@@ -1364,9 +1427,13 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                 body = parse_expression(ctx, l);
             }
 
+            exit_scope(ctx);
+
             ASTNode *c = ast_create(NODE_MATCH_CASE);
             c->match_case.pattern = pattern;
-            c->match_case.binding_name = binding;
+            c->match_case.binding_names = bindings;      // New multi-binding field
+            c->match_case.binding_count = binding_count; // New binding count field
+            c->match_case.binding_refs = binding_refs;
             c->match_case.is_destructuring = is_destructure;
             c->match_case.guard = guard;
             c->match_case.body = body;
@@ -1579,13 +1646,59 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                                                 v = v->next;
                                             }
                                         }
+                                        int resolved = 0;
                                         if (is_variant)
                                         {
                                             sprintf(tmp, "%s_%.*s", acc, suffix.len, suffix.start);
+                                            resolved = 1;
                                         }
                                         else
                                         {
-                                            sprintf(tmp, "%s__%.*s", acc, suffix.len, suffix.start);
+                                            char inherent_name[256];
+                                            sprintf(inherent_name, "%s__%.*s", acc, suffix.len,
+                                                    suffix.start);
+
+                                            if (find_func(ctx, inherent_name))
+                                            {
+                                                strcpy(tmp, inherent_name);
+                                                resolved = 1;
+                                            }
+                                            else
+                                            {
+                                                GenericImplTemplate *it = ctx->impl_templates;
+                                                while (it)
+                                                {
+                                                    if (strcmp(it->struct_name, gname) == 0)
+                                                    {
+                                                        char *tname = NULL;
+                                                        if (it->impl_node &&
+                                                            it->impl_node->type == NODE_IMPL_TRAIT)
+                                                        {
+                                                            tname = it->impl_node->impl_trait
+                                                                        .trait_name;
+                                                        }
+                                                        if (tname)
+                                                        {
+                                                            char cand[512];
+                                                            sprintf(cand, "%s__%s_%.*s", acc, tname,
+                                                                    suffix.len, suffix.start);
+
+                                                            if (find_func(ctx, cand))
+                                                            {
+                                                                strcpy(tmp, cand);
+                                                                resolved = 1;
+                                                                break;
+                                                            }
+                                                        }
+                                                    }
+                                                    it = it->next;
+                                                }
+                                            }
+                                            if (!resolved)
+                                            {
+                                                sprintf(tmp, "%s__%.*s", acc, suffix.len,
+                                                        suffix.start);
+                                            }
                                         }
                                         handled_as_generic = 1;
                                         break; // Found and handled
@@ -2046,7 +2159,7 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                 node->type_info = type_new(TYPE_INT); // Returns count
 
                 ASTNode *fmt_node = ast_create(NODE_EXPR_LITERAL);
-                fmt_node->literal.type_kind = 2; // string
+                fmt_node->literal.type_kind = LITERAL_STRING; // string
                 fmt_node->literal.string_val = xstrdup(fmt);
 
                 ASTNode *head = fmt_node, *tail = fmt_node;
@@ -2540,7 +2653,7 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                 // Constant Folding for 'def', emits literal
                 node = ast_create(NODE_EXPR_LITERAL);
                 node->token = t;
-                node->literal.type_kind = 0; // INT (assumed for now from const_int_val)
+                node->literal.type_kind = LITERAL_INT; // INT (assumed for now from const_int_val)
                 node->literal.int_val = sym->const_int_val;
                 node->type_info = type_new(TYPE_INT);
                 // No need for resolution
@@ -2740,15 +2853,15 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
                 if (elements[i]->type == NODE_EXPR_LITERAL)
                 {
                     char buf[256];
-                    if (elements[i]->literal.type_kind == 0) // int
+                    if (elements[i]->literal.type_kind == LITERAL_INT) // int
                     {
                         sprintf(buf, "%lld", elements[i]->literal.int_val);
                     }
-                    else if (elements[i]->literal.type_kind == 1) // float
+                    else if (elements[i]->literal.type_kind == LITERAL_FLOAT) // float
                     {
                         sprintf(buf, "%f", elements[i]->literal.float_val);
                     }
-                    else if (elements[i]->literal.type_kind == 2) // string
+                    else if (elements[i]->literal.type_kind == LITERAL_STRING) // string
                     {
                         sprintf(buf, "\"%s\"", elements[i]->literal.string_val);
                     }
@@ -2969,7 +3082,7 @@ ASTNode *parse_primary(ParserContext *ctx, Lexer *l)
             if (node->type_info && node->type_info->kind == TYPE_ARRAY &&
                 node->type_info->array_size > 0)
             {
-                if (index->type == NODE_EXPR_LITERAL && index->literal.type_kind == 0)
+                if (index->type == NODE_EXPR_LITERAL && index->literal.type_kind == LITERAL_INT)
                 {
                     int idx = index->literal.int_val;
                     if (idx < 0 || idx >= node->type_info->array_size)
@@ -3403,7 +3516,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
             }
 
             // Reuse printf sugar to generate the prompt print
-            char *print_code = process_printf_sugar(ctx, inner, 0, "stdout", NULL, NULL);
+            char *print_code = process_printf_sugar(ctx, inner, 0, "stdout", NULL, NULL, 1);
             free(inner);
 
             // Checks for (args...) suffix for SCAN mode
@@ -3498,7 +3611,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 call->type_info = type_new(TYPE_INT);
 
                 ASTNode *fmt_node = ast_create(NODE_EXPR_LITERAL);
-                fmt_node->literal.type_kind = TOK_STRING;
+                fmt_node->literal.type_kind = LITERAL_STRING;
                 fmt_node->literal.string_val = xstrdup(fmt);
                 ASTNode *head = fmt_node, *tail = fmt_node;
 
@@ -3532,7 +3645,10 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 free(print_code);
 
                 ASTNode *n = ast_create(NODE_RAW_STMT);
-                n->raw_stmt.content = final_code;
+                char *stmt_code = xmalloc(strlen(final_code) + 2);
+                sprintf(stmt_code, "%s;", final_code);
+                free(final_code);
+                n->raw_stmt.content = stmt_code;
                 return n;
             }
         }
@@ -3568,11 +3684,14 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 newline = 0;
             }
 
-            char *code = process_printf_sugar(ctx, inner, newline, "stderr", NULL, NULL);
+            char *code = process_printf_sugar(ctx, inner, newline, "stderr", NULL, NULL, 1);
             free(inner);
 
             ASTNode *n = ast_create(NODE_RAW_STMT);
-            n->raw_stmt.content = code;
+            char *stmt_code = xmalloc(strlen(code) + 2);
+            sprintf(stmt_code, "%s;", code);
+            free(code);
+            n->raw_stmt.content = stmt_code;
             return n;
         }
     }
@@ -4441,7 +4560,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                 if (lhs->type_info && lhs->type_info->kind == TYPE_ARRAY &&
                     lhs->type_info->array_size > 0)
                 {
-                    if (start->type == NODE_EXPR_LITERAL && start->literal.type_kind == 0)
+                    if (start->type == NODE_EXPR_LITERAL && start->literal.type_kind == LITERAL_INT)
                     {
                         int idx = start->literal.int_val;
                         if (idx < 0 || idx >= lhs->type_info->array_size)
@@ -4564,7 +4683,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                                     if (find_func(ctx, trait_mangled))
                                     {
                                         strcpy(mangled, trait_mangled); // Update mangled name
-                                        sig = find_func(ctx, mangled);
+                                        sig = find_func(ctx, trait_mangled);
                                         break;
                                     }
                                 }
@@ -4819,7 +4938,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
 
         if (strcmp(bin->binary.op, "/") == 0 || strcmp(bin->binary.op, "%") == 0)
         {
-            if (rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == 0 &&
+            if (rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == LITERAL_INT &&
                 rhs->literal.int_val == 0)
             {
                 warn_division_by_zero(op);
@@ -4845,8 +4964,8 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                     }
                 }
             }
-            else if (lhs->type == NODE_EXPR_LITERAL && lhs->literal.type_kind == 0 &&
-                     rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == 0)
+            else if (lhs->type == NODE_EXPR_LITERAL && lhs->literal.type_kind == LITERAL_INT &&
+                     rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == LITERAL_INT)
             {
                 // Check if literals make sense (e.g. 5 > 5)
                 if (lhs->literal.int_val == rhs->literal.int_val)
@@ -4865,7 +4984,7 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
 
             if (lhs->type_info && type_is_unsigned(lhs->type_info))
             {
-                if (rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == 0 &&
+                if (rhs->type == NODE_EXPR_LITERAL && rhs->literal.type_kind == LITERAL_INT &&
                     rhs->literal.int_val == 0)
                 {
                     if (strcmp(bin->binary.op, ">=") == 0)
@@ -5245,6 +5364,63 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
         // Standard Type Checking (if no overload found)
         if (lhs->type_info && rhs->type_info)
         {
+            // Ensure type_info is set for variables (critical for inference)
+            if (lhs->type == NODE_EXPR_VAR && !lhs->type_info)
+            {
+                Symbol *s = find_symbol_entry(ctx, lhs->var_ref.name);
+                if (s)
+                {
+                    lhs->type_info = s->type_info;
+                }
+            }
+            if (rhs->type == NODE_EXPR_VAR && !rhs->type_info)
+            {
+                Symbol *s = find_symbol_entry(ctx, rhs->var_ref.name);
+                if (s)
+                {
+                    rhs->type_info = s->type_info;
+                }
+            }
+
+            // Backward Inference for Lambda Params
+            // LHS is Unknown Var, RHS is Known
+            if (lhs->type == NODE_EXPR_VAR && lhs->type_info &&
+                lhs->type_info->kind == TYPE_UNKNOWN && rhs->type_info &&
+                rhs->type_info->kind != TYPE_UNKNOWN)
+            {
+                // Infer LHS type from RHS
+                Symbol *sym = find_symbol_entry(ctx, lhs->var_ref.name);
+                if (sym)
+                {
+                    // Update Symbol
+                    sym->type_info = rhs->type_info;
+                    sym->type_name = type_to_string(rhs->type_info);
+
+                    // Update AST Node
+                    lhs->type_info = rhs->type_info;
+                    lhs->resolved_type = xstrdup(sym->type_name);
+                }
+            }
+
+            // RHS is Unknown Var, LHS is Known
+            if (rhs->type == NODE_EXPR_VAR && rhs->type_info &&
+                rhs->type_info->kind == TYPE_UNKNOWN && lhs->type_info &&
+                lhs->type_info->kind != TYPE_UNKNOWN)
+            {
+                // Infer RHS type from LHS
+                Symbol *sym = find_symbol_entry(ctx, rhs->var_ref.name);
+                if (sym)
+                {
+                    // Update Symbol
+                    sym->type_info = lhs->type_info;
+                    sym->type_name = type_to_string(lhs->type_info);
+
+                    // Update AST Node
+                    rhs->type_info = lhs->type_info;
+                    rhs->resolved_type = xstrdup(sym->type_name);
+                }
+            }
+
             if (is_comparison_op(bin->binary.op))
             {
                 bin->type_info = type_new(TYPE_INT); // bool
@@ -5366,6 +5542,51 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
 
                     if (!is_ptr_arith && !alias_match)
                     {
+                        // ** Backward Inference for Binary Ops **
+                        // Case 1: LHS is Unknown Var, RHS is Known
+                        if (lhs->type == NODE_EXPR_VAR && lhs->type_info &&
+                            lhs->type_info->kind == TYPE_UNKNOWN && rhs->type_info &&
+                            rhs->type_info->kind != TYPE_UNKNOWN)
+                        {
+                            // Infer LHS type from RHS
+                            Symbol *sym = find_symbol_entry(ctx, lhs->var_ref.name);
+                            if (sym)
+                            {
+                                // Update Symbol
+                                sym->type_info = rhs->type_info;
+                                sym->type_name = type_to_string(rhs->type_info);
+
+                                // Update AST Node
+                                lhs->type_info = rhs->type_info;
+                                lhs->resolved_type = xstrdup(sym->type_name);
+
+                                bin->type_info = rhs->type_info;
+                                goto bin_inference_success;
+                            }
+                        }
+
+                        // Case 2: RHS is Unknown Var, LHS is Known
+                        if (rhs->type == NODE_EXPR_VAR && rhs->type_info &&
+                            rhs->type_info->kind == TYPE_UNKNOWN && lhs->type_info &&
+                            lhs->type_info->kind != TYPE_UNKNOWN)
+                        {
+                            // Infer RHS type from LHS
+                            Symbol *sym = find_symbol_entry(ctx, rhs->var_ref.name);
+                            if (sym)
+                            {
+                                // Update Symbol
+                                sym->type_info = lhs->type_info;
+                                sym->type_name = type_to_string(lhs->type_info);
+
+                                // Update AST Node
+                                rhs->type_info = lhs->type_info;
+                                rhs->resolved_type = xstrdup(sym->type_name);
+
+                                bin->type_info = lhs->type_info;
+                                goto bin_inference_success;
+                            }
+                        }
+
                         // Allow assigning 0 to pointer (NULL)
                         int is_null_assign = 0;
                         if (strcmp(bin->binary.op, "=") == 0)
@@ -5452,6 +5673,8 @@ ASTNode *parse_expr_prec(ParserContext *ctx, Lexer *l, Precedence min_prec)
                                 zpanic_with_suggestion(op, msg, suggestion);
                             }
                         }
+
+                    bin_inference_success:;
                     }
                 }
             }
@@ -5469,21 +5692,21 @@ ASTNode *parse_arrow_lambda_single(ParserContext *ctx, Lexer *l, char *param_nam
     lambda->lambda.param_names[0] = param_name;
     lambda->lambda.num_params = 1;
 
-    // Default param type: int
+    // Default param type: unknown (to be inferred)
     lambda->lambda.param_types = xmalloc(sizeof(char *));
-    lambda->lambda.param_types[0] = xstrdup("int");
+    lambda->lambda.param_types[0] = NULL;
 
-    // Create Type Info: int -> int
+    // Create Type Info: unknown -> unknown
     Type *t = type_new(TYPE_FUNCTION);
-    t->inner = type_new(TYPE_INT); // Return
+    t->inner = type_new(TYPE_INT); // Return (default to int)
     t->args = xmalloc(sizeof(Type *));
-    t->args[0] = type_new(TYPE_INT); // Arg
+    t->args[0] = type_new(TYPE_UNKNOWN); // Arg
     t->arg_count = 1;
     lambda->type_info = t;
 
     // Register parameter in scope for body parsing
     enter_scope(ctx);
-    add_symbol(ctx, param_name, "int", type_new(TYPE_INT));
+    add_symbol(ctx, param_name, NULL, t->args[0]);
 
     // Body parsing...
     ASTNode *body_block = NULL;
@@ -5500,7 +5723,55 @@ ASTNode *parse_arrow_lambda_single(ParserContext *ctx, Lexer *l, char *param_nam
         body_block->block.statements = ret;
     }
     lambda->lambda.body = body_block;
-    lambda->lambda.return_type = xstrdup("int");
+
+    // Attempt to infer return type from body if it's a simple return
+    if (lambda->lambda.body->block.statements &&
+        lambda->lambda.body->block.statements->type == NODE_RETURN &&
+        !lambda->lambda.body->block.statements->next)
+    {
+        ASTNode *ret_val = lambda->lambda.body->block.statements->ret.value;
+        if (ret_val->type_info && ret_val->type_info->kind != TYPE_UNKNOWN)
+        {
+            if (param_name[0] == 'x')
+            {
+                // fprintf(stderr, "DEBUG: Updating return type to %d\n", ret_val->type_info->kind);
+            }
+            // Update return type
+            if (t->inner)
+            {
+                free(t->inner);
+            }
+            t->inner = ret_val->type_info;
+        }
+    }
+
+    // Update parameter types from symbol table (in case inference happened)
+    Symbol *sym = find_symbol_entry(ctx, param_name);
+    if (sym && sym->type_info && sym->type_info->kind != TYPE_UNKNOWN)
+    {
+        free(lambda->lambda.param_types[0]);
+        lambda->lambda.param_types[0] = type_to_string(sym->type_info);
+        t->args[0] = sym->type_info;
+    }
+    else
+    {
+        // Fallback to int if still unknown
+        if (lambda->lambda.param_types[0])
+        {
+            free(lambda->lambda.param_types[0]);
+        }
+        lambda->lambda.param_types[0] = xstrdup("int");
+        t->args[0] = type_new(TYPE_INT); // FIX: Update AST type info too!
+
+        // Update symbol to match fallback
+        if (sym)
+        {
+            sym->type_name = xstrdup("int");
+            sym->type_info = type_new(TYPE_INT);
+        }
+    }
+
+    lambda->lambda.return_type = type_to_string(t->inner);
     lambda->lambda.lambda_id = ctx->lambda_counter++;
     lambda->lambda.is_expression = 1;
     register_lambda(ctx, lambda);