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 inputs(&self) -> Vec<(DataType, String)> {
15        vec![(DataType::U32, "num_scalars".to_string())]
16    }
17
18    fn outputs(&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        ) -> Vec<BFieldElement> {
101            let sponge = sponge.as_mut().expect("sponge must be initialized");
102            let num_scalars = stack.pop().unwrap().value() as usize;
103            let num_squeezes = (num_scalars * 3 + 9) / tip5::RATE;
104            let pseudorandomness = (0..num_squeezes)
105                .flat_map(|_| sponge.squeeze().to_vec())
106                .collect_vec();
107            let scalars = pseudorandomness
108                .chunks(3)
109                .take(num_scalars)
110                .map(|ch| XFieldElement::new(ch.try_into().unwrap()))
111                .collect_vec();
112            let scalars_pointer = DYN_MALLOC_FIRST_ADDRESS;
113
114            encode_to_memory(memory, scalars_pointer, &scalars);
115
116            // store all pseudorandomness (not just sampled scalars) to memory
117            let safety_offset = BFieldElement::new(1);
118            for (i, pr) in pseudorandomness.iter().enumerate() {
119                memory.insert(
120                    BFieldElement::new(i as u64) + scalars_pointer + safety_offset,
121                    *pr,
122                );
123            }
124
125            // the list of scalars was allocated properly; reflect that fact
126            memory.insert(scalars_pointer, BFieldElement::new(num_scalars as u64));
127
128            stack.push(scalars_pointer);
129            vec![]
130        }
131
132        fn pseudorandom_initial_state(
133            &self,
134            seed: [u8; 32],
135            bench_case: Option<BenchmarkCase>,
136        ) -> ProcedureInitialState {
137            let mut rng = StdRng::from_seed(seed);
138            let num_scalars = match bench_case {
139                Some(BenchmarkCase::CommonCase) => 10,
140                Some(BenchmarkCase::WorstCase) => 100,
141                None => rng.random_range(0..40),
142            };
143            let mut stack = empty_stack();
144            stack.push(BFieldElement::new(num_scalars as u64));
145            let sponge = Tip5 {
146                state: rng.random(),
147            };
148
149            ProcedureInitialState {
150                stack,
151                sponge: Some(sponge),
152                ..Default::default()
153            }
154        }
155
156        fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
157            let zero_to_twenty_scalars_empty_sponge = (0..20)
158                .map(|num_scalars| {
159                    let mut stack = empty_stack();
160                    stack.push(BFieldElement::new(num_scalars as u64));
161                    let sponge = Tip5::init();
162
163                    ProcedureInitialState {
164                        stack,
165                        sponge: Some(sponge),
166                        ..Default::default()
167                    }
168                })
169                .collect_vec();
170            let zero_to_twenty_scalars_non_empty_sponge = (0..20)
171                .map(|num_scalars| {
172                    let mut stack = empty_stack();
173                    stack.push(BFieldElement::new(num_scalars as u64));
174                    let mut sponge = Tip5::init();
175                    sponge.absorb([BFieldElement::new(42); Tip5::RATE]);
176
177                    ProcedureInitialState {
178                        stack,
179                        sponge: Some(sponge),
180                        ..Default::default()
181                    }
182                })
183                .collect_vec();
184
185            [
186                zero_to_twenty_scalars_empty_sponge,
187                zero_to_twenty_scalars_non_empty_sponge,
188            ]
189            .concat()
190        }
191    }
192
193    #[test]
194    fn test() {
195        ShadowedProcedure::new(SampleScalars).test();
196    }
197
198    /// This is a regression test that verifies that this implementation of `sample_scalars`
199    /// agrees with Tip5's version from twenty-first. For the bugfix, see:
200    /// <https://github.com/Neptune-Crypto/twenty-first/commit/e708b305>
201    #[test]
202    fn verify_agreement_with_tip5_sample_scalars() {
203        let empty_sponge = Tip5::init();
204        let mut non_empty_sponge = Tip5::init();
205        non_empty_sponge.absorb([BFieldElement::new(100); Tip5::RATE]);
206
207        for init_sponge in [empty_sponge, non_empty_sponge] {
208            for num_scalars in 0..30 {
209                let init_stack = [
210                    SampleScalars.init_stack_for_isolated_run(),
211                    vec![BFieldElement::new(num_scalars as u64)],
212                ]
213                .concat();
214                let tasm = tasm_final_state(
215                    &ShadowedProcedure::new(SampleScalars),
216                    &init_stack,
217                    &[],
218                    NonDeterminism::default(),
219                    &Some(init_sponge.clone()),
220                );
221
222                let final_ram = tasm.ram;
223                let snippet_output_scalar_pointer =
224                    tasm.op_stack.stack[tasm.op_stack.stack.len() - 1];
225
226                let scalars_from_tip5 = Tip5::sample_scalars(&mut init_sponge.clone(), num_scalars);
227
228                for (i, expected_scalar) in scalars_from_tip5.into_iter().enumerate() {
229                    assert_eq!(
230                        expected_scalar.coefficients.to_vec(),
231                        rust_shadowing_helper_functions::list::list_get(
232                            snippet_output_scalar_pointer,
233                            i,
234                            &final_ram,
235                            EXTENSION_DEGREE
236                        )
237                    );
238                }
239            }
240        }
241    }
242}
243
244#[cfg(test)]
245mod bench {
246    use super::*;
247    use crate::test_prelude::*;
248
249    #[test]
250    fn benchmark() {
251        ShadowedProcedure::new(SampleScalars).bench();
252    }
253}
tasm_lib/hashing/algebraic_hasher/sample_scalars.rs

tasm_lib/hashing/algebraic_hasher/
sample_scalars.rs
