tasm_lib/hashing/algebraic_hasher/
sample_scalars.rs

1use triton_vm::prelude::*;
2use twenty_first::math::x_field_element::EXTENSION_DEGREE;
3
4use crate::hashing::squeeze_repeatedly::SqueezeRepeatedly;
5use crate::list::new::New;
6use crate::list::set_length::SetLength;
7use crate::prelude::*;
8
9/// Squeeze the sponge to sample a given number of `XFieldElement`s.
10#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
11pub struct SampleScalars;
12
13impl BasicSnippet for SampleScalars {
14    fn parameters(&self) -> Vec<(DataType, String)> {
15        vec![(DataType::U32, "num_scalars".to_string())]
16    }
17
18    fn return_values(&self) -> Vec<(DataType, String)> {
19        vec![(
20            DataType::List(Box::new(DataType::Xfe)),
21            "*scalars".to_string(),
22        )]
23    }
24
25    fn entrypoint(&self) -> String {
26        "tasmlib_hashing_algebraic_hasher_sample_scalars".to_string()
27    }
28
29    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
30        assert_eq!(10, tip5::RATE, "Code assumes Tip5's RATE is 10");
31        assert_eq!(3, EXTENSION_DEGREE, "Code assumes extension degree 3");
32
33        let entrypoint = self.entrypoint();
34        let set_length = library.import(Box::new(SetLength));
35        let new_list = library.import(Box::new(New));
36        let safety_offset = 1;
37        let squeeze_repeatedly = library.import(Box::new(SqueezeRepeatedly));
38        triton_asm! {
39            // BEFORE: _ num_scalars
40            // AFTER:  _ *scalars
41            {entrypoint}:
42                call {new_list}
43                                // _ num_scalars *scalars
44
45                // set length
46                dup 1           // _ num_scalars *scalars num_scalars
47                call {set_length}
48                                // _ num_scalars *scalars
49
50                // calculate number of squeezes
51                dup 1           // _ num_scalars *scalars num_scalars
52                push {EXTENSION_DEGREE} mul
53
54                // _ num_scalars *scalars num_bfes
55                push 9 add      // _ num_scalars *scalars (num_bfes+9)
56                push {tip5::RATE} swap 1
57                                // _ num_scalars *scalars rate (num_bfes+9)
58                div_mod pop 1   // _ num_scalars *scalars floor((num_bfes+9)/rate)
59                                // _ num_scalars *scalars num_squeezes
60
61                // prepare stack for call to squeeze_repeatedly
62                dup 1
63                push {safety_offset}
64                add
65                swap 1          // _ num_scalars *scalars (*scalars+so) num_squeezes
66
67                // squeeze
68                call {squeeze_repeatedly}
69                                // _ num_scalars *scalars *scalars' 0
70
71                // clean up stack
72                pop 2
73                swap 1
74                pop 1           // _ *scalars
75                return
76
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use twenty_first::prelude::*;
84
85    use super::*;
86    use crate::empty_stack;
87    use crate::memory::dyn_malloc::DYN_MALLOC_FIRST_ADDRESS;
88    use crate::rust_shadowing_helper_functions;
89    use crate::test_helpers::tasm_final_state;
90    use crate::test_prelude::*;
91
92    impl Procedure for SampleScalars {
93        fn rust_shadow(
94            &self,
95            stack: &mut Vec<BFieldElement>,
96            memory: &mut HashMap<BFieldElement, BFieldElement>,
97            _: &NonDeterminism,
98            _: &[BFieldElement],
99            sponge: &mut Option<Tip5>,
100        ) -> Result<Vec<BFieldElement>, RustShadowError> {
101            let Some(sponge) = sponge.as_mut() else {
102                return Err(RustShadowError::SpongeUninitialized);
103            };
104            let num_scalars = stack.pop().ok_or(RustShadowError::StackUnderflow)?.value() as usize;
105            let num_squeezes = (num_scalars * 3 + 9) / tip5::RATE;
106            let pseudorandomness = (0..num_squeezes)
107                .flat_map(|_| sponge.squeeze().to_vec())
108                .collect_vec();
109            let scalars = pseudorandomness
110                .chunks(3)
111                .take(num_scalars)
112                .map(|ch| XFieldElement::new(ch.try_into().unwrap()))
113                .collect_vec();
114            let scalars_pointer = DYN_MALLOC_FIRST_ADDRESS;
115
116            encode_to_memory(memory, scalars_pointer, &scalars);
117
118            // store all pseudorandomness (not just sampled scalars) to memory
119            let safety_offset = BFieldElement::new(1);
120            for (i, pr) in pseudorandomness.iter().enumerate() {
121                memory.insert(
122                    BFieldElement::new(i as u64) + scalars_pointer + safety_offset,
123                    *pr,
124                );
125            }
126
127            // the list of scalars was allocated properly; reflect that fact
128            memory.insert(scalars_pointer, BFieldElement::new(num_scalars as u64));
129            stack.push(scalars_pointer);
130
131            Ok(Vec::new())
132        }
133
134        fn pseudorandom_initial_state(
135            &self,
136            seed: [u8; 32],
137            bench_case: Option<BenchmarkCase>,
138        ) -> ProcedureInitialState {
139            let mut rng = StdRng::from_seed(seed);
140            let num_scalars = match bench_case {
141                Some(BenchmarkCase::CommonCase) => 10,
142                Some(BenchmarkCase::WorstCase) => 100,
143                None => rng.random_range(0..40),
144            };
145            let mut stack = empty_stack();
146            stack.push(BFieldElement::new(num_scalars as u64));
147            let sponge = Tip5 {
148                state: rng.random(),
149            };
150
151            ProcedureInitialState {
152                stack,
153                sponge: Some(sponge),
154                ..Default::default()
155            }
156        }
157
158        fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
159            let zero_to_twenty_scalars_empty_sponge = (0..20)
160                .map(|num_scalars| {
161                    let mut stack = empty_stack();
162                    stack.push(BFieldElement::new(num_scalars as u64));
163                    let sponge = Tip5::init();
164
165                    ProcedureInitialState {
166                        stack,
167                        sponge: Some(sponge),
168                        ..Default::default()
169                    }
170                })
171                .collect_vec();
172            let zero_to_twenty_scalars_non_empty_sponge = (0..20)
173                .map(|num_scalars| {
174                    let mut stack = empty_stack();
175                    stack.push(BFieldElement::new(num_scalars as u64));
176                    let mut sponge = Tip5::init();
177                    sponge.absorb([BFieldElement::new(42); Tip5::RATE]);
178
179                    ProcedureInitialState {
180                        stack,
181                        sponge: Some(sponge),
182                        ..Default::default()
183                    }
184                })
185                .collect_vec();
186
187            [
188                zero_to_twenty_scalars_empty_sponge,
189                zero_to_twenty_scalars_non_empty_sponge,
190            ]
191            .concat()
192        }
193    }
194
195    #[macro_rules_attr::apply(test)]
196    fn test() {
197        ShadowedProcedure::new(SampleScalars).test();
198    }
199
200    /// This is a regression test that verifies that this implementation of `sample_scalars`
201    /// agrees with Tip5's version from twenty-first. For the bugfix, see:
202    /// <https://github.com/Neptune-Crypto/twenty-first/commit/e708b305>
203    #[macro_rules_attr::apply(test)]
204    fn verify_agreement_with_tip5_sample_scalars() {
205        let empty_sponge = Tip5::init();
206        let mut non_empty_sponge = Tip5::init();
207        non_empty_sponge.absorb([BFieldElement::new(100); Tip5::RATE]);
208
209        for init_sponge in [empty_sponge, non_empty_sponge] {
210            for num_scalars in 0..30 {
211                let init_stack = [
212                    SampleScalars.init_stack_for_isolated_run(),
213                    vec![BFieldElement::new(num_scalars as u64)],
214                ]
215                .concat();
216                let tasm = tasm_final_state(
217                    &ShadowedProcedure::new(SampleScalars),
218                    &init_stack,
219                    &[],
220                    NonDeterminism::default(),
221                    &Some(init_sponge.clone()),
222                )
223                .unwrap();
224
225                let final_ram = tasm.ram;
226                let snippet_output_scalar_pointer =
227                    tasm.op_stack.stack[tasm.op_stack.stack.len() - 1];
228
229                let scalars_from_tip5 = Tip5::sample_scalars(&mut init_sponge.clone(), num_scalars);
230
231                for (i, expected_scalar) in scalars_from_tip5.into_iter().enumerate() {
232                    assert_eq!(
233                        expected_scalar.coefficients.to_vec(),
234                        rust_shadowing_helper_functions::list::list_get(
235                            snippet_output_scalar_pointer,
236                            i,
237                            &final_ram,
238                            EXTENSION_DEGREE
239                        )
240                        .unwrap()
241                    );
242                }
243            }
244        }
245    }
246}
247
248#[cfg(test)]
249mod bench {
250    use super::*;
251    use crate::test_prelude::*;
252
253    #[macro_rules_attr::apply(test)]
254    fn benchmark() {
255        ShadowedProcedure::new(SampleScalars).bench();
256    }
257}
tasm_lib/hashing/algebraic_hasher/sample_scalars.rs

tasm_lib/hashing/algebraic_hasher/
sample_scalars.rs
