feat(tazjin/rlox): Implement early return from functions r/2108

In the book this is implemented via exceptions as control flow, and
I'm sticking somewhat closely to that by doing it via an error
variant.

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

3 files changed, 47 insertions, 1 deletions

diff --git a/users/tazjin/rlox/src/errors.rs b/users/tazjin/rlox/src/errors.rs
index 4f173d016039..875f268ca64d 100644
--- a/users/tazjin/rlox/src/errors.rs
+++ b/users/tazjin/rlox/src/errors.rs
@@ -1,3 +1,5 @@
+use crate::interpreter::Value;
+
 #[derive(Debug)]
 pub enum ErrorKind {
     UnexpectedChar(char),
@@ -12,6 +14,16 @@ pub enum ErrorKind {
     InternalError(String),
     InvalidAssignmentTarget(String),
     RuntimeError(String),
+
+    // This variant is not an error, rather it is used for
+    // short-circuiting out of a function body that hits a `return`
+    // statement.
+    //
+    // It's implemented this way because in the original book the
+    // author uses exceptions for control flow, and this is the
+    // closest equivalent that I had available without diverging too
+    // much.
+    FunctionReturn(Value),
 }
 
 #[derive(Debug)]
diff --git a/users/tazjin/rlox/src/interpreter.rs b/users/tazjin/rlox/src/interpreter.rs
index 5fb5283bc791..5fdde9adac68 100644
--- a/users/tazjin/rlox/src/interpreter.rs
+++ b/users/tazjin/rlox/src/interpreter.rs
@@ -40,7 +40,18 @@ impl Callable {
                     fn_env.define(param, value)?;
                 }
 
-                lox.interpret_block(Rc::new(RwLock::new(fn_env)), &func.body)
+                let result = lox.interpret_block(Rc::new(RwLock::new(fn_env)), &func.body);
+
+                match result {
+                    // extract returned values if applicable
+                    Err(Error {
+                        kind: ErrorKind::FunctionReturn(value),
+                        ..
+                    }) => Ok(value),
+
+                    // otherwise just return the result itself
+                    _ => result,
+                }
             }
         }
     }
@@ -221,6 +232,12 @@ impl Interpreter {
             Statement::If(if_stmt) => return self.interpret_if(if_stmt),
             Statement::While(while_stmt) => return self.interpret_while(while_stmt),
             Statement::Function(func) => return self.interpret_function(func.clone()),
+            Statement::Return(ret) => {
+                return Err(Error {
+                    line: 0,
+                    kind: ErrorKind::FunctionReturn(self.eval(&ret.value)?),
+                })
+            }
         };
 
         Ok(value)
diff --git a/users/tazjin/rlox/src/parser.rs b/users/tazjin/rlox/src/parser.rs
index ed5e670ecb79..495304686b2f 100644
--- a/users/tazjin/rlox/src/parser.rs
+++ b/users/tazjin/rlox/src/parser.rs
@@ -80,6 +80,11 @@ pub struct Var {
 }
 
 #[derive(Debug)]
+pub struct Return {
+    pub value: Expr,
+}
+
+#[derive(Debug)]
 pub struct If {
     pub condition: Expr,
     pub then_branch: Box<Statement>,
@@ -110,6 +115,7 @@ pub enum Statement {
     If(If),
     While(While),
     Function(Rc<Function>),
+    Return(Return),
 }
 
 // Parser
@@ -130,6 +136,7 @@ statement      → exprStmt
                | forStmt
                | ifStmt
                | printStmt
+               | returnStmt
                | whileStmt
                | block ;
 
@@ -137,6 +144,8 @@ forStmt        → "for" "(" ( varDecl | exprStmt | ";" )
                  expression? ";"
                  expression? ")" statement ;
 
+returnStmt     → "return" expression? ";" ;
+
 whileStmt      → "while" "(" expression ")" statement ;
 
 exprStmt       → expression ";" ;
@@ -256,6 +265,8 @@ impl Parser {
             self.while_statement()
         } else if self.match_token(&TokenKind::For) {
             self.for_statement()
+        } else if self.match_token(&TokenKind::Return) {
+            self.return_statement()
         } else {
             self.expr_statement()
         }
@@ -379,6 +390,12 @@ impl Parser {
         Ok(body)
     }
 
+    fn return_statement(&mut self) -> StmtResult {
+        let value = self.expression()?;
+        self.consume(&TokenKind::Semicolon, ErrorKind::ExpectedSemicolon)?;
+        Ok(Statement::Return(Return { value }))
+    }
+
     fn expr_statement(&mut self) -> StmtResult {
         let expr = self.expression()?;
         self.consume(&TokenKind::Semicolon, ErrorKind::ExpectedSemicolon)?;