Skip to main content

tasm_lib/hashing/sponge_hasher/
absorb.rs

1use triton_vm::prelude::*;
2use twenty_first::tip5::RATE;
3
4use crate::data_type::ArrayType;
5use crate::prelude::*;
6
7#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
8pub struct Absorb;
9
10impl BasicSnippet for Absorb {
11    fn parameters(&self) -> Vec<(DataType, String)> {
12        vec![(
13            DataType::Array(Box::new(ArrayType {
14                element_type: DataType::Bfe,
15                length: RATE,
16            })),
17            "input".to_owned(),
18        )]
19    }
20
21    fn return_values(&self) -> Vec<(DataType, String)> {
22        vec![]
23    }
24
25    fn entrypoint(&self) -> String {
26        "tasmlib_hashing_sponge_hasher_absorb".to_string()
27    }
28
29    fn code(&self, _library: &mut Library) -> Vec<LabelledInstruction> {
30        assert_eq!(10, RATE, "Code assumes RATE is 10");
31
32        triton_asm!(
33            {self.entrypoint()}:
34                // _ *input
35
36                push 0
37                push 0
38                push 0
39                push 0
40                pick 4
41                // _ 0 0 0 0 *input
42
43                sponge_absorb_mem
44                // _ g0 g1 g2 g3 (*input+10)
45
46                pop 5
47                // _
48
49                return
50        )
51
52        // Can also be implemented without the use of `sponge_absorb_mem`:
53        // triton_asm!(
54        //     {entrypoint}:
55        //         // _ *input
56
57        //         push 9 add
58        //         // _ *last_bfe
59
60        //         read_mem 5
61        //         read_mem 5
62        //         // _ [word_9..word_0] (*first_elem - 1)
63
64        //         pop 1
65        //         // _ [word_9..word_0]
66
67        //         sponge_absorb
68        //         // _
69
70        //         return
71        // )
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use arbitrary::Arbitrary;
78    use arbitrary::Unstructured;
79    use twenty_first::math::other::random_elements;
80    use twenty_first::prelude::Sponge;
81
82    use super::*;
83    use crate::empty_stack;
84    use crate::test_prelude::*;
85
86    fn init_memory_and_stack() -> (HashMap<BFieldElement, BFieldElement>, Vec<BFieldElement>) {
87        let array_pointer: BFieldElement = rand::random();
88        let init_stack = [empty_stack(), vec![array_pointer]].concat();
89        let init_memory: HashMap<BFieldElement, BFieldElement> = (0..RATE)
90            .map(|i| array_pointer + BFieldElement::new(i as u64))
91            .zip(random_elements(RATE))
92            .collect();
93
94        (init_memory, init_stack)
95    }
96
97    impl Procedure for Absorb {
98        fn rust_shadow(
99            &self,
100            stack: &mut Vec<BFieldElement>,
101            memory: &mut HashMap<BFieldElement, BFieldElement>,
102            _nondeterminism: &NonDeterminism,
103            _public_input: &[BFieldElement],
104            sponge: &mut Option<Tip5>,
105        ) -> Result<Vec<BFieldElement>, RustShadowError> {
106            let Some(sponge) = sponge.as_mut() else {
107                return Err(RustShadowError::SpongeUninitialized);
108            };
109            let input_pointer: BFieldElement =
110                stack.pop().ok_or(RustShadowError::StackUnderflow)?;
111            let input: [BFieldElement; 10] = (0..10)
112                .map(|i| memory[&(BFieldElement::new(i) + input_pointer)])
113                .collect_vec()
114                .try_into()
115                .unwrap();
116            sponge.absorb(input);
117
118            Ok(Vec::new())
119        }
120
121        fn pseudorandom_initial_state(
122            &self,
123            seed: [u8; 32],
124            _bench_case: Option<BenchmarkCase>,
125        ) -> ProcedureInitialState {
126            let mut rng = StdRng::from_seed(seed);
127            let mut bytes = [0u8; 400];
128            rng.fill_bytes(&mut bytes);
129            let mut unstructured = Unstructured::new(&bytes);
130            let (memory, stack) = init_memory_and_stack();
131
132            ProcedureInitialState {
133                stack,
134                nondeterminism: NonDeterminism::default().with_ram(memory),
135                public_input: Vec::default(),
136                sponge: Some(Tip5::arbitrary(&mut unstructured).unwrap()),
137            }
138        }
139
140        fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
141            let empty_sponge = {
142                let (memory, stack) = init_memory_and_stack();
143                ProcedureInitialState {
144                    stack,
145                    nondeterminism: NonDeterminism::default().with_ram(memory),
146                    public_input: Vec::default(),
147                    sponge: Some(Tip5::init()),
148                }
149            };
150
151            vec![empty_sponge]
152        }
153    }
154
155    #[macro_rules_attr::apply(test)]
156    fn absorb_test() {
157        ShadowedProcedure::new(Absorb).test();
158    }
159}
160
161#[cfg(test)]
162mod benches {
163    use super::*;
164    use crate::test_prelude::*;
165
166    #[macro_rules_attr::apply(test)]
167    fn benchmark() {
168        ShadowedProcedure::new(Absorb).bench();
169    }
170}