tasm_lib/traits/
procedure.rs

1use std::collections::HashMap;
2
3use rand::prelude::*;
4use triton_vm::prelude::*;
5
6use crate::InitVmState;
7use crate::linker::execute_bench;
8use crate::prelude::Tip5;
9use crate::snippet_bencher::BenchmarkCase;
10use crate::snippet_bencher::NamedBenchmarkResult;
11use crate::snippet_bencher::write_benchmarks;
12use crate::test_helpers::test_rust_equivalence_given_complete_state;
13use crate::traits::basic_snippet::BasicSnippet;
14use crate::traits::rust_shadow::RustShadow;
15use crate::traits::rust_shadow::RustShadowError;
16
17/// A trait that can modify all parts of the VM state.
18///
19/// A Procedure is a piece of tasm code that can do almost anything: modify stack, read
20/// from and write to memory, take in nondeterminism, and read and write from standard
21/// input/output. What it cannot do is stand alone. In other words, it is still wrapped
22/// in a function (lower case f, as in 'labelled scope'); and cannot be proved as
23/// a standalone program.
24///
25/// See also: [closure], [function], [algorithm], [read_only_algorithm],
26///           [accessor], [mem_preserver]
27///
28/// [closure]: crate::traits::closure::Closure
29/// [function]: crate::traits::function::Function
30/// [algorithm]: crate::traits::algorithm::Algorithm
31/// [read_only_algorithm]: crate::traits::read_only_algorithm::ReadOnlyAlgorithm
32/// [accessor]: crate::traits::accessor::Accessor
33/// [mem_preserver]: crate::traits::mem_preserver::MemPreserver
34pub trait Procedure: BasicSnippet {
35    /// Returns standard output
36    fn rust_shadow(
37        &self,
38        stack: &mut Vec<BFieldElement>,
39        memory: &mut HashMap<BFieldElement, BFieldElement>,
40        nondeterminism: &NonDeterminism,
41        public_input: &[BFieldElement],
42        sponge: &mut Option<Tip5>,
43    ) -> Result<Vec<BFieldElement>, RustShadowError>;
44
45    fn preprocess<T: BFieldCodec>(_meta_input: T, _nondeterminism: &mut NonDeterminism) {}
46
47    fn pseudorandom_initial_state(
48        &self,
49        seed: [u8; 32],
50        bench_case: Option<BenchmarkCase>,
51    ) -> ProcedureInitialState;
52
53    fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
54        Vec::new()
55    }
56}
57
58#[derive(Debug, Clone, Default)]
59pub struct ProcedureInitialState {
60    pub stack: Vec<BFieldElement>,
61    pub nondeterminism: NonDeterminism,
62    pub public_input: Vec<BFieldElement>,
63    pub sponge: Option<Tip5>,
64}
65
66impl From<ProcedureInitialState> for InitVmState {
67    fn from(value: ProcedureInitialState) -> Self {
68        Self {
69            stack: value.stack,
70            public_input: value.public_input,
71            nondeterminism: value.nondeterminism,
72            sponge: value.sponge,
73        }
74    }
75}
76
77pub struct ShadowedProcedure<P: Procedure> {
78    procedure: P,
79}
80
81impl<P: Procedure> ShadowedProcedure<P> {
82    pub fn new(procedure: P) -> Self {
83        Self { procedure }
84    }
85}
86
87impl<P: Procedure> RustShadow for ShadowedProcedure<P> {
88    fn inner(&self) -> &dyn BasicSnippet {
89        &self.procedure
90    }
91
92    fn rust_shadow_wrapper(
93        &self,
94        stdin: &[BFieldElement],
95        nondeterminism: &NonDeterminism,
96        stack: &mut Vec<BFieldElement>,
97        memory: &mut HashMap<BFieldElement, BFieldElement>,
98        sponge: &mut Option<Tip5>,
99    ) -> Result<Vec<BFieldElement>, RustShadowError> {
100        self.procedure
101            .rust_shadow(stack, memory, nondeterminism, stdin, sponge)
102    }
103
104    fn test(&self) {
105        let num_states = 5;
106        let seed: [u8; 32] = rand::rng().random();
107        let mut rng = StdRng::from_seed(seed);
108        let procedure = &self.procedure;
109
110        for corner_case in procedure.corner_case_initial_states().into_iter() {
111            self.test_initial_state(corner_case);
112        }
113
114        for _ in 0..num_states {
115            let seed: [u8; 32] = rng.random();
116            let state = procedure.pseudorandom_initial_state(seed, None);
117
118            self.test_initial_state(state);
119        }
120    }
121
122    fn bench(&self) {
123        let mut rng = StdRng::from_seed(
124            hex::decode("73a24b6b8b32e4d7d563a4d9a85f476573a24b6b8b32e4d7d563a4d9a85f4765")
125                .unwrap()
126                .try_into()
127                .unwrap(),
128        );
129        let mut benchmarks = Vec::with_capacity(2);
130
131        for bench_case in [BenchmarkCase::CommonCase, BenchmarkCase::WorstCase] {
132            let ProcedureInitialState {
133                stack,
134                nondeterminism,
135                public_input,
136                sponge,
137            } = self
138                .procedure
139                .pseudorandom_initial_state(rng.random(), Some(bench_case));
140            let program = self.procedure.link_for_isolated_run();
141            let benchmark = execute_bench(&program, &stack, public_input, nondeterminism, sponge);
142            let benchmark = NamedBenchmarkResult {
143                name: self.procedure.entrypoint(),
144                benchmark_result: benchmark,
145                case: bench_case,
146            };
147            benchmarks.push(benchmark);
148        }
149
150        write_benchmarks(benchmarks);
151    }
152}
153
154impl<P: Procedure> ShadowedProcedure<P> {
155    fn test_initial_state(&self, state: ProcedureInitialState) {
156        test_rust_equivalence_given_complete_state(
157            self,
158            &state.stack,
159            &state.public_input,
160            &state.nondeterminism,
161            &state.sponge,
162            None,
163        );
164    }
165}
tasm_lib/traits/procedure.rs

tasm_lib/traits/
procedure.rs
