tasm_lib/traits/
algorithm.rs1use std::collections::HashMap;
2
3use rand::prelude::*;
4use triton_vm::prelude::*;
5
6use super::basic_snippet::BasicSnippet;
7use super::rust_shadow::RustShadow;
8use super::rust_shadow::RustShadowError;
9use crate::InitVmState;
10use crate::linker::execute_bench;
11use crate::prelude::Tip5;
12use crate::snippet_bencher::BenchmarkCase;
13use crate::snippet_bencher::NamedBenchmarkResult;
14use crate::snippet_bencher::write_benchmarks;
15use crate::test_helpers::test_rust_equivalence_given_complete_state;
16
17pub trait Algorithm: BasicSnippet {
33 fn rust_shadow(
34 &self,
35 stack: &mut Vec<BFieldElement>,
36 memory: &mut HashMap<BFieldElement, BFieldElement>,
37 nondeterminism: &NonDeterminism,
38 ) -> Result<(), RustShadowError>;
39
40 fn preprocess<T: BFieldCodec>(_meta_input: T, _nondeterminism: &mut NonDeterminism) {}
60
61 fn pseudorandom_initial_state(
62 &self,
63 seed: [u8; 32],
64 bench_case: Option<BenchmarkCase>,
65 ) -> AlgorithmInitialState;
66
67 fn corner_case_initial_states(&self) -> Vec<AlgorithmInitialState> {
68 vec![]
69 }
70}
71
72#[derive(Debug, Clone, Default)]
73pub struct AlgorithmInitialState {
74 pub stack: Vec<BFieldElement>,
75 pub nondeterminism: NonDeterminism,
76}
77
78impl From<AlgorithmInitialState> for InitVmState {
79 fn from(value: AlgorithmInitialState) -> Self {
80 Self {
81 stack: value.stack,
82 nondeterminism: value.nondeterminism,
83 ..Default::default()
84 }
85 }
86}
87
88pub struct ShadowedAlgorithm<T: Algorithm> {
89 algorithm: T,
90}
91
92impl<T: Algorithm> ShadowedAlgorithm<T> {
93 pub fn new(algorithm: T) -> Self {
94 Self { algorithm }
95 }
96}
97
98impl<T> RustShadow for ShadowedAlgorithm<T>
99where
100 T: Algorithm,
101{
102 fn inner(&self) -> &dyn BasicSnippet {
103 &self.algorithm
104 }
105
106 fn rust_shadow_wrapper(
107 &self,
108 _stdin: &[BFieldElement],
109 nondeterminism: &NonDeterminism,
110 stack: &mut Vec<BFieldElement>,
111 memory: &mut HashMap<BFieldElement, BFieldElement>,
112 _sponge: &mut Option<Tip5>,
113 ) -> Result<Vec<BFieldElement>, RustShadowError> {
114 self.algorithm.rust_shadow(stack, memory, nondeterminism)?;
115
116 Ok(Vec::new())
117 }
118
119 fn test(&self) {
120 for corner_case in self.algorithm.corner_case_initial_states() {
121 let stdin = vec![];
122 test_rust_equivalence_given_complete_state(
123 self,
124 &corner_case.stack,
125 &stdin,
126 &corner_case.nondeterminism,
127 &None,
128 None,
129 );
130 }
131
132 let num_states = 10;
133 let seed: [u8; 32] = rand::rng().random();
134 let mut rng = StdRng::from_seed(seed);
135 for _ in 0..num_states {
136 let seed: [u8; 32] = rng.random();
137 let AlgorithmInitialState {
138 stack,
139 nondeterminism,
140 } = self.algorithm.pseudorandom_initial_state(seed, None);
141
142 let stdin = vec![];
143 test_rust_equivalence_given_complete_state(
144 self,
145 &stack,
146 &stdin,
147 &nondeterminism,
148 &None,
149 None,
150 );
151 }
152 }
153
154 fn bench(&self) {
155 let mut rng = StdRng::from_seed(
156 hex::decode("73a24b6b8b32e4d7d563a4d9a85f476573a24b6b8b32e4d7d563a4d9a85f4765")
157 .unwrap()
158 .try_into()
159 .unwrap(),
160 );
161 let mut benchmarks = Vec::with_capacity(2);
162
163 for bench_case in [BenchmarkCase::CommonCase, BenchmarkCase::WorstCase] {
164 let AlgorithmInitialState {
165 stack,
166 nondeterminism,
167 } = self
168 .algorithm
169 .pseudorandom_initial_state(rng.random(), Some(bench_case));
170 let program = self.algorithm.link_for_isolated_run();
171 let benchmark = execute_bench(&program, &stack, vec![], nondeterminism, None);
172 let benchmark = NamedBenchmarkResult {
173 name: self.algorithm.entrypoint(),
174 benchmark_result: benchmark,
175 case: bench_case,
176 };
177 benchmarks.push(benchmark);
178 }
179
180 write_benchmarks(benchmarks);
181 }
182}