feat: include debug information when printing Program

triton-vm/src/parser.rs

@@ -1085,4 +1085,35 @@ pub(crate) mod tests {
         let program = triton_program! { break halt break };
         assert_eq!(1, program.instructions.len());
     }
+
+    #[test]
+    fn printing_program_includes_debug_information() {
+        let source_code = "\
+            call foo\n\
+            break\n\
+            call bar\n\
+            halt\n\
+            foo:\n\
+            break\n\
+            call baz\n\
+            push 1\n\
+            nop\n\
+            return\n\
+            baz:\n\
+            hash\n\
+            return\n\
+            nop\n\
+            pop\n\
+            bar:\n\
+            divine\n\
+            skiz\n\
+            split\n\
+            break\n\
+            return\n\
+        ";
+        let program = Program::from_code(source_code).unwrap();
+        let printed_program = format!("{program}");
+        assert_eq!(source_code, &printed_program);
+        println!("{program}");
+    }
 }

triton-vm/src/program.rs

@@ -1,5 +1,6 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::fmt::Display;
 use std::fmt::Formatter;
 use std::fmt::Result as FmtResult;
@@ -40,13 +41,8 @@ pub struct Program {
 
 impl Display for Program {
     fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
-        let mut stream = self.instructions.iter();
-        while let Some(instruction) = stream.next() {
+        for instruction in self.labelled_instructions() {
             writeln!(f, "{instruction}")?;
-            // 2-word instructions already print their arguments
-            for _ in 1..instruction.size() {
-                stream.next();
-            }
         }
         Ok(())
     }
@@ -236,6 +232,46 @@ impl Program {
             .map_err(|err| anyhow!("{err}"))
     }
 
+    pub fn labelled_instructions(&self) -> Vec<LabelledInstruction> {
+        let call_targets = self.call_targets();
+        let instructions_with_labels = self.instructions.iter().map(|instruction| {
+            instruction.map_call_address(|&address| self.label_for_address(address))
+        });
+
+        let mut labelled_instructions = vec![];
+        let mut address = 0;
+        let mut instruction_stream = instructions_with_labels.into_iter();
+        while let Some(instruction) = instruction_stream.next() {
+            let instruction_size = instruction.size() as u64;
+            if call_targets.contains(&address) {
+                let address = address.into();
+                let label = self.label_for_address(address);
+                let label = LabelledInstruction::Label(label);
+                labelled_instructions.push(label);
+            }
+            if self.breakpoints[address as usize] {
+                labelled_instructions.push(LabelledInstruction::Breakpoint);
+            }
+            labelled_instructions.push(LabelledInstruction::Instruction(instruction));
+
+            for _ in 1..instruction_size {
+                instruction_stream.next();
+            }
+            address += instruction_size;
+        }
+        labelled_instructions
+    }
+
+    fn call_targets(&self) -> HashSet<u64> {
+        self.instructions
+            .iter()
+            .filter_map(|instruction| match instruction {
+                Instruction::Call(address) => Some(address.value()),
+                _ => None,
+            })
+            .collect()
+    }
+
     /// Convert a `Program` to a `Vec<BFieldElement>`.
     ///
     /// Every single-word instruction is converted to a single word.