tasm_lib/traits/
closure.rs

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

tasm_lib/traits/
closure.rs
