From 0f074c5ebcd6a1ae9cb778d18be1771f56ae262c Mon Sep 17 00:00:00 2001
From: Alex van de Sandt <asandt@fieldwire.com>
Date: Fri, 7 Jun 2024 18:32:28 -0400
Subject: [PATCH] Implement variable assignment

---
 src/environment.rs |  8 ++++++++
 src/interpreter.rs | 42 +++++++++++++++++++++++++++++++++++++-----
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/environment.rs b/src/environment.rs
index c8788db..7c91a52 100644
--- a/src/environment.rs
+++ b/src/environment.rs
@@ -1,4 +1,5 @@
 use crate::value::Value;
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 
 pub struct Env(HashMap<String, Value>);
@@ -12,6 +13,13 @@ impl Env {
         self.0.insert(var, val);
     }
 
+    pub fn assign(&mut self, var: String, new_val: Value) -> bool {
+        matches!(
+            self.0.entry(var).and_modify(|entry| *entry = new_val),
+            Entry::Occupied(_)
+        )
+    }
+
     pub fn get(&self, var: &str) -> Option<Value> {
         self.0.get(var).cloned()
     }
diff --git a/src/interpreter.rs b/src/interpreter.rs
index 7510dc1..95326ad 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -2,6 +2,7 @@ use std::io::Write;
 
 use miette::SourceSpan;
 
+use crate::ast::Lval;
 use crate::{
     ast::{Ast, BinaryOp, Expr, Stmt, UnaryOp},
     environment::Env,
@@ -81,7 +82,7 @@ impl<Output: Write> Interpreter<Output> {
     }
 
     #[tracing::instrument(name = "expr", skip_all)]
-    fn interpret_expr(&self, expr: &Spanned<Expr>) -> Result<Value> {
+    fn interpret_expr(&mut self, expr: &Spanned<Expr>) -> Result<Value> {
         match expr.as_ref() {
             Expr::Binary { left, op, right } => self.interpret_binary_op(left, op, right),
             Expr::Grouping { inner } => self.interpret_grouping(inner),
@@ -94,13 +95,26 @@ impl<Output: Write> Interpreter<Output> {
                 .env
                 .get(name)
                 .ok_or_else(|| RuntimeError::undefined_var(name)),
-            Expr::Assignment { .. } => todo!("assignment interpreting"),
+            Expr::Assignment { target, value } => {
+                let name = match target.as_ref() {
+                    Lval::Ident { name } => name.clone(),
+                };
+                let value = self.interpret_expr(value)?;
+                if !self.env.assign(name.clone(), value.clone()) {
+                    Err(RuntimeError::undefined_var(&Spanned::new(
+                        name,
+                        target.span(),
+                    )))
+                } else {
+                    Ok(value)
+                }
+            }
         }
     }
 
     #[tracing::instrument(name = "binary", skip_all)]
     fn interpret_binary_op(
-        &self,
+        &mut self,
         left: &Spanned<Expr>,
         op: &Spanned<BinaryOp>,
         right: &Spanned<Expr>,
@@ -204,14 +218,14 @@ impl<Output: Write> Interpreter<Output> {
     }
 
     #[tracing::instrument(name = "grouping", skip_all)]
-    fn interpret_grouping(&self, inner: &Spanned<Expr>) -> Result<Value> {
+    fn interpret_grouping(&mut self, inner: &Spanned<Expr>) -> Result<Value> {
         tracing::trace!("descending to group");
         // TODO: investigate using more span features here
         self.interpret_expr(inner)
     }
 
     #[tracing::instrument(name = "unary", skip_all)]
-    fn interpret_unary(&self, op: &Spanned<UnaryOp>, expr: &Spanned<Expr>) -> Result<Value> {
+    fn interpret_unary(&mut self, op: &Spanned<UnaryOp>, expr: &Spanned<Expr>) -> Result<Value> {
         tracing::trace!("descending into unary");
         let val = self.interpret_expr(expr)?;
         match op.as_ref() {
@@ -334,6 +348,7 @@ mod test {
     fn executes_expr_stmt() {
         assert_eq!(execute_stmts("1 + 1;").1, "");
     }
+
     #[test]
     fn executes_print_stmt() {
         assert_eq!(execute_stmts("print 1;").1, "1\n");
@@ -357,6 +372,15 @@ mod test {
         assert_some_eq!(e.get("bar"), Value::Number(1.0));
     }
 
+    #[test]
+    fn assigns_var() {
+        let (e, _) = execute_stmts("var foo; foo = 1;");
+        assert_some_eq!(e.get("foo"), Value::Number(1.0));
+
+        let (e, _) = execute_stmts("var foo = true; foo = !foo;");
+        assert_some_eq!(e.get("foo"), Value::Bool(false));
+    }
+
     #[test]
     fn interprets_arithmetic() {
         for (input, expected_value) in [
@@ -462,6 +486,14 @@ mod test {
         );
     }
 
+    #[test]
+    fn errs_on_undefined_var_assignment() {
+        assert_matches!(
+            interpret_expr_to_err("foo = 1"),
+            RuntimeError::UndefinedVar { .. }
+        );
+    }
+
     #[test]
     fn errs_on_illegal_arithmetic() {
         for input in [
-- 
2.51.2