diff --git a/triton-vm/src/parser.rs b/triton-vm/src/parser.rs index 7f28eac7..d848a5b5 100644 --- a/triton-vm/src/parser.rs +++ b/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}"); + } } diff --git a/triton-vm/src/program.rs b/triton-vm/src/program.rs index 95a1d001..0df21362 100644 --- a/triton-vm/src/program.rs +++ b/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 { + 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 { + self.instructions + .iter() + .filter_map(|instruction| match instruction { + Instruction::Call(address) => Some(address.value()), + _ => None, + }) + .collect() + } + /// Convert a `Program` to a `Vec`. /// /// Every single-word instruction is converted to a single word.