From 01cc38f31d3a61c7470f028d903b5ff95386f556 Mon Sep 17 00:00:00 2001
From: Toni <antpiasecki@gmail.com>
Date: Thu, 29 May 2025 19:59:34 +0200
Subject: [PATCH] assignment

---
 src/codegen.rs | 24 +++++++++++++++++++-----
 src/parser.rs  | 25 ++++++++++++++++++++++++-
 test.mot       |  6 ++++++
 3 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 test.mot

diff --git a/src/codegen.rs b/src/codegen.rs
index bf3f788..46c3d4c 100644
--- a/src/codegen.rs
+++ b/src/codegen.rs
@@ -2,11 +2,11 @@ use std::{collections::HashMap, error::Error, fmt::Write};
 
 use crate::{
     parser::{Expr, Stmt},
-    tokenizer::TokenType,
+    tokenizer::{MotError, TokenType, error},
 };
 
 pub struct Env {
-    locals: HashMap<String, i32>,
+    locals: HashMap<String, usize>,
 }
 
 impl Env {
@@ -79,7 +79,7 @@ section .note.GNU-stack
             }
             Stmt::Var { name, initializer } => {
                 self.compile_expr(env, initializer)?;
-                let offset = (env.locals.len() as i32 + 1) * 8;
+                let offset = (env.locals.len() + 1) * 8;
                 env.locals.insert(name.lexeme, offset);
                 writeln!(&mut self.output, "    mov QWORD [rbp-{}], rax", offset)?;
             }
@@ -136,10 +136,24 @@ section .note.GNU-stack
                 }
             }
             Expr::Variable(name) => {
-                // TODO: handle error
-                let offset = env.locals.get(&name.lexeme).unwrap();
+                let offset = match env.locals.get(&name.lexeme) {
+                    Some(x) => x,
+                    None => {
+                        return error!(name.loc, format!("undefined variable: {}", &name.lexeme));
+                    }
+                };
                 writeln!(&mut self.output, "    mov rax, QWORD [rbp-{}]", offset)?
             }
+            Expr::Assign { name, value } => {
+                self.compile_expr(env, *value)?;
+                let offset = match env.locals.get(&name.lexeme) {
+                    Some(x) => x,
+                    None => {
+                        return error!(name.loc, format!("undefined variable: {}", &name.lexeme));
+                    }
+                };
+                writeln!(&mut self.output, "    mov QWORD [rbp-{}], rax", offset)?;
+            }
         }
         Ok(())
     }
diff --git a/src/parser.rs b/src/parser.rs
index c21523c..869436a 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -21,6 +21,10 @@ pub enum Expr {
         right: Box<Expr>,
     },
     Variable(Token),
+    Assign {
+        name: Token,
+        value: Box<Expr>,
+    },
 }
 
 pub struct Parser {
@@ -66,7 +70,26 @@ impl Parser {
     }
 
     fn expression(&mut self) -> Result<Expr, MotError> {
-        self.equality()
+        self.assignment()
+    }
+
+    fn assignment(&mut self) -> Result<Expr, MotError> {
+        let expr = self.equality()?;
+
+        if self.match_token(&[TokenType::Equal]) {
+            let equals = self.previous().clone();
+            let value = self.assignment()?;
+
+            return match expr {
+                Expr::Variable(name) => Ok(Expr::Assign {
+                    name,
+                    value: Box::new(value),
+                }),
+                _ => return error!(equals.loc, "invalid assignment target"),
+            };
+        }
+
+        Ok(expr)
     }
 
     fn equality(&mut self) -> Result<Expr, MotError> {
diff --git a/test.mot b/test.mot
new file mode 100644
index 0000000..4013d61
--- /dev/null
+++ b/test.mot
@@ -0,0 +1,6 @@
+let a = 5
+let b = 7
+print(b - a)
+a = 4
+print(b - a)
+print(b)
\ No newline at end of file