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

tasm_lib/traits/
compiled_program.rs
