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 inputs(&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 outputs(&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        ) {
154            let new_leaf = pop_encodable::<Digest>(stack);
155            let peaks_pointer = stack.pop().unwrap();
156            let old_num_leafs = pop_encodable::<u64>(stack);
157            let old_peaks = *Vec::decode_from_memory(memory, peaks_pointer).unwrap();
158
159            // Mimic all potential artifacts of the snippet.
160            // This is _not_ shadowing the actual behavior, only intermediate results.
161            list_push(peaks_pointer, new_leaf.encode(), memory);
162            let old_num_peaks = old_num_leafs.count_ones();
163            for _ in 0..old_num_peaks {
164                let left = list_pop(peaks_pointer, memory, Digest::LEN);
165                let right = list_pop(peaks_pointer, memory, Digest::LEN);
166                let new = Tip5::hash_pair(right.try_into().unwrap(), left.try_into().unwrap());
167                list_push(peaks_pointer, new.encode(), memory);
168            }
169
170            // actually shadow the snippet
171            let (new_peaks, proof) =
172                calculate_new_peaks_from_append(old_num_leafs, old_peaks, new_leaf);
173
174            let auth_path_pointer = DYN_MALLOC_FIRST_ADDRESS;
175            encode_to_memory(memory, peaks_pointer, &new_peaks);
176            encode_to_memory(memory, auth_path_pointer, &proof.authentication_path);
177            stack.push(peaks_pointer);
178            stack.push(auth_path_pointer);
179        }
180
181        fn pseudorandom_initial_state(
182            &self,
183            seed: [u8; 32],
184            bench_case: Option<BenchmarkCase>,
185        ) -> FunctionInitialState {
186            let mut rng = StdRng::from_seed(seed);
187
188            let old_num_leafs: u64 = match bench_case {
189                Some(BenchmarkCase::CommonCase) => (1 << 31) - 1,
190                Some(BenchmarkCase::WorstCase) => (1 << 62) - 1,
191                None => rng.random_range(0..1 << 63),
192            };
193            let peaks = (0..old_num_leafs.count_ones())
194                .map(|_| rng.random())
195                .collect();
196            let mmr = MmrAccumulator::init(peaks, old_num_leafs);
197
198            self.set_up_initial_state(mmr, rng.random())
199        }
200
201        fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
202            let mut states = vec![];
203            for num_leafs in (0_u64..=17).chain([100, 1000, 1 << 31, 1 << 32, 1 << 33]) {
204                let num_peaks = num_leafs.count_ones();
205                let peaks = (0..num_peaks).map(|i| Tip5::hash(&i)).collect();
206                let mmr = MmrAccumulator::init(peaks, num_leafs);
207                let new_leaf = Tip5::hash(&num_leafs);
208                states.push(self.set_up_initial_state(mmr, new_leaf));
209            }
210
211            states
212        }
213    }
214
215    #[test]
216    fn rust_shadow() {
217        ShadowedFunction::new(CalculateNewPeaksFromAppend).test();
218    }
219}
220
221#[cfg(test)]
222mod benches {
223    use super::*;
224    use crate::test_prelude::*;
225
226    #[test]
227    fn calculate_new_peaks_from_append_benchmark() {
228        ShadowedFunction::new(CalculateNewPeaksFromAppend).bench();
229    }
230}
tasm_lib/mmr/calculate_new_peaks_from_append.rs

tasm_lib/mmr/
calculate_new_peaks_from_append.rs
