tasm_lib/hashing/
squeeze_repeatedly_static_number.rs

1use triton_vm::prelude::*;
2use twenty_first::tip5::RATE;
3
4use crate::prelude::*;
5
6/// Squeeze the sponge a statically-known number of times.
7///
8/// Squeeze the sponge `num_squeezes` times, storing all the produced pseudorandom `BFieldElement`s
9/// contiguously in memory. It is the caller's responsibility to allocate enough memory.
10/// Number of squeezes must be statically known.
11#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
12pub struct SqueezeRepeatedlyStaticNumber {
13    pub num_squeezes: usize,
14}
15
16impl BasicSnippet for SqueezeRepeatedlyStaticNumber {
17    fn inputs(&self) -> Vec<(DataType, String)> {
18        vec![(DataType::VoidPointer, "address".to_string())]
19    }
20
21    fn outputs(&self) -> Vec<(DataType, String)> {
22        vec![]
23    }
24
25    fn entrypoint(&self) -> String {
26        format!(
27            "tasmlib_hashing_squeeze_repeatedly_static_number_{}",
28            self.num_squeezes
29        )
30    }
31
32    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
33        assert_eq!(10, RATE, "Code assumes RATE is 10");
34
35        let entrypoint = self.entrypoint();
36        let one_squeeze = triton_asm!(
37            // _ *address
38
39            sponge_squeeze
40            // _ *address r9 r8 r7 r6 r5 r4 r3 r2 r1 r0
41
42            dup 10
43            // _ *address r9 r8 r7 r6 r5 r4 r3 r2 r1 r0 *address
44
45            write_mem 5
46            write_mem 5
47            // _ *address (*address + 10)
48
49            swap 1
50            pop 1
51            // _ (*address + 10)
52        );
53
54        let all_squeezes = vec![one_squeeze; self.num_squeezes].concat();
55
56        triton_asm!(
57            {entrypoint}:
58                // _ *address
59
60                {&all_squeezes}
61                 // _ *address'
62
63                pop 1
64
65                return
66        )
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    use crate::hashing::squeeze_repeatedly::SqueezeRepeatedly;
74    use crate::test_helpers::tasm_final_state;
75    use crate::test_helpers::verify_memory_equivalence;
76    use crate::test_helpers::verify_sponge_equivalence;
77    use crate::test_helpers::verify_stack_equivalence;
78    use crate::test_prelude::*;
79
80    impl Procedure for SqueezeRepeatedlyStaticNumber {
81        fn rust_shadow(
82            &self,
83            stack: &mut Vec<BFieldElement>,
84            memory: &mut HashMap<BFieldElement, BFieldElement>,
85            nondeterminism: &NonDeterminism,
86            public_input: &[BFieldElement],
87            sponge: &mut Option<Tip5>,
88        ) -> Vec<BFieldElement> {
89            stack.push(BFieldElement::new(self.num_squeezes as u64));
90            let ret =
91                SqueezeRepeatedly.rust_shadow(stack, memory, nondeterminism, public_input, sponge);
92            stack.pop();
93            stack.pop();
94
95            ret
96        }
97
98        fn pseudorandom_initial_state(
99            &self,
100            seed: [u8; 32],
101            bench_case: Option<BenchmarkCase>,
102        ) -> ProcedureInitialState {
103            let mut init_state = SqueezeRepeatedly.pseudorandom_initial_state(seed, bench_case);
104            init_state.stack.pop();
105            init_state
106        }
107    }
108
109    #[test]
110    fn squeeze_repeatedly_static_number_pbt() {
111        for num_squeezes in 0..25 {
112            ShadowedProcedure::new(SqueezeRepeatedlyStaticNumber { num_squeezes }).test();
113        }
114    }
115
116    #[test]
117    fn test_dyn_equivalence() {
118        // Verify that the snippets for the dynamically known and statically known
119        // number of squeezes agree.
120        fn dyn_output(seed: [u8; 32]) -> (VMState, usize) {
121            let dyn_snippet = ShadowedProcedure::new(SqueezeRepeatedly);
122            let ProcedureInitialState {
123                stack,
124                nondeterminism,
125                public_input: stdin,
126                sponge,
127            } = SqueezeRepeatedly.pseudorandom_initial_state(seed, None);
128            let num_squeeze_count = stack.last().unwrap().value();
129
130            (
131                tasm_final_state(&dyn_snippet, &stack, &stdin, nondeterminism, &sponge),
132                num_squeeze_count as usize,
133            )
134        }
135
136        fn stat_output(seed: [u8; 32], num_squeezes: usize) -> VMState {
137            let snippet = SqueezeRepeatedlyStaticNumber { num_squeezes };
138            let ProcedureInitialState {
139                stack,
140                nondeterminism,
141                public_input: stdin,
142                sponge,
143            } = snippet.pseudorandom_initial_state(seed, None);
144
145            tasm_final_state(
146                &ShadowedProcedure::new(snippet),
147                &stack,
148                &stdin,
149                nondeterminism,
150                &sponge,
151            )
152        }
153
154        let mut seed = [0u8; 32];
155        rand::rng().fill_bytes(&mut seed);
156        let (mut dyn_output, num_squeezes) = dyn_output(seed);
157        dyn_output.op_stack.stack.pop();
158        dyn_output.op_stack.stack.pop();
159
160        let stat_output = stat_output(seed, num_squeezes);
161
162        verify_memory_equivalence(
163            "Snippet with dynamic symbols",
164            &dyn_output.ram,
165            "Snippet with static symbols",
166            &stat_output.ram,
167        );
168        verify_sponge_equivalence(&dyn_output.sponge, &stat_output.sponge);
169        verify_stack_equivalence(
170            "Snippet with dynamic symbols",
171            &dyn_output.op_stack.stack,
172            "Snippet with static symbols",
173            &stat_output.op_stack.stack,
174        );
175    }
176}
177
178#[cfg(test)]
179mod benches {
180    use super::*;
181    use crate::test_prelude::*;
182
183    #[test]
184    fn bench_10() {
185        ShadowedProcedure::new(SqueezeRepeatedlyStaticNumber { num_squeezes: 10 }).bench();
186    }
187
188    #[test]
189    fn bench_200() {
190        ShadowedProcedure::new(SqueezeRepeatedlyStaticNumber { num_squeezes: 200 }).bench();
191    }
192}
tasm_lib/hashing/squeeze_repeatedly_static_number.rs

tasm_lib/hashing/
squeeze_repeatedly_static_number.rs
