tasm_lib/array/
sum_of_bfes.rs1use num::Zero;
2use triton_vm::prelude::*;
3
4use crate::data_type::ArrayType;
5use crate::memory::load_words_from_memory_pop_pointer;
6use crate::prelude::*;
7
8#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
9pub struct SumOfBfes {
10 length: usize,
11}
12
13impl BasicSnippet for SumOfBfes {
14 fn inputs(&self) -> Vec<(DataType, String)> {
15 vec![(
16 DataType::Array(Box::new(ArrayType {
17 element_type: DataType::Bfe,
18 length: self.length,
19 })),
20 "*array".to_owned(),
21 )]
22 }
23
24 fn outputs(&self) -> Vec<(DataType, String)> {
25 vec![(DataType::Bfe, "sum".to_owned())]
26 }
27
28 fn entrypoint(&self) -> String {
29 format!("tasmlib_array_sum_of_{}_bfes", self.length)
30 }
31
32 fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
33 let move_pointer_to_last_word = match self.length {
34 0 | 1 => triton_asm!(),
35 n => triton_asm!( addi {n - 1} ),
36 };
37
38 let load_all_elements_to_stack = load_words_from_memory_pop_pointer(self.length);
39
40 let sum = if self.length.is_zero() {
41 triton_asm!(push 0)
42 } else {
43 triton_asm![add; self.length - 1]
44 };
45
46 triton_asm!(
47 {self.entrypoint()}:
48 {&move_pointer_to_last_word}
51 {&load_all_elements_to_stack}
54 {&sum}
57 return
60 )
61 }
62}
63
64#[cfg(test)]
65mod tests {
66 use num_traits::ConstZero;
67
68 use super::*;
69 use crate::rust_shadowing_helper_functions::array::insert_random_array;
70 use crate::test_prelude::*;
71
72 impl Function for SumOfBfes {
73 fn rust_shadow(
74 &self,
75 stack: &mut Vec<BFieldElement>,
76 memory: &mut HashMap<BFieldElement, BFieldElement>,
77 ) {
78 let mut array_pointer = stack.pop().unwrap();
79 let mut array_quote_unquote = bfe_vec![0; self.length];
80 for array_elem in array_quote_unquote.iter_mut() {
81 memory
82 .get(&array_pointer)
83 .unwrap_or(&BFieldElement::ZERO)
84 .clone_into(array_elem);
85 array_pointer.increment();
86 }
87
88 let sum = array_quote_unquote.into_iter().sum();
89
90 stack.push(sum);
91 }
92
93 fn pseudorandom_initial_state(
94 &self,
95 seed: [u8; 32],
96 _: Option<BenchmarkCase>,
97 ) -> FunctionInitialState {
98 self.prepare_state(StdRng::from_seed(seed).random())
99 }
100
101 fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
102 let all_zeros = {
103 let mut init_stack = self.init_stack_for_isolated_run();
104 init_stack.push(BFieldElement::new(500));
105 FunctionInitialState {
106 stack: init_stack,
107 memory: HashMap::default(),
108 }
109 };
110
111 vec![all_zeros]
112 }
113 }
114
115 impl SumOfBfes {
116 fn prepare_state(&self, array_pointer: BFieldElement) -> FunctionInitialState {
117 let mut memory = HashMap::default();
118 insert_random_array(&DataType::Bfe, array_pointer, self.length, &mut memory);
119
120 let mut stack = self.init_stack_for_isolated_run();
121 stack.push(array_pointer);
122
123 FunctionInitialState { stack, memory }
124 }
125 }
126
127 #[test]
128 fn sum_bfes_pbt() {
129 let snippets = (0..20).chain(100..110).map(|x| SumOfBfes { length: x });
130 for test_case in snippets {
131 ShadowedFunction::new(test_case).test()
132 }
133 }
134}
135
136#[cfg(test)]
137mod benches {
138 use super::*;
139 use crate::test_prelude::*;
140
141 #[test]
142 fn sum_bfes_bench_100() {
143 ShadowedFunction::new(SumOfBfes { length: 100 }).bench();
144 }
145
146 #[test]
147 fn sum_bfes_bench_200() {
148 ShadowedFunction::new(SumOfBfes { length: 200 }).bench();
149 }
150}