Skip to content

Commit

Permalink
refactor: store debug information address_to_label in Program
Browse files Browse the repository at this point in the history
  • Loading branch information
@jan-ferdinand
jan-ferdinand committed Oct 16, 2023
1 parent 9bc7a5e commit d857e83
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 82 deletions.
53 changes: 2 additions & 51 deletions triton-vm/src/instruction.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt::Display;
use std::fmt::Formatter;
Expand Down Expand Up @@ -260,9 +259,10 @@ impl<Dest: PartialEq + Default> AnInstruction<Dest> {
((opcode >> bit_number) & 1).into()
}

fn map_call_address<F, NewDest: PartialEq + Default>(&self, f: F) -> AnInstruction<NewDest>
pub(crate) fn map_call_address<F, NewDest>(&self, f: F) -> AnInstruction<NewDest>
where
F: Fn(&Dest) -> NewDest,
NewDest: PartialEq + Default,
{
match self {
Pop => Pop,
Expand Down Expand Up @@ -448,55 +448,6 @@ impl TryFrom<usize> for Instruction {
}
}

/// Convert a program with labels to a program with absolute addresses.
pub fn convert_all_labels_to_addresses(program: &[LabelledInstruction]) -> Vec<Instruction> {
let label_map = build_label_to_address_map(program);
program
.iter()
.flat_map(|instruction| convert_label_to_address_for_instruction(instruction, &label_map))
.collect()
}

pub(crate) fn build_label_to_address_map(program: &[LabelledInstruction]) -> HashMap<String, u64> {
use LabelledInstruction::*;

let mut label_map = HashMap::new();
let mut instruction_pointer = 0;

for labelled_instruction in program.iter() {
match labelled_instruction {
Label(label) => match label_map.entry(label.clone()) {
Entry::Occupied(_) => panic!("Duplicate label: {label}"),
Entry::Vacant(entry) => {
entry.insert(instruction_pointer);
}
},
Instruction(instruction) => instruction_pointer += instruction.size() as u64,
Breakpoint => (),
}
}
label_map
}

/// Convert a single instruction with a labelled call target to an instruction with an absolute
/// address as the call target. Discards all labels.
fn convert_label_to_address_for_instruction(
labelled_instruction: &LabelledInstruction,
label_map: &HashMap<String, u64>,
) -> Option<Instruction> {
let LabelledInstruction::Instruction(instruction) = labelled_instruction else {
return None;
};

let instruction_with_absolute_address = instruction.map_call_address(|label| {
label_map
.get(label)
.map(|&address| BFieldElement::new(address))
.unwrap_or_else(|| panic!("Label not found: {label}"))
});
Some(instruction_with_absolute_address)
}

const fn all_instructions_without_args() -> [AnInstruction<BFieldElement>; Instruction::COUNT] {
[
Pop,
Expand Down
110 changes: 79 additions & 31 deletions triton-vm/src/program.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt::Display;
use std::fmt::Formatter;
Expand All @@ -20,8 +21,6 @@ use twenty_first::util_types::algebraic_hasher::AlgebraicHasher;
use crate::aet::AlgebraicExecutionTrace;
use crate::ensure_eq;
use crate::error::InstructionError::InstructionPointerOverflow;
use crate::instruction::build_label_to_address_map;
use crate::instruction::convert_all_labels_to_addresses;
use crate::instruction::Instruction;
use crate::instruction::LabelledInstruction;
use crate::parser::parse;
Expand All @@ -35,6 +34,7 @@ use crate::vm::VMState;
#[derive(Debug, Clone, Default, PartialEq, Eq, GetSize, Serialize, Deserialize)]
pub struct Program {
pub instructions: Vec<Instruction>,
pub address_to_label: HashMap<u64, String>,
}

impl Display for Program {
Expand Down Expand Up @@ -87,7 +87,11 @@ impl BFieldCodec for Program {
}

ensure_eq!(idx, program_length);
Ok(Box::new(Program { instructions }))

Ok(Box::new(Program {
instructions,
address_to_label: HashMap::new(),
}))
}

fn encode(&self) -> Vec<BFieldElement> {
Expand Down Expand Up @@ -134,12 +138,61 @@ impl IntoIterator for Program {
impl Program {
/// Create a `Program` from a slice of `Instruction`.
pub fn new(labelled_instructions: &[LabelledInstruction]) -> Self {
let instructions = convert_all_labels_to_addresses(labelled_instructions)
let label_to_address = Self::build_label_to_address_map(labelled_instructions);
let instructions = labelled_instructions
.iter()
.flat_map(|&instr| vec![instr; instr.size()])
.flat_map(|instr| Self::turn_label_to_address_for_instruction(instr, &label_to_address))
.flat_map(|instr| vec![instr; instr.size()])
.collect();
let address_to_label = label_to_address
.into_iter()
.map(|(label, address)| (address, label))
.collect();

Program { instructions }
Program {
instructions,
address_to_label,
}
}

fn build_label_to_address_map(program: &[LabelledInstruction]) -> HashMap<String, u64> {
use LabelledInstruction::*;

let mut label_map = HashMap::new();
let mut instruction_pointer = 0;

for labelled_instruction in program.iter() {
match labelled_instruction {
Label(label) => match label_map.entry(label.clone()) {
Entry::Occupied(_) => panic!("Duplicate label: {label}"),
Entry::Vacant(entry) => {
entry.insert(instruction_pointer);
}
},
Instruction(instruction) => instruction_pointer += instruction.size() as u64,
Breakpoint => (),
}
}
label_map
}

/// Convert a single instruction with a labelled call target to an instruction with an absolute
/// address as the call target. Discards all labels.
fn turn_label_to_address_for_instruction(
labelled_instruction: &LabelledInstruction,
label_map: &HashMap<String, u64>,
) -> Option<Instruction> {
let LabelledInstruction::Instruction(instruction) = labelled_instruction else {
return None;
};

let instruction_with_absolute_address = instruction.map_call_address(|label| {
label_map
.get(label)
.map(|&address| BFieldElement::new(address))
.unwrap_or_else(|| panic!("Label not found: {label}"))
});
Some(instruction_with_absolute_address)
}

/// Create a `Program` by parsing source code.
Expand Down Expand Up @@ -317,35 +370,24 @@ impl Program {
}
}

/// Run Triton VM on the given program with the given public and secret input,
/// but record the number of cycles spent in each callable block of instructions.
/// This function returns a Result wrapping a program profiler report, which is a
/// Vec of [`ProfileLine`]s.
///
/// Note that the program is given as a list of [`LabelledInstruction`]s rather
/// than as a [`Program`] because the labels are needed to build a meaningful profiler
/// report.
/// Run Triton VM with the given public and secret input, but record the number of cycles spent
/// in each callable block of instructions. This function returns a Result wrapping a program
/// profiler report, which is a Vec of [`ProfileLine`]s.
///
/// See also [`run`](Self::run), [`trace_execution`](Self::trace_execution) and
/// [`debug`](Self::debug).
pub fn profile(
labelled_instructions: &[LabelledInstruction],
&self,
public_input: PublicInput,
non_determinism: NonDeterminism<BFieldElement>,
) -> Result<(Vec<BFieldElement>, Vec<ProfileLine>)> {
let address_to_label_map = build_label_to_address_map(labelled_instructions)
.into_iter()
.map(|(label, address)| (address, label))
.collect::<HashMap<_, _>>();
let mut call_stack = vec![];
let mut profile = vec![];

let program = Self::new(labelled_instructions);
let mut state = VMState::new(&program, public_input, non_determinism);
let mut state = VMState::new(self, public_input, non_determinism);
while !state.halting {
if let Instruction::Call(address) = state.current_instruction()? {
let address = address.value();
let label = address_to_label_map[&address].to_owned();
let label = self.label_for_address(address);
let profile_line = ProfileLine::new(call_stack.len(), label, 0);
let profile_line_number = profile.len();
profile.push(profile_line);
Expand All @@ -368,6 +410,15 @@ impl Program {

Ok((state.public_output, profile))
}

fn label_for_address(&self, address: BFieldElement) -> String {
let address = address.value();
let substitute_for_unknown_label = format! {"address_{address}"};
self.address_to_label
.get(&address)
.unwrap_or(&substitute_for_unknown_label)
.to_owned()
}
}

/// A single line in a profile report for profiling Triton Assembly programs.
Expand Down Expand Up @@ -559,9 +610,8 @@ mod tests {
use rand::Rng;
use twenty_first::shared_math::tip5::Tip5;

use crate::example_programs::calculate_new_mmr_peaks_from_append_with_safe_lists;
use crate::example_programs::CALCULATE_NEW_MMR_PEAKS_FROM_APPEND_WITH_SAFE_LISTS;
use crate::parser::tests::program_gen;
use crate::triton_asm;
use crate::triton_program;

use super::*;
Expand Down Expand Up @@ -705,10 +755,8 @@ mod tests {

#[test]
fn test_profile() {
let labelled_instructions = calculate_new_mmr_peaks_from_append_with_safe_lists();
let (profile_output, profile) =
Program::profile(&labelled_instructions, [].into(), [].into()).unwrap();
let program = Program::new(&labelled_instructions);
let program = CALCULATE_NEW_MMR_PEAKS_FROM_APPEND_WITH_SAFE_LISTS.clone();
let (profile_output, profile) = program.profile([].into(), [].into()).unwrap();
let debug_terminal_state = program
.debug_terminal_state([].into(), [].into(), None, None)
.unwrap();
Expand All @@ -727,15 +775,15 @@ mod tests {

#[test]
fn test_profile_with_open_calls() {
let labelled_instructions = triton_asm! {
let program = triton_program! {
push 2 call outer_fn
outer_fn:
call inner_fn
dup 0 skiz recurse halt
inner_fn:
push -1 add return
};
let (_, profile) = Program::profile(&labelled_instructions, [].into(), [].into()).unwrap();
let (_, profile) = program.profile([].into(), [].into()).unwrap();

println!();
for line in profile.iter() {
Expand Down

0 comments on commit d857e83

Please sign in to comment.