Implement while loop

froth · Mar 7, 2024 · a70dca9 · a70dca9
1 parent 3b8d8ba
commit a70dca9
Show file tree

Hide file tree

Showing 7 changed files with 117 additions and 34 deletions.
diff --git a/src/ast/stmt.rs b/src/ast/stmt.rs
@@ -55,6 +55,15 @@ impl Stmt {
             location,
         }
     }
+
+    pub fn while_stmt(condition: Expr, body: Stmt, location: SourceSpan) -> Self {
+        let src = condition.src.clone();
+        Stmt {
+            stmt_type: StmtType::While(condition, body.into()),
+            src,
+            location,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -64,6 +73,7 @@ pub enum StmtType {
     Var(Name, Option<Expr>),
     Block(Vec<Stmt>),
     If(Expr, Box<Stmt>, Option<Box<Stmt>>),
+    While(Expr, Box<Stmt>),
 }
 
 impl Display for Stmt {
@@ -91,6 +101,11 @@ impl Display for Stmt {
                 write!(f, "{}", then_branch)?;
                 writeln!(f, "endif")
             }
+            While(condition, body) => {
+                writeln!(f, "while {} {{", condition)?;
+                write!(f, "{}", body)?;
+                writeln!(f, "}}")
+            }
         }
     }
 }
