tasm_lib/mmr/
leaf_index_to_mt_index_and_peak_index.rs

1use std::collections::HashMap;
2
3use triton_vm::prelude::*;
4
5use crate::arithmetic::u64::add::Add;
6use crate::arithmetic::u64::and::And;
7use crate::arithmetic::u64::decr::Decr;
8use crate::arithmetic::u64::log_2_floor::Log2Floor;
9use crate::arithmetic::u64::lt_preserve_args::LtPreserveArgs;
10use crate::arithmetic::u64::popcount::PopCount;
11use crate::arithmetic::u64::pow2::Pow2;
12use crate::prelude::*;
13use crate::traits::basic_snippet::Reviewer;
14use crate::traits::basic_snippet::SignOffFingerprint;
15
16/// Compute both, merkle tree and peak index, for the given leaf index in an MMR
17/// with the given number of leafs. [See also][rust].
18///
19/// ### Behavior
20///
21/// ```text
22/// BEFORE: _ [num_leafs: u64] [leaf_index: u64]
23/// AFTER:  _ [merkle_tree_index: 64] [peak_index: u32]
24/// ```
25///
26/// ### Preconditions
27///
28/// - the `leaf_index` is smaller than `num_leafs`
29/// - all input arguments are properly [`BFieldCodec`] encoded
30///
31/// ### Postconditions
32///
33/// - all output arguments are properly [`BFieldCodec`] encoded
34///
35/// [rust]: twenty_first::util_types::mmr::shared_basic::leaf_index_to_mt_index_and_peak_index
36#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
37pub struct MmrLeafIndexToMtIndexAndPeakIndex;
38
39impl MmrLeafIndexToMtIndexAndPeakIndex {
40    pub const LEAF_INDEX_GE_NUM_LEAFS_ERROR_ID: i128 = 350;
41}
42
43impl BasicSnippet for MmrLeafIndexToMtIndexAndPeakIndex {
44    fn parameters(&self) -> Vec<(DataType, String)> {
45        ["num_leafs", "leaf_index"]
46            .map(|s| (DataType::U64, s.to_string()))
47            .to_vec()
48    }
49
50    fn return_values(&self) -> Vec<(DataType, String)> {
51        vec![
52            (DataType::U64, "mt_index".to_string()),
53            (DataType::U32, "peak_index".to_string()),
54        ]
55    }
56
57    fn entrypoint(&self) -> String {
58        "tasmlib_mmr_leaf_index_to_mt_index_and_peak_index".to_string()
59    }
60
61    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
62        let entrypoint = self.entrypoint();
63        let log_2_floor_u64 = library.import(Box::new(Log2Floor));
64        let lt_u64 = library.import(Box::new(LtPreserveArgs));
65        let add_u64 = library.import(Box::new(Add));
66        let and_u64 = library.import(Box::new(And));
67        let pow2_u64 = library.import(Box::new(Pow2));
68        let decr_u64 = library.import(Box::new(Decr));
69        let popcount_u64 = library.import(Box::new(PopCount));
70
71        triton_asm!(
72        // BEFORE: _ [num_leafs: u64] [leaf_index: u64]
73        // AFTER:  _ [mt_index: u64] peak_index
74        {entrypoint}:
75            /* assert leaf_index < num_leafs */
76            call {lt_u64}
77            assert error_id {Self::LEAF_INDEX_GE_NUM_LEAFS_ERROR_ID}
78            // _ [num_leafs: u64] [leaf_index: u64]
79
80            dup 3
81            dup 2
82            xor
83            // _ [num_leafs: u64] [leaf_index: u64] discrepancies_hi
84
85            dup 3
86            dup 2
87            xor
88            // _ [num_leafs: u64] [leaf_index: u64] [discrepancies: u64]
89
90            call {log_2_floor_u64}
91            // _ [num_leafs: u64] [leaf_index: u64] local_mt_height
92
93            call {pow2_u64}
94            // _ [num_leafs: u64] [leaf_index: u64] [local_mt_num_leafs: u64]
95
96            dup 1 dup 1
97            call {decr_u64}
98            // _ [num_leafs: u64] [leaf_index: u64] [local_mt_num_leafs: u64] [remainder_bitmask: u64]
99
100            dup 1 dup 1
101            pick 7 pick 7
102            call {and_u64}
103            // _ [num_leafs: u64] [local_mt_num_leafs: u64] [remainder_bitmask: u64] [local_leaf_index: u64]
104
105            pick 5 pick 5
106            // _ [num_leafs: u64] [remainder_bitmask: u64] [local_leaf_index: u64] [local_mt_num_leafs: u64]
107
108            call {add_u64}
109            // _ [num_leafs: u64] [remainder_bitmask: u64] [merkle_tree_index: u64]
110
111            place 5 place 5
112            // _ [merkle_tree_index: u64] [num_leafs: u64] [remainder_bitmask: u64]
113
114            dup 3 dup 3
115            call {popcount_u64}
116            // _ [merkle_tree_index: u64] [num_leafs: u64] [remainder_bitmask: u64] all_the_ones
117
118            place 4
119            // _ [merkle_tree_index: u64] all_the_ones [num_leafs: u64] [remainder_bitmask: u64]
120
121            call {and_u64}
122            call {popcount_u64}
123            // _ [merkle_tree_index: u64] all_the_ones ones_to_subtract
124
125            push -1
126            mul
127            add
128            addi -1
129            // _ [merkle_tree_index: u64] peak_index
130
131            return
132        )
133    }
134
135    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
136        let mut sign_offs = HashMap::new();
137        sign_offs.insert(Reviewer("ferdinand"), 0xfaeccefe8b0b120.into());
138        sign_offs
139    }
140}
141
142#[cfg(test)]
143pub(crate) mod tests {
144    use twenty_first::util_types::mmr::shared_basic::leaf_index_to_mt_index_and_peak_index;
145
146    use super::*;
147    use crate::test_prelude::*;
148
149    impl MmrLeafIndexToMtIndexAndPeakIndex {
150        pub fn assert_expected_behavior(&self, num_leafs: u64, leaf_index: u64) {
151            let initial_stack = self.set_up_test_stack((num_leafs, leaf_index));
152
153            let mut expected_stack = initial_stack.clone();
154            self.rust_shadow(&mut expected_stack).unwrap();
155
156            test_rust_equivalence_given_complete_state(
157                &ShadowedClosure::new(Self),
158                &initial_stack,
159                &[],
160                &NonDeterminism::default(),
161                &None,
162                Some(&expected_stack),
163            );
164        }
165    }
166
167    impl Closure for MmrLeafIndexToMtIndexAndPeakIndex {
168        type Args = (u64, u64);
169
170        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
171            let (num_leafs, leaf_index) = pop_encodable::<Self::Args>(stack)?;
172
173            if leaf_index >= num_leafs {
174                return Err(RustShadowError::Other);
175            }
176            let (mt_index, peak_index) =
177                leaf_index_to_mt_index_and_peak_index(leaf_index, num_leafs);
178
179            push_encodable(stack, &mt_index);
180            push_encodable(stack, &peak_index);
181            Ok(())
182        }
183
184        fn pseudorandom_args(
185            &self,
186            seed: [u8; 32],
187            bench_case: Option<BenchmarkCase>,
188        ) -> Self::Args {
189            let Some(bench_case) = bench_case else {
190                let mut rng = StdRng::from_seed(seed);
191                let num_leafs = rng.random_range(1..1 << 63);
192                let leaf_index = rng.random_range(0..num_leafs);
193                return (num_leafs, leaf_index);
194            };
195
196            match bench_case {
197                BenchmarkCase::WorstCase => ((1 << 63) - 1, (1 << 63) - 63),
198                BenchmarkCase::CommonCase => ((1 << 32) - 1, (3 << 30) + 100_000),
199            }
200        }
201
202        fn corner_case_args(&self) -> Vec<Self::Args> {
203            // todo: can the collection of cases below be less
204            //   “stab-in-the-dark” and more actual edge cases?
205
206            let mut states = vec![];
207            for num_leafs in (1..=5).chain([14]).chain(32..=37) {
208                for leaf_index in 0..num_leafs {
209                    states.push((num_leafs, leaf_index));
210                }
211            }
212
213            let more_states = (10..20)
214                .map(|pow| 1 << pow)
215                .flat_map(|n| [(14, n), (n + 9, n + 11), (n + 10, n + 11)])
216                .map(|(leaf_index, num_leafs)| (num_leafs, leaf_index));
217            states.extend(more_states);
218
219            states.push((1 << 31, 5_550_001));
220            states.push(((1 << 31) + (1 << 20), 5_550_001));
221
222            for num_leafs in [(1 << 31) + (1 << 20) - 1, (1 << 63) + (1 << 62) - 1] {
223                states.push((num_leafs, num_leafs - 1));
224            }
225
226            states
227        }
228    }
229
230    #[macro_rules_attr::apply(test)]
231    fn rust_shadow() {
232        ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex).test();
233    }
234
235    #[macro_rules_attr::apply(proptest)]
236    fn property_test(
237        #[strategy(1_u64..)] num_leafs: u64,
238        #[strategy(0_u64..#num_leafs)] leaf_index: u64,
239    ) {
240        MmrLeafIndexToMtIndexAndPeakIndex.assert_expected_behavior(num_leafs, leaf_index);
241    }
242
243    #[macro_rules_attr::apply(proptest)]
244    fn negative_property_test(num_leafs: u64, #[strategy(#num_leafs..)] leaf_index: u64) {
245        let initial_stack =
246            MmrLeafIndexToMtIndexAndPeakIndex.set_up_test_stack((num_leafs, leaf_index));
247
248        test_assertion_failure(
249            &ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex),
250            InitVmState::with_stack(initial_stack),
251            &[MmrLeafIndexToMtIndexAndPeakIndex::LEAF_INDEX_GE_NUM_LEAFS_ERROR_ID],
252        );
253    }
254}
255
256#[cfg(test)]
257mod benches {
258    use super::*;
259    use crate::test_prelude::*;
260
261    #[macro_rules_attr::apply(test)]
262    fn benchmark() {
263        ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex).bench();
264    }
265}
tasm_lib/mmr/leaf_index_to_mt_index_and_peak_index.rs

tasm_lib/mmr/
leaf_index_to_mt_index_and_peak_index.rs
