feat(tazjin/rlox): Support local variables r/2978

WIP

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

6 files changed, 193 insertions, 12 deletions

diff --git a/users/tazjin/rlox/examples/var.lox b/users/tazjin/rlox/examples/var.lox
index 133906ad1b2d..7af90b3f0bee 100644
--- a/users/tazjin/rlox/examples/var.lox
+++ b/users/tazjin/rlox/examples/var.lox
@@ -1,8 +1,8 @@
-{
+var a = 10;
+var b = 5;
 
 {
-  var a = 5;
-  print a;
-}
-
+  var b = 10;
+  var c = 2;
+  a * b * c;
 }
diff --git a/users/tazjin/rlox/src/bytecode/compiler.rs b/users/tazjin/rlox/src/bytecode/compiler.rs
index a6675d1c20b6..28b00cbb6ed8 100644
--- a/users/tazjin/rlox/src/bytecode/compiler.rs
+++ b/users/tazjin/rlox/src/bytecode/compiler.rs
@@ -8,12 +8,47 @@ use crate::scanner::{self, Token, TokenKind};
 #[cfg(feature = "disassemble")]
 use super::chunk;
 
+#[derive(Debug)]
+enum Depth {
+    Unitialised,
+    At(usize),
+}
+
+impl Depth {
+    fn above(&self, theirs: usize) -> bool {
+        match self {
+            Depth::Unitialised => false,
+            Depth::At(ours) => *ours > theirs,
+        }
+    }
+
+    fn below(&self, theirs: usize) -> bool {
+        match self {
+            Depth::Unitialised => false,
+            Depth::At(ours) => *ours < theirs,
+        }
+    }
+}
+
+#[derive(Debug)]
+struct Local {
+    name: Token,
+    depth: Depth,
+}
+
+#[derive(Debug, Default)]
+struct Locals {
+    locals: Vec<Local>,
+    scope_depth: usize,
+}
+
 struct Compiler<T: Iterator<Item = Token>> {
     tokens: T,
     chunk: Chunk,
     panic: bool,
     errors: Vec<Error>,
     strings: Interner,
+    locals: Locals,
 
     current: Option<Token>,
     previous: Option<Token>,
@@ -200,7 +235,14 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
     }
 
     fn define_variable(&mut self, var: usize) -> LoxResult<()> {
-        self.emit_op(OpCode::OpDefineGlobal(var));
+        if self.locals.scope_depth == 0 {
+            self.emit_op(OpCode::OpDefineGlobal(var));
+        } else {
+            self.locals.locals.last_mut()
+                .expect("fatal: variable not yet added at definition")
+                .depth = Depth::At(self.locals.scope_depth);
+        }
+
         Ok(())
     }
 
@@ -221,6 +263,11 @@ 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::LeftBrace) {
+            self.begin_scope();
+            self.block()?;
+            self.end_scope();
+            Ok(())
         } else {
             self.expression_statement()
         }
@@ -233,6 +280,37 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         Ok(())
     }
 
+    fn begin_scope(&mut self) {
+        self.locals.scope_depth += 1;
+    }
+
+    fn end_scope(&mut self) {
+        debug_assert!(self.locals.scope_depth > 0, "tried to end global scope");
+        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.emit_op(OpCode::OpPop);
+            self.locals.locals.remove(self.locals.locals.len() - 1);
+        }
+    }
+
+    fn block(&mut self) -> LoxResult<()> {
+        while !self.check(&TokenKind::RightBrace)
+            && !self.check(&TokenKind::Eof)
+        {
+            self.declaration()?;
+        }
+
+        consume!(
+            self,
+            TokenKind::RightBrace,
+            ErrorKind::ExpectedToken("Expected '}' after block.")
+        );
+        Ok(())
+    }
+
     fn expression_statement(&mut self) -> LoxResult<()> {
         self.expression()?;
         self.expect_semicolon("expect ';' after expression")?;
@@ -341,15 +419,29 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
     }
 
     fn named_variable(&mut self) -> LoxResult<()> {
-        let ident = self.identifier_str(Self::previous)?;
-        let constant_id =
-            self.emit_constant(Value::String(ident.into()), false);
+        let local_idx = self.resolve_local();
+
+        let ident = if local_idx.is_some() {
+            None
+        } else {
+            Some(self.identifier_constant()?)
+        };
 
         if self.match_token(&TokenKind::Equal) {
             self.expression()?;
-            self.emit_op(OpCode::OpSetGlobal(constant_id));
+            match local_idx {
+                Some(idx) => self.emit_op(OpCode::OpSetLocal(idx)),
+                None => {
+                    self.emit_op(OpCode::OpSetGlobal(ident.unwrap()));
+                }
+            }
         } else {
-            self.emit_op(OpCode::OpGetGlobal(constant_id));
+            match local_idx {
+                Some(idx) => self.emit_op(OpCode::OpGetLocal(idx)),
+                None => {
+                    self.emit_op(OpCode::OpGetGlobal(ident.unwrap()))
+                }
+            }
         }
 
         Ok(())
@@ -401,6 +493,59 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
         Ok(self.strings.intern(ident))
     }
 
