Add block expressions and if expressions

If and block statements are now expressions. The last expression, if
any, is used as the return value of this expression.

Also fixed a (major) bug in the if statement generation that was causing
the wrong jump address to be generated.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

src/ast.rs

@@ -1,5 +1,5 @@
 // This is an auto-generated file. Any changes made to this file may be overwritten.
-// This file was created at: 2024-10-16 10:33:48
+// This file was created at: 2024-10-21 15:09:43
 #![allow(dead_code)]
 use std::fmt::Debug;
 use std::any::Any;
@@ -18,6 +18,8 @@ pub trait ExprVisitor {
     fn visit_function_expr(&mut self, expr: &FunctionExpr) -> Result<(), Error>;
     fn visit_list_expr(&mut self, expr: &ListExpr) -> Result<(), Error>;
     fn visit_map_expr(&mut self, expr: &MapExpr) -> Result<(), Error>;
+    fn visit_if_expr(&mut self, expr: &IfExpr) -> Result<(), Error>;
+    fn visit_block_expr(&mut self, expr: &BlockExpr) -> Result<(), Error>;
 }
 
 pub trait Expr: Debug + Any {
@@ -170,15 +172,47 @@ impl Expr for MapExpr {
     fn as_any_ref(&self) -> &dyn Any { self }
 }
 
+#[derive(Debug)]
+pub struct IfExpr {
+    pub if_kw: Token,
+    pub condition: ExprP,
+    pub then_branch: BlockExpr,
+    pub else_branch: Option<BlockExpr>,
+}
+
+impl Expr for IfExpr {
+    fn accept(&self, visitor: &mut dyn ExprVisitor) -> Result<(), Error> {
+        visitor.visit_if_expr(self)
+    }
+    
+    fn as_any(self: Box<Self>) -> Box<dyn Any> { self }
+    fn as_any_ref(&self) -> &dyn Any { self }
+}
+
+#[derive(Debug)]
+pub struct BlockExpr {
+    pub lbrace: Token,
+    pub stmts: Vec<StmtP>,
+    pub return_expr: Option<ExprP>,
+    pub rbrace: Token,
+}
+
+impl Expr for BlockExpr {
+    fn accept(&self, visitor: &mut dyn ExprVisitor) -> Result<(), Error> {
+        visitor.visit_block_expr(self)
+    }
+    
+    fn as_any(self: Box<Self>) -> Box<dyn Any> { self }
+    fn as_any_ref(&self) -> &dyn Any { self }
+}
+
 pub trait StmtVisitor {
     fn visit_import_stmt(&mut self, stmt: &ImportStmt) -> Result<(), Error>;
     fn visit_expr_stmt(&mut self, stmt: &ExprStmt) -> Result<(), Error>;
     fn visit_assign_stmt(&mut self, stmt: &AssignStmt) -> Result<(), Error>;
     fn visit_index_assign_stmt(&mut self, stmt: &IndexAssignStmt) -> Result<(), Error>;
     fn visit_set_stmt(&mut self, stmt: &SetStmt) -> Result<(), Error>;
-    fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<(), Error>;
     fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<(), Error>;
-    fn visit_if_stmt(&mut self, stmt: &IfStmt) -> Result<(), Error>;
 }
 
 pub trait Stmt: Debug + Any {
@@ -269,22 +303,6 @@ impl Stmt for SetStmt {
     fn as_any_ref(&self) -> &dyn Any { self }
 }
 
-#[derive(Debug)]
-pub struct BlockStmt {
-    pub lbrace: Token,
-    pub stmts: Vec<StmtP>,
-    pub rbrace: Token,
-}
-
-impl Stmt for BlockStmt {
-    fn accept(&self, visitor: &mut dyn StmtVisitor) -> Result<(), Error> {
-        visitor.visit_block_stmt(self)
-    }
-    
-    fn as_any(self: Box<Self>) -> Box<dyn Any> { self }
-    fn as_any_ref(&self) -> &dyn Any { self }
-}
-
 #[derive(Debug)]
 pub struct ReturnStmt {
     pub return_kw: Token,
@@ -296,23 +314,6 @@ impl Stmt for ReturnStmt {
         visitor.visit_return_stmt(self)
     }
     
-    fn as_any(self: Box<Self>) -> Box<dyn Any> { self }
-    fn as_any_ref(&self) -> &dyn Any { self }
-}
-
-#[derive(Debug)]
-pub struct IfStmt {
-    pub if_kw: Token,
-    pub condition: ExprP,
-    pub then_branch: BlockStmt,
-    pub else_branch: Vec<StmtP>,
-}
-
-impl Stmt for IfStmt {
-    fn accept(&self, visitor: &mut dyn StmtVisitor) -> Result<(), Error> {
-        visitor.visit_if_stmt(self)
-    }
-    
     fn as_any(self: Box<Self>) -> Box<dyn Any> { self }
     fn as_any_ref(&self) -> &dyn Any { self }
 }

src/compiler.rs

@@ -747,15 +747,6 @@ impl StmtVisitor for Compiler<'_> {
         Ok(())
     }
 
-    fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<()> {
-        self.begin_scope(ScopeKind::Local);
-        for s in &stmt.stmts {
-            self.compile_stmt(s)?;
-        }
-        self.end_scope((stmt.rbrace.line, stmt.rbrace.line));
-        Ok(())
-    }
-
     fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<()> {
         if let Some(expr) = &stmt.expr {
             self.compile_expr(expr)?;
@@ -766,55 +757,6 @@ impl StmtVisitor for Compiler<'_> {
         self.emit(stmt_line_number(stmt), Op::Return);
         Ok(())
     }
-    fn visit_if_stmt(&mut self, stmt: &IfStmt) -> Result<()> {
-        // condition
-        self.compile_expr(&stmt.condition)?;
-        // call obj.to_bool()
-        let bool_attr = self.insert_constant(Str::create("to_bool"))?;
-        self.emit(expr_line_number(&*stmt.condition), Op::GetAttr(bool_attr));
-        self.emit(expr_line_number(&*stmt.condition), Op::Call(0));
-        let condition_patch_index = self.chunk().code.len();
-        self.emit(expr_line_number(&*stmt.condition), Op::JumpFalse(0));
-
-        // then branch
-        // pop the condition on top of the stack (no jump taken)
-        self.emit(expr_line_number(&*stmt.condition), Op::Pop);
-        // not using  compile_stmt because then_branch isn't a pointer, it's an honest-to-goodness
-        // value
-        stmt.then_branch.accept(self)?;
-        let exit_patch_index = self.chunk().code.len();
-        self.emit(stmt_line_number(&stmt.then_branch), Op::Jump(0));
-
-        // else branch
-        // patch the condition index - this is where the JUMP_FALSE will jump to
-        assert_matches!(self.chunk().code[condition_patch_index], Op::JumpFalse(_));
-        let offset = self.chunk().code.len() - condition_patch_index;
-        assert!(
-            offset <= (JumpOpArg::MAX as usize),
-            "jump offset too large between lines {:?} - this is a compiler limitation, sorry",
-            stmt_line_number(&stmt.then_branch)
-        );
-        self.chunk_mut().code[condition_patch_index] = Op::JumpFalse(offset as JumpOpArg);
-
-        // pop the condition on top of the stack (jump taken)
-        self.emit(expr_line_number(&*stmt.condition), Op::Pop);
-        for s in &stmt.else_branch {
-            self.compile_stmt(s)?;
-        }
-
-        // patch the "then" branch exit jump address - this is where Op::Jump will jump to.
-        // TODO : see if we can eliminate duplicates by checking the last two instructions
-        assert_matches!(self.chunk().code[exit_patch_index], Op::Jump(_));
-        let offset = self.chunk().code.len() - condition_patch_index;
-        assert!(
-            offset <= (JumpOpArg::MAX as usize),
-            "jump offset too large between lines {:?} - this is a compiler limitation, sorry",
-            stmt_line_number(&stmt.then_branch)
-        );
-        self.chunk_mut().code[exit_patch_index] = Op::Jump(offset as JumpOpArg);
-
-        Ok(())
-    }
 }
 
 impl ExprVisitor for Compiler<'_> {
@@ -1071,4 +1013,132 @@ impl ExprVisitor for Compiler<'_> {
 
         Ok(())
     }
