tasm_lib/traits/
closure.rs

1use rand::prelude::*;
2use triton_vm::prelude::*;
3
4use super::basic_snippet::BasicSnippet;
5use super::rust_shadow::RustShadow;
6use crate::linker::execute_bench;
7use crate::prelude::Tip5;
8use crate::push_encodable;
9use crate::snippet_bencher::write_benchmarks;
10use crate::snippet_bencher::BenchmarkCase;
11use crate::snippet_bencher::NamedBenchmarkResult;
12use crate::test_helpers::test_rust_equivalence_given_complete_state;
13
14/// A Closure is a piece of tasm code that modifies the top of the stack without access to
15/// memory or nondeterminism or standard input/output.
16///
17/// See also: [function], [algorithm], [read_only_algorithm], [procedure],
18///           [accessor], [mem_preserver]
19///
20/// [function]: crate::traits::function::Function
21/// [algorithm]: crate::traits::algorithm::Algorithm
22/// [procedure]: crate::traits::procedure::Procedure
23/// [accessor]: crate::traits::accessor::Accessor
24/// [mem_preserver]: crate::traits::mem_preserver::MemPreserver
25/// [read_only_algorithm]: crate::traits::read_only_algorithm::ReadOnlyAlgorithm
26pub trait Closure: BasicSnippet {
27    /// The arguments of this closure.
28    ///
29    /// Because rust does not support variadic types (yet?), it is recommended to
30    /// use a tuple if the closure takes more than one argument. Because of the way
31    /// [`BFieldCodec`] works for tuples, specify the element the closure expects on
32    /// top of the stack as the last element of the tuple.
33    type Args: BFieldCodec;
34
35    fn rust_shadow(&self, stack: &mut Vec<BFieldElement>);
36
37    /// Given a [seed](StdRng::Seed), generate pseudorandom [arguments](Self::Args),
38    /// from which an [initial stack](Self::set_up_test_stack) can be constructed.
39    fn pseudorandom_args(&self, seed: [u8; 32], bench_case: Option<BenchmarkCase>) -> Self::Args;
40
41    fn corner_case_args(&self) -> Vec<Self::Args> {
42        vec![]
43    }
44
45    fn set_up_test_stack(&self, args: Self::Args) -> Vec<BFieldElement> {
46        let mut stack = self.init_stack_for_isolated_run();
47        push_encodable(&mut stack, &args);
48        stack
49    }
50}
51
52pub struct ShadowedClosure<C: Closure> {
53    closure: C,
54}
55
56impl<C: Closure> ShadowedClosure<C> {
57    pub fn new(closure: C) -> Self {
58        Self { closure }
59    }
60}
61
62impl<C: Closure> RustShadow for ShadowedClosure<C> {
63    fn inner(&self) -> &dyn BasicSnippet {
64        &self.closure
65    }
66
67    fn rust_shadow_wrapper(
68        &self,
69        _stdin: &[BFieldElement],
70        _nondeterminism: &NonDeterminism,
71        stack: &mut Vec<BFieldElement>,
72        _memory: &mut std::collections::HashMap<BFieldElement, BFieldElement>,
73        _sponge: &mut Option<Tip5>,
74    ) -> Vec<BFieldElement> {
75        self.closure.rust_shadow(stack);
76        vec![]
77    }
78
79    fn test(&self) {
80        let num_states = 5;
81        let mut rng = rand::rng();
82
83        // First test corner-cases as they're easier to debug on failure
84        for args in self.closure.corner_case_args() {
85            test_rust_equivalence_given_complete_state(
86                self,
87                &self.closure.set_up_test_stack(args),
88                &[],
89                &NonDeterminism::default(),
90                &None,
91                None,
92            );
93        }
94
95        for _ in 0..num_states {
96            let seed: [u8; 32] = rng.random();
97            let args = self.closure.pseudorandom_args(seed, None);
98            let stack = self.closure.set_up_test_stack(args);
99
100            let stdin = vec![];
101            test_rust_equivalence_given_complete_state(
102                self,
103                &stack,
104                &stdin,
105                &NonDeterminism::default(),
106                &None,
107                None,
108            );
109        }
110    }
111
112    fn bench(&self) {
113        let mut rng = StdRng::from_seed(
114            hex::decode("73a24b6b8b32e4d7d563a4d9a85f476573a24b6b8b32e4d7d563a4d9a85f4765")
115                .unwrap()
116                .try_into()
117                .unwrap(),
118        );
119        let mut benchmarks = Vec::with_capacity(2);
120
121        for bench_case in [BenchmarkCase::CommonCase, BenchmarkCase::WorstCase] {
122            let args = self
123                .closure
124                .pseudorandom_args(rng.random(), Some(bench_case));
125            let stack = self.closure.set_up_test_stack(args);
126            let program = self.closure.link_for_isolated_run();
127            let benchmark =
128                execute_bench(&program, &stack, vec![], NonDeterminism::new(vec![]), None);
129            let benchmark = NamedBenchmarkResult {
130                name: self.closure.entrypoint(),
131                benchmark_result: benchmark,
132                case: bench_case,
133            };
134            benchmarks.push(benchmark);
135        }
136
137        write_benchmarks(benchmarks);
138    }
139}
tasm_lib/traits/closure.rs

tasm_lib/traits/
closure.rs
