tasm_lib/traits/
compiled_program.rs1use 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
49pub 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 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 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 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#[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 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}