+
+    fn visit_if_expr(&mut self, expr: &IfExpr) -> Result<()> {
+        // condition
+        self.compile_expr(&expr.condition)?;
+
+        // call obj.to_bool()
+        let bool_attr = self.insert_constant(Str::create("to_bool"))?;
+        self.emit(expr_line_number(&*expr.condition), Op::GetAttr(bool_attr));
+        self.emit(expr_line_number(&*expr.condition), Op::Call(0));
+        let condition_patch_index = self.chunk().code.len();
+        self.emit(expr_line_number(&*expr.condition), Op::JumpFalse(0));
+
+        // then branch
+        // pop the condition on top of the stack (no jump taken)
+        self.emit(expr_line_number(&*expr.condition), Op::Pop);
+        // not using compile_expr because then_branch isn't a pointer, it's an honest-to-goodness
+        // value
+        expr.then_branch.accept(self)?;
+        let exit_patch_index = self.chunk().code.len();
+        self.emit(expr_line_number(&expr.then_branch), Op::Jump(0));
+
+        // else branch
+        // patch the condition index - this is where the JUMP_FALSE will jump to
+        assert_matches!(self.chunk().code[condition_patch_index], Op::JumpFalse(_));
+        let offset = self.chunk().code.len() - condition_patch_index;
+        assert!(
+            offset <= (JumpOpArg::MAX as usize),
+            "jump offset too large between lines {:?} - this is a compiler limitation, sorry",
+            expr_line_number(&expr.then_branch)
+        );
+        self.chunk_mut().code[condition_patch_index] = Op::JumpFalse(offset as JumpOpArg);
+
+        // pop the condition on top of the stack (jump taken)
+        self.emit(expr_line_number(&*expr.condition), Op::Pop);
+
+        if let Some(branch) = expr.else_branch.as_ref() {
+            branch.accept(self)?;
+        } else {
+            // push a nil value as the expression return value
+            let line = expr.then_branch.rbrace.line;
+            let nil_constant = self.insert_constant(Nil::create())?;
+            self.emit((line, line), Op::PushConstant(nil_constant));
+        }
+
+        // patch the "then" branch exit jump address - this is where Op::Jump will jump to.
+        // TODO : see if we can eliminate duplicates by checking the last two instructions
+        assert_matches!(self.chunk().code[exit_patch_index], Op::Jump(_));
+        let offset = self.chunk().code.len() - exit_patch_index;
+        assert!(
+            offset <= (JumpOpArg::MAX as usize),
+            "jump offset too large between lines {:?} - this is a compiler limitation, sorry",
+            expr_line_number(&expr.then_branch)
+        );
+        self.chunk_mut().code[exit_patch_index] = Op::Jump(offset as JumpOpArg);
+
+        Ok(())
+    }
+
+    fn visit_block_expr(&mut self, expr: &BlockExpr) -> Result<()> {
+        // Block expr scope problem:
+        //
+        // Let's say we have an expression that looks like this:
+        // foo = 1 + {
+        //     a = 2
+        //     a * 3  # this value is what gets returned by the block
+        // }
+        //
+        // We need some way to "return" the 'a * 3' value. Just leaving the value on top of the
+        // stack will not work - when the scope exits, it will try to pop the 'a' value off the
+        // stack (when in actuality it's the 'a * 3' value) and we will have a misaligned stack.
+        //
+        // Instead, we can create a "dummy" scope that only will contain the "return" value. Then,
+        // we take the last expression and assign it to that "return" value (or nil if there isn't
+        // any). We push another new local scope (call this the "real" scope), and compile the
+        // statements in there. At the very end, we fill in the "return" local inside of the dummy
+        // scope, and then exit the "real" scope like normal. Then we can just pop the dummy scope
+        // frame (without calling end_scope(), because that would pop our return value) and then
+        // leave that return value on top of the stack.
+        //
+        // # enter "dummy" scope
+        // PUSH_CONSTANT nil
+        // # enter "real" scope
+        // PUSH_CONSTANT 2
+        // # do 'a * 3'
+        // GET_LOCAL "a"
+        // GET_ATTR "__mul__"
+        // PUSH_CONSTANT 3
+        // CALL 1
+        // # fill in the value in the "dummy" scope
+        // SET_LOCAL "<block result value>"
+        // # exit "real" scope
+        // POP
+        // # exit "dummy" scope (don't pop anything)
+
+        // create a "dummy" scope, and the slot for storing the result value
+        self.begin_scope(ScopeKind::Local);
+        let nil_constant = self.insert_constant(Nil::create())?;
+        self.emit(
+            (expr.lbrace.line, expr.lbrace.line),
+            Op::PushConstant(nil_constant),
+        );
+        let result_local = self
+            .insert_local("<block result value>".to_string())?
+            .clone();
+
+        // create the "real" scope
+        self.begin_scope(ScopeKind::Local);
+        for s in &expr.stmts {
+            self.compile_stmt(s)?;
+        }
+
+        // compile the last expression, if any.
+        // if there is not an expression, then there will just be a "nil" value in the local slot.
+        if let Some(expr) = &expr.return_expr {
+            self.compile_expr(expr)?;
+            self.emit(
+                expr_line_number(expr.as_ref()),
+                Op::SetLocal(result_local.index),
+            );
+        }
+
+        // exit "real" scope
+        self.end_scope((expr.rbrace.line, expr.rbrace.line));
+
+        // dummy value is left on top of the stack
+        self.scopes.pop().expect("no scope");
+        Ok(())
+    }
 }

