Block-level scoping and variable manager

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,
@@ -379,6 +382,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 {
@@ -419,6 +446,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,16 +16,18 @@ 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());
 
     // println!("Function Locations: \n");
     // for (loc, name) in function_locations.iter() {
     //     println!("{name}: {loc}");
     // }
+    /* let compiled = */
+    compile_to_low_level_bytecode(intermediate_bytecode, program, function_locations).unwrap() //;
     // println!("\n\nCompiled Program:\n\n{compiled}");
-    compile_to_low_level_bytecode(intermediate_bytecode, program, function_locations).unwrap()
+    // compiled
 }
 
 pub fn compile_and_run(

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

@@ -382,30 +382,34 @@ fn test_inlined() {
 }
 
 #[test]
-fn test_match() {
+fn test_inlined_2() {
     let program = r#"
     fn main() {
-        for x in 0..3 unroll {
-            func_match(x);
-        }
-        for x in 0..2 unroll {
-            match x {
-                0 => {
-                    y = 10 * (x + 8);
-                    z = 10 * y;
-                    print(z);
-                }
-                1 => {
-                    y = 10 * x;
-                    z = func_2(y);
-                    print(z);
-                }
-            }
+        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#"
+    fn main() {
+        func_match(1);
         return;
     }
 
-    fn func_match(x) inline {
+    fn func_match(x) {
         match x {
             0 => {
                 print(41);
@@ -415,8 +419,6 @@ fn test_match() {
                 print(y + 1);
             }
             2 => {
-                y = 10 * x;
-                print(y);
             }
         }
         return;
@@ -425,10 +427,6 @@ fn test_match() {
     fn func_1(x) -> 1 {
         return x * x * x * x;
     }
-
-    fn func_2(x) inline -> 1 {
-        return x * x * x * x * x * x;
-    }
    "#;
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }

crates/lean_prover/tests/test_zkvm.rs

@@ -125,6 +125,6 @@ fn test_zk_vm_helper(program_str: &str, (public_input, private_input): (&[F], &[
     );
     let proof_time = time.elapsed();
     verify_execution(&bytecode, public_input, proof_data, whir_config_builder()).unwrap();
-    println!("{}", summary);
+    println!("{summary}");
     println!("Proof time: {:.3} s", proof_time.as_secs_f32());
 }

crates/lean_vm/src/diagnostics/error.rs

@@ -9,8 +9,12 @@ pub enum RunnerError {
     #[error("Out of memory")]
     OutOfMemory,
 
-    #[error("Memory already set")]
-    MemoryAlreadySet,
+    #[error("Memory already set: m[{address}] = {prev_value} != {new_value}")]
+    MemoryAlreadySet {
+        address: usize,
+        prev_value: F,
+        new_value: F,
+    },
 
     #[error("Not a pointer")]
     NotAPointer,
@@ -21,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/diagnostics/stack_trace.rs

@@ -10,6 +10,7 @@ pub(crate) fn pretty_stack_trace(
     source_code: &str,
     instructions: &[SourceLineNumber], // SourceLineNumber = usize
     function_locations: &BTreeMap<usize, String>,
+    last_pc: usize,
 ) -> String {
     let source_lines: Vec<&str> = source_code.lines().collect();
     let mut result = String::new();
@@ -97,7 +98,11 @@ pub(crate) fn pretty_stack_trace(
     if !call_stack.is_empty() {
         result.push_str("\nCall stack:\n");
         for (i, (line, func)) in call_stack.iter().enumerate() {
-            result.push_str(&format!("  {}. {} (line {})\n", i + 1, func, line));
+            if i + 1 == call_stack.len() {
+                result.push_str(&format!("  {}. {} (line {}, pc {})\n", i + 1, func, line, last_pc));
+            } else {
+                result.push_str(&format!("  {}. {} (line {})\n", i + 1, func, line));
+            }
         }
     }
 

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
@@ -41,7 +45,11 @@ impl Memory {
         }
         if let Some(existing) = &mut self.0[index] {
             if *existing != value {
-                return Err(RunnerError::MemoryAlreadySet);
+                return Err(RunnerError::MemoryAlreadySet {
+                    address: index,
+                    prev_value: *existing,
+                    new_value: value,
+                });
             }
         } else {
             self.0[index] = Some(value);

crates/lean_vm/src/execution/runner.rs

@@ -70,15 +70,16 @@ pub fn execute_bytecode(
         profiling,
         (poseidons_16_precomputed, poseidons_24_precomputed),
     )
-    .unwrap_or_else(|err| {
+    .unwrap_or_else(|(last_pc, err)| {
         let lines_history = &instruction_history.lines;
         let latest_instructions = &lines_history[lines_history.len().saturating_sub(STACK_TRACE_INSTRUCTIONS)..];
         println!(
             "\n{}",
             crate::diagnostics::pretty_stack_trace(
                 &bytecode.program,
                 latest_instructions,
-                &bytecode.function_locations
+                &bytecode.function_locations,
+                last_pc
             )
         );
         if !std_out.is_empty() {
@@ -146,15 +147,15 @@ fn execute_bytecode_helper(
     no_vec_runtime_memory: usize,
     profiling: bool,
     (poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History),
-) -> Result<ExecutionResult, RunnerError> {
+) -> Result<ExecutionResult, (CodeAddress, RunnerError)> {
     // set public memory
     let mut memory = Memory::new(build_public_memory(public_input));
 
     let public_memory_size = (NONRESERVED_PROGRAM_INPUT_START + public_input.len()).next_power_of_two();
     let mut fp = public_memory_size;
 
     for (i, value) in private_input.iter().enumerate() {
-        memory.set(fp + i, *value)?;
+        memory.set(fp + i, *value).expect("to set private input in memory");
     }
 
     let mut mem_profile = MemoryProfile {
@@ -199,7 +200,7 @@ fn execute_bytecode_helper(
 
     while pc != ENDING_PC {
         if pc >= bytecode.instructions.len() {
-            return Err(RunnerError::PCOutOfBounds);
+            return Err((pc, RunnerError::PCOutOfBounds));
         }
 
         pcs.push(pc);
@@ -225,7 +226,7 @@ fn execute_bytecode_helper(
                 profiling,
                 memory_profile: &mut mem_profile,
             };
-            hint.execute_hint(&mut hint_ctx)?;
+            hint.execute_hint(&mut hint_ctx).map_err(|e| (pc, e))?;
         }
 
         let instruction = &bytecode.instructions[pc];
@@ -244,7 +245,9 @@ fn execute_bytecode_helper(
             n_poseidon16_precomputed_used: &mut n_poseidon16_precomputed_used,
             n_poseidon24_precomputed_used: &mut n_poseidon24_precomputed_used,
         };
-        instruction.execute_instruction(&mut instruction_ctx)?;
+        instruction
+            .execute_instruction(&mut instruction_ctx)
+            .map_err(|e| (pc, e))?;
     }
 
     assert_eq!(

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]
@@ -25,7 +25,14 @@ fn test_memory_already_set_error() {
     memory.set(0, F::ONE).unwrap();
 
     // Setting different value should fail
-    assert!(matches!(memory.set(0, F::ZERO), Err(RunnerError::MemoryAlreadySet)));
+    assert!(matches!(
+        memory.set(0, F::ZERO),
+        Err(RunnerError::MemoryAlreadySet {
+            address: 0,
+            prev_value: F::ONE,
+            new_value: F::ZERO
+        })
+    ));
 }
 
 #[test]