Update object code pretty printing with labels

crates/lean_compiler/src/a_simplify_lang.rs

@@ -7,7 +7,7 @@ use crate::{
     },
     precompiles::Precompile,
 };
-use lean_vm::LocationInSourceCode;
+use lean_vm::SourceLineNumber;
 use std::{
     collections::{BTreeMap, BTreeSet},
     fmt::{Display, Formatter},
@@ -137,7 +137,7 @@ pub enum SimpleLine {
     },
     // noop, debug purpose only
     LocationReport {
-        location: LocationInSourceCode,
+        location: SourceLineNumber,
     },
 }
 

crates/lean_compiler/src/c_compile_final.rs

@@ -1,6 +1,6 @@
 use crate::{F, PUBLIC_INPUT_START, ZERO_VEC_PTR, intermediate_bytecode::*, lang::*};
 use lean_vm::*;
-use p3_field::PrimeCharacteristicRing;
+use p3_field::{PrimeCharacteristicRing, PrimeField32};
 use std::collections::BTreeMap;
 use utils::ToUsize;
 
@@ -28,9 +28,9 @@ impl IntermediateInstruction {
 
 struct Compiler {
     memory_size_per_function: BTreeMap<String, usize>,
-    label_to_pc: BTreeMap<Label, usize>,
+    label_to_pc: BTreeMap<Label, CodeAddress>,
     match_block_sizes: Vec<usize>,
-    match_first_block_starts: Vec<usize>,
+    match_first_block_starts: Vec<CodeAddress>,
 }
 
 pub fn compile_to_low_level_bytecode(
@@ -91,6 +91,14 @@ pub fn compile_to_low_level_bytecode(
     let mut low_level_bytecode = Vec::new();
     let mut hints = BTreeMap::new();
 
+    for (label, pc) in label_to_pc.clone() {
+        let pc_hints: &mut Vec<Hint> = match hints.try_insert(pc, vec![]) {
+            Ok(pc_hints) => pc_hints,
+            Err(_) => hints.get_mut(&pc).unwrap(),
+        };
+        pc_hints.push(Hint::Label { label });
+    }
+
     let compiler = Compiler {
         memory_size_per_function: intermediate_bytecode.memory_size_per_function,
         label_to_pc,
@@ -119,9 +127,9 @@ pub fn compile_to_low_level_bytecode(
 fn compile_block(
     compiler: &Compiler,
     block: &[IntermediateInstruction],
-    pc_start: usize,
+    pc_start: CodeAddress,
     low_level_bytecode: &mut Vec<Instruction>,
-    hints: &mut BTreeMap<usize, Vec<Hint>>,
+    hints: &mut BTreeMap<CodeAddress, Vec<Hint>>,
 ) {
     let try_as_mem_or_constant = |value: &IntermediateValue| {
         if let Some(cst) = try_as_constant(value, compiler) {
@@ -143,6 +151,39 @@ fn compile_block(
         _ => None,
     };
 
+    let codegen_jump = |hints: &BTreeMap<CodeAddress, Vec<Hint>>,
+                        low_level_bytecode: &mut Vec<Instruction>,
+                        condition: IntermediateValue,
+                        dest: IntermediateValue,
+                        updated_fp: Option<IntermediateValue>| {
+        let dest =
+            try_as_mem_or_constant(&dest).expect("Fatal: Could not materialize jump destination");
+        let label = match dest {
+            MemOrConstant::Constant(dest) => hints
+                .get(&usize::try_from(dest.as_canonical_u32()).unwrap())
+                .and_then(|hints: &Vec<Hint>| {
+                    hints.iter().find_map(|x| match x {
+                        Hint::Label { label } => Some(label),
+                        _ => None,
+                    })
+                })
+                .expect("Fatal: Unlabeled jump destination")
+                .clone(),
+            MemOrConstant::MemoryAfterFp { offset } => {
+                format!("fp+{offset}")
+            }
+        };
+        let updated_fp = updated_fp
+            .map(|fp| try_as_mem_or_fp(&fp).unwrap())
+            .unwrap_or(MemOrFp::Fp);
+        low_level_bytecode.push(Instruction::Jump {
+            condition: try_as_mem_or_constant(&condition).unwrap(),
+            label,
+            dest,
+            updated_fp,
+        });
+    };
+
     let mut pc = pc_start;
     for instruction in block {
         match instruction.clone() {
@@ -216,24 +257,11 @@ fn compile_block(
                 condition,
                 dest,
                 updated_fp,
-            } => {
-                let updated_fp = updated_fp
-                    .map(|fp| try_as_mem_or_fp(&fp).unwrap())
-                    .unwrap_or(MemOrFp::Fp);
-                low_level_bytecode.push(Instruction::Jump {
-                    condition: try_as_mem_or_constant(&condition).unwrap(),
-                    dest: try_as_mem_or_constant(&dest).unwrap(),
-                    updated_fp,
-                });
-            }
+            } => codegen_jump(hints, low_level_bytecode, condition, dest, updated_fp),
             IntermediateInstruction::Jump { dest, updated_fp } => {
-                low_level_bytecode.push(Instruction::Jump {
-                    condition: MemOrConstant::one(),
-                    dest: try_as_mem_or_constant(&dest).unwrap(),
-                    updated_fp: updated_fp
-                        .map(|fp| try_as_mem_or_fp(&fp).unwrap())
-                        .unwrap_or(MemOrFp::Fp),
-                });
+                let one =
+                    IntermediateValue::Constant(ConstExpression::Value(ConstantValue::Scalar(1)));
+                codegen_jump(hints, low_level_bytecode, one, dest, updated_fp)
             }
             IntermediateInstruction::Poseidon2_16 { arg_a, arg_b, res } => {
                 low_level_bytecode.push(Instruction::Poseidon2_16 {

crates/lean_compiler/src/intermediate_bytecode.rs

@@ -150,7 +150,7 @@ pub enum IntermediateInstruction {
     },
     // noop, debug purpose only
     LocationReport {
-        location: LocationInSourceCode,
+        location: SourceLineNumber,
     },
 }
 

crates/lean_compiler/src/lang.rs

@@ -356,7 +356,7 @@ pub enum Line {
     },
     // noop, debug purpose only
     LocationReport {
-        location: LocationInSourceCode,
+        location: SourceLineNumber,
     },
 }
 impl Display for Expression {

crates/lean_compiler/src/lib.rs

@@ -1,3 +1,5 @@
+#![feature(map_try_insert)]
+
 use std::collections::BTreeMap;
 
 use lean_vm::*;
@@ -24,7 +26,11 @@ pub fn compile_program(program: &str) -> (Bytecode, BTreeMap<usize, String>) {
     let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
     // println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
     let compiled = compile_to_low_level_bytecode(intermediate_bytecode).unwrap();
-    println!("Compiled Program:\n\n{compiled}");
+    println!("Function Locations: \n");
+    for (loc, name) in function_locations.iter() {
+        println!("{name}: {loc}");
+    }
+    println!("\n\nCompiled Program:\n\n{compiled}");
     (compiled, function_locations)
 }
 

crates/lean_compiler/src/parser.rs

@@ -175,7 +175,7 @@ fn parse_parameter(pair: Pair<'_, Rule>) -> Result<(String, bool), ParseError> {
     Ok((first.as_str().to_string(), false))
 }
 
-fn on_new_located_line(lines: &mut Vec<Line>, line_and_location: (Line, LocationInSourceCode)) {
+fn on_new_located_line(lines: &mut Vec<Line>, line_and_location: (Line, SourceLineNumber)) {
     let (line, location) = line_and_location;
     lines.push(Line::LocationReport { location });
     lines.push(line);
@@ -185,7 +185,7 @@ fn parse_statement(
     pair: Pair<'_, Rule>,
     constants: &BTreeMap<String, usize>,
     trash_var_count: &mut usize,
-) -> Result<(Line, LocationInSourceCode), ParseError> {
+) -> Result<(Line, SourceLineNumber), ParseError> {
     let location = pair.line_col().0;
     let inner = pair.into_inner().next().unwrap();
     let line = match inner.as_rule() {

crates/lean_prover/witness_generation/src/instruction_encoder.rs

@@ -53,6 +53,7 @@ pub fn field_representation(instr: &Instruction) -> [F; N_INSTRUCTION_COLUMNS] {
         }
         Instruction::Jump {
             condition,
+            label: _,
             dest,
             updated_fp,
         } => {

crates/lean_vm/src/core/types.rs

@@ -7,7 +7,10 @@ pub type F = KoalaBear;
 pub type EF = QuinticExtensionFieldKB;
 
 /// Location in source code for debugging
-pub type LocationInSourceCode = usize;
+pub type SourceLineNumber = usize;
 
 /// String label for bytecode locations
 pub type Label = String;
+
+/// Bytecode address (i.e., a value of the program counter)
+pub type CodeAddress = usize;

crates/lean_vm/src/diagnostics/stack_trace.rs

@@ -2,13 +2,13 @@ use std::collections::BTreeMap;
 
 use colored::Colorize;
 
-use crate::LocationInSourceCode;
+use crate::SourceLineNumber;
 
 const STACK_TRACE_MAX_LINES_PER_FUNCTION: usize = 5;
 
 pub(crate) fn pretty_stack_trace(
     source_code: &str,
-    instructions: &[LocationInSourceCode], // LocationInSourceCode = usize
+    instructions: &[SourceLineNumber], // SourceLineNumber = usize
     function_locations: &BTreeMap<usize, String>,
 ) -> String {
     let source_lines: Vec<&str> = source_code.lines().collect();
@@ -131,7 +131,7 @@ pub(crate) fn find_function_for_line(
 
 fn count_remaining_lines_in_function(
     current_idx: usize,
-    instructions: &[LocationInSourceCode],
+    instructions: &[SourceLineNumber],
     function_locations: &BTreeMap<usize, String>,
     current_function_line: usize,
 ) -> usize {

crates/lean_vm/src/execution/context.rs

@@ -1,9 +1,9 @@
-use crate::core::LocationInSourceCode;
+use crate::core::SourceLineNumber;
 use std::collections::BTreeMap;
 
 #[derive(Debug, Clone, Default)]
 pub struct ExecutionHistory {
-    pub lines: Vec<LocationInSourceCode>,
+    pub lines: Vec<SourceLineNumber>,
     pub cycles: Vec<usize>, // for each line, how many cycles it took
 }
 
@@ -12,7 +12,7 @@ impl ExecutionHistory {
         Self::default()
     }
 
-    pub fn add_line(&mut self, location: LocationInSourceCode, cycles: usize) {
+    pub fn add_line(&mut self, location: SourceLineNumber, cycles: usize) {
         self.lines.push(location);
         self.cycles.push(cycles);
     }

crates/lean_vm/src/isa/bytecode.rs

@@ -1,6 +1,6 @@
 //! Bytecode representation and management
 
-use crate::Hint;
+use crate::{CodeAddress, Hint};
 
 use super::Instruction;
 use std::collections::BTreeMap;
@@ -10,9 +10,9 @@ use std::fmt::{Display, Formatter};
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct Bytecode {
     pub instructions: Vec<Instruction>,
-    pub hints: BTreeMap<usize, Vec<Hint>>, // pc -> hints
+    pub hints: BTreeMap<CodeAddress, Vec<Hint>>, // pc -> hints
     pub starting_frame_memory: usize,
-    pub ending_pc: usize,
+    pub ending_pc: CodeAddress,
 }
 
 impl Display for Bytecode {

crates/lean_vm/src/isa/hint.rs

@@ -1,4 +1,4 @@
-use crate::core::{DIMENSION, F, LOG_VECTOR_LEN, LocationInSourceCode, VECTOR_LEN};
+use crate::core::{DIMENSION, F, LOG_VECTOR_LEN, Label, SourceLineNumber, VECTOR_LEN};
 use crate::diagnostics::RunnerError;
 use crate::execution::{ExecutionHistory, Memory};
 use crate::isa::operands::MemOrConstant;
@@ -50,8 +50,10 @@ pub enum Hint {
     /// Report source code location for debugging
     LocationReport {
         /// Source code location
-        location: LocationInSourceCode,
+        location: SourceLineNumber,
     },
+    /// Jump destination label (for debugging purposes)
+    Label { label: Label },
 }
 
 /// Execution state for hint processing
@@ -169,6 +171,7 @@ impl Hint {
                 *ctx.cpu_cycles_before_new_line = 0;
                 Ok(false)
             }
+            Self::Label { label: _ } => Ok(false),
         }
     }
 }
@@ -220,7 +223,14 @@ impl Display for Hint {
             Self::Inverse { arg, res_offset } => {
                 write!(f, "m[fp + {res_offset}] = inverse({arg})")
             }
-            Self::LocationReport { .. } => Ok(()),
+            Self::LocationReport {
+                location: line_number,
+            } => {
+                write!(f, "source line number: {line_number}")
+            }
+            Self::Label { label } => {
+                write!(f, "label: {label}")
+            }
         }
     }
 }

crates/lean_vm/src/isa/instruction.rs

@@ -2,7 +2,7 @@
 
 use super::Operation;
 use super::operands::{MemOrConstant, MemOrFp, MemOrFpOrConstant};
-use crate::core::{DIMENSION, EF, F, VECTOR_LEN};
+use crate::core::{DIMENSION, EF, F, Label, VECTOR_LEN};
 use crate::diagnostics::RunnerError;
 use crate::execution::Memory;
 use crate::witness::{
@@ -46,6 +46,8 @@ pub enum Instruction {
     Jump {
         /// Jump condition (jump if non-zero)
         condition: MemOrConstant,
+        /// Jump destination label (for debugging purposes)
+        label: Label,
         /// Jump destination address
         dest: MemOrConstant,
         /// New frame pointer value after jump
@@ -187,6 +189,7 @@ impl Instruction {
             }
             Self::Jump {
                 condition,
+                label: _,
                 dest,
                 updated_fp,
             } => {
@@ -433,12 +436,13 @@ impl Display for Instruction {
             }
             Self::Jump {
                 condition,
+                label,
                 dest,
                 updated_fp,
             } => {
                 write!(
                     f,
-                    "if {condition} != 0 jump to {dest} with next(fp) = {updated_fp}"
+                    "if {condition} != 0 jump to {label} = {dest} with next(fp) = {updated_fp}"
                 )
             }
             Self::Poseidon2_16 { arg_a, arg_b, res } => {

crates/packed_pcs/src/lib.rs

@@ -321,8 +321,7 @@ pub fn packed_pcs_global_statements_for_prover<
                 assert_eq!(
                     chunks[0].n_vars,
                     statement.point.0.len(),
-                    "poly: {}",
-                    poly_index
+                    "poly: {poly_index}"
                 );
                 assert!(chunks[0].offset_in_packed.unwrap() % (1 << chunks[0].n_vars) == 0);
 

crates/utils/src/multilinear.rs

@@ -174,9 +174,7 @@ pub fn multilinear_eval_constants_at_right<F: Field>(limit: usize, point: &[F])
 
     assert!(
         limit <= (1 << n_vars),
-        "limit {} is too large for n_vars {}",
-        limit,
-        n_vars
+        "limit {limit} is too large for n_vars {n_vars}"
     );
 
     if limit == 1 << n_vars {