From 08769c7a6d4195fcd3402a9e7b07542e6e4d405f Mon Sep 17 00:00:00 2001
From: Judah Caruso <judahrcr@gmail.com>
Date: Tue, 3 Jun 2025 21:43:11 -0600
Subject: [PATCH] [vm] return statements and start of procedure parsing

---
 vm/module.jai |   3 ++
 vm/parser.jai | 132 ++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 131 insertions(+), 4 deletions(-)

diff --git a/vm/module.jai b/vm/module.jai
index 267b041..4140cfb 100644
--- a/vm/module.jai
+++ b/vm/module.jai
@@ -21,6 +21,9 @@ strings :: #import "String"; // @future
    init(*parser, context.allocator);
 
    ok := parse_string(*parser, #string END
+      fn add(x, y) do return x + y end
+      fn sub(x, y) do return x - y end
+
       var x = 10.0
       var y = 20.0
       var z = x + y * 2.0 / 3.0
diff --git a/vm/parser.jai b/vm/parser.jai
index 98f11a3..f0b01c7 100644
--- a/vm/parser.jai
+++ b/vm/parser.jai
@@ -13,6 +13,7 @@ Token :: struct {
       kw_var;
       kw_def;
       kw_type;
+      kw_fn;
       kw_do;
       kw_end;
       kw_if;
@@ -61,12 +62,16 @@ Node :: struct {
    Kind :: enum {
       invalid;
 
+      block;
+
       stmt_start;
          print;
+         return_;
       stmt_end;
 
       decl_start;
          variable;
+         procedure;
       decl_end;
 
       expr_start;
@@ -184,6 +189,50 @@ Node_Type :: struct {
    }
 }
 
+/*
+fn add(x int, y int) int do return x + y end
+*/
+
+Node_Procedure :: struct {
+   #as using n: Node;
+   n.kind = .procedure;
+
+   header: *Node_Procedure_Header;
+
+   args:  [..]Node_Argument;
+   rets:  [..]Node_Argument;
+   body:  *Node_Block;
+   flags: Flag;
+
+   Flag :: enum_flags {
+      inline_;
+   }
+}
+
+Node_Argument :: struct {
+   symbol: *Node_Symbol;
+   type:   *Node_Type;
+   value:  *Node; // always an expression, can be null
+}
+
+Node_Procedure_Header :: struct {
+   args: [..]*Node_Type;
+   rets: [..]*Node_Type;
+}
+
+Node_Block :: struct {
+   #as using n: Node;
+   n.kind = .block;
+   body: [..]*Node;
+}
+
+Node_Return :: struct {
+   #as using n: Node;
+   n.kind = .return_;
+
+   values: [..]*Node;
+}
+
 Parser :: struct {
    allocator: Allocator;
    toplevel:  [..]*Node;
@@ -210,9 +259,7 @@ parse_string :: (p: *Parser, source: string) -> bool {
       }
 
       node := parse_toplevel(p);
-      if node == null break;
-
-      array.append(*p.toplevel, node);
+      if node != null array.append(*p.toplevel, node);
    }
 
    return false;
@@ -222,7 +269,7 @@ parse_string :: (p: *Parser, source: string) -> bool {
 #scope_file;
 
 parse_toplevel :: (p: *Parser) -> *Node {
-   t, ok := expect_token(p, .kw_var, .kw_def, .kw_print);
+   t, ok := expect_token(p, .kw_var, .kw_def, .kw_fn, .kw_print, .kw_do, .kw_return);
    basic.assert(ok, "var, def, print, found '%'", t.str); // @errors
 
    if t.kind == {
@@ -269,6 +316,22 @@ parse_toplevel :: (p: *Parser) -> *Node {
 
          return node;
 
+      // return
+      // return expr0, ..exprN
+      case .kw_return;
+         node := make_node(p, Node_Return);
+         array.init(*node.values, p.allocator);
+
+         prev_offset := p.offset;
+         expr := parse_expression(p);
+         if expr == null {
+            p.offset = prev_offset;
+            return node;
+         }
+
+         array.append(*node.values, expr);
+         return node;
+
       // print(expr)
       // print expr
       case .kw_print;
@@ -278,11 +341,71 @@ parse_toplevel :: (p: *Parser) -> *Node {
          node := make_node(p, Node_Print);
          node.expr = expr;
          return node;
+
+      // fn symbol(arg0, ..argN) do end
+      case .kw_fn;
+         symbol, ok := expect_token(p, .symbol);
+         basic.assert(ok, "expected name for procedure"); // @errors @todo(judah): lambdas
+
+         t, ok = expect_token(p, .l_paren);
+         basic.assert(ok, "expected '(' but found '%'", t.str); // @errors
+
+         while !at_end(p) {
+            t = peek_token(p);
+            if t.kind == .r_paren break;
+
+            expr := parse_expression(p);
+            basic.assert(expr != null); // @errors
+
+            t = peek_token(p);
+            if t.kind == {
+               case .comma;
+                  consume_token(p);
+                  continue;
+               case .r_paren;
+                  break;
+               case;
+                  basic.assert(false, "expected ',' or ')' but found '%'", t.str);
+            }
+         }
+
+         _, ok = expect_token(p, .r_paren);
+         basic.assert(ok, "expected ')'"); // @errors
+
+         block := parse_block(p);
+         basic.assert(block != null, "expected block"); // @errors
+
+         node := make_node(p, Node_Procedure);
+
+         return null;
    }
 
    return null;
 }
 
+parse_block :: (p: *Parser) -> *Node_Block {
+   t, ok := expect_token(p, .kw_do);
+   basic.assert(ok, "expected 'do' found '%'", t.str); // @errors
+
+   block := make_node(p, Node_Block);
+   array.init(*block.body, p.allocator);
+
+   while !at_end(p) {
+      t = peek_token(p);
+      if t.kind == .kw_end break;
+
+      node := parse_toplevel(p);
+      basic.assert(node != null); // @errors
+
+      array.append(*block.body, node);
+   }
+
+   t, ok = expect_token(p, .kw_end);
+   basic.assert(ok, "expected 'end' found '%'", t.str); // @errors
+
+   return block;
+}
+
 parse_type_expression :: (p: *Parser) -> *Node_Type {
    t, ok := expect_token(p, .symbol, .star, .l_square);
    basic.assert(ok, "type expression"); // @errors
@@ -596,6 +719,7 @@ consume_token :: (p: *Parser) -> Token {
          case "var";      t.kind = .kw_var;
          case "def";      t.kind = .kw_def;
          case "type";     t.kind = .kw_type;
+         case "fn";       t.kind = .kw_fn;
          case "do";       t.kind = .kw_do;
          case "end";      t.kind = .kw_end;
          case "if";       t.kind = .kw_if;