+    fn identifier_constant(&mut self) -> LoxResult<usize> {
+        let ident = self.identifier_str(Self::previous)?;
+        Ok(self.emit_constant(Value::String(ident.into()), false))
+    }
+
+    fn resolve_local(&mut self) -> Option<usize> {
+        dbg!(&self.locals);
+        for (idx, local) in self.locals.locals.iter().enumerate().rev() {
+            if self.previous().lexeme == local.name.lexeme {
+                if let Depth::Unitialised = local.depth {
+                    // TODO(tazjin): *return* err
+                    panic!("can't read variable in its own initialiser");
+                }
+                return Some(idx);
+            }
+        }
+
+        None
+    }
+
+    fn add_local(&mut self, name: Token) -> LoxResult<()> {
+        let local = Local {
+            name,
+            depth: Depth::Unitialised,
+        };
+
+        self.locals.locals.push(local);
+        Ok(()) // TODO(tazjin): needed?
+    }
+
+    fn declare_variable(&mut self) -> LoxResult<()> {
+        if self.locals.scope_depth == 0 {
+            return Ok(());
+        }
+
+        let name = self.previous().clone();
+
+        for local in self.locals.locals.iter().rev() {
+            if local.depth.below(self.locals.scope_depth) {
+                break;
+            }
+
+            if name.lexeme == local.name.lexeme {
+                return Err(Error {
+                    kind: ErrorKind::VariableShadowed(name.lexeme.into()),
+                    line: name.line,
+                });
+            }
+        }
+
+        self.add_local(name)
+    }
+
     fn parse_variable(&mut self) -> LoxResult<usize> {
         consume!(
             self,
@@ -408,6 +553,11 @@ impl<T: Iterator<Item = Token>> Compiler<T> {
             ErrorKind::ExpectedToken("expected identifier")
         );
 
+        self.declare_variable()?;
+        if self.locals.scope_depth > 0 {
+            return Ok(0); // TODO(tazjin): grr sentinel
+        }
+
         let id = self.identifier_str(Self::previous)?;
         Ok(self.emit_constant(Value::String(id.into()), false))
     }
@@ -520,6 +670,7 @@ pub fn compile(code: &str) -> Result<(Interner, Chunk), Vec<Error>> {
         panic: false,
         errors: vec![],
         strings: Interner::with_capacity(1024),
+        locals: Default::default(),
         current: None,
         previous: None,
     };
diff --git a/users/tazjin/rlox/src/bytecode/errors.rs b/users/tazjin/rlox/src/bytecode/errors.rs
index c6b86172f86d..988031f763cf 100644
--- a/users/tazjin/rlox/src/bytecode/errors.rs
+++ b/users/tazjin/rlox/src/bytecode/errors.rs
@@ -9,6 +9,7 @@ pub enum ErrorKind {
     ExpectedToken(&'static str),
     InternalError(&'static str),
     TypeError(String),
+    VariableShadowed(String),
 }
 
 #[derive(Debug)]
diff --git a/users/tazjin/rlox/src/bytecode/opcode.rs b/users/tazjin/rlox/src/bytecode/opcode.rs
index c7ccdd89e85f..143a79f41916 100644
--- a/users/tazjin/rlox/src/bytecode/opcode.rs
+++ b/users/tazjin/rlox/src/bytecode/opcode.rs
@@ -34,4 +34,6 @@ pub enum OpCode {
     OpDefineGlobal(usize),
     OpGetGlobal(usize),
     OpSetGlobal(usize),
+    OpGetLocal(usize),
+    OpSetLocal(usize),
 }
diff --git a/users/tazjin/rlox/src/bytecode/tests.rs b/users/tazjin/rlox/src/bytecode/tests.rs
index 13e64400ade5..de482275ecfd 100644
--- a/users/tazjin/rlox/src/bytecode/tests.rs
+++ b/users/tazjin/rlox/src/bytecode/tests.rs
@@ -128,6 +128,23 @@ fn global_assignment() {
           breakfast = "beignets with " + beverage;
           breakfast;
         "#,
-        "beignets with cafe au lait"
+        "beignets with cafe au lait",
+    );
+}
+
+#[test]
+fn local_variables() {
+    expect_num(
+        r#"
+          var a = 10;
+          var b = 5;
+
+          {
+            var b = 10;
+            var c = 2;
+            a * b * c;
+          }
+        "#,
+        200.0,
     );
 }
diff --git a/users/tazjin/rlox/src/bytecode/vm.rs b/users/tazjin/rlox/src/bytecode/vm.rs
index 4432473ee649..10d0f595d6d4 100644
--- a/users/tazjin/rlox/src/bytecode/vm.rs
+++ b/users/tazjin/rlox/src/bytecode/vm.rs
@@ -194,6 +194,16 @@ impl VM {
                         }
                     });
                 }
+
+                OpCode::OpGetLocal(local_idx) => {
+                    let value = self.stack[*local_idx].clone();
+                    self.push(value);
+                }
+
+                OpCode::OpSetLocal(local_idx) => {
+                    debug_assert!(self.stack.len() > *local_idx, "stack is not currently large enough for local");
+                    self.stack[*local_idx] = self.stack.last().unwrap().clone();
+                }
             }
 
             #[cfg(feature = "disassemble")]