Tracing executed debug locations.

riscv/src/disambiguator.rs

@@ -88,11 +88,14 @@ fn disambiguate_file_ids(
         .iter()
         .flat_map(|(name, statements)| extract_file_ids(name, statements))
         .collect::<Vec<_>>();
-    let debug_file_id_mapping = debug_file_ids
-        .iter()
-        .enumerate()
-        .map(|(i, (asm_name, file_id, ..))| ((asm_name.to_string(), *file_id), i as i64 + 1))
-        .collect::<HashMap<_, _>>();
+    // ensure the ids are densely packed:
+    let debug_file_id_mapping = {
+        let mut map = HashMap::new();
+        for (asm_name, file_id, ..) in debug_file_ids.iter() {
+            map.insert((asm_name.to_string(), *file_id), map.len() as i64 + 1);
+        }
+        map
+    };
     let new_debug_file_ids = debug_file_ids
         .into_iter()
         .map(|(asm_file, id, dir, file)| {

riscv_executor/src/lib.rs

@@ -15,7 +15,10 @@ use std::{
 
 use ast::{
     asm_analysis::{AnalysisASMFile, CallableSymbol, FunctionStatement, LabelStatement, Machine},
-    parsed::{asm::AssignmentRegister, Expression},
+    parsed::{
+        asm::{AssignmentRegister, DebugDirective},
+        Expression,
+    },
 };
 use builder::{MemoryBuilder, TraceBuilder};
 use number::{BigInt, FieldElement};
@@ -132,21 +135,25 @@ mod builder {
             self.trace.values[self.curr_idx + idx]
         }
 
-        /// set next value of register
+        /// set next value of register, accounting to x0 or pc writes
         pub(crate) fn s(&mut self, idx: &str, value: impl Into<Elem>) {
-            self.s_idx(self.trace.reg_map[idx], value.into());
+            self.s_impl(idx, value.into())
         }
 
-        /// set next value of register by register index instead of name
-        fn s_idx(&mut self, idx: usize, value: Elem) {
+        fn s_impl(&mut self, idx: &str, value: Elem) {
+            let idx = self.trace.reg_map[idx];
             if idx == self.x0_idx {
                 return;
             } else if idx == self.pc_idx {
                 // PC has been written, so we must update our statement-based
                 // program counter accordingly:
                 self.next_statement_line = self.batch_to_line_map[value.u() as usize];
             }
+            self.s_idx(idx, value.into());
+        }
 
+        /// raw set next value of register by register index instead of name
+        fn s_idx(&mut self, idx: usize, value: Elem) {
             let final_idx = self.curr_idx + self.reg_len() + idx;
             self.trace.values[final_idx] = value;
         }
@@ -179,6 +186,7 @@ mod builder {
             let curr_pc = self.g_idx(self.pc_idx).u();
 
             let line_of_next_batch = self.batch_to_line_map[curr_pc as usize + 1];
+
             self.s_idx(
                 self.pc_idx,
                 if self.next_statement_line >= line_of_next_batch {
@@ -229,7 +237,12 @@ fn get_main_machine<T: FieldElement>(program: &AnalysisASMFile<T>) -> &Machine<T
 /// labels to indices into that list, and the list with the start of each batch.
 fn preprocess_main_function<T: FieldElement>(
     machine: &Machine<T>,
-) -> (Vec<&FunctionStatement<T>>, HashMap<&str, Elem>, Vec<u32>) {
+) -> (
+    Vec<&FunctionStatement<T>>,
+    HashMap<&str, Elem>,
+    Vec<u32>,
+    Vec<(&str, &str)>,
+) {
     let CallableSymbol::Function(main_function) = &machine.callable.0["main"] else {
         panic!("main function missing")
     };
@@ -239,6 +252,7 @@ fn preprocess_main_function<T: FieldElement>(
     let mut statements = Vec::new();
     let mut label_map = HashMap::new();
     let mut batch_to_line_map = Vec::new();
+    let mut debug_files = Vec::new();
 
     for (batch_idx, batch) in orig_statements.iter_batches().enumerate() {
         batch_to_line_map.push(statements.len() as u32);
@@ -251,15 +265,25 @@ fn preprocess_main_function<T: FieldElement>(
                     statement_seen = true;
                     statements.push(s)
                 }
+                FunctionStatement::DebugDirective(d) => {
+                    match &d.directive {
+                        ast::parsed::asm::DebugDirective::File(idx, dir, file) => {
+                            // debug files should be densely packed starting
+                            // from 1, so the idx should match vec size + 1:
+                            assert_eq!(*idx, debug_files.len() + 1);
+                            debug_files.push((dir.as_str(), file.as_str()));
+                        }
+                        ast::parsed::asm::DebugDirective::Loc(_, _, _) => {
+                            // keep debug locs for debugging purposes
+                            statements.push(s);
+                        }
+                    }
+                }
                 FunctionStatement::Label(LabelStatement { start: _, name }) => {
                     // assert there are no statements in the middle of a block
                     assert!(!statement_seen);
                     label_map.insert(name.as_str(), (batch_idx as i64).into());
                 }
-                FunctionStatement::DebugDirective(_) => {
-                    // discard debug directives
-                    // TODO: use them to debug
-                }
             }
         }
     }
@@ -268,7 +292,7 @@ fn preprocess_main_function<T: FieldElement>(
     // add a final element to the map so the queries don't overflow:
     batch_to_line_map.push(statements.len() as u32);
 
-    (statements, label_map, batch_to_line_map)
+    (statements, label_map, batch_to_line_map, debug_files)
 }
 
 struct Executor<'a, 'b, F: FieldElement> {
@@ -688,7 +712,8 @@ pub fn execute_ast<'a, T: FieldElement>(
     inputs: &[T],
 ) -> ExecutionTrace<'a> {
     let main_machine = get_main_machine(program);
-    let (statements, label_map, batch_to_line_map) = preprocess_main_function(main_machine);
+    let (statements, label_map, batch_to_line_map, debug_files) =
+        preprocess_main_function(main_machine);
 
     let mut e = Executor {
         proc: TraceBuilder::new(main_machine, &batch_to_line_map),
@@ -720,7 +745,14 @@ pub fn execute_ast<'a, T: FieldElement>(
                 e.exec_instruction(&i.instruction, &i.inputs);
             }
             FunctionStatement::Return(_) => break,
-            FunctionStatement::DebugDirective(_) | FunctionStatement::Label(_) => {
+            FunctionStatement::DebugDirective(dd) => {
+                let DebugDirective::Loc(file, line, column) = &dd.directive else {
+                    unreachable!()
+                };
+                let (dir, file) = debug_files[file - 1];
+                println!("Executed {dir}/{file}:{line}:{column}");
+            }
+            FunctionStatement::Label(_) => {
                 unreachable!()
             }
         }