tasm_lib/traits/
algorithm.rs

1use std::collections::HashMap;
2
3use rand::prelude::*;
4use triton_vm::prelude::*;
5
6use super::basic_snippet::BasicSnippet;
7use super::rust_shadow::RustShadow;
8use crate::linker::execute_bench;
9use crate::prelude::Tip5;
10use crate::snippet_bencher::write_benchmarks;
11use crate::snippet_bencher::BenchmarkCase;
12use crate::snippet_bencher::NamedBenchmarkResult;
13use crate::test_helpers::test_rust_equivalence_given_complete_state;
14use crate::InitVmState;
15
16/// An Algorithm can modify memory and take ND_input.
17///
18/// An Algorithm is a piece of tasm code that can modify memory even at addresses below
19/// the dynamic memory allocator, and can take nondeterministic input. It cannot read from
20/// standard in or write to standard out.
21///
22/// See also: [closure], [function], [procedure], [accessor], [mem_preserver],
23///           [read_only_algorithm]
24///
25/// [closure]: crate::traits::closure::Closure
26/// [function]: crate::traits::function::Function
27/// [procedure]: crate::traits::procedure::Procedure
28/// [accessor]: crate::traits::accessor::Accessor
29/// [mem_preserver]: crate::traits::mem_preserver::MemPreserver
30/// [read_only_algorithm]: crate::traits::read_only_algorithm::ReadOnlyAlgorithm
31pub trait Algorithm: BasicSnippet {
32    fn rust_shadow(
33        &self,
34        stack: &mut Vec<BFieldElement>,
35        memory: &mut HashMap<BFieldElement, BFieldElement>,
36        nondeterminism: &NonDeterminism,
37    );
38
39    /// Take a object about which something is being proven in order to extract out the
40    /// right nondeterminism. Update the mutably referenced non-determism argument.
41    ///
42    /// For example:
43    ///  - When proving the correct verification of a proof, you might want to pull all
44    ///    digests out of the authentication structures and put them in the `digests`
45    ///    field of the non-determinism. This way the VM can avoid the processing
46    ///    required by authentication structures and just divine in the right digests
47    ///    as it walks up the Merkle trees.
48    ///  - When proving the correct sorting of a list, the VM ought to avoid running a
49    ///    sorting algorithm; instead it should divine the sorted list and then prove
50    ///    that the two lists are equal. The preprocessing step in this case would take
51    ///    the unsorted list, sort it, and use the sorted list to populate the non-
52    ///    determinism.
53    ///  - When verifying a Falcon signature, at some point the NTT of a vector is needed.
54    ///    The VM should not compute the NTT because expensive; instead it should divine
55    ///    the NTT-transformed vector and verify that it satisfies the right relation. In
56    ///    this case the preprocessor calculates the NTT and populates the non-determinism
57    ///    with the transformed vector.
58    fn preprocess<T: BFieldCodec>(_meta_input: T, _nondeterminism: &mut NonDeterminism) {}
59
60    fn pseudorandom_initial_state(
61        &self,
62        seed: [u8; 32],
63        bench_case: Option<BenchmarkCase>,
64    ) -> AlgorithmInitialState;
65
66    fn corner_case_initial_states(&self) -> Vec<AlgorithmInitialState> {
67        vec![]
68    }
69}
70
71#[derive(Debug, Clone, Default)]
72pub struct AlgorithmInitialState {
73    pub stack: Vec<BFieldElement>,
74    pub nondeterminism: NonDeterminism,
75}
76
77impl From<AlgorithmInitialState> for InitVmState {
78    fn from(value: AlgorithmInitialState) -> Self {
79        Self {
80            stack: value.stack,
81            nondeterminism: value.nondeterminism,
82            ..Default::default()
83        }
84    }
85}
86
87pub struct ShadowedAlgorithm<T: Algorithm> {
88    algorithm: T,
89}
90
91impl<T: Algorithm> ShadowedAlgorithm<T> {
92    pub fn new(algorithm: T) -> Self {
93        Self { algorithm }
94    }
95}
96
97impl<T> RustShadow for ShadowedAlgorithm<T>
98where
99    T: Algorithm,
100{
101    fn inner(&self) -> &dyn BasicSnippet {
102        &self.algorithm
103    }
104
105    fn rust_shadow_wrapper(
106        &self,
107        _stdin: &[BFieldElement],
108        nondeterminism: &NonDeterminism,
109        stack: &mut Vec<BFieldElement>,
110        memory: &mut HashMap<BFieldElement, BFieldElement>,
111        _sponge: &mut Option<Tip5>,
112    ) -> Vec<BFieldElement> {
113        self.algorithm.rust_shadow(stack, memory, nondeterminism);
114        vec![]
115    }
116
117    fn test(&self) {
118        for corner_case in self.algorithm.corner_case_initial_states() {
119            let stdin = vec![];
120            test_rust_equivalence_given_complete_state(
121                self,
122                &corner_case.stack,
123                &stdin,
124                &corner_case.nondeterminism,
125                &None,
126                None,
127            );
128        }
129
130        let num_states = 10;
131        let seed: [u8; 32] = rand::rng().random();
132        let mut rng = StdRng::from_seed(seed);
133        for _ in 0..num_states {
134            let seed: [u8; 32] = rng.random();
135            let AlgorithmInitialState {
136                stack,
137                nondeterminism,
138            } = self.algorithm.pseudorandom_initial_state(seed, None);
139
140            let stdin = vec![];
141            test_rust_equivalence_given_complete_state(
142                self,
143                &stack,
144                &stdin,
145                &nondeterminism,
146                &None,
147                None,
148            );
149        }
150    }
151
152    fn bench(&self) {
153        let mut rng = StdRng::from_seed(
154            hex::decode("73a24b6b8b32e4d7d563a4d9a85f476573a24b6b8b32e4d7d563a4d9a85f4765")
155                .unwrap()
156                .try_into()
157                .unwrap(),
158        );
159        let mut benchmarks = Vec::with_capacity(2);
160
161        for bench_case in [BenchmarkCase::CommonCase, BenchmarkCase::WorstCase] {
162            let AlgorithmInitialState {
163                stack,
164                nondeterminism,
165            } = self
166                .algorithm
167                .pseudorandom_initial_state(rng.random(), Some(bench_case));
168            let program = self.algorithm.link_for_isolated_run();
169            let benchmark = execute_bench(&program, &stack, vec![], nondeterminism, None);
170            let benchmark = NamedBenchmarkResult {
171                name: self.algorithm.entrypoint(),
172                benchmark_result: benchmark,
173                case: bench_case,
174            };
175            benchmarks.push(benchmark);
176        }
177
178        write_benchmarks(benchmarks);
179    }
180}
tasm_lib/traits/algorithm.rs

tasm_lib/traits/
algorithm.rs
