From f98ca8075d8128893eaa389d9890f055bb93b809 Mon Sep 17 00:00:00 2001
From: Toni <antpiasecki@gmail.com>
Date: Sat, 31 May 2025 16:12:18 +0200
Subject: [PATCH] parse return types, disallow nested functions

---
 examples/fib.zr       |  2 +-
 src/codegen.rs        |  4 ++--
 src/codegen_x86_64.rs | 14 ++++++++++----
 src/parser.rs         | 33 +++++++++++++++++++++++++++++----
 4 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/examples/fib.zr b/examples/fib.zr
index 37344da..a199d24 100644
--- a/examples/fib.zr
+++ b/examples/fib.zr
@@ -1,4 +1,4 @@
-func main[]
+func main[] : I64
     let a: I64 = 0
     let b: I64 = 1
 
diff --git a/src/codegen.rs b/src/codegen.rs
index 778a5ac..73111a1 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -16,7 +16,7 @@ impl Env {
     pub fn new() -> Env {
         Env {
             scopes: vec![HashMap::new()],
-            next_offset: 1,
+            next_offset: 8,
         }
     }
 
@@ -30,7 +30,7 @@ impl Env {
 
     pub fn define_var(&mut self, name: String, var_type: String) -> usize {
         let offset = self.next_offset;
-        self.next_offset += 1;
+        self.next_offset += 8;
         self.scopes.last_mut().unwrap().insert(name, Var {
             var_type,
             stack_offset: offset,
diff --git a/src/codegen_x86_64.rs b/src/codegen_x86_64.rs
index 7a4e719..8b4539c 100644
--- a/src/codegen_x86_64.rs
+++ b/src/codegen_x86_64.rs
@@ -72,7 +72,7 @@ print:
 
                 self.compile_expr(env, initializer)?;
                 let offset = env.define_var(name.lexeme.clone(), var_type.lexeme);
-                writeln!(&mut self.output, "    mov QWORD [rbp-{}], rax", offset * 8)?;
+                writeln!(&mut self.output, "    mov QWORD [rbp-{}], rax", offset)?;
             }
             Stmt::Block(statements) => {
                 env.push_scope();
@@ -110,8 +110,14 @@ print:
                 writeln!(&mut self.output, "    jmp {}", begin_label)?;
                 writeln!(&mut self.output, "{}:", end_label)?;
             }
-            Stmt::Function { name, params, body } => {
+            Stmt::Function {
+                name,
+                params,
+                return_type,
+                body,
+            } => {
                 assert!(params.is_empty()); // TODO
+                assert!(return_type.lexeme == "I64");
 
                 writeln!(&mut self.output, "global {}", name.lexeme)?;
                 writeln!(&mut self.output, "{}:", name.lexeme)?;
@@ -222,7 +228,7 @@ print:
                 writeln!(
                     &mut self.output,
                     "    mov rax, QWORD [rbp-{}]",
-                    var.stack_offset * 8
+                    var.stack_offset
                 )?
             }
             Expr::Assign { name, value } => {
@@ -237,7 +243,7 @@ print:
                 writeln!(
                     &mut self.output,
                     "    mov QWORD [rbp-{}], rax",
-                    var.stack_offset * 8
+                    var.stack_offset
                 )?;
             }
             Expr::Call {
diff --git a/src/parser.rs b/src/parser.rs
index 1d1e879..5f4b406 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -29,6 +29,7 @@ pub enum Stmt {
     Function {
         name: Token,
         params: Vec<Param>,
+        return_type: Token,
         body: Box<Stmt>,
     },
     Return(Expr),
@@ -62,11 +63,16 @@ pub enum Expr {
 pub struct Parser {
     tokens: Vec<Token>,
     current: usize,
+    is_inside_function: bool,
 }
 
 impl Parser {
     pub fn new(tokens: Vec<Token>) -> Parser {
-        Parser { tokens, current: 0 }
+        Parser {
+            tokens,
+            current: 0,
+            is_inside_function: false,
+        }
     }
 
     pub fn parse(mut self) -> Result<Vec<Stmt>, Box<dyn Error>> {
@@ -78,11 +84,19 @@ impl Parser {
     }
 
     fn declaration(&mut self) -> Result<Stmt, Box<dyn Error>> {
+        if !self.is_inside_function {
+            if self.match_token(&[TokenType::KeywordFunc]) {
+                return self.func_declaration();
+            }
+            return error!(
+                self.peek().loc,
+                "statements not allowed outside function body"
+            );
+        }
+
         // TODO: synchronization after parse error
         if self.match_token(&[TokenType::KeywordLet]) {
             self.let_declaration()
-        } else if self.match_token(&[TokenType::KeywordFunc]) {
-            self.func_declaration()
         } else {
             self.statement()
         }
@@ -107,8 +121,19 @@ impl Parser {
         }
 
         self.consume(TokenType::RightBracket, "expected ']' after arguments")?;
+        self.consume(TokenType::Colon, "expected ':' after '['")?;
+        let return_type = self.consume(TokenType::Identifier, "expected return type")?;
+
+        self.is_inside_function = true;
         let body = Box::new(self.block()?);
-        Ok(Stmt::Function { name, params, body })
+        self.is_inside_function = false;
+
+        Ok(Stmt::Function {
+            name,
+            params,
+            return_type,
+            body,
+        })
     }
 
     fn let_declaration(&mut self) -> Result<Stmt, Box<dyn Error>> {