Block-level scoping (#92) and variable manager (#88)

crates/lean_compiler/src/grammar.pest

@@ -16,6 +16,7 @@ return_count = { "->" ~ number }
 
 // Statements
 statement = {
+    forward_declaration |
     single_assignment |
     array_assign |
     if_statement |
@@ -34,6 +35,8 @@ return_statement = { "return" ~ (tuple_expression)? ~ ";" }
 break_statement = { "break" ~ ";" }
 continue_statement = { "continue" ~ ";" }
 
+forward_declaration = { "var" ~ identifier ~ ";" }
+
 single_assignment = { identifier ~ "=" ~ expression ~ ";" }
 
 array_assign = { identifier ~ "[" ~ expression ~ "]" ~ "=" ~ expression ~ ";" }

crates/lean_compiler/src/lang.rs

@@ -1,7 +1,7 @@
 use lean_vm::*;
 use multilinear_toolkit::prelude::*;
 use p3_util::log2_ceil_usize;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::{Display, Formatter};
 use utils::ToUsize;
 
@@ -307,6 +307,9 @@ pub enum Line {
         value: Expression,
         arms: Vec<(usize, Vec<Self>)>,
     },
+    ForwardDeclaration {
+        var: Var,
+    },
     Assignment {
         var: Var,
         value: Expression,
@@ -378,6 +381,30 @@ pub enum Line {
         location: SourceLineNumber,
     },
 }
+
+/// A context specifying which variables are in scope.
+pub struct Context {
+    /// A list of lexical scopes, innermost scope last.
+    pub scopes: Vec<Scope>,
+}
+
+impl Context {
+    pub fn defines(&self, var: &Var) -> bool {
+        for scope in self.scopes.iter() {
+            if scope.vars.contains(var) {
+                return true;
+            }
+        }
+        false
+    }
+}
+
+#[derive(Default)]
+pub struct Scope {
+    /// A set of declared variables.
+    pub vars: BTreeSet<Var>,
+}
+
 impl Display for Expression {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -418,6 +445,9 @@ impl Line {
                     .join("\n");
                 format!("match {value} {{\n{arms_str}\n{spaces}}}")
             }
+            Self::ForwardDeclaration { var } => {
+                format!("var {var}")
+            }
             Self::Assignment { var, value } => {
                 format!("{var} = {value}")
             }

crates/lean_compiler/src/lib.rs

@@ -16,7 +16,7 @@ pub fn compile_program(program: String) -> Bytecode {
     let (parsed_program, function_locations) = parse_program(&program).unwrap();
     // println!("Parsed program: {}", parsed_program.to_string());
     let simple_program = simplify_program(parsed_program);
-    // println!("Simplified program: {}", simple_program.to_string());
+    // println!("Simplified program: {}", simple_program);
     let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
     // println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
 

crates/lean_compiler/src/parser/parsers/statement.rs

@@ -19,6 +19,7 @@ impl Parse<Line> for StatementParser {
         let inner = next_inner_pair(&mut pair.into_inner(), "statement body")?;
 
         match inner.as_rule() {
+            Rule::forward_declaration => ForwardDeclarationParser::parse(inner, ctx),
             Rule::single_assignment => AssignmentParser::parse(inner, ctx),
             Rule::array_assign => ArrayAssignParser::parse(inner, ctx),
             Rule::if_statement => IfStatementParser::parse(inner, ctx),
@@ -35,6 +36,17 @@ impl Parse<Line> for StatementParser {
     }
 }
 
+/// Parser for forward declarations of variables.
+pub struct ForwardDeclarationParser;
+
+impl Parse<Line> for ForwardDeclarationParser {
+    fn parse(pair: ParsePair<'_>, _ctx: &mut ParseContext) -> ParseResult<Line> {
+        let mut inner = pair.into_inner();
+        let var = next_inner_pair(&mut inner, "variable name")?.as_str().to_string();
+        Ok(Line::ForwardDeclaration { var })
+    }
+}
+
 /// Parser for variable assignments.
 pub struct AssignmentParser;
 

crates/lean_compiler/tests/test_compiler.rs

@@ -381,6 +381,26 @@ fn test_inlined() {
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
+#[test]
+fn test_inlined_2() {
+    let program = r#"
+    fn main() {
+        b = is_one();
+        c = b;
+        return;
+    }
+
+    fn is_one() inline -> 1 {
+        if 1 {
+            return 1;
+        } else {
+            return 0;
+        }
+    }
+   "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}
+
 #[test]
 fn test_match() {
     let program = r#"
@@ -433,6 +453,29 @@ fn test_match() {
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
+#[test]
+fn test_match_shrink() {
+    let program = r#"
+    fn main() {
+        match 1 {
+            0 => {
+                y = 90;
+            }
+            1 => {
+                y = 10;
+                z = func_2(y);
+            }
+        }
+        return;
+    }
+
+    fn func_2(x) inline -> 1 {
+        return x * x;
+    }
+   "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}
+
 // #[test]
 // fn inline_bug_mre() {
 //     let program = r#"
@@ -523,3 +566,26 @@ fn test_nested_inline_functions() {
 
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
+
+#[test]
+fn test_const_and_nonconst_malloc_sharing_name() {
+    let program = r#"
+    fn main() {
+        f(1);
+        return;
+    }
+
+    fn f(n) {
+        if 0 == 0 {
+            res = malloc(2);
+            res[1] = 0;
+            return;
+        } else {
+            res = malloc(n * 1);
+            return;
+        }
+    }
+    "#;
+
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}

crates/lean_vm/src/diagnostics/error.rs

@@ -25,8 +25,8 @@ pub enum RunnerError {
     #[error("Computation invalid: {0} != {1}")]
     NotEqual(F, F),
 
-    #[error("Undefined memory access")]
-    UndefinedMemory,
+    #[error("Undefined memory access: {0}")]
+    UndefinedMemory(usize),
 
     #[error("Program counter out of bounds")]
     PCOutOfBounds,

crates/lean_vm/src/execution/memory.rs

@@ -25,7 +25,11 @@ impl Memory {
     ///
     /// Returns an error if the address is uninitialized
     pub fn get(&self, index: usize) -> Result<F, RunnerError> {
-        self.0.get(index).copied().flatten().ok_or(RunnerError::UndefinedMemory)
+        self.0
+            .get(index)
+            .copied()
+            .flatten()
+            .ok_or(RunnerError::UndefinedMemory(index))
     }
 
     /// Sets a value at a memory address

crates/lean_vm/src/execution/tests.rs

@@ -13,7 +13,7 @@ fn test_basic_memory_operations() {
     assert_eq!(memory.get(5).unwrap(), F::from_usize(42));
 
     // Test undefined memory access
-    assert!(matches!(memory.get(1), Err(RunnerError::UndefinedMemory)));
+    assert!(matches!(memory.get(1), Err(RunnerError::UndefinedMemory(1))));
 }
 
 #[test]

crates/lookup/src/quotient_gkr.rs

@@ -220,7 +220,7 @@ fn sum_quotients_helper<F: PrimeCharacteristicRing + Sync + Send + Copy>(
 ) -> Vec<Vec<F>> {
     assert_eq!(numerators_and_denominators.len(), n_groups);
     let n = numerators_and_denominators[0].len();
-    assert!(n.is_power_of_two() && n >= 2, "n = {}", n);
+    assert!(n.is_power_of_two() && n >= 2, "n = {n}");
     let mut new_numerators = Vec::new();
     let mut new_denominators = Vec::new();
     let (prev_numerators, prev_denominators) = numerators_and_denominators.split_at(n_groups / 2);