diff --git a/src/interpreter/environment.rs b/src/interpreter/environment.rs
@@ -9,8 +9,8 @@ pub struct Environment {
 }
 
 impl Environment {
-    pub fn define(&mut self, key: Name, value: Value) {
-        self.values.insert(key, value);
+    pub fn define(&mut self, key: &Name, value: Value) {
+        self.values.insert(key.clone(), value);
     }
 
     pub fn get(&self, key: &Name) -> Option<Value> {
@@ -54,7 +54,7 @@ mod environment_tests {
     fn define_get() {
         let mut env = Environment::default();
         let name = Name::new("x".to_string());
-        env.define(name.clone(), Value::Boolean(true));
+        env.define(&name, Value::Boolean(true));
         let returned = env.get(&name);
         assert_eq!(returned, Some(Value::Boolean(true)))
     }
@@ -63,7 +63,7 @@ mod environment_tests {
     fn define_assign_get() {
         let mut env = Environment::default();
         let name = Name::new("x".to_string());
-        env.define(name.clone(), Value::Boolean(true));
+        env.define(&name, Value::Boolean(true));
         let assigned = env.assign(&name, &Value::Boolean(false));
         assert!(assigned);
         let returned = env.get(&name);
@@ -84,7 +84,7 @@ mod environment_tests {
     fn define_get_from_parent() {
         let mut global = Environment::default();
         let name = Name::new("x".to_string());
-        global.define(name.clone(), Value::Boolean(true));
+        global.define(&name, Value::Boolean(true));
         let local = local(global);
         let returned = local.get(&name);
         assert_eq!(returned, Some(Value::Boolean(true)))
@@ -94,7 +94,7 @@ mod environment_tests {
     fn define_assign_to_parent() {
         let mut global = Environment::default();
         let name = Name::new("x".to_string());
-        global.define(name.clone(), Value::Nil);
+        global.define(&name, Value::Nil);
         let mut local = local(global);
         let assigned = local.assign(&name, &Value::Boolean(false));
         assert!(assigned);
@@ -107,9 +107,9 @@ mod environment_tests {
     fn shadowing() {
         let mut global = Environment::default();
         let name = Name::new("x".to_string());
-        global.define(name.clone(), Value::Nil);
+        global.define(&name, Value::Nil);
         let mut local = local(global);
-        local.define(name.clone(), Value::Boolean(false));
+        local.define(&name, Value::Boolean(false));
         let loc_return = local.get(&name);
         global = *local.parent.unwrap();
         let glob_return = global.get(&name);

diff --git a/src/interpreter/expression.rs b/src/interpreter/expression.rs
@@ -356,7 +356,7 @@ mod value_interpreter_tests {
         let right = literal(false.into());
         let expr = Expr::assign(name_expr(name.clone()), right);
         let mut env = Environment::default();
-        env.define(name, Value::Nil);
+        env.define(&name, Value::Nil);
         let mut under_test = Interpreter::with_env(Box::new(VecPrinter::new()), env);
         assert_matches!(
             under_test.interpret_expr(&expr).unwrap(),

diff --git a/src/interpreter/mod.rs b/src/interpreter/mod.rs
@@ -20,9 +20,7 @@ pub struct Interpreter {
 
 impl Interpreter {
     pub fn interpret(&mut self, statements: Vec<Stmt>) -> Result<()> {
-        statements
-            .into_iter()
-            .try_for_each(|s| self.interpret_stmt(s))
+        statements.iter().try_for_each(|s| self.interpret_stmt(s))
     }
 
     fn push_environment(&mut self) {

diff --git a/src/interpreter/statement.rs b/src/interpreter/statement.rs
@@ -9,43 +9,51 @@ use crate::{
 use super::{Interpreter, Result};
 
 impl Interpreter {
-    pub(super) fn interpret_stmt(&mut self, statement: Stmt) -> Result<()> {
-        match statement.stmt_type {
-            Expression(expr) => self.interpret_expr(&expr).map(|_| ()),
+    pub(super) fn interpret_stmt(&mut self, statement: &Stmt) -> Result<()> {
+        match &statement.stmt_type {
+            Expression(expr) => self.interpret_expr(expr).map(|_| ()),
             Print(expr) => self
-                .interpret_expr(&expr)
+                .interpret_expr(expr)
                 .map(|value| self.printer.print(value)),
             Var(key, initializer) => self.define_var(key, initializer),
             Block(stmts) => self.execute_block(stmts),
-            If(condition, then_stmt, else_stmt) => {
-                self.execute_if(condition, *then_stmt, else_stmt.map(|x| *x))
-            }
+            If(condition, then_stmt, else_stmt) => self.execute_if(condition, then_stmt, else_stmt),
+            While(condition, body) => self.execute_while(condition, body.as_ref()),
         }
     }
 
-    fn define_var(&mut self, key: Name, initializer: Option<Expr>) -> Result<()> {
-        let initializer = initializer.map_or(Ok(Value::Nil), |expr| self.interpret_expr(&expr))?;
+    fn define_var(&mut self, key: &Name, initializer: &Option<Expr>) -> Result<()> {
+        let initializer = initializer
+            .as_ref()
+            .map_or(Ok(Value::Nil), |expr| self.interpret_expr(expr))?;
         self.environment.define(key, initializer);
         Ok(())
     }
 
-    fn execute_block(&mut self, stmts: Vec<Stmt>) -> Result<()> {
+    fn execute_block(&mut self, stmts: &[Stmt]) -> Result<()> {
         self.push_environment();
-        let result = stmts.into_iter().try_for_each(|s| self.interpret_stmt(s));
+        let result = stmts.iter().try_for_each(|s| self.interpret_stmt(s));
         self.pop_environment();
         result
     }
 
     fn execute_if(
         &mut self,
-        condition: Expr,
-        then_stmt: Stmt,
-        else_stmt: Option<Stmt>,
+        condition: &Expr,
+        then_stmt: &Stmt,
+        else_stmt: &Option<Box<Stmt>>,
     ) -> Result<()> {
-        if self.interpret_expr(&condition)?.is_truthy() {
+        if self.interpret_expr(condition)?.is_truthy() {
             self.interpret_stmt(then_stmt)?;
         } else if let Some(else_stmt) = else_stmt {
-            self.interpret_stmt(else_stmt)?;
+            self.interpret_stmt(else_stmt.as_ref())?;
+        }
+        Ok(())
+    }
+
+    fn execute_while(&mut self, condition: &Expr, body: &Stmt) -> Result<()> {
+        while self.interpret_expr(condition)?.is_truthy() {
+            self.interpret_stmt(body)?;
         }
         Ok(())
     }
@@ -72,7 +80,7 @@ mod stmt_interpreter_tests {
             (0, 1).into(),
         );
         let mut interpreter = Interpreter::new(Box::new(printer.clone()));
-        interpreter.interpret_stmt(stmt).unwrap();
+        interpreter.interpret_stmt(&stmt).unwrap();
         assert_eq!(printer.get_lines(), vec!["string".into()])
     }
 
@@ -81,12 +89,12 @@ mod stmt_interpreter_tests {
         let printer = VecPrinter::new();
         let stmt = block(vec![var("a")]);
         let mut interpreter = Interpreter::new(Box::new(printer.clone()));
-        interpreter.interpret_stmt(stmt).unwrap();
+        interpreter.interpret_stmt(&stmt).unwrap();
         let stmt = Stmt::expr(
             Expr::variable("a".to_string(), token(TokenType::Eof)),
             (0, 1).into(),
         );
-        let err = interpreter.interpret_stmt(stmt).unwrap_err();
+        let err = interpreter.interpret_stmt(&stmt).unwrap_err();
         assert_matches!(err, RuntimeError::UndefinedVariable { .. })
     }
 
@@ -99,12 +107,12 @@ mod stmt_interpreter_tests {
         );
         let stmt = block(vec![var("a"), read_undefined_var]);
         let mut interpreter = Interpreter::new(Box::new(printer.clone()));
-        let _ = interpreter.interpret_stmt(stmt).unwrap_err();
+        let _ = interpreter.interpret_stmt(&stmt).unwrap_err();
         let stmt = Stmt::expr(
             Expr::variable("a".to_string(), token(TokenType::Eof)),
             (0, 1).into(),
         );
-        let err = interpreter.interpret_stmt(stmt).unwrap_err();
+        let err = interpreter.interpret_stmt(&stmt).unwrap_err();
         assert_matches!(err, RuntimeError::UndefinedVariable { .. })
     }
 

diff --git a/src/parser/statement.rs b/src/parser/statement.rs
@@ -64,6 +64,7 @@ impl Parser {
             Print => self.print_statement(),
             LeftBrace => self.block_statement(),
             If => self.if_statement(),
+            While => self.while_statement(),
             _ => self.expression_statement(),
         }
     }
@@ -103,8 +104,12 @@ impl Parser {
             .is_some()
             .then(|| self.statement())
             .transpose()?;
-        let location = if_location.until(then_statement.location);
 
+        let end_location = else_statement
+            .as_ref()
+            .map(|s| s.location)
+            .unwrap_or(then_statement.location);
+        let location = if_location.until(end_location);
         Ok(Stmt::if_stmt(
             condition,
             then_statement,
@@ -113,6 +118,31 @@ impl Parser {
         ))
     }
 
+    fn while_statement(&mut self) -> Result<Stmt> {
+        use TokenType::*;
+        let while_location = self.advance().location;
+
+        consume!(self, LeftParen, |t: &Token| {
+            ExpectedLeftParen {
+                src: t.src.clone(),
+                location: self.previous_if_eof(t.location),
+            }
+        });
+
+        let condition = self.expression()?;
+
+        consume!(self, RightParen, |t: &Token| {
+            ExpectedRightParen {
+                src: t.src.clone(),
+                location: self.previous_if_eof(t.location),
+            }
+        });
+
+        let body = self.statement()?;
+        let location = while_location.until(body.location);
+        Ok(Stmt::while_stmt(condition, body, location))
+    }
+
     fn block(&mut self) -> Result<InternalBlock> {
         let left_brace_location = self.advance().location;
         let mut stmts = vec![];
@@ -350,4 +380,19 @@ mod tests {
             }
         )
     }
+
+    #[test]
+    fn parse_while() {
+        let tokens = vec![
+            token(TokenType::While),
+            token(TokenType::LeftParen),
+            token(TokenType::Nil),
+            token(TokenType::RightParen),
+            token(TokenType::Nil),
+            token(TokenType::Semicolon),
+            token(TokenType::Eof),
+        ];
+        let stmt = parse_stmt(tokens).unwrap();
+        assert_eq!(stmt.to_string().trim_end(), "while (nil) {\nExpr(nil)\n}")
+    }
 }
diff --git a/tests/while.lox b/tests/while.lox
@@ -0,0 +1,17 @@
+interpret
+var a = 0;
+var b = 0;
+while (a <= 10) {
+    if(b==2) {
+        b=0;
+        print a;
+    }
+    a = a+1;
+    b = b+1;
+}
+----
+2
+4
+6
+8
+10