Skip to content

Commit

Permalink
temp
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpraynaud
jpraynaud committed Apr 30, 2024
1 parent 09000ea commit 3ce9542
Showing 1 changed file with 124 additions and 36 deletions.
160 changes: 124 additions & 36 deletions mithril-aggregator/src/services/prover.rs
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
use std::{collections::HashMap, sync::Arc};
use std::{
collections::{BTreeMap, HashMap},
sync::Arc,
};

use async_trait::async_trait;

use mithril_common::{
crypto_helper::MKTree,
crypto_helper::{MKMapValue, MKTree},
entities::{
BlockRange, CardanoDbBeacon, CardanoTransaction, CardanoTransactionsSetProof,
TransactionHash,
Expand Down Expand Up @@ -56,6 +59,7 @@ impl MithrilProverService {
}
}

/// Get all the transactions of the block ranges in which the given transaction hashes are included
async fn get_transactions_grouped_by_block_range(
&self,
hashes: &[TransactionHash],
Expand Down Expand Up @@ -87,9 +91,13 @@ impl ProverService for MithrilProverService {
let block_range_transactions = self
.get_transactions_grouped_by_block_range(transaction_hashes)
.await?;
println!(
"block_range_transactions: {:?}",
block_range_transactions.keys().collect::<Vec<_>>()
);

// 2 - Compute block ranges sub Merkle trees
let mut mk_trees = HashMap::new();
let mut mk_trees = BTreeMap::new();
for (block_range, transactions) in block_range_transactions {
let mk_tree = MKTree::new(
&transactions
Expand All @@ -105,19 +113,38 @@ impl ProverService for MithrilProverService {
.block_range_root_retriever
.compute_merkle_map_from_block_range_roots(up_to.immutable_file_number)
.await?;
println!(
"keys={:?}",
mk_map.iter().map(|(k, _)| k).collect::<Vec<_>>()
);

// 4 - Enrich the Merkle map with the block ranges Merkle trees
for (block_range, mk_tree) in mk_trees {
println!(
"block_range={:?}, mk_tree_len: {}, mk_tree_root: {:?}",
block_range,
mk_tree.total_leaves(),
mk_tree.compute_root().unwrap()
);
let mk_map_node = mk_map.get(&block_range).unwrap().to_owned();
println!(
"mk_map[{:?}]={:?}",
block_range,
mk_map_node.compute_root().unwrap()
);
mk_map.insert(block_range, mk_tree.into())?;
}

// 5 - Compute the proof for each transaction
let mk_proof = mk_map.compute_proof(transaction_hashes)?;

Ok(vec![CardanoTransactionsSetProof::new(
transaction_hashes.to_vec(),
mk_proof,
)])
// TODO: return STdResult<Option<MKProof>> to handle proof not computed beacause no transaction is valid?
if let Ok(mk_proof) = mk_map.compute_proof(transaction_hashes) {
Ok(vec![CardanoTransactionsSetProof::new(
transaction_hashes.to_vec(),
mk_proof,
)])
} else {
Ok(vec![])
}
}
}

Expand Down Expand Up @@ -161,7 +188,7 @@ mod tests {
let hash = format!("tx-{i}");
transactions.push(CardanoTransaction::new(
&hash,
max(0, 10 * i - 1) as u64,
max(0, 5 * i - 1) as u64,
100 * i as u64,
format!("block_hash-{i}"),
i as u64,
Expand All @@ -172,6 +199,18 @@ mod tests {
(hashes, transactions)
}

fn compute_mk_map_from_transactions(
transactions: Vec<CardanoTransaction>,
) -> MKMap<BlockRange, MKMapNode<BlockRange>> {
MKMap::new_from_iter(transactions.iter().map(|tx| {
(
BlockRange::from_block_number(tx.block_number),
MKMapNode::TreeNode(tx.transaction_hash.clone().into()),
)
}))
.unwrap()
}

fn build_prover<F, G>(
transaction_retriever_mock_config: F,
block_range_root_retriever_mock_config: G,
Expand All @@ -193,18 +232,14 @@ mod tests {

#[tokio::test]
async fn compute_proof_for_one_set_with_multiple_transactions() {
let (transaction_hashes, transactions) = generate_transactions(3);
let (transaction_hashes, transactions) = generate_transactions(30);
let prover = build_prover(
|retriever_mock| {
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
Expand All @@ -227,27 +262,38 @@ mod tests {

#[tokio::test]
async fn cant_compute_proof_for_unknown_transaction() {
let (transaction_hashes, transactions) = generate_transactions(3);
let (transaction_hashes, transactions) = generate_transactions(30);
let transaction_hash_unknown = "tx-unknown-123".to_string();
let beacon = CardanoDbBeacon {
immutable_file_number: transactions.last().unwrap().immutable_file_number,
..CardanoDbBeacon::default()
};
let prover = build_prover(
|retriever_mock| {
let transactions = transactions[..].to_vec();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.returning(|_| Ok(vec![]));
.expect_get_by_hashes()
.with(eq(vec![transaction_hash_unknown.clone()]))
.return_once(move |_| Ok(vec![]));
},
|block_range_root_retriever_mock| {
let transactions_to_compute_mk_map = transactions[..].to_vec();
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
.return_once(|_| {
Ok(compute_mk_map_from_transactions(
transactions_to_compute_mk_map,
))
});
},
);

let transactions_set_proof = prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
.compute_transactions_proofs(&beacon, &[transaction_hash_unknown])
.await
.unwrap();

Expand All @@ -256,44 +302,63 @@ mod tests {

#[tokio::test]
async fn compute_proof_for_one_set_of_three_known_transactions_and_two_unknowns() {
let (transaction_hashes, transactions) = generate_transactions(5);
let (transaction_hashes, transactions) = generate_transactions(3);
let transaction_hashes_unknown =
vec!["tx-unknown-123".to_string(), "tx-unknown-456".to_string()];
let transaction_hashes_unknown = vec![];
let transaction_hashes_to_verify = transaction_hashes
.clone()
.into_iter()
.chain(transaction_hashes_unknown.into_iter())
.collect::<Vec<_>>();
let beacon = CardanoDbBeacon {
immutable_file_number: transactions.last().unwrap().immutable_file_number,
..CardanoDbBeacon::default()
};
let prover = build_prover(
|retriever_mock| {
// The last two are not in the "store"
let transactions = transactions[0..=2].to_vec();
let transactions = transactions[..].to_vec();
let transactions_by_hashes_res = transactions.clone();
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.with(eq(transaction_hashes_to_verify.clone()))
.return_once(move |_| Ok(transactions_by_hashes_res));
retriever_mock
.expect_get_up_to()
.with(eq(fake_data::beacon()))
.return_once(move |_| Ok(transactions));
},
|block_range_root_retriever_mock| {
let transactions_to_compute_mk_map = transactions[..].to_vec();
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| MKMap::new(&[]));
.return_once(|_| {
println!(
"123 transactions_to_compute_mk_map={:?}",
transactions_to_compute_mk_map
);
Ok(compute_mk_map_from_transactions(
transactions_to_compute_mk_map,
))
});
},
);

println!(
"transaction_hashes_to_verify: {:?}",
transaction_hashes_to_verify
);
let transactions_set_proof = prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
.compute_transactions_proofs(&beacon, &transaction_hashes_to_verify)
.await
.unwrap();

assert_eq!(transactions_set_proof.len(), 1);
assert_eq!(
transactions_set_proof[0].transactions_hashes(),
&transaction_hashes[0..=2].to_vec()
transaction_hashes
);
transactions_set_proof[0].verify().unwrap();
}

#[tokio::test]
async fn cant_compute_proof_if_retriever_fail() {
let (transaction_hashes, _transactions) = generate_transactions(3);
async fn cant_compute_proof_if_transaction_retriever_fail() {
let (transaction_hashes, _transactions) = generate_transactions(30);
let prover = build_prover(
|retriever_mock| {
retriever_mock
Expand All @@ -311,6 +376,29 @@ mod tests {
prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
.await
.expect_err("Should have failed because of its retriever");
.expect_err("Should have failed because of its transaction retriever");
}

#[tokio::test]
async fn cant_compute_proof_if_block_range_root_retriever_fail() {
let (transaction_hashes, transactions) = generate_transactions(3);
let prover = build_prover(
|retriever_mock| {
retriever_mock
.expect_get_by_hashes()
.with(eq(transaction_hashes.clone()))
.return_once(move |_| Ok(transactions));
},
|block_range_root_retriever_mock| {
block_range_root_retriever_mock
.expect_compute_merkle_map_from_block_range_roots()
.return_once(|_| Err(anyhow!("Error")));
},
);

prover
.compute_transactions_proofs(&fake_data::beacon(), &transaction_hashes)
.await
.expect_err("Should have failed because of its transaction retriever");
}
}

0 comments on commit 3ce9542

Please sign in to comment.