feat(tazjin/rlox): Implement simple conditionals r/2986

... basically just optional blocks (no else).

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

4 files changed, 84 insertions, 16 deletions

diff --git a/users/tazjin/rlox/examples/if.lox b/users/tazjin/rlox/examples/if.lox
new file mode 100644
index 000000000000..b59f00b20161
--- /dev/null
+++ b/users/tazjin/rlox/examples/if.lox
@@ -0,0 +1,5 @@
+if (true) {
+  print "yes";
+}
+
+print "afterwards";
diff --git a/users/tazjin/rlox/src/bytecode/compiler.rs b/users/tazjin/rlox/src/bytecode/compiler.rs
index 392fc9b72d6e..1b87e94a5512 100644
--- a/users/tazjin/rlox/src/bytecode/compiler.rs
+++ b/users/tazjin/rlox/src/bytecode/compiler.rs
@@ -1,7 +1,7 @@
 use super::chunk::Chunk;
 use super::errors::{Error, ErrorKind, LoxResult};
 use super::interner::{InternedStr, Interner};
-use super::opcode::{CodeIdx, ConstantIdx, OpCode, StackIdx};
+use super::opcode::{CodeIdx, CodeOffset, ConstantIdx, OpCode, StackIdx};
 use super::value::Value;
 use crate::scanner::{self, Token, TokenKind};
 
@@ -236,9 +236,13 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
 
     fn define_variable(&mut self, var: Option<ConstantIdx>) -> LoxResult<()> {
         if self.locals.scope_depth == 0 {
-            self.emit_op(OpCode::OpDefineGlobal(var.expect("should be global")));
+            self.emit_op(OpCode::OpDefineGlobal(
+                var.expect("should be global"),
+            ));
         } else {
-            self.locals.locals.last_mut()
+            self.locals
+                .locals
+                .last_mut()
                 .expect("fatal: variable not yet added at definition")
                 .depth = Depth::At(self.locals.scope_depth);
         }
@@ -263,6 +267,8 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
     fn statement(&mut self) -> LoxResult<()> {
         if self.match_token(&TokenKind::Print) {
             self.print_statement()
+        } else if self.match_token(&TokenKind::If) {
+            self.if_statement()
         } else if self.match_token(&TokenKind::LeftBrace) {
             self.begin_scope();
             self.block()?;
@@ -289,7 +295,9 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         self.locals.scope_depth -= 1;
 
         while self.locals.locals.len() > 0
-            && self.locals.locals[self.locals.locals.len() - 1].depth.above(self.locals.scope_depth)
+            && self.locals.locals[self.locals.locals.len() - 1]
+                .depth
+                .above(self.locals.scope_depth)
         {
             self.emit_op(OpCode::OpPop);
             self.locals.locals.remove(self.locals.locals.len() - 1);
@@ -319,6 +327,28 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         Ok(())
     }
 
+    fn if_statement(&mut self) -> LoxResult<()> {
+        consume!(
+            self,
+            TokenKind::LeftParen,
+            ErrorKind::ExpectedToken("Expected '(' after 'if'")
+        );
+
+        self.expression()?;
+
+        consume!(
+            self,
+            TokenKind::RightParen,
+            ErrorKind::ExpectedToken("Expected ')' after condition")
+        );
+
+        let then_jump = self.emit_op(OpCode::OpJumpPlaceholder(false));
+        self.statement()?;
+        self.patch_jump(then_jump);
+
+        Ok(())
+    }
+
     fn number(&mut self) -> LoxResult<()> {
         if let TokenKind::Number(num) = self.previous().kind {
             self.emit_constant(Value::Number(num), true);
@@ -431,16 +461,12 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
             self.expression()?;
             match local_idx {
                 Some(idx) => self.emit_op(OpCode::OpSetLocal(idx)),
-                None => {
-                    self.emit_op(OpCode::OpSetGlobal(ident.unwrap()))
-                }
+                None => self.emit_op(OpCode::OpSetGlobal(ident.unwrap())),
             };
         } else {
             match local_idx {
                 Some(idx) => self.emit_op(OpCode::OpGetLocal(idx)),
-                None => {
-                    self.emit_op(OpCode::OpGetGlobal(ident.unwrap()))
-                }
+                None => self.emit_op(OpCode::OpGetGlobal(ident.unwrap())),
             };
         }
 
@@ -477,10 +503,7 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         Ok(())
     }
 
-    fn identifier_str(
-        &mut self,
-        token: &Token,
-    ) -> LoxResult<InternedStr> {
+    fn identifier_str(&mut self, token: &Token) -> LoxResult<InternedStr> {
         let ident = match &token.kind {
             TokenKind::Identifier(ident) => ident.to_string(),
             _ => {
@@ -594,6 +617,20 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         idx
     }
 
+    fn patch_jump(&mut self, idx: CodeIdx) {
+        let offset = CodeOffset(self.chunk.code.len() - idx.0 - 1);
+
+        if let OpCode::OpJumpPlaceholder(false) = self.chunk.code[idx.0] {
+            self.chunk.code[idx.0] = OpCode::OpJumpIfFalse(offset);
+            return;
+        }
+
+        panic!(
+            "attempted to patch unsupported op: {:?}",
+            self.chunk.code[idx.0]
+        );
+    }
+
     fn previous(&self) -> &Token {
         self.previous
             .as_ref()
diff --git a/users/tazjin/rlox/src/bytecode/opcode.rs b/users/tazjin/rlox/src/bytecode/opcode.rs
index 13e2f23939ed..accc6b3bac10 100644
--- a/users/tazjin/rlox/src/bytecode/opcode.rs
+++ b/users/tazjin/rlox/src/bytecode/opcode.rs
@@ -7,6 +7,9 @@ pub struct StackIdx(pub usize);
 #[derive(Clone, Copy, Debug)]
 pub struct CodeIdx(pub usize);
 
+#[derive(Clone, Copy, Debug)]
+pub struct CodeOffset(pub usize);
+
 #[derive(Debug)]
 pub enum OpCode {
     /// Push a constant onto the stack.
@@ -45,4 +48,8 @@ pub enum OpCode {
     OpSetGlobal(ConstantIdx),
     OpGetLocal(StackIdx),
     OpSetLocal(StackIdx),
+
+    // Control flow
+    OpJumpPlaceholder(bool),
+    OpJumpIfFalse(CodeOffset),
 }
diff --git a/users/tazjin/rlox/src/bytecode/vm.rs b/users/tazjin/rlox/src/bytecode/vm.rs
index bbdf70d8dc2f..e0db97abe5a2 100644
--- a/users/tazjin/rlox/src/bytecode/vm.rs
+++ b/users/tazjin/rlox/src/bytecode/vm.rs
@@ -201,8 +201,27 @@ impl VM {
                 }
 
                 OpCode::OpSetLocal(local_idx) => {
-                    debug_assert!(self.stack.len() > local_idx.0, "stack is not currently large enough for local");
-                    self.stack[local_idx.0] = self.stack.last().unwrap().clone();
+                    debug_assert!(
+                        self.stack.len() > local_idx.0,
+                        "stack is not currently large enough for local"
+                    );
+                    self.stack[local_idx.0] =
+                        self.stack.last().unwrap().clone();
+                }
+
+                OpCode::OpJumpPlaceholder(_) => {
+                    panic!("unpatched jump detected - this is a fatal compiler error!");
+                }
+
+                OpCode::OpJumpIfFalse(offset) => {
+                    if self
+                        .stack
+                        .last()
+                        .expect("condition should leave a value on the stack")
+                        .is_falsey()
+                    {
+                        self.ip += offset.0;
+                    }
                 }
             }