feat(tazjin/rlox): Implement variable depth resolver r/2113

Implements the first part of the resolver from
https://craftinginterpreters.com/resolving-and-binding.html

This is wired up to the execution paths in main, but not yet in the
tests. The resolved depth is also not actually used for variable
lookups (yet).

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

5 files changed, 217 insertions, 16 deletions

diff --git a/users/tazjin/rlox/src/errors.rs b/users/tazjin/rlox/src/errors.rs
index 875f268ca64d..27ccb962f736 100644
--- a/users/tazjin/rlox/src/errors.rs
+++ b/users/tazjin/rlox/src/errors.rs
@@ -14,6 +14,7 @@ pub enum ErrorKind {
     InternalError(String),
     InvalidAssignmentTarget(String),
     RuntimeError(String),
+    StaticError(String),
 
     // This variant is not an error, rather it is used for
     // short-circuiting out of a function body that hits a `return`
diff --git a/users/tazjin/rlox/src/interpreter.rs b/users/tazjin/rlox/src/interpreter.rs
index a1d246ba2eb2..6c61288d5a7d 100644
--- a/users/tazjin/rlox/src/interpreter.rs
+++ b/users/tazjin/rlox/src/interpreter.rs
@@ -107,13 +107,13 @@ impl Environment {
     }
 
     fn get(&self, name: &parser::Variable) -> Result<Value, Error> {
-        let ident = identifier_str(&name.0)?;
+        let ident = identifier_str(&name.name)?;
 
         self.values
             .get(ident)
             .map(Clone::clone)
             .ok_or_else(|| Error {
-                line: name.0.line,
+                line: name.name.line,
                 kind: ErrorKind::UndefinedVariable(ident.into()),
             })
             .or_else(|err| {
diff --git a/users/tazjin/rlox/src/main.rs b/users/tazjin/rlox/src/main.rs
index 8970349bfa69..24ebe503b692 100644
--- a/users/tazjin/rlox/src/main.rs
+++ b/users/tazjin/rlox/src/main.rs
@@ -7,6 +7,7 @@ use std::process;
 mod errors;
 mod interpreter;
 mod parser;
+mod resolver;
 mod scanner;
 
 fn main() {
@@ -48,17 +49,13 @@ fn run_prompt() {
 fn run(lox: &mut interpreter::Interpreter, code: &str) {
     let chars: Vec<char> = code.chars().collect();
 
-    match scanner::scan(&chars) {
-        Ok(tokens) => match parser::parse(tokens) {
-            Ok(program) => {
-                println!("Program:\n{:?}", program);
-                if let Err(err) = lox.interpret(&program) {
-                    println!("Error in program: {:?}", err);
-                }
-            }
-            Err(errors) => report_errors(errors),
-        },
-        Err(errors) => report_errors(errors),
+    let result = scanner::scan(&chars)
+        .and_then(|tokens| parser::parse(tokens))
+        .and_then(|program| resolver::resolve(program).map_err(|e| vec![e]))
+        .and_then(|program| lox.interpret(&program).map_err(|e| vec![e]));
+
+    if let Err(errors) = result {
+        report_errors(errors);
     }
 }
 
diff --git a/users/tazjin/rlox/src/parser.rs b/users/tazjin/rlox/src/parser.rs
index 495304686b2f..e28404aa9fce 100644
--- a/users/tazjin/rlox/src/parser.rs
+++ b/users/tazjin/rlox/src/parser.rs
@@ -15,6 +15,7 @@ use std::rc::Rc;
 pub struct Assign {
     pub name: Token,
     pub value: Box<Expr>,
+    pub depth: Option<usize>,
 }
 
 #[derive(Debug)]
@@ -57,7 +58,10 @@ pub struct Call {
 
 // Not to be confused with `Var`, which is for assignment.
 #[derive(Debug)]
-pub struct Variable(pub Token);
+pub struct Variable {
+    pub name: Token,
+    pub depth: Option<usize>,
+}
 
 #[derive(Debug)]
 pub enum Expr {
@@ -413,10 +417,11 @@ impl Parser {
             let equals = self.previous().clone();
             let value = self.assignment()?;
 
-            if let Expr::Variable(Variable(name)) = expr {
+            if let Expr::Variable(Variable { name, .. }) = expr {
                 return Ok(Expr::Assign(Assign {
                     name,
                     value: Box::new(value),
+                    depth: None,
                 }));
             }
 
@@ -549,7 +554,12 @@ impl Parser {
                 return Ok(Expr::Grouping(Grouping(Box::new(expr))));
             }
 
-            TokenKind::Identifier(_) => return Ok(Expr::Variable(Variable(next))),
+            TokenKind::Identifier(_) => {
+                return Ok(Expr::Variable(Variable {
+                    name: next,
+                    depth: None,
+                }))
+            }
 
             unexpected => {
                 eprintln!("encountered {:?}", unexpected);
diff --git a/users/tazjin/rlox/src/resolver.rs b/users/tazjin/rlox/src/resolver.rs
new file mode 100644
index 000000000000..f5ef91b99a19
--- /dev/null
+++ b/users/tazjin/rlox/src/resolver.rs
@@ -0,0 +1,193 @@
+// Resolves variable access to their specific instances in the
+// environment chain.
+//
+// https://craftinginterpreters.com/resolving-and-binding.html
+
+use std::collections::HashMap;
+use std::rc::Rc;
+
+use crate::errors::{Error, ErrorKind};
+use crate::parser::{self, Expr, Statement};
+use crate::scanner::Token;
+
+#[derive(Default)]
+struct Resolver<'a> {
+    scopes: Vec<HashMap<&'a str, bool>>,
+}
+
+impl<'a> Resolver<'a> {
+    // AST traversal
+    fn resolve(&mut self, program: &'a mut parser::Block) -> Result<(), Error> {
+        self.begin_scope();
+        for stmt in program {
+            self.resolve_stmt(stmt)?;
+        }
+        self.end_scope();
+
+        Ok(())
+    }
+
+    fn resolve_stmt(&mut self, stmt: &'a mut Statement) -> Result<(), Error> {
+        match stmt {
+            Statement::Expr(expr) => self.resolve_expr(expr),
+            Statement::Print(expr) => self.resolve_expr(expr),
+            Statement::Var(var) => self.resolve_var(var),
+            Statement::Return(ret) => self.resolve_expr(&mut ret.value),
+            Statement::Block(block) => self.resolve(block),
+
+            Statement::If(if_stmt) => {
+                self.resolve_expr(&mut if_stmt.condition)?;
+                self.resolve_stmt(&mut if_stmt.then_branch)?;
+
+                if let Some(branch) = if_stmt.else_branch.as_mut() {
+                    self.resolve_stmt(branch)?;
+                }
+
+                Ok(())
+            }
+
+            Statement::While(while_stmt) => {
+                self.resolve_expr(&mut while_stmt.condition)?;
+                self.resolve_stmt(&mut while_stmt.body)
+            }
+
+            Statement::Function(func) => match Rc::get_mut(func) {
+                Some(func) => self.resolve_function(func),
+                // The resolver does not clone references, so unless
+                // the interpreter is called before the resolver this
+                // case should never happen.
+                None => {
+                    return Err(Error {
+                        line: 0,
+                        kind: ErrorKind::InternalError(
+                            "multiple function references before interpretation".into(),
+                        ),
+                    })
+                }
+            },
+        }
+    }
+
+    fn resolve_var(&mut self, var: &'a mut parser::Var) -> Result<(), Error> {
+        self.declare(&var.name);
+
+        if let Some(init) = &mut var.initialiser {
+            self.resolve_expr(init)?;
+        }
+
+        self.define(&var.name);
+
+        Ok(())
+    }
+
+    fn resolve_function(&mut self, func: &'a mut parser::Function) -> Result<(), Error> {
+        self.declare(&func.name);
+        self.define(&func.name);
+
+        self.begin_scope();
+
+        for param in &func.params {
+            self.declare(param);
+            self.define(param);
+        }
+
+        for stmt in &mut func.body {
+            self.resolve_stmt(stmt)?;
+        }
+
+        self.end_scope();
+
+        Ok(())
+    }
+
+    fn resolve_expr(&mut self, expr: &'a mut Expr) -> Result<(), Error> {
+        match expr {
+            Expr::Variable(var) => self.resolve_variable(var),
+            Expr::Assign(assign) => self.resolve_assign(assign),
+            Expr::Grouping(grouping) => self.resolve_expr(&mut grouping.0),
+            Expr::Call(call) => self.resolve_call(call),
+            Expr::Literal(_) => Ok(()),
+            Expr::Unary(unary) => self.resolve_expr(&mut unary.right),
+
+            Expr::Logical(log) => {
+                self.resolve_expr(&mut log.left)?;
+                self.resolve_expr(&mut log.right)
+            }
+
+            Expr::Binary(binary) => {
+                self.resolve_expr(&mut binary.left)?;
+                self.resolve_expr(&mut binary.right)
+            }
+        }
+    }
+
+    fn resolve_variable(&mut self, var: &'a mut parser::Variable) -> Result<(), Error> {
+        if let Some(scope) = self.scopes.last_mut() {
+            if let Some(false) = scope.get(var.name.lexeme.as_str()) {
+                return Err(Error {
+                    line: var.name.line,
+                    kind: ErrorKind::StaticError(
+                        "can't read local variable in its own initialiser".into(),
+                    ),
+                });
+            }
+        }
+
+        var.depth = self.resolve_local(&var.name);
+        Ok(())
+    }
+
+    fn resolve_assign(&mut self, assign: &'a mut parser::Assign) -> Result<(), Error> {
+        self.resolve_expr(&mut assign.value)?;
+        assign.depth = self.resolve_local(&assign.name);
+        Ok(())
+    }
+
+    fn resolve_local(&mut self, name: &'a Token) -> Option<usize> {
+        for (c, scope) in self.scopes.iter().rev().enumerate() {
+            if scope.contains_key(name.lexeme.as_str()) {
+                return Some(c);
+            }
+        }
+
+        None
+    }
+
+    fn resolve_call(&mut self, call: &'a mut parser::Call) -> Result<(), Error> {
+        self.resolve_expr(&mut call.callee)?;
+
+        for arg in call.args.iter_mut() {
+            self.resolve_expr(arg)?;
+        }
+
+        Ok(())
+    }
+
+    // Internal helpers
+
+    fn declare(&mut self, name: &'a Token) {
+        if let Some(scope) = self.scopes.last_mut() {
+            scope.insert(&name.lexeme, false);
+        }
+    }
+
+    fn define(&mut self, name: &'a Token) {
+        if let Some(scope) = self.scopes.last_mut() {
+            scope.insert(&name.lexeme, true);
+        }
+    }
+
+    fn begin_scope(&mut self) {
+        self.scopes.push(Default::default());
+    }
+
+    fn end_scope(&mut self) {
+        self.scopes.pop();
+    }
+}
+
+pub fn resolve(mut block: parser::Block) -> Result<parser::Block, Error> {
+    let mut resolver: Resolver = Default::default();
+    resolver.resolve(&mut block)?;
+    Ok(block)
+}