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

tasm_lib/traits/
procedure.rs
