Skip to main content

tasm_lib/hashing/sponge_hasher/
squeeze.rs

1use triton_vm::prelude::*;
2use twenty_first::tip5::RATE;
3
4use crate::data_type::ArrayType;
5use crate::prelude::*;
6
7/// Squeeze the sponge and return an array of `[RATE]` elements
8///
9/// Snippet that emulates the Tip5 implementation of twenty-first's
10/// `sponge_hasher` trait function `squeeze`. You probably don't want to use
11/// this snippet for whatever cryptography you're doing and instead use the
12/// instruction `sponge_squeeze` directly, or some version of
13/// `squeeze_repeatedly`.
14#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
15pub struct Squeeze;
16
17impl BasicSnippet for Squeeze {
18    fn parameters(&self) -> Vec<(DataType, String)> {
19        vec![]
20    }
21
22    fn return_values(&self) -> Vec<(DataType, String)> {
23        let produce_type = DataType::Array(Box::new(ArrayType {
24            element_type: DataType::Bfe,
25            length: RATE,
26        }));
27
28        vec![(produce_type, "produce".to_string())]
29    }
30
31    fn entrypoint(&self) -> String {
32        "tasmlib_hashing_sponge_hasher_squeeze".to_string()
33    }
34
35    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
36        assert_eq!(10, RATE, "Code assumes RATE is 10");
37        let entrypoint = self.entrypoint();
38        let dyn_malloc_label = library.import(Box::new(DynMalloc));
39
40        triton_asm!(
41            {entrypoint}:
42                // _
43                sponge_squeeze
44                // _ [word_9..word_0]
45
46
47                // Allocate memory for the returned array
48                call {dyn_malloc_label}
49                // _ [word_9..word_0] *array
50
51                // Write words to array
52                write_mem 5
53                write_mem 5
54                // _ (*array + 10)
55
56                push -10
57                add
58                // _ *array
59
60                return
61        )
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use arbitrary::Arbitrary;
68    use arbitrary::Unstructured;
69    use twenty_first::prelude::Sponge;
70
71    use super::*;
72    use crate::empty_stack;
73    use crate::memory::dyn_malloc;
74    use crate::memory::dyn_malloc::DYN_MALLOC_ADDRESS;
75    use crate::rust_shadowing_helper_functions;
76    use crate::test_prelude::*;
77
78    impl Procedure for Squeeze {
79        fn rust_shadow(
80            &self,
81            stack: &mut Vec<BFieldElement>,
82            memory: &mut HashMap<BFieldElement, BFieldElement>,
83            _nondeterminism: &NonDeterminism,
84            _public_input: &[BFieldElement],
85            sponge: &mut Option<Tip5>,
86        ) -> Result<Vec<BFieldElement>, RustShadowError> {
87            let Some(sponge) = sponge.as_mut() else {
88                return Err(RustShadowError::SpongeUninitialized);
89            };
90            let mut array_pointer =
91                rust_shadowing_helper_functions::dyn_malloc::dynamic_allocator(memory);
92            stack.push(array_pointer);
93            let produce = sponge.squeeze();
94            for elem in produce.into_iter() {
95                memory.insert(array_pointer, elem);
96                array_pointer.increment();
97            }
98
99            Ok(Vec::new())
100        }
101
102        fn pseudorandom_initial_state(
103            &self,
104            seed: [u8; 32],
105            _bench_case: Option<BenchmarkCase>,
106        ) -> ProcedureInitialState {
107            let mut rng = StdRng::from_seed(seed);
108            let mut init_memory: HashMap<BFieldElement, BFieldElement> = HashMap::default();
109            let random_dynmalloc_init_page_counter =
110                rng.random_range(0..dyn_malloc::NUM_ALLOCATABLE_PAGES);
111            init_memory.insert(DYN_MALLOC_ADDRESS, bfe!(random_dynmalloc_init_page_counter));
112
113            let mut rng = StdRng::from_seed(seed);
114            let mut bytes = [0u8; 400];
115            rng.fill_bytes(&mut bytes);
116            let mut unstructured = Unstructured::new(&bytes);
117
118            ProcedureInitialState {
119                stack: empty_stack(),
120                nondeterminism: NonDeterminism::default().with_ram(init_memory),
121                public_input: Vec::default(),
122                sponge: Some(Tip5::arbitrary(&mut unstructured).unwrap()),
123            }
124        }
125    }
126
127    #[macro_rules_attr::apply(test)]
128    fn squeeze_test() {
129        ShadowedProcedure::new(Squeeze).test();
130    }
131}
132
133#[cfg(test)]
134mod benches {
135    use super::*;
136    use crate::test_prelude::*;
137
138    #[macro_rules_attr::apply(test)]
139    fn benchmark() {
140        ShadowedProcedure::new(Squeeze).bench();
141    }
142}