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 inputs(&self) -> Vec<(DataType, String)> {
19        vec![]
20    }
21
22    fn outputs(&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        ) -> Vec<BFieldElement> {
87            let sponge = sponge.as_mut().expect("sponge must be initialized");
88            let mut array_pointer =
89                rust_shadowing_helper_functions::dyn_malloc::dynamic_allocator(memory);
90            stack.push(array_pointer);
91            let produce = sponge.squeeze();
92            for elem in produce.into_iter() {
93                memory.insert(array_pointer, elem);
94                array_pointer.increment();
95            }
96
97            Vec::default()
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 mut init_memory: HashMap<BFieldElement, BFieldElement> = HashMap::default();
107            let random_dynmalloc_init_page_counter =
108                rng.random_range(0..dyn_malloc::NUM_ALLOCATABLE_PAGES);
109            init_memory.insert(DYN_MALLOC_ADDRESS, bfe!(random_dynmalloc_init_page_counter));
110
111            let mut rng = StdRng::from_seed(seed);
112            let mut bytes = [0u8; 400];
113            rng.fill_bytes(&mut bytes);
114            let mut unstructured = Unstructured::new(&bytes);
115
116            ProcedureInitialState {
117                stack: empty_stack(),
118                nondeterminism: NonDeterminism::default().with_ram(init_memory),
119                public_input: Vec::default(),
120                sponge: Some(Tip5::arbitrary(&mut unstructured).unwrap()),
121            }
122        }
123    }
124
125    #[test]
126    fn squeeze_test() {
127        ShadowedProcedure::new(Squeeze).test();
128    }
129}
130
131#[cfg(test)]
132mod benches {
133    use super::*;
134    use crate::test_prelude::*;
135
136    #[test]
137    fn benchmark() {
138        ShadowedProcedure::new(Squeeze).bench();
139    }
140}