src/compiler/visitors.rs

@@ -59,11 +59,6 @@ impl StmtVisitor for LineNumber {
         Ok(())
     }
 
-    fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<()> {
-        self.update_start(stmt.lbrace.line);
-        self.update_end(stmt.rbrace.line);
-        Ok(())
-    }
     fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<()> {
         self.update_start(stmt.return_kw.line);
         self.update_end(stmt.return_kw.line);
@@ -72,15 +67,6 @@ impl StmtVisitor for LineNumber {
         }
         Ok(())
     }
-    fn visit_if_stmt(&mut self, stmt: &IfStmt) -> Result<()> {
-        self.update_start(stmt.if_kw.line);
-        stmt.condition.accept(self).unwrap();
-        stmt.then_branch.accept(self).unwrap();
-        for stmt in &stmt.else_branch {
-            stmt.accept(self).unwrap();
-        }
-        Ok(())
-    }
 }
 
 impl ExprVisitor for LineNumber {
@@ -137,6 +123,22 @@ impl ExprVisitor for LineNumber {
         self.update_end(expr.rbracket.line);
         Ok(())
     }
+
+    fn visit_if_expr(&mut self, expr: &IfExpr) -> Result<()> {
+        self.update_start(expr.if_kw.line);
+        expr.condition.accept(self).unwrap();
+        expr.then_branch.accept(self).unwrap();
+        if let Some(block) = &expr.else_branch {
+            block.accept(self).unwrap();
+        }
+        Ok(())
+    }
+
+    fn visit_block_expr(&mut self, stmt: &BlockExpr) -> Result<()> {
+        self.update_start(stmt.lbrace.line);
+        self.update_end(stmt.rbrace.line);
+        Ok(())
+    }
 }
 
 pub(super) fn expr_line_number(expr: &dyn Expr) -> LineRange {
@@ -215,51 +217,12 @@ impl StmtVisitor for LocalAssignCollector {
         Ok(())
     }
 
-    fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<()> {
-        // we visit the block statement because even though it goes below the current "local"
-        // scope, we're ultimately trying to get a list of ALL local names that are assigned to in
-        // this scope.
-        // TODO FIXME BUG this does create some weirdness, for example take this:
-        // outer_function = () {
-        //     some_value = 1234
-        //     inner_function = () {
-        //         {
-        //             # this is a local value because we're assigning to it
-        //             some_value = 5678
-        //         }
-        //         # our local named "some_value" has gone out of scope, so hypothetically we
-        //         # should be using the "some_value" that was defined in the scope above us.
-        //         # however, since we're collecting local assignments in all blocks, this should
-        //         # error out as "unknown local 'some_value'"
-        //         println(some_value)
-        //     }
-        //     return inner_function
-        // }
-        //
-        // Ideally, we would be checking nonlocals with every new scope layer, and every new block.
-        // This is a pretty tough bug to solve with how things are set up right now. not sure how
-        // we'll go about solving this one.
-        for stmt in &stmt.stmts {
-            stmt.accept(self)?;
-        }
-        Ok(())
-    }
-
     fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<()> {
         if let Some(expr) = stmt.expr.as_ref() {
             expr.accept(self)?;
         }
         Ok(())
     }
-
-    fn visit_if_stmt(&mut self, stmt: &IfStmt) -> Result<()> {
-        stmt.condition.accept(self)?;
-        stmt.then_branch.accept(self)?;
-        for stmt in &stmt.else_branch {
-            stmt.accept(self)?;
-        }
-        Ok(())
-    }
 }
 
 impl ExprVisitor for LocalAssignCollector {
@@ -306,6 +269,48 @@ impl ExprVisitor for LocalAssignCollector {
         // there shouldn't be any local assignments inside of a map expression
         Ok(())
     }
+
+    fn visit_if_expr(&mut self, expr: &IfExpr) -> Result<()> {
+        expr.condition.accept(self)?;
+        expr.then_branch.accept(self)?;
+        if let Some(block) = &expr.else_branch {
+            block.accept(self)?;
+        }
+        Ok(())
+    }
+
+    fn visit_block_expr(&mut self, expr: &BlockExpr) -> Result<()> {
+        // we visit the block statement because even though it goes below the current "local"
+        // scope, we're ultimately trying to get a list of ALL local names that are assigned to in
+        // this scope.
+        // TODO FIXME BUG this does create some weirdness, for example take this:
+        // outer_function = () {
+        //     some_value = 1234
+        //     inner_function = () {
+        //         {
+        //             # this is a local value because we're assigning to it
+        //             some_value = 5678
+        //         }
+        //         # our local named "some_value" has gone out of scope, so hypothetically we
+        //         # should be using the "some_value" that was defined in the scope above us.
+        //         # however, since we're collecting local assignments in all blocks, this should
+        //         # error out as "unknown local 'some_value'"
+        //         println(some_value)
+        //     }
+        //     return inner_function
+        // }
+        //
+        // Ideally, we would be checking nonlocals with every new scope layer, and every new block.
+        // This is a pretty tough bug to solve with how things are set up right now. not sure how
+        // we'll go about solving this one.
+        for stmt in &expr.stmts {
+            stmt.accept(self)?;
+        }
+        if let Some(expr) = &expr.return_expr {
+            expr.accept(self)?;
+        }
+        Ok(())
+    }
 }
 
 #[derive(Default)]
