tasm_lib/hashing/algebraic_hasher/
sample_scalars_static_length_kmalloc.rs

1use triton_vm::prelude::*;
2use twenty_first::math::x_field_element::EXTENSION_DEGREE;
3use twenty_first::tip5::RATE;
4
5use crate::hashing::algebraic_hasher::sample_scalars_static_length_dyn_malloc::SampleScalarsStaticLengthDynMalloc;
6use crate::hashing::squeeze_repeatedly_static_number::SqueezeRepeatedlyStaticNumber;
7use crate::prelude::*;
8
9/// Squeeze the sponge to sample a given number of [`XFieldElement`]s. Puts the scalars into
10/// statically allocated memory.
11///
12/// # Panics
13///
14/// Panics if both fields are 0 because the static allocator will be unhappy. :)
15#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
16pub struct SampleScalarsStaticLengthKMalloc {
17    pub num_elements_to_sample: usize,
18
19    /// Number of additional elements to statically allocate, in number of
20    /// [extension field element][xfe]s.
21    /// Necessary for [`Challenges`][chall].
22    ///
23    /// [chall]: crate::verifier::challenges::new_empty_input_and_output::NewEmptyInputAndOutput
24    /// [xfe]: XFieldElement
25    pub extra_capacity: usize,
26}
27
28impl SampleScalarsStaticLengthKMalloc {
29    fn num_words_to_allocate(&self) -> u32 {
30        ((self.num_elements_to_sample + self.extra_capacity) * EXTENSION_DEGREE)
31            .try_into()
32            .unwrap()
33    }
34}
35
36impl BasicSnippet for SampleScalarsStaticLengthKMalloc {
37    fn parameters(&self) -> Vec<(DataType, String)> {
38        vec![]
39    }
40
41    fn return_values(&self) -> Vec<(DataType, String)> {
42        vec![]
43    }
44
45    fn entrypoint(&self) -> String {
46        format!(
47            "tasmlib_hashing_algebraic_hasher_sample_scalars_static_length_k_malloc{}",
48            self.num_elements_to_sample
49        )
50    }
51
52    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
53        assert_eq!(10, RATE, "Code assumes Tip5's RATE is 10");
54        assert_eq!(3, EXTENSION_DEGREE, "Code assumes extension degree 3");
55        assert!(
56            self.extra_capacity + self.num_elements_to_sample > 0,
57            "Must allocate positive number of words"
58        );
59        let num_squeezes =
60            SampleScalarsStaticLengthDynMalloc::num_squeezes(self.num_elements_to_sample);
61
62        let num_squeezed_words = num_squeezes * RATE;
63        debug_assert!(
64            self.num_elements_to_sample * EXTENSION_DEGREE <= num_squeezed_words,
65            "need {} elements but getting {num_squeezed_words}",
66            self.num_elements_to_sample * EXTENSION_DEGREE,
67        );
68
69        let entrypoint = self.entrypoint();
70        let squeeze_repeatedly_static_number =
71            library.import(Box::new(SqueezeRepeatedlyStaticNumber { num_squeezes }));
72        let scalars_pointer_alloc = library.kmalloc(self.num_words_to_allocate());
73
74        triton_asm!(
75            {entrypoint}:
76                push {scalars_pointer_alloc.write_address()}
77                call {squeeze_repeatedly_static_number}
78                return
79        )
80    }
81}
82
83#[cfg(test)]
84pub(crate) mod tests {
85    use std::ops::Neg;
86
87    use twenty_first::util_types::sponge::Sponge;
88
89    use super::*;
90    use crate::rust_shadowing_helper_functions::array::array_get;
91    use crate::rust_shadowing_helper_functions::array::insert_as_array;
92    use crate::test_helpers::tasm_final_state;
93    use crate::test_prelude::*;
94
95    impl SampleScalarsStaticLengthKMalloc {
96        /// For testing purposes only.
97        pub fn k_malloc_address_isolated_run(&self) -> BFieldElement {
98            let dyn_allocator_state_size = 1;
99            BFieldElement::from(self.num_words_to_allocate() + dyn_allocator_state_size).neg()
100        }
101    }
102
103    impl Procedure for SampleScalarsStaticLengthKMalloc {
104        fn rust_shadow(
105            &self,
106            _stack: &mut Vec<BFieldElement>,
107            memory: &mut HashMap<BFieldElement, BFieldElement>,
108            _nondeterminism: &NonDeterminism,
109            _public_input: &[BFieldElement],
110            sponge: &mut Option<Tip5>,
111        ) -> Result<Vec<BFieldElement>, RustShadowError> {
112            let Some(sponge) = sponge.as_mut() else {
113                return Err(RustShadowError::SpongeUninitialized);
114            };
115            let num_squeezes =
116                SampleScalarsStaticLengthDynMalloc::num_squeezes(self.num_elements_to_sample);
117            let pseudorandomness = (0..num_squeezes)
118                .flat_map(|_| sponge.squeeze().to_vec())
119                .collect_vec();
120            let scalars_pointer = self.k_malloc_address_isolated_run();
121            insert_as_array(scalars_pointer, memory, pseudorandomness);
122
123            Ok(Vec::new())
124        }
125
126        fn pseudorandom_initial_state(
127            &self,
128            seed: [u8; 32],
129            _: Option<BenchmarkCase>,
130        ) -> ProcedureInitialState {
131            let mut rng = StdRng::from_seed(seed);
132            let stack = self.init_stack_for_isolated_run();
133            let sponge = Tip5 {
134                state: rng.random(),
135            };
136
137            ProcedureInitialState {
138                stack,
139                sponge: Some(sponge),
140                ..Default::default()
141            }
142        }
143
144        fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
145            let freshly_initialized_sponge = ProcedureInitialState {
146                stack: self.init_stack_for_isolated_run(),
147                sponge: Some(Tip5::init()),
148                ..Default::default()
149            };
150
151            vec![freshly_initialized_sponge]
152        }
153    }
154
155    #[macro_rules_attr::apply(test)]
156    fn sample_scalars_static_length_pbt() {
157        for num_elements_to_sample in 0..11 {
158            for extra_capacity in 0..11 {
159                if num_elements_to_sample + extra_capacity == 0 {
160                    continue;
161                }
162                ShadowedProcedure::new(SampleScalarsStaticLengthKMalloc {
163                    num_elements_to_sample,
164                    extra_capacity,
165                })
166                .test();
167            }
168        }
169    }
170
171    #[macro_rules_attr::apply(proptest)]
172    fn verify_agreement_with_tip5_sample_scalars(
173        #[strategy(1_usize..500)] num_elements_to_sample: usize,
174        #[strategy(1_usize..500)] extra_capacity: usize,
175        #[strategy(arb())] mut sponge: Tip5,
176    ) {
177        let snippet = SampleScalarsStaticLengthKMalloc {
178            num_elements_to_sample,
179            extra_capacity,
180        };
181        let init_stack = snippet.init_stack_for_isolated_run();
182        let tasm = tasm_final_state(
183            &ShadowedProcedure::new(snippet),
184            &init_stack,
185            &[],
186            NonDeterminism::default(),
187            &Some(sponge.clone()),
188        )
189        .unwrap();
190
191        let scalar_pointer = snippet.k_malloc_address_isolated_run();
192        let read_scalar = |i| array_get(scalar_pointer, i, &tasm.ram, EXTENSION_DEGREE);
193
194        let scalars_from_tip5 = sponge.sample_scalars(num_elements_to_sample);
195        for (i, expected_scalar) in scalars_from_tip5.into_iter().enumerate() {
196            assert_eq!(expected_scalar.coefficients.to_vec(), read_scalar(i));
197        }
198    }
199}
200
201#[cfg(test)]
202mod bench {
203    use super::*;
204    use crate::test_prelude::*;
205
206    #[macro_rules_attr::apply(test)]
207    fn bench_10() {
208        ShadowedProcedure::new(SampleScalarsStaticLengthKMalloc {
209            num_elements_to_sample: 10,
210            extra_capacity: 4,
211        })
212        .bench();
213    }
214
215    #[macro_rules_attr::apply(test)]
216    fn bench_100() {
217        ShadowedProcedure::new(SampleScalarsStaticLengthKMalloc {
218            num_elements_to_sample: 100,
219            extra_capacity: 4,
220        })
221        .bench();
222    }
223
224    #[macro_rules_attr::apply(test)]
225    fn bench_63() {
226        ShadowedProcedure::new(SampleScalarsStaticLengthKMalloc {
227            num_elements_to_sample: 63,
228            extra_capacity: 4,
229        })
230        .bench();
231    }
232}
tasm_lib/hashing/algebraic_hasher/sample_scalars_static_length_kmalloc.rs

tasm_lib/hashing/algebraic_hasher/
sample_scalars_static_length_kmalloc.rs
