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 inputs(&self) -> Vec<(DataType, String)> {
45        ["num_leafs", "leaf_index"]
46            .map(|s| (DataType::U64, s.to_string()))
47            .to_vec()
48    }
49
50    fn outputs(&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"), 0xdbc9fd10de0f2dbd.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);
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>) {
171            let (num_leafs, leaf_index) = pop_encodable::<Self::Args>(stack);
172
173            let (mt_index, peak_index) =
174                leaf_index_to_mt_index_and_peak_index(leaf_index, num_leafs);
175
176            push_encodable(stack, &mt_index);
177            push_encodable(stack, &peak_index);
178        }
179
180        fn pseudorandom_args(
181            &self,
182            seed: [u8; 32],
183            bench_case: Option<BenchmarkCase>,
184        ) -> Self::Args {
185            let Some(bench_case) = bench_case else {
186                let mut rng = StdRng::from_seed(seed);
187                let num_leafs = rng.random_range(1..1 << 63);
188                let leaf_index = rng.random_range(0..num_leafs);
189                return (num_leafs, leaf_index);
190            };
191
192            match bench_case {
193                BenchmarkCase::WorstCase => ((1 << 63) - 1, (1 << 63) - 63),
194                BenchmarkCase::CommonCase => ((1 << 32) - 1, (3 << 30) + 100_000),
195            }
196        }
197
198        fn corner_case_args(&self) -> Vec<Self::Args> {
199            // todo: can the collection of cases below be less
200            //   “stab-in-the-dark” and more actual edge cases?
201
202            let mut states = vec![];
203            for num_leafs in (1..=5).chain([14]).chain(32..=37) {
204                for leaf_index in 0..num_leafs {
205                    states.push((num_leafs, leaf_index));
206                }
207            }
208
209            let more_states = (10..20)
210                .map(|pow| 1 << pow)
211                .flat_map(|n| [(14, n), (n + 9, n + 11), (n + 10, n + 11)])
212                .map(|(leaf_index, num_leafs)| (num_leafs, leaf_index));
213            states.extend(more_states);
214
215            states.push((1 << 31, 5_550_001));
216            states.push(((1 << 31) + (1 << 20), 5_550_001));
217
218            for num_leafs in [(1 << 31) + (1 << 20) - 1, (1 << 63) + (1 << 62) - 1] {
219                states.push((num_leafs, num_leafs - 1));
220            }
221
222            states
223        }
224    }
225
226    #[test]
227    fn rust_shadow() {
228        ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex).test();
229    }
230
231    #[proptest]
232    fn property_test(
233        #[strategy(1_u64..)] num_leafs: u64,
234        #[strategy(0_u64..#num_leafs)] leaf_index: u64,
235    ) {
236        MmrLeafIndexToMtIndexAndPeakIndex.assert_expected_behavior(num_leafs, leaf_index);
237    }
238
239    #[proptest]
240    fn negative_property_test(num_leafs: u64, #[strategy(#num_leafs..)] leaf_index: u64) {
241        let initial_stack =
242            MmrLeafIndexToMtIndexAndPeakIndex.set_up_test_stack((num_leafs, leaf_index));
243
244        test_assertion_failure(
245            &ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex),
246            InitVmState::with_stack(initial_stack),
247            &[MmrLeafIndexToMtIndexAndPeakIndex::LEAF_INDEX_GE_NUM_LEAFS_ERROR_ID],
248        );
249    }
250}
251
252#[cfg(test)]
253mod benches {
254    use super::*;
255    use crate::test_prelude::*;
256
257    #[test]
258    fn benchmark() {
259        ShadowedClosure::new(MmrLeafIndexToMtIndexAndPeakIndex).bench();
260    }
261}
tasm_lib/mmr/leaf_index_to_mt_index_and_peak_index.rs

tasm_lib/mmr/
leaf_index_to_mt_index_and_peak_index.rs
