Skip to main content

tasm_lib/hashing/sponge_hasher/
pad_and_absorb_all.rs

1use triton_vm::prelude::*;
2
3use crate::list::LIST_METADATA_SIZE;
4use crate::prelude::*;
5
6#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
7pub struct PadAndAbsorbAll;
8
9impl BasicSnippet for PadAndAbsorbAll {
10    fn parameters(&self) -> Vec<(DataType, String)> {
11        vec![(DataType::List(Box::new(DataType::Bfe)), "input".to_string())]
12    }
13
14    fn return_values(&self) -> Vec<(DataType, String)> {
15        vec![]
16    }
17
18    fn entrypoint(&self) -> String {
19        "tasmlib_hashing_sponge_hasher_pad_and_absorb_all".into()
20    }
21
22    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
23        let entrypoint = self.entrypoint();
24        let hash_absorb_snippet_subroutine =
25            library.import(Box::new(crate::hashing::absorb_multiple::AbsorbMultiple));
26        triton_asm!(
27            {entrypoint}:
28                // _ *input_list
29
30                read_mem 1
31                // _ length (*input_list - 1)
32
33                addi {LIST_METADATA_SIZE + 1}
34                // _ length *first_element
35
36                pick 1
37                // _ *first_element length
38
39                call {hash_absorb_snippet_subroutine}
40                // _
41
42                return
43        )
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use arbitrary::Arbitrary;
50    use arbitrary::Unstructured;
51    use twenty_first::prelude::Sponge;
52
53    use super::*;
54    use crate::empty_stack;
55    use crate::rust_shadowing_helper_functions::list::insert_random_list;
56    use crate::test_prelude::*;
57
58    impl PadAndAbsorbAll {
59        fn init_memory_and_stack(
60            &self,
61            input_length: usize,
62        ) -> (HashMap<BFieldElement, BFieldElement>, Vec<BFieldElement>) {
63            let list_pointer: BFieldElement = rand::random();
64            let init_stack = [empty_stack(), vec![list_pointer]].concat();
65            let mut init_memory = HashMap::default();
66            insert_random_list(&DataType::Bfe, list_pointer, input_length, &mut init_memory);
67
68            (init_memory, init_stack)
69        }
70    }
71
72    impl Procedure for PadAndAbsorbAll {
73        fn rust_shadow(
74            &self,
75            stack: &mut Vec<BFieldElement>,
76            memory: &mut HashMap<BFieldElement, BFieldElement>,
77            _nondeterminism: &NonDeterminism,
78            _public_input: &[BFieldElement],
79            sponge: &mut Option<Tip5>,
80        ) -> Result<Vec<BFieldElement>, RustShadowError> {
81            let Some(sponge) = sponge.as_mut() else {
82                return Err(RustShadowError::SpongeUninitialized);
83            };
84            let input_pointer: BFieldElement =
85                stack.pop().ok_or(RustShadowError::StackUnderflow)?;
86            let first_word = input_pointer + BFieldElement::new(LIST_METADATA_SIZE as u64);
87            let input_length = memory[&input_pointer].value();
88            let input = (0..input_length)
89                .map(|i| memory[&(BFieldElement::new(i) + first_word)])
90                .collect_vec();
91            sponge.pad_and_absorb_all(&input);
92
93            Ok(Vec::new())
94        }
95
96        fn pseudorandom_initial_state(
97            &self,
98            seed: [u8; 32],
99            bench_case: Option<BenchmarkCase>,
100        ) -> ProcedureInitialState {
101            let mut rng = StdRng::from_seed(seed);
102            let mut bytes = [0u8; 400];
103            rng.fill_bytes(&mut bytes);
104            let mut unstructured = Unstructured::new(&bytes);
105            let input_length = match bench_case {
106                Some(BenchmarkCase::CommonCase) => 100,
107                Some(BenchmarkCase::WorstCase) => 10000,
108                None => rng.random_range(0..200),
109            };
110            let (memory, stack) = self.init_memory_and_stack(input_length);
111
112            ProcedureInitialState {
113                stack,
114                nondeterminism: NonDeterminism::default().with_ram(memory),
115                public_input: Vec::default(),
116                sponge: Some(Tip5::arbitrary(&mut unstructured).unwrap()),
117            }
118        }
119
120        fn corner_case_initial_states(&self) -> Vec<ProcedureInitialState> {
121            let empty_input_empty_sponge = {
122                let (memory, stack) = self.init_memory_and_stack(0);
123                ProcedureInitialState {
124                    stack,
125                    nondeterminism: NonDeterminism::default().with_ram(memory),
126                    public_input: Vec::default(),
127                    sponge: Some(Tip5::init()),
128                }
129            };
130            let empty_input_random_sponge = {
131                let (memory, stack) = self.init_memory_and_stack(0);
132                let mut rng = StdRng::from_seed([12u8; 32]);
133                let mut bytes = [0u8; 400];
134                rng.fill_bytes(&mut bytes);
135                let mut unstructured = Unstructured::new(&bytes);
136                ProcedureInitialState {
137                    stack,
138                    nondeterminism: NonDeterminism::default().with_ram(memory),
139                    public_input: Vec::default(),
140                    sponge: Some(Tip5::arbitrary(&mut unstructured).unwrap()),
141                }
142            };
143            let length_one_input_empty_sponge = {
144                let (memory, stack) = self.init_memory_and_stack(1);
145                ProcedureInitialState {
146                    stack,
147                    nondeterminism: NonDeterminism::default().with_ram(memory),
148                    public_input: Vec::default(),
149                    sponge: Some(Tip5::init()),
150                }
151            };
152            vec![
153                empty_input_empty_sponge,
154                empty_input_random_sponge,
155                length_one_input_empty_sponge,
156            ]
157        }
158    }
159
160    #[macro_rules_attr::apply(test)]
161    fn pad_and_absorb_all_test() {
162        ShadowedProcedure::new(PadAndAbsorbAll).test();
163    }
164}