tasm_lib/mmr/
calculate_new_peaks_from_append.rs

1use triton_vm::prelude::*;
2
3use crate::arithmetic::u64::incr::Incr;
4use crate::arithmetic::u64::trailing_zeros::TrailingZeros;
5use crate::list::new::New;
6use crate::list::pop::Pop;
7use crate::list::push::Push;
8use crate::prelude::*;
9
10#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
11pub struct CalculateNewPeaksFromAppend;
12
13impl BasicSnippet for CalculateNewPeaksFromAppend {
14    fn parameters(&self) -> Vec<(DataType, String)> {
15        let list_type = DataType::List(Box::new(DataType::Digest));
16
17        vec![
18            (DataType::U64, "old_num_leafs".to_string()),
19            (list_type, "peaks".to_string()),
20            (DataType::Digest, "new_leaf".to_string()),
21        ]
22    }
23
24    fn return_values(&self) -> Vec<(DataType, String)> {
25        let list_type = DataType::List(Box::new(DataType::Digest));
26
27        vec![
28            (list_type.clone(), "*new_peaks".to_string()),
29            (list_type, "*auth_path".to_string()),
30        ]
31    }
32
33    fn entrypoint(&self) -> String {
34        "tasmlib_mmr_calculate_new_peaks_from_append".into()
35    }
36
37    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
38        let entrypoint = self.entrypoint();
39        let while_loop_label = format!("{entrypoint}_while");
40
41        let new_list = library.import(Box::new(New));
42        let push = library.import(Box::new(Push::new(DataType::Digest)));
43        let pop = library.import(Box::new(Pop::new(DataType::Digest)));
44        let u64incr = library.import(Box::new(Incr));
45        let right_lineage_count = library.import(Box::new(TrailingZeros));
46
47        triton_asm!(
48            // BEFORE: _ [old_num_leafs: u64] *peaks [new_leaf: Digest]
49            // AFTER:  _ *new_peaks *auth_path
50            {entrypoint}:
51                dup 5
52                place 5
53                call {push}
54                // _ [old_num_leafs: u64] *peaks
55
56                /* create auth_path return value */
57                call {new_list}
58                // _ [old_num_leafs: u64] *peaks *auth_path
59
60                pick 3 pick 3
61                // _ *peaks *auth_path [old_num_leafs: u64]
62
63                call {u64incr}
64                call {right_lineage_count}
65                // _ *peaks *auth_path right_lineage_count
66
67                call {while_loop_label}
68                // _ *peaks *auth_path 0
69
70                pop 1
71                // _ *peaks *auth_path
72
73                return
74
75            // INVARIANT: _ *peaks *auth_path rlc
76            {while_loop_label}:
77                dup 0
78                push 0
79                eq
80                skiz
81                    return
82                // _ *peaks *auth_path rlc
83
84                dup 2
85                dup 0
86                call {pop}
87                // _ *peaks *auth_path rlc *peaks [new_hash: Digest]
88
89                dup 5
90                // _ *peaks *auth_path rlc *peaks [new_hash: Digest] *peaks
91
92                call {pop}
93                // _ *peaks *auth_path rlc *peaks [new_hash: Digest] [previous_peak: Digest]
94
95                /* update authentication path with latest previous_peak */
96                dup 12
97                dup 5 dup 5 dup 5 dup 5 dup 5
98                call {push}
99                // _ *peaks *auth_path rlc *peaks [new_hash: Digest] [previous_peak: Digest]
100
101                hash
102                // _ *peaks *auth_path rlc *peaks [new_peak: Digest]
103
104                call {push}
105                // _ *peaks *auth_path rlc
106
107                addi -1
108                // _ *auth_path *peaks (rlc - 1)
109
110                recurse
111        )
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator;
118
119    use super::*;
120    use crate::memory::FIRST_NON_DETERMINISTICALLY_INITIALIZED_MEMORY_ADDRESS;
121    use crate::memory::dyn_malloc::DYN_MALLOC_FIRST_ADDRESS;
122    use crate::rust_shadowing_helper_functions::list::list_pop;
123    use crate::rust_shadowing_helper_functions::list::list_push;
124    use crate::test_prelude::*;
125    use crate::twenty_first::prelude::Mmr;
126    use crate::twenty_first::util_types::mmr::shared_basic::calculate_new_peaks_from_append;
127
128    impl CalculateNewPeaksFromAppend {
129        fn set_up_initial_state(
130            &self,
131            mmr: MmrAccumulator,
132            new_leaf: Digest,
133        ) -> FunctionInitialState {
134            let peaks_pointer = FIRST_NON_DETERMINISTICALLY_INITIALIZED_MEMORY_ADDRESS;
135
136            let mut stack = self.init_stack_for_isolated_run();
137            push_encodable(&mut stack, &mmr.num_leafs());
138            push_encodable(&mut stack, &peaks_pointer);
139            push_encodable(&mut stack, &new_leaf.reversed());
140
141            let mut memory = HashMap::default();
142            encode_to_memory(&mut memory, peaks_pointer, &mmr.peaks());
143
144            FunctionInitialState { stack, memory }
145        }
146    }
147
148    impl Function for CalculateNewPeaksFromAppend {
149        fn rust_shadow(
150            &self,
151            stack: &mut Vec<BFieldElement>,
152            memory: &mut HashMap<BFieldElement, BFieldElement>,
153        ) -> Result<(), RustShadowError> {
154            let new_leaf = pop_encodable::<Digest>(stack)?;
155            let peaks_pointer = stack.pop().ok_or(RustShadowError::StackUnderflow)?;
156            let old_num_leafs = pop_encodable::<u64>(stack)?;
157            let old_peaks = *Vec::decode_from_memory(memory, peaks_pointer)
158                .map_err(|_| RustShadowError::DecodingError)?;
159
160            // Mimic all potential artifacts of the snippet.
161            // This is _not_ shadowing the actual behavior, only intermediate results.
162            list_push(peaks_pointer, new_leaf.encode(), memory)?;
163            let old_num_peaks = old_num_leafs.count_ones();
164            for _ in 0..old_num_peaks {
165                let left = list_pop(peaks_pointer, memory, Digest::LEN)?;
166                let right = list_pop(peaks_pointer, memory, Digest::LEN)?;
167                let right: Digest = right
168                    .try_into()
169                    .map_err(|_| RustShadowError::DecodingError)?;
170                let left: Digest = left
171                    .try_into()
172                    .map_err(|_| RustShadowError::DecodingError)?;
173                let new = Tip5::hash_pair(right, left);
174                list_push(peaks_pointer, new.encode(), memory)?;
175            }
176
177            // actually shadow the snippet
178            let (new_peaks, proof) =
179                calculate_new_peaks_from_append(old_num_leafs, old_peaks, new_leaf);
180
181            let auth_path_pointer = DYN_MALLOC_FIRST_ADDRESS;
182            encode_to_memory(memory, peaks_pointer, &new_peaks);
183            encode_to_memory(memory, auth_path_pointer, &proof.authentication_path);
184            stack.push(peaks_pointer);
185            stack.push(auth_path_pointer);
186            Ok(())
187        }
188
189        fn pseudorandom_initial_state(
190            &self,
191            seed: [u8; 32],
192            bench_case: Option<BenchmarkCase>,
193        ) -> FunctionInitialState {
194            let mut rng = StdRng::from_seed(seed);
195
196            let old_num_leafs: u64 = match bench_case {
197                Some(BenchmarkCase::CommonCase) => (1 << 31) - 1,
198                Some(BenchmarkCase::WorstCase) => (1 << 62) - 1,
199                None => rng.random_range(0..1 << 63),
200            };
201            let peaks = (0..old_num_leafs.count_ones())
202                .map(|_| rng.random())
203                .collect();
204            let mmr = MmrAccumulator::init(peaks, old_num_leafs);
205
206            self.set_up_initial_state(mmr, rng.random())
207        }
208
209        fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
210            let mut states = vec![];
211            for num_leafs in (0_u64..=17).chain([100, 1000, 1 << 31, 1 << 32, 1 << 33]) {
212                let num_peaks = num_leafs.count_ones();
213                let peaks = (0..num_peaks).map(|i| Tip5::hash(&i)).collect();
214                let mmr = MmrAccumulator::init(peaks, num_leafs);
215                let new_leaf = Tip5::hash(&num_leafs);
216                states.push(self.set_up_initial_state(mmr, new_leaf));
217            }
218
219            states
220        }
221    }
222
223    #[macro_rules_attr::apply(test)]
224    fn rust_shadow() {
225        ShadowedFunction::new(CalculateNewPeaksFromAppend).test();
226    }
227}
228
229#[cfg(test)]
230mod benches {
231    use super::*;
232    use crate::test_prelude::*;
233
234    #[macro_rules_attr::apply(test)]
235    fn calculate_new_peaks_from_append_benchmark() {
236        ShadowedFunction::new(CalculateNewPeaksFromAppend).bench();
237    }
238}
tasm_lib/mmr/calculate_new_peaks_from_append.rs

tasm_lib/mmr/
calculate_new_peaks_from_append.rs