@@ -366,28 +371,12 @@ impl StmtVisitor for LocalNameCollector {
         Ok(())
     }
 
-    fn visit_block_stmt(&mut self, stmt: &BlockStmt) -> Result<()> {
-        for stmt in &stmt.stmts {
-            stmt.accept(self)?;
-        }
-        Ok(())
-    }
-
     fn visit_return_stmt(&mut self, stmt: &ReturnStmt) -> Result<()> {
         if let Some(expr) = stmt.expr.as_ref() {
             expr.accept(self)?;
         }
         Ok(())
     }
-
-    fn visit_if_stmt(&mut self, stmt: &IfStmt) -> Result<()> {
-        stmt.condition.accept(self)?;
-        stmt.then_branch.accept(self)?;
-        for stmt in &stmt.else_branch {
-            stmt.accept(self)?;
-        }
-        Ok(())
-    }
 }
 
 impl ExprVisitor for LocalNameCollector {
@@ -444,4 +433,23 @@ impl ExprVisitor for LocalNameCollector {
         }
         Ok(())
     }
+
+    fn visit_if_expr(&mut self, expr: &IfExpr) -> Result<()> {
+        expr.condition.accept(self)?;
+        expr.then_branch.accept(self)?;
+        if let Some(block) = &expr.else_branch {
+            block.accept(self)?;
+        }
+        Ok(())
+    }
+
+    fn visit_block_expr(&mut self, expr: &BlockExpr) -> Result<()> {
+        for stmt in &expr.stmts {
+            stmt.accept(self)?;
+        }
+        if let Some(expr) = &expr.return_expr {
+            expr.accept(self)?;
+        }
+        Ok(())
+    }
 }

