mithril_common/test/crypto_helper/
mk_extensions.rs

1use anyhow::Context;
2use std::collections::BTreeSet;
3
4use crate::StdResult;
5use crate::crypto_helper::{
6    MKMap, MKMapNode, MKProof, MKTree, MKTreeNode, MKTreeStoreInMemory, MKTreeStorer,
7};
8use crate::entities::{
9    BlockNumber, BlockNumberOffset, BlockRange, CardanoBlock, CardanoBlockWithTransactions,
10    CardanoTransaction, IntoMKTreeNode, MkSetProof,
11};
12use crate::messages::{CardanoBlocksProofsMessage, CardanoTransactionsProofsV2Message};
13use crate::test::crypto_helper::mkmap_helpers;
14use crate::test::entities_extensions::BlockNumberTestExtension;
15
16/// Extension trait adding test utilities to [MKTree]
17pub trait MKTreeTestExtension {
18    /// `TEST ONLY` - Generate root of the Merkle tree built from the given leaves
19    ///
20    /// Shortcut for `MKTree::new_from_iter(leaves)?.compute_root()`
21    fn compute_root_from_iter<T: IntoIterator<Item = U>, U: Into<MKTreeNode>>(
22        leaves: T,
23    ) -> StdResult<MKTreeNode>;
24}
25
26impl<S: MKTreeStorer> MKTreeTestExtension for MKTree<S> {
27    fn compute_root_from_iter<T: IntoIterator<Item = U>, U: Into<MKTreeNode>>(
28        leaves: T,
29    ) -> StdResult<MKTreeNode> {
30        let mk_tree = Self::new_from_iter(leaves)?;
31        mk_tree.compute_root()
32    }
33}
34
35/// Extension trait adding test utilities to [MKProof]
36pub trait MKProofTestExtension {
37    /// `TEST ONLY` - Build a [MKProof] based on the given leaves
38    fn from_leaves<T: Into<MKTreeNode> + Clone>(leaves: &[T]) -> StdResult<MKProof>;
39
40    /// `TEST ONLY` - Build a [MKProof] based on the given leaves
41    fn from_subset_of_leaves<T: Into<MKTreeNode> + Clone>(
42        leaves: &[T],
43        leaves_to_verify: &[T],
44    ) -> StdResult<MKProof>;
45}
46
47impl MKProofTestExtension for MKProof {
48    fn from_leaves<T: Into<MKTreeNode> + Clone>(leaves: &[T]) -> StdResult<MKProof> {
49        Self::from_subset_of_leaves(leaves, leaves)
50    }
51
52    fn from_subset_of_leaves<T: Into<MKTreeNode> + Clone>(
53        leaves: &[T],
54        leaves_to_verify: &[T],
55    ) -> StdResult<MKProof> {
56        fn list_to_mknode<T: Into<MKTreeNode> + Clone>(hashes: &[T]) -> Vec<MKTreeNode> {
57            hashes.iter().map(|h| h.clone().into()).collect()
58        }
59
60        let leaves = list_to_mknode(leaves);
61        let leaves_to_verify = list_to_mknode(leaves_to_verify);
62
63        let mktree = MKTree::<MKTreeStoreInMemory>::new(&leaves)
64            .with_context(|| "MKTree creation should not fail")?;
65        mktree.compute_proof(&leaves_to_verify)
66    }
67}
68
69/// Extension trait adding test utilities to [MKMap]
70pub trait MKMapTestExtension<K, V, S: MKTreeStorer> {
71    /// `TEST ONLY` - Get the root of the merkle tree of a merkelized map built from an iterator
72    ///
73    /// Shortcut for `MKMap::new_from_iter(entries)?.compute_root()`
74    fn compute_root_from_iter<T: IntoIterator<Item = (K, V)>>(entries: T) -> StdResult<MKTreeNode>;
75
76    /// `TEST ONLY` - Helper to create a MKMap from a list of cardano blocks with transactions
77    fn from_blocks_with_transactions(
78        leaves: &[CardanoBlockWithTransactions],
79    ) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>;
80
81    /// `TEST ONLY` - Helper to create a proof for a list of transactions hashes
82    ///
83    /// IMPORTANT: the `all_leaves` list must contains all the leaves used to create the MKMap
84    fn compute_proof_for_transactions_hashes<H: AsRef<str>>(
85        &self,
86        transactions_hashes_to_prove: &[H],
87        all_leaves: &[CardanoBlockWithTransactions],
88    ) -> StdResult<MkSetProof<CardanoTransaction>>;
89
90    /// `TEST ONLY` - Helper to create a proof message for a list of transactions hashes
91    ///
92    /// IMPORTANT: the `all_leaves` list must contains all the leaves used to create the MKMap
93    fn compute_proof_message_for_transactions_hashes<H: AsRef<str>>(
94        &self,
95        certificate_hash: &str,
96        transactions_hashes_to_prove: &[H],
97        all_leaves: &[CardanoBlockWithTransactions],
98    ) -> StdResult<CardanoTransactionsProofsV2Message> {
99        let proof =
100            self.compute_proof_for_transactions_hashes(transactions_hashes_to_prove, all_leaves)?;
101
102        let message = CardanoTransactionsProofsV2Message::new(
103            certificate_hash,
104            Some(proof.try_into()?),
105            Vec::new(),
106            BlockNumber(9999),
107            BlockNumberOffset(15),
108        );
109
110        Ok(message)
111    }
112
113    /// `TEST ONLY` - Helper to create a proof for a list of blocks hashes
114    ///
115    /// IMPORTANT: the `all_leaves` list must contains all the leaves used to create the MKMap
116    fn compute_proof_for_blocks_hashes<H: AsRef<str>>(
117        &self,
118        blocks_hashes_to_prove: &[H],
119        all_leaves: &[CardanoBlockWithTransactions],
120    ) -> StdResult<MkSetProof<CardanoBlock>>;
121
122    /// `TEST ONLY` - Helper to create a proof message for a list of blocks hashes
123    ///
124    /// IMPORTANT: the `all_leaves` list must contains all the leaves used to create the MKMap
125    fn compute_proof_message_for_blocks_hashes<H: AsRef<str>>(
126        &self,
127        certificate_hash: &str,
128        blocks_hashes_to_prove: &[H],
129        all_leaves: &[CardanoBlockWithTransactions],
130    ) -> StdResult<CardanoBlocksProofsMessage> {
131        let proof = self.compute_proof_for_blocks_hashes(blocks_hashes_to_prove, all_leaves)?;
132
133        let message = CardanoBlocksProofsMessage::new(
134            certificate_hash,
135            Some(proof.try_into()?),
136            Vec::new(),
137            BlockNumber(9999),
138            BlockNumberOffset(15),
139        );
140
141        Ok(message)
142    }
143}
144
145impl<S: MKTreeStorer> MKMapTestExtension<BlockRange, MKMapNode<BlockRange, S>, S>
146    for MKMap<BlockRange, MKMapNode<BlockRange, S>, S>
147{
148    fn compute_root_from_iter<T: IntoIterator<Item = (BlockRange, MKMapNode<BlockRange, S>)>>(
149        entries: T,
150    ) -> StdResult<MKTreeNode> {
151        let mk_map = Self::new_from_iter(entries)?;
152        mk_map.compute_root()
153    }
154
155    fn from_blocks_with_transactions(
156        leaves: &[CardanoBlockWithTransactions],
157    ) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>> {
158        let ordered_leaves: BTreeSet<_> =
159            leaves.iter().flat_map(|l| l.clone().into_mk_tree_node()).collect();
160        let node_per_block_range = BlockNumber::group_items_by_block_range(
161            ordered_leaves.into_iter().map(|n| (n.block_number(), n)),
162        );
163
164        mkmap_helpers::fold_nodes_per_block_range_into_mkmap::<_, _, S>(node_per_block_range)
165    }
166
167    fn compute_proof_for_transactions_hashes<H: AsRef<str>>(
168        &self,
169        transactions_hashes_to_prove: &[H],
170        all_leaves: &[CardanoBlockWithTransactions],
171    ) -> StdResult<MkSetProof<CardanoTransaction>> {
172        let hashes_to_prove: Vec<_> =
173            transactions_hashes_to_prove.iter().map(AsRef::as_ref).collect();
174        let leaves_to_prove: Vec<CardanoTransaction> = all_leaves
175            .iter()
176            .flat_map(|l| l.clone().into_transactions())
177            .filter(|tx| hashes_to_prove.contains(&tx.transaction_hash.as_str()))
178            .collect();
179        let mk_tree_nodes_to_prove: Vec<MKTreeNode> = leaves_to_prove
180            .iter()
181            .cloned()
182            .map(|l| l.into_mk_tree_node())
183            .collect();
184
185        let proof = self.compute_proof(&mk_tree_nodes_to_prove)?;
186        Ok(MkSetProof::new(leaves_to_prove, proof))
187    }
188
189    fn compute_proof_for_blocks_hashes<H: AsRef<str>>(
190        &self,
191        blocks_hashes_to_prove: &[H],
192        all_leaves: &[CardanoBlockWithTransactions],
193    ) -> StdResult<MkSetProof<CardanoBlock>> {
194        let hashes_to_prove: Vec<_> = blocks_hashes_to_prove.iter().map(AsRef::as_ref).collect();
195        let leaves_to_prove: Vec<CardanoBlock> = all_leaves
196            .iter()
197            .filter(|b| hashes_to_prove.contains(&b.block_hash.as_ref()))
198            .cloned()
199            .map(Into::into)
200            .collect();
201        let mk_tree_nodes_to_prove: Vec<MKTreeNode> = leaves_to_prove
202            .iter()
203            .cloned()
204            .map(|l| l.into_mk_tree_node())
205            .collect();
206
207        let proof = self.compute_proof(&mk_tree_nodes_to_prove)?;
208        Ok(MkSetProof::new(leaves_to_prove, proof))
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use crate::entities::SlotNumber;
215
216    use super::*;
217
218    #[test]
219    fn mk_map_from_blocks_with_txs_order_the_leaves() {
220        let ordered_blocks_with_txs = [
221            CardanoBlockWithTransactions::new(
222                "block_hash-10",
223                BlockNumber(10),
224                SlotNumber(100),
225                vec!["tx_hash-1", "tx_hash-2"],
226            ),
227            CardanoBlockWithTransactions::new(
228                "block_hash-15",
229                BlockNumber(15),
230                SlotNumber(150),
231                vec!["tx_hash-4"],
232            ),
233            CardanoBlockWithTransactions::new(
234                "block_hash-16",
235                BlockNumber(16),
236                SlotNumber(160),
237                vec!["tx_hash-5", "tx_hash-6", "tx_hash-7"],
238            ),
239        ];
240        let unordered_blocks = [
241            CardanoBlockWithTransactions::new(
242                "block_hash-16",
243                BlockNumber(16),
244                SlotNumber(160),
245                vec!["tx_hash-5", "tx_hash-6", "tx_hash-7"],
246            ),
247            CardanoBlockWithTransactions::new(
248                "block_hash-10",
249                BlockNumber(10),
250                SlotNumber(100),
251                vec!["tx_hash-1", "tx_hash-2"],
252            ),
253            CardanoBlockWithTransactions::new(
254                "block_hash-15",
255                BlockNumber(15),
256                SlotNumber(150),
257                vec!["tx_hash-4"],
258            ),
259        ];
260        let unordered_txs = [
261            CardanoBlockWithTransactions::new(
262                "block_hash-10",
263                BlockNumber(10),
264                SlotNumber(100),
265                vec!["tx_hash-2", "tx_hash-1"],
266            ),
267            CardanoBlockWithTransactions::new(
268                "block_hash-15",
269                BlockNumber(15),
270                SlotNumber(150),
271                vec!["tx_hash-4"],
272            ),
273            CardanoBlockWithTransactions::new(
274                "block_hash-16",
275                BlockNumber(16),
276                SlotNumber(160),
277                vec!["tx_hash-6", "tx_hash-7", "tx_hash-5"],
278            ),
279        ];
280
281        let ordered_blocks_with_txs_mk_map_root =
282            MKMap::<_, _, MKTreeStoreInMemory>::from_blocks_with_transactions(
283                &ordered_blocks_with_txs,
284            )
285            .unwrap()
286            .compute_root()
287            .unwrap();
288        let unordered_blocks_mk_map_root =
289            MKMap::<_, _, MKTreeStoreInMemory>::from_blocks_with_transactions(&unordered_blocks)
290                .unwrap()
291                .compute_root()
292                .unwrap();
293        let unordered_txs_mk_map_root =
294            MKMap::<_, _, MKTreeStoreInMemory>::from_blocks_with_transactions(&unordered_txs)
295                .unwrap()
296                .compute_root()
297                .unwrap();
298
299        assert_eq!(
300            ordered_blocks_with_txs_mk_map_root,
301            unordered_blocks_mk_map_root
302        );
303        assert_eq!(
304            ordered_blocks_with_txs_mk_map_root,
305            unordered_txs_mk_map_root
306        );
307    }
308
309    #[test]
310    fn compute_proofs_for_blocks_and_txs_from_the_same_mkmap() {
311        let blocks_with_txs = [
312            CardanoBlockWithTransactions::new(
313                "block_hash-10",
314                BlockNumber(10),
315                SlotNumber(100),
316                vec!["tx_hash-1", "tx_hash-2"],
317            ),
318            CardanoBlockWithTransactions::new(
319                "block_hash-15",
320                BlockNumber(15),
321                SlotNumber(150),
322                vec!["tx_hash-4"],
323            ),
324            CardanoBlockWithTransactions::new(
325                "block_hash-16",
326                BlockNumber(16),
327                SlotNumber(160),
328                vec!["tx_hash-5", "tx_hash-6", "tx_hash-7"],
329            ),
330        ];
331
332        let mk_map =
333            MKMap::<_, _, MKTreeStoreInMemory>::from_blocks_with_transactions(&blocks_with_txs)
334                .unwrap();
335
336        let proof_for_blocks_subset = mk_map
337            .compute_proof_for_blocks_hashes(&["block_hash-10", "block_hash-16"], &blocks_with_txs)
338            .unwrap();
339        proof_for_blocks_subset.verify().unwrap();
340
341        let proof_for_txs_subset = mk_map
342            .compute_proof_for_transactions_hashes(&["tx_hash-1", "tx_hash-4"], &blocks_with_txs)
343            .unwrap();
344        proof_for_txs_subset.verify().unwrap();
345
346        assert_eq!(
347            proof_for_blocks_subset.merkle_root(),
348            proof_for_txs_subset.merkle_root()
349        );
350    }
351}
mithril_common/test/crypto_helper/mk_extensions.rs

mithril_common/test/crypto_helper/
mk_extensions.rs
