Skip to main content

tasm_lib/traits/
compiled_program.rs

1use anyhow::Result;
2use anyhow::anyhow;
3use triton_vm::prelude::*;
4
5use crate::library::Library;
6use crate::snippet_bencher::BenchmarkResult;
7
8pub trait CompiledProgram {
9    fn rust_shadow(
10        public_input: &PublicInput,
11        nondeterminism: &NonDeterminism,
12    ) -> Result<Vec<BFieldElement>>;
13
14    fn program() -> Program {
15        let (program_instructions, library) = Self::code();
16        let library_instructions = library.all_imports();
17
18        Program::new(&[program_instructions, library_instructions].concat())
19    }
20
21    fn run(
22        public_input: &PublicInput,
23        nondeterminism: &NonDeterminism,
24    ) -> Result<Vec<BFieldElement>> {
25        VM::run(
26            Self::program(),
27            public_input.clone(),
28            nondeterminism.clone(),
29        )
30        .map_err(|err| anyhow!(err))
31    }
32
33    fn code() -> (Vec<LabelledInstruction>, Library);
34
35    fn crash_conditions() -> Vec<String> {
36        vec![]
37    }
38}
39
40pub fn test_rust_shadow<P: CompiledProgram>(
41    public_input: &PublicInput,
42    nondeterminism: &NonDeterminism,
43) {
44    let rust_output = P::rust_shadow(public_input, nondeterminism).unwrap();
45    let tasm_output = P::run(public_input, nondeterminism).unwrap();
46    assert_eq!(rust_output, tasm_output);
47}
48
49/// Run the program, collect benchmarkable performance statistics (including a profile),
50/// and write them to disk.
51pub fn bench_and_profile_program<P: CompiledProgram>(
52    name: &str,
53    case: crate::snippet_bencher::BenchmarkCase,
54    public_input: &PublicInput,
55    nondeterminism: &NonDeterminism,
56) {
57    use crate::snippet_bencher::NamedBenchmarkResult;
58
59    let (program_instructions, library) = P::code();
60    let library_instructions = library.all_imports();
61    let all_instructions = [program_instructions, library_instructions].concat();
62    let program = Program::new(&all_instructions);
63
64    // run in trace mode to get table heights
65    let (aet, _output) = VM::trace_execution(
66        program.clone(),
67        public_input.clone(),
68        nondeterminism.clone(),
69    )
70    .unwrap();
71    let benchmark_result = BenchmarkResult::new(&aet);
72    let benchmark = NamedBenchmarkResult {
73        name: name.to_owned(),
74        benchmark_result,
75        case,
76    };
77
78    crate::snippet_bencher::write_benchmarks(vec![benchmark]);
79
80    // write profile to standard output in case someone is watching
81    let profile = crate::generate_full_profile(name, program, public_input, nondeterminism);
82    write_profile(name, &profile);
83    println!("{profile}");
84}
85
86#[cfg(not(target_arch = "wasm32"))]
87fn write_profile(name: &str, profile: &str) {
88    use std::io::Write;
89
90    // write profile to profile file
91    let mut path = std::path::PathBuf::new();
92    path.push("profiles");
93    std::fs::create_dir_all(&path).expect("profiles directory should exist");
94
95    path.push(std::path::Path::new(name).with_extension("profile"));
96    let mut file = std::fs::File::create(path).expect("open file for writing");
97    write!(file, "{profile}").unwrap();
98}
99
100// file access is not possible on `wasm32` architectures; ignore attempts
101#[cfg(target_arch = "wasm32")]
102fn write_profile(_: &str, _: &str) {}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use crate::test_prelude::*;
108
109    pub(super) struct FiboTest;
110
111    impl CompiledProgram for FiboTest {
112        fn rust_shadow(
113            public_input: &PublicInput,
114            _secret_input: &NonDeterminism,
115        ) -> Result<Vec<BFieldElement>> {
116            let num_iterations = public_input.individual_tokens[0].value() as usize;
117            let mut a = BFieldElement::new(0);
118            let mut b = BFieldElement::new(1);
119            for _ in 0..num_iterations {
120                let c = a + b;
121                a = b;
122                b = c;
123            }
124
125            Ok(vec![b])
126        }
127
128        fn code() -> (Vec<LabelledInstruction>, Library) {
129            let code = triton_asm!(
130                push 0
131                push 1
132                read_io 1
133                call fibo_test_loop
134                pop 1
135                write_io 1
136                halt
137
138                // INVARIANT: _ a b itr
139                fibo_test_loop:
140                    dup 0 push 0 eq
141                    skiz return
142
143                    push -1 add
144
145                    dup 2
146                    dup 2
147                    add
148                    swap 1
149                    recurse
150            );
151
152            (code, Library::default())
153        }
154    }
155
156    #[macro_rules_attr::apply(test)]
157    fn test_fibo_shadow() {
158        let public_input = PublicInput::new(vec![BFieldElement::new(501)]);
159        let nondeterminism = NonDeterminism::new(vec![]);
160        test_rust_shadow::<FiboTest>(&public_input, &nondeterminism);
161    }
162}
163
164#[cfg(test)]
165mod benches {
166    use super::tests::FiboTest;
167    use super::*;
168    use crate::test_prelude::*;
169
170    #[macro_rules_attr::apply(test)]
171    fn bench_fibo() {
172        let public_input = PublicInput::new(vec![BFieldElement::new(501)]);
173        let secret_input = NonDeterminism::new(vec![]);
174        bench_and_profile_program::<FiboTest>(
175            "fibo_test",
176            BenchmarkCase::CommonCase,
177            &public_input,
178            &secret_input,
179        );
180    }
181}