Skip to main content

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 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
17/// An Algorithm can modify memory and take ND_input.
18///
19/// An Algorithm is a piece of tasm code that can modify memory even at addresses below
20/// the dynamic memory allocator, and can take nondeterministic input. It cannot read from
21/// standard in or write to standard out.
22///
23/// See also: [closure], [function], [procedure], [accessor], [mem_preserver],
24///           [read_only_algorithm]
25///
26/// [closure]: crate::traits::closure::Closure
27/// [function]: crate::traits::function::Function
28/// [procedure]: crate::traits::procedure::Procedure
29/// [accessor]: crate::traits::accessor::Accessor
30/// [mem_preserver]: crate::traits::mem_preserver::MemPreserver
31/// [read_only_algorithm]: crate::traits::read_only_algorithm::ReadOnlyAlgorithm
32pub 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    /// Take a object about which something is being proven in order to extract out the
41    /// right nondeterminism. Update the mutably referenced non-determism argument.
42    ///
43    /// For example:
44    ///  - When proving the correct verification of a proof, you might want to pull all
45    ///    digests out of the authentication structures and put them in the `digests`
46    ///    field of the non-determinism. This way the VM can avoid the processing
47    ///    required by authentication structures and just divine in the right digests
48    ///    as it walks up the Merkle trees.
49    ///  - When proving the correct sorting of a list, the VM ought to avoid running a
50    ///    sorting algorithm; instead it should divine the sorted list and then prove
51    ///    that the two lists are equal. The preprocessing step in this case would take
52    ///    the unsorted list, sort it, and use the sorted list to populate the non-
53    ///    determinism.
54    ///  - When verifying a Falcon signature, at some point the NTT of a vector is needed.
55    ///    The VM should not compute the NTT because expensive; instead it should divine
56    ///    the NTT-transformed vector and verify that it satisfies the right relation. In
57    ///    this case the preprocessor calculates the NTT and populates the non-determinism
58    ///    with the transformed vector.
59    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}