feat(tazjin/rlox): Implement mutable variable assignment r/2033

Change-Id: I56db10a5ac26958ae27a8d4c8fa7b8f8285bc7e1
Reviewed-on: https://cl.tvl.fyi/c/depot/+/2299
Reviewed-by: tazjin <mail@tazj.in>
Tested-by: BuildkiteCI

3 files changed, 74 insertions, 33 deletions

diff --git a/users/tazjin/rlox/src/errors.rs b/users/tazjin/rlox/src/errors.rs
index 2c29aa4c2f0a..162b80dc00d4 100644
--- a/users/tazjin/rlox/src/errors.rs
+++ b/users/tazjin/rlox/src/errors.rs
@@ -9,6 +9,7 @@ pub enum ErrorKind {
     TypeError(String),
     UndefinedVariable(String),
     InternalError(String),
+    InvalidAssignmentTarget(String),
 }
 
 #[derive(Debug)]
diff --git a/users/tazjin/rlox/src/interpreter.rs b/users/tazjin/rlox/src/interpreter.rs
index 0b936e89d2e1..0855068b6387 100644
--- a/users/tazjin/rlox/src/interpreter.rs
+++ b/users/tazjin/rlox/src/interpreter.rs
@@ -1,6 +1,6 @@
 use crate::errors::{Error, ErrorKind};
 use crate::parser::{self, Declaration, Expr, Literal, Program, Statement};
-use crate::scanner::TokenKind;
+use crate::scanner::{self, TokenKind};
 use std::collections::HashMap;
 
 // Tree-walk interpreter
@@ -11,24 +11,31 @@ struct Environment {
 }
 
 impl Environment {
-    fn define(&mut self, name: &str, value: Literal) {
-        self.values.insert(name.into(), value);
+    fn define(&mut self, name: &scanner::Token, value: Literal) -> Result<(), Error> {
+        let ident = identifier_str(name)?;
+        self.values.insert(ident.into(), value);
+        Ok(())
     }
 
     fn get(&self, name: &parser::Variable) -> Result<Literal, Error> {
-        if let TokenKind::Identifier(ident) = &name.0.kind {
-            return self
-                .values
-                .get(ident)
-                .map(Clone::clone)
-                .ok_or_else(|| Error {
-                    line: name.0.line,
-                    kind: ErrorKind::UndefinedVariable(ident.into()),
-                });
-        }
+        let ident = identifier_str(&name.0)?;
+
+        self.values
+            .get(ident)
+            .map(Clone::clone)
+            .ok_or_else(|| Error {
+                line: name.0.line,
+                kind: ErrorKind::UndefinedVariable(ident.into()),
+            })
+    }
+}
 
+fn identifier_str<'a>(name: &'a scanner::Token) -> Result<&'a str, Error> {
+    if let TokenKind::Identifier(ident) = &name.kind {
+        Ok(ident)
+    } else {
         Err(Error {
-            line: name.0.line,
+            line: name.line,
             kind: ErrorKind::InternalError("unexpected identifier kind".into()),
         })
     }