src/parser.rs

@@ -562,19 +562,8 @@ impl Parser {
     fn stmt_wrapped(&mut self) -> Result<StmtP> {
         if self.mat(TokenKind::Return)? {
             self.return_stmt()
-        } else if self.mat(TokenKind::If)? {
-            self.if_stmt()
         } else if self.mat(TokenKind::Import)? {
             self.import_stmt()
-        } else if self.mat(TokenKind::LBrace)? {
-            let lbrace = self.prev.clone().unwrap();
-            let stmts = self.block()?;
-            let rbrace = self.prev.clone().unwrap();
-            Ok(Box::new(BlockStmt {
-                lbrace,
-                stmts,
-                rbrace,
-            }) as Box<dyn Stmt + 'static>)
         } else if self.current.kind == TokenKind::Name
             && (self.next.kind == TokenKind::Eq
                 || self.next.kind == TokenKind::PlusEq
@@ -657,21 +646,30 @@ impl Parser {
         Ok(Box::new(ReturnStmt { return_kw, expr }))
     }
 
-    fn if_stmt(&mut self) -> Result<StmtP> {
+    fn if_expr(&mut self) -> Result<ExprP> {
         let if_kw = self.prev.clone().unwrap();
         let condition = self.expr()?;
         self.expect("expect '{' after 'if' condition", TokenKind::LBrace)?;
-        let then_branch = self.block_stmt()?;
-        let mut else_branch = Vec::new();
+        let then_branch = self.block_expr()?;
+        let mut else_branch = None;
+
         if self.mat(TokenKind::Else)? {
             if self.mat(TokenKind::If)? {
-                else_branch.push(self.if_stmt()?);
+                let lbrace = self.current.clone();
+                let if_expr = self.if_expr()?;
+                let rbrace = self.prev.clone().unwrap();
+                else_branch = Some(BlockExpr {
+                    lbrace,
+                    stmts: Default::default(),
+                    return_expr: Some(if_expr),
+                    rbrace,
+                });
             } else {
                 self.expect("expect '{' after else statement", TokenKind::LBrace)?;
-                else_branch = self.block()?;
+                else_branch = Some(self.block_expr()?);
             }
         }
-        Ok(Box::new(IfStmt {
+        Ok(Box::new(IfExpr {
             if_kw,
             condition,
             then_branch,
@@ -754,16 +752,28 @@ impl Parser {
         Ok(Box::new(import_stmt))
     }
 
-    fn block_stmt(&mut self) -> Result<BlockStmt> {
+    fn block_expr(&mut self) -> Result<BlockExpr> {
         let lbrace = self.prev.clone().unwrap();
         assert_eq!(lbrace.kind, TokenKind::LBrace);
-        let stmts = self.block()?;
+        let mut stmts = self.block()?;
         let rbrace = self.prev.clone().unwrap();
         assert_eq!(rbrace.kind, TokenKind::RBrace);
-        Ok(BlockStmt {
+
+        let return_expr = if let Some(last) = stmts.pop() {
+            if last.as_any_ref().downcast_ref::<ExprStmt>().is_some() {
+                let stmt = last.as_any().downcast::<ExprStmt>().unwrap();
+                Some(stmt.expr)
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+        Ok(BlockExpr {
             lbrace,
             stmts,
             rbrace,
+            return_expr,
         })
     }
 
@@ -951,6 +961,10 @@ impl Parser {
             Ok(expr)
         } else if self.mat(TokenKind::LBracket)? {
             self.list_or_map()
+        } else if self.mat(TokenKind::LBrace)? {
+            Ok(Box::new(self.block_expr()?))
+        } else if self.mat(TokenKind::If)? {
+            self.if_expr()
         } else {
             Err(self.error(format!("unexpected token {:?}", self.current.kind)))
         }