tasm_lib/array/
sum_of_xfes.rs

1use num::Zero;
2use triton_vm::prelude::*;
3use twenty_first::math::x_field_element::EXTENSION_DEGREE;
4
5use crate::data_type::ArrayType;
6use crate::memory::load_words_from_memory_pop_pointer;
7use crate::prelude::*;
8
9#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
10pub struct SumOfXfes {
11    length: usize,
12}
13
14impl SumOfXfes {
15    fn input_type(&self) -> DataType {
16        DataType::Array(Box::new(ArrayType {
17            element_type: DataType::Xfe,
18            length: self.length,
19        }))
20    }
21}
22
23impl BasicSnippet for SumOfXfes {
24    fn inputs(&self) -> Vec<(DataType, String)> {
25        vec![(self.input_type(), "*array".to_owned())]
26    }
27
28    fn outputs(&self) -> Vec<(DataType, String)> {
29        vec![(DataType::Xfe, "sum".to_owned())]
30    }
31
32    fn entrypoint(&self) -> String {
33        format!("tasmlib_array_sum_of_{}_xfes", self.length)
34    }
35
36    fn code(&self, _library: &mut Library) -> Vec<LabelledInstruction> {
37        let entrypoint = self.entrypoint();
38
39        let move_pointer_to_last_word = match self.length {
40            0 => triton_asm!(),
41            n => {
42                let word_length_minus_one = EXTENSION_DEGREE * n - 1;
43                triton_asm!(
44                    push {word_length_minus_one}
45                    add
46                )
47            }
48        };
49
50        let load_all_elements_to_stack =
51            load_words_from_memory_pop_pointer(EXTENSION_DEGREE * self.length);
52
53        let sum = if self.length.is_zero() {
54            triton_asm!(
55                push 0
56                push 0
57                push 0
58            )
59        } else {
60            vec![triton_asm!(xx_add); self.length - 1].concat()
61        };
62
63        triton_asm!(
64            {entrypoint}:
65                // _ *array
66
67                {&move_pointer_to_last_word}
68                // _ *last_word
69
70                {&load_all_elements_to_stack}
71                // _ [elements]
72
73                {&sum}
74                // _ sum
75
76                return
77        )
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use num::One;
84    use num::Zero;
85    use num_traits::ConstZero;
86
87    use super::*;
88    use crate::rust_shadowing_helper_functions::array::insert_as_array;
89    use crate::rust_shadowing_helper_functions::array::insert_random_array;
90    use crate::test_prelude::*;
91
92    impl Function for SumOfXfes {
93        fn rust_shadow(
94            &self,
95            stack: &mut Vec<BFieldElement>,
96            memory: &mut HashMap<BFieldElement, BFieldElement>,
97        ) {
98            let mut array_pointer = stack.pop().unwrap();
99            let mut array_quote_unquote = vec![XFieldElement::zero(); self.length];
100            for array_elem in array_quote_unquote.iter_mut() {
101                for word in array_elem.coefficients.iter_mut() {
102                    memory
103                        .get(&array_pointer)
104                        .unwrap_or(&BFieldElement::ZERO)
105                        .clone_into(word);
106                    array_pointer.increment();
107                }
108            }
109
110            let sum: XFieldElement = array_quote_unquote.into_iter().sum();
111
112            for word in sum.coefficients.into_iter().rev() {
113                stack.push(word);
114            }
115        }
116
117        fn pseudorandom_initial_state(
118            &self,
119            seed: [u8; 32],
120            _bench_case: Option<BenchmarkCase>,
121        ) -> crate::traits::function::FunctionInitialState {
122            let mut rng = StdRng::from_seed(seed);
123            let list_pointer = BFieldElement::new(rng.random());
124            self.prepare_state(list_pointer)
125        }
126
127        fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
128            let all_zeros = {
129                let mut init_stack = self.init_stack_for_isolated_run();
130                init_stack.push(BFieldElement::new(500));
131                FunctionInitialState {
132                    stack: init_stack,
133                    memory: HashMap::default(),
134                }
135            };
136            let all_ones = {
137                let mut init_stack = self.init_stack_for_isolated_run();
138                let array_pointer = BFieldElement::new(500);
139                init_stack.push(array_pointer);
140                let all_ones = vec![XFieldElement::one(); self.length];
141                let mut memory = HashMap::default();
142                insert_as_array(array_pointer, &mut memory, all_ones);
143
144                FunctionInitialState {
145                    stack: init_stack,
146                    memory,
147                }
148            };
149
150            vec![all_zeros, all_ones]
151        }
152    }
153
154    impl SumOfXfes {
155        fn prepare_state(&self, array_pointer: BFieldElement) -> FunctionInitialState {
156            let mut memory = HashMap::default();
157            insert_random_array(&DataType::Xfe, array_pointer, self.length, &mut memory);
158
159            let mut init_stack = self.init_stack_for_isolated_run();
160            init_stack.push(array_pointer);
161            FunctionInitialState {
162                stack: init_stack,
163                memory,
164            }
165        }
166    }
167
168    #[test]
169    fn sum_xfes_pbt() {
170        let snippets = (0..20).chain(100..110).map(|x| SumOfXfes { length: x });
171        for test_case in snippets {
172            ShadowedFunction::new(test_case).test()
173        }
174    }
175
176    #[test]
177    fn xfe_array_sum_unit_test() {
178        let xfes = vec![
179            XFieldElement::new([
180                BFieldElement::new(100),
181                BFieldElement::zero(),
182                BFieldElement::new(10),
183            ]),
184            XFieldElement::new([
185                BFieldElement::new(200),
186                BFieldElement::zero(),
187                BFieldElement::new(4),
188            ]),
189        ];
190
191        let expected_sum = XFieldElement::new([
192            BFieldElement::new(300),
193            BFieldElement::zero(),
194            BFieldElement::new(14),
195        ]);
196        assert_eq!(xfes.iter().cloned().sum::<XFieldElement>(), expected_sum);
197
198        let mut memory = HashMap::default();
199        let array_pointer = BFieldElement::new(1u64 << 44);
200        insert_as_array(array_pointer, &mut memory, xfes);
201
202        let snippet = SumOfXfes { length: 2 };
203        let expected_final_stack = [
204            snippet.init_stack_for_isolated_run(),
205            expected_sum.coefficients.into_iter().rev().collect_vec(),
206        ]
207        .concat();
208        let init_stack = [snippet.init_stack_for_isolated_run(), vec![array_pointer]].concat();
209        test_rust_equivalence_given_complete_state(
210            &ShadowedFunction::new(snippet),
211            &init_stack,
212            &[],
213            &NonDeterminism::default().with_ram(memory),
214            &None,
215            Some(&expected_final_stack),
216        );
217    }
218}
219
220#[cfg(test)]
221mod benches {
222    use super::*;
223    use crate::test_prelude::*;
224
225    #[test]
226    fn sum_xfes_bench_100() {
227        ShadowedFunction::new(SumOfXfes { length: 100 }).bench();
228    }
229
230    #[test]
231    fn sum_xfes_bench_200() {
232        ShadowedFunction::new(SumOfXfes { length: 200 }).bench();
233    }
234}
tasm_lib/array/sum_of_xfes.rs

tasm_lib/array/
sum_of_xfes.rs