@@ -40,7 +47,7 @@ pub struct Interpreter {
 }
 
 impl Interpreter {
-    fn interpret_stmt<'a>(&self, stmt: &Statement<'a>) -> Result<(), Error> {
+    fn interpret_stmt<'a>(&mut self, stmt: &Statement<'a>) -> Result<(), Error> {
         match stmt {
             Statement::Expr(expr) => {
                 self.eval(expr)?;
@@ -55,20 +62,13 @@ impl Interpreter {
     }
 
     fn interpret_var<'a>(&mut self, var: &parser::Var<'a>) -> Result<(), Error> {
-        if let TokenKind::Identifier(ident) = &var.name.kind {
-            let init = var.initialiser.as_ref().ok_or_else(|| Error {
-                line: var.name.line,
-                kind: ErrorKind::InternalError("missing variable initialiser".into()),
-            })?;
-
-            self.globals.define(ident, self.eval(init)?);
-            return Ok(());
-        }
-
-        Err(Error {
+        let init = var.initialiser.as_ref().ok_or_else(|| Error {
             line: var.name.line,
-            kind: ErrorKind::InternalError("unexpected identifier kind".into()),
-        })
+            kind: ErrorKind::InternalError("missing variable initialiser".into()),
+        })?;
+        let value = self.eval(init)?;
+        self.globals.define(&var.name, value)?;
+        return Ok(());
     }
 
     pub fn interpret<'a>(&mut self, program: &Program<'a>) -> Result<(), Error> {
@@ -82,8 +82,9 @@ impl Interpreter {
         Ok(())
     }
 
-    fn eval<'a>(&self, expr: &Expr<'a>) -> Result<Literal, Error> {
+    fn eval<'a>(&mut self, expr: &Expr<'a>) -> Result<Literal, Error> {
         match expr {
+            Expr::Assign(assign) => self.eval_assign(assign),
             Expr::Literal(lit) => Ok(lit.clone()),
             Expr::Grouping(grouping) => self.eval(&*grouping.0),
             Expr::Unary(unary) => self.eval_unary(unary),
@@ -92,7 +93,7 @@ impl Interpreter {
         }
     }
 
-    fn eval_unary<'a>(&self, expr: &parser::Unary<'a>) -> Result<Literal, Error> {
+    fn eval_unary<'a>(&mut self, expr: &parser::Unary<'a>) -> Result<Literal, Error> {
         let right = self.eval(&*expr.right)?;
 
         match (&expr.operator.kind, right) {
@@ -109,7 +110,7 @@ impl Interpreter {
         }
     }
 
-    fn eval_binary<'a>(&self, expr: &parser::Binary<'a>) -> Result<Literal, Error> {
+    fn eval_binary<'a>(&mut self, expr: &parser::Binary<'a>) -> Result<Literal, Error> {
         let left = self.eval(&*expr.left)?;
         let right = self.eval(&*expr.right)?;
 
@@ -152,6 +153,12 @@ impl Interpreter {
 
         Ok(result)
     }
+
+    fn eval_assign<'a>(&mut self, assign: &parser::Assign<'a>) -> Result<Literal, Error> {
+        let value = self.eval(&assign.value)?;
+        self.globals.define(&assign.name, value.clone())?;
+        Ok(value)
+    }
 }
 
 // Interpreter functions not dependent on interpreter-state.
diff --git a/users/tazjin/rlox/src/parser.rs b/users/tazjin/rlox/src/parser.rs
index 3f1dda6e39cf..b46e570994c4 100644
--- a/users/tazjin/rlox/src/parser.rs
+++ b/users/tazjin/rlox/src/parser.rs
@@ -11,6 +11,12 @@ use crate::scanner::{Token, TokenKind};
 // AST
 
 #[derive(Debug)]
+pub struct Assign<'a> {
+    pub name: Token<'a>,
+    pub value: Box<Expr<'a>>,
+}
+
+#[derive(Debug)]
 pub struct Binary<'a> {
     pub left: Box<Expr<'a>>,
     pub operator: Token<'a>,
@@ -34,11 +40,13 @@ pub struct Unary<'a> {
     pub right: Box<Expr<'a>>,
 }
 
+// Not to be confused with `Var`, which is for assignment.
 #[derive(Debug)]
 pub struct Variable<'a>(pub Token<'a>);
 
 #[derive(Debug)]
 pub enum Expr<'a> {
+    Assign(Assign<'a>),
     Binary(Binary<'a>),
     Grouping(Grouping<'a>),
     Literal(Literal),
@@ -81,7 +89,9 @@ statement      → exprStmt
 exprStmt       → expression ";" ;
 printStmt      → "print" expression ";" ;
 
-expression     → equality ;
+expression     → assignment ;
+assignment     → IDENTIFIER "=" assignment
+               | equality ;
 equality       → comparison ( ( "!=" | "==" ) comparison )* ;
 comparison     → term ( ( ">" | ">=" | "<" | "<=" ) term )* ;
 term           → factor ( ( "-" | "+" ) factor )* ;
@@ -156,7 +166,30 @@ impl<'a> Parser<'a> {
     }
 
     fn expression(&mut self) -> ExprResult<'a> {
-        self.equality()
+        self.assignment()
+    }
+
+    fn assignment(&mut self) -> ExprResult<'a> {
+        let expr = self.equality()?;
+
+        if self.match_token(&[TokenKind::Equal]) {
+            let equals = self.previous().clone();
+            let value = self.assignment()?;
+
+            if let Expr::Variable(Variable(name)) = expr {
+                return Ok(Expr::Assign(Assign {
+                    name,
+                    value: Box::new(value),
+                }));
+            }
+
+            return Err(Error {
+                line: equals.line,
+                kind: ErrorKind::InvalidAssignmentTarget(format!("{:?}", equals)),
+            });
+        }
+
+        Ok(expr)
     }
 
     fn equality(&mut self) -> ExprResult<'a> {