tasm_lib/hashing/
squeeze_repeatedly.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5/// Squeeze the sponge n times, storing all the produced pseudorandom `BFieldElement`s
6/// contiguously in memory. It is the caller's responsibility to allocate enough memory.
7#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
8pub struct SqueezeRepeatedly;
9
10impl BasicSnippet for SqueezeRepeatedly {
11    fn inputs(&self) -> Vec<(DataType, String)> {
12        vec![
13            (DataType::VoidPointer, "address".to_string()),
14            (DataType::U32, "num_squeezes".to_string()),
15        ]
16    }
17
18    fn outputs(&self) -> Vec<(DataType, String)> {
19        vec![
20            (DataType::VoidPointer, "address".to_string()),
21            (DataType::U32, "num_squeezes".to_string()),
22        ]
23    }
24
25    fn entrypoint(&self) -> String {
26        "tasmlib_hashing_squeeze_repeatedly".to_string()
27    }
28
29    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
30        triton_asm! {
31            // BEFORE: _ address num_squeezes
32            // AFTER:  _ address' 0
33            {self.entrypoint()}:
34
35                // test termination condition
36                dup 0
37                push 0 eq       // _ address num_squeezes num_squeezes==0
38                skiz return
39
40                addi -1
41
42                sponge_squeeze  // _ address (num_squeezes-1) r9 r8 r7 r6 r5 r4 r3 r2 r1 r0
43
44                pick 11         // _ (num_squeezes-1) r9 r8 r7 r6 r5 r4 r3 r2 r1 r0 address
45                write_mem 5
46                write_mem 5
47                                // _ (num_squeezes-1) (address + 10)
48
49                place 1         // _ (address + 10) (num_squeezes-1)
50
51                recurse
52        }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use twenty_first::prelude::Sponge;
59
60    use super::*;
61    use crate::empty_stack;
62    use crate::test_helpers::rust_final_state;
63    use crate::test_helpers::tasm_final_state;
64    use crate::test_helpers::verify_memory_equivalence;
65    use crate::test_helpers::verify_sponge_equivalence;
66    use crate::test_helpers::verify_stack_equivalence;
67    use crate::test_helpers::verify_stack_growth;
68    use crate::test_prelude::*;
69
70    impl Procedure for SqueezeRepeatedly {
71        fn rust_shadow(
72            &self,
73            stack: &mut Vec<BFieldElement>,
74            memory: &mut HashMap<BFieldElement, BFieldElement>,
75            _: &NonDeterminism,
76            _: &[BFieldElement],
77            sponge: &mut Option<Tip5>,
78        ) -> Vec<BFieldElement> {
79            let num_squeezes = stack.pop().unwrap().value() as usize;
80            let address = stack.pop().unwrap();
81
82            let sponge = sponge.as_mut().expect("sponge must be initialized");
83            let sequence = (0..num_squeezes)
84                .flat_map(|_| sponge.squeeze().to_vec())
85                .collect_vec();
86
87            for (i, s) in sequence.into_iter().enumerate() {
88                memory.insert(address + BFieldElement::new(i as u64), s);
89            }
90
91            let new_address = address + BFieldElement::new(tip5::RATE as u64 * num_squeezes as u64);
92            stack.push(new_address);
93            stack.push(BFieldElement::new(0));
94
95            vec![]
96        }
97
98        fn pseudorandom_initial_state(
99            &self,
100            seed: [u8; 32],
101            bench_case: Option<BenchmarkCase>,
102        ) -> ProcedureInitialState {
103            let mut rng = StdRng::from_seed(seed);
104            let num_squeezes = match bench_case {
105                Some(BenchmarkCase::CommonCase) => 10,
106                Some(BenchmarkCase::WorstCase) => 200,
107                None => rng.random_range(0..10),
108            };
109
110            let sponge = Tip5 {
111                state: rng.random(),
112            };
113            let mut stack = empty_stack();
114            let address = BFieldElement::new(rng.next_u64() % (1 << 20));
115            stack.push(address);
116            stack.push(BFieldElement::new(num_squeezes as u64));
117
118            ProcedureInitialState {
119                stack,
120                nondeterminism: NonDeterminism::default(),
121                public_input: vec![],
122                sponge: Some(sponge),
123            }
124        }
125    }
126
127    #[test]
128    fn test() {
129        // custom test procedure because it is a procedure but we do want to test memory equivalence
130
131        let shadow = ShadowedProcedure::new(SqueezeRepeatedly);
132        let num_states = 15;
133        let mut rng = rand::rng();
134        let entrypoint = shadow.inner().entrypoint();
135
136        for _ in 0..num_states {
137            let seed: [u8; 32] = rng.random();
138            println!("testing {entrypoint} common case with seed: {seed:x?}");
139            let ProcedureInitialState {
140                stack,
141                nondeterminism,
142                public_input: stdin,
143                sponge,
144            } = SqueezeRepeatedly.pseudorandom_initial_state(seed, None);
145
146            let init_stack = stack.to_vec();
147
148            let rust = rust_final_state(&shadow, &stack, &stdin, &nondeterminism, &sponge);
149
150            // run tvm
151            let tasm = tasm_final_state(&shadow, &stack, &stdin, nondeterminism, &sponge);
152
153            assert_eq!(
154                rust.public_output, tasm.public_output,
155                "Rust shadowing and VM std out must agree"
156            );
157
158            verify_stack_equivalence(
159                "Rust-shadow",
160                &rust.stack,
161                "TVM execution",
162                &tasm.op_stack.stack,
163            );
164            verify_memory_equivalence("Rust-shadow", &rust.ram, "TVM execution", &tasm.ram);
165            verify_stack_growth(&shadow, &init_stack, &tasm.op_stack.stack);
166            verify_sponge_equivalence(&rust.sponge, &tasm.sponge);
167        }
168    }
169}
170
171#[cfg(test)]
172mod benches {
173    use super::*;
174    use crate::test_prelude::*;
175
176    #[test]
177    fn benchmark() {
178        ShadowedProcedure::new(SqueezeRepeatedly).bench();
179    }
180}
tasm_lib/hashing/squeeze_repeatedly.rs

tasm_lib/hashing/
squeeze_repeatedly.rs
