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 parameters(&self) -> Vec<(DataType, String)> {
12        vec![
13            (DataType::VoidPointer, "address".to_string()),
14            (DataType::U32, "num_squeezes".to_string()),
15        ]
16    }
17
18    fn return_values(&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        ) -> Result<Vec<BFieldElement>, RustShadowError> {
79            let num_squeezes = stack.pop().ok_or(RustShadowError::StackUnderflow)?.value() as usize;
80            let address = stack.pop().ok_or(RustShadowError::StackUnderflow)?;
81
82            let Some(sponge) = sponge.as_mut() else {
83                return Err(RustShadowError::SpongeUninitialized);
84            };
85            let sequence = (0..num_squeezes)
86                .flat_map(|_| sponge.squeeze().to_vec())
87                .collect_vec();
88
89            for (i, s) in sequence.into_iter().enumerate() {
90                memory.insert(address + BFieldElement::new(i as u64), s);
91            }
92
93            let new_address = address + BFieldElement::new(tip5::RATE as u64 * num_squeezes as u64);
94            stack.push(new_address);
95            stack.push(BFieldElement::new(0));
96
97            Ok(Vec::new())
98        }
99
100        fn pseudorandom_initial_state(
101            &self,
102            seed: [u8; 32],
103            bench_case: Option<BenchmarkCase>,
104        ) -> ProcedureInitialState {
105            let mut rng = StdRng::from_seed(seed);
106            let num_squeezes = match bench_case {
107                Some(BenchmarkCase::CommonCase) => 10,
108                Some(BenchmarkCase::WorstCase) => 200,
109                None => rng.random_range(0..10),
110            };
111
112            let sponge = Tip5 {
113                state: rng.random(),
114            };
115            let mut stack = empty_stack();
116            let address = BFieldElement::new(rng.next_u64() % (1 << 20));
117            stack.push(address);
118            stack.push(BFieldElement::new(num_squeezes as u64));
119
120            ProcedureInitialState {
121                stack,
122                nondeterminism: NonDeterminism::default(),
123                public_input: vec![],
124                sponge: Some(sponge),
125            }
126        }
127    }
128
129    #[macro_rules_attr::apply(test)]
130    fn test() {
131        // custom test procedure because it is a procedure but we do want to test memory equivalence
132
133        let shadow = ShadowedProcedure::new(SqueezeRepeatedly);
134        let num_states = 15;
135        let mut rng = rand::rng();
136        let entrypoint = shadow.inner().entrypoint();
137
138        for _ in 0..num_states {
139            let seed: [u8; 32] = rng.random();
140            println!("testing {entrypoint} common case with seed: {seed:x?}");
141            let ProcedureInitialState {
142                stack,
143                nondeterminism,
144                public_input: stdin,
145                sponge,
146            } = SqueezeRepeatedly.pseudorandom_initial_state(seed, None);
147
148            let init_stack = stack.to_vec();
149
150            let rust = rust_final_state(&shadow, &stack, &stdin, &nondeterminism, &sponge);
151
152            // run tvm
153            let tasm = tasm_final_state(&shadow, &stack, &stdin, nondeterminism, &sponge).unwrap();
154
155            assert_eq!(
156                rust.public_output, tasm.public_output,
157                "Rust shadowing and VM std out must agree"
158            );
159
160            verify_stack_equivalence(
161                "Rust-shadow",
162                &rust.stack,
163                "TVM execution",
164                &tasm.op_stack.stack,
165            );
166            verify_memory_equivalence("Rust-shadow", &rust.ram, "TVM execution", &tasm.ram);
167            verify_stack_growth(&shadow, &init_stack, &tasm.op_stack.stack);
168            verify_sponge_equivalence(&rust.sponge, &tasm.sponge);
169        }
170    }
171}
172
173#[cfg(test)]
174mod benches {
175    use super::*;
176    use crate::test_prelude::*;
177
178    #[macro_rules_attr::apply(test)]
179    fn benchmark() {
180        ShadowedProcedure::new(SqueezeRepeatedly).bench();
181    }
182}
tasm_lib/hashing/squeeze_repeatedly.rs

tasm_lib/hashing/
squeeze_repeatedly.rs
