use crate::{BlockBuilder, SnapshotVaultIndex};
use ex3_blockchain_public_types::{ActivatedWalletRegistry, SnapshotCursor};
use ex3_canister_types::range::ReportRange;
use ex3_crypto::sha256;
use ex3_node_types::number::Ex3Uint;
use ex3_node_types::range::CandidRange;
use ex3_node_types::transaction::EncodedTransaction;
use ex3_node_types::{BlockHeight, WalletRegistrySeqId};
use ex3_node_types::{MerkleNode, WalletRegisterId};
use ex3_serde::bincode::serialize;
use ex3_wallet_registry_public_types::WalletTxShardingSubReport;
use num_traits::{CheckedSub, One, ToPrimitive, Zero};
use rs_merkle::algorithms::Sha256;
use rs_merkle::MerkleTree;
use std::cmp::max;
use std::collections::{BTreeMap, VecDeque};
impl BlockBuilder {
pub(crate) fn generate_wallet_tx_sharding_sub_reports(
&self,
tx_map: BTreeMap<WalletRegistrySeqId, Vec<(EncodedTransaction, usize)>>,
txs_merkle_tree: &MerkleTree<Sha256>,
snapshot_vault_cursor_of_pre_height: &SnapshotCursor,
) -> BTreeMap<WalletRegistrySeqId, Vec<WalletTxShardingSubReport>> {
let mut wallet_registry_tx_sharding_sub_reports = BTreeMap::new();
let max_wallet_registry_seq_id = max(
snapshot_vault_cursor_of_pre_height.max_activated_wallet_registry_seq_id,
tx_map.keys().max().unwrap_or(&0).clone(),
);
let mut next_wallet_registry_seq_id: WalletRegistrySeqId = 0;
let mut tx_map_mut = tx_map;
let total_tx_count = txs_merkle_tree.leaves_len();
while next_wallet_registry_seq_id <= max_wallet_registry_seq_id {
let mut wallet_register_txs_sub_reports = vec![];
if !tx_map_mut.contains_key(&next_wallet_registry_seq_id) {
continue;
}
let mut wallet_register_txs = tx_map_mut
.remove(&next_wallet_registry_seq_id)
.unwrap()
.into_iter()
.collect::<VecDeque<_>>();
let total_tx_count_in_sharding = wallet_register_txs.len();
loop {
let mut txs = vec![];
for _ in 0..self.max_wallet_register_txs_per_sharding_sub_report {
if let Some(tx) = wallet_register_txs.pop_front() {
txs.push(tx);
} else {
break;
}
}
if txs.len() == 0 {
break;
}
let (txs, indexes) =
txs.into_iter()
.map(|(tx, index)| (tx, index.clone()))
.unzip::<EncodedTransaction, usize, Vec<EncodedTransaction>, Vec<usize>>();
let merkle_proof = txs_merkle_tree.proof(&indexes);
let report_start = wallet_register_txs_sub_reports.len()
* self.max_wallet_register_txs_per_sharding_sub_report.clone();
let report_end = report_start + indexes.len();
let range = ReportRange {
total_count: total_tx_count.try_into().unwrap(),
sharding_range: CandidRange {
start: 0u32,
end: total_tx_count_in_sharding.try_into().unwrap(),
},
report_range: CandidRange {
start: report_start.try_into().unwrap(),
end: report_end.try_into().unwrap(),
},
};
let report = WalletTxShardingSubReport {
merkle_proof: merkle_proof.to_bytes(),
range,
txs,
tx_indexes: indexes
.into_iter()
.map(|index| index.try_into().unwrap())
.collect(),
};
wallet_register_txs_sub_reports.push(report);
}
wallet_registry_tx_sharding_sub_reports
.insert(next_wallet_registry_seq_id, wallet_register_txs_sub_reports);
next_wallet_registry_seq_id += 1;
}
wallet_registry_tx_sharding_sub_reports
}
pub(crate) fn generate_wallet_registry_data_integrity(
&self,
current_block_height: &BlockHeight,
wallet_register_tx_sharding_sub_reports: &BTreeMap<
WalletRegistrySeqId,
Vec<WalletTxShardingSubReport>,
>,
snapshot_vault_cursor: &SnapshotCursor,
snapshot_vault_index: &SnapshotVaultIndex,
) -> (
BTreeMap<WalletRegistrySeqId, MerkleNode>,
Vec<ActivatedWalletRegistry>,
) {
let mut wallet_registry_data_integrity = BTreeMap::new();
let mut new_activated_wallet_registries = vec![];
wallet_register_tx_sharding_sub_reports.iter().for_each(
|(wallet_registry_seq_id, reports)| {
let mut pre_block_height_processed: Option<BlockHeight> = None;
if wallet_registry_seq_id
<= &snapshot_vault_cursor.max_activated_wallet_registry_seq_id
&& *current_block_height > BlockHeight::zero()
{
pre_block_height_processed = Some(
current_block_height
.checked_sub(&BlockHeight::one())
.unwrap(),
);
};
let total_vaults_count = snapshot_vault_index.get_total_vaults_count();
let total_merkle_node_count = snapshot_vault_index.get_total_merkle_node_count();
let wallet_registry_vault_index =
snapshot_vault_index.get_wallet_registry_index(wallet_registry_seq_id);
let total_register_tx_count: u32 = reports
.iter()
.map(|report| report.txs.len())
.sum::<usize>()
.try_into()
.unwrap();
let mut data_buffer = vec![];
data_buffer.extend(
serialize(&pre_block_height_processed).expect("serialize should not fail"),
);
data_buffer.extend(current_block_height.0.to_bytes_le());
data_buffer.extend(wallet_registry_seq_id.to_le_bytes());
data_buffer.extend(total_register_tx_count.to_le_bytes());
data_buffer.extend(total_merkle_node_count.to_le_bytes());
data_buffer.extend(wallet_registry_vault_index.to_le_bytes());
reports.iter().for_each(|report| {
data_buffer.extend(&report.merkle_proof);
});
let data_integrity: MerkleNode = sha256(data_buffer);
wallet_registry_data_integrity
.insert(wallet_registry_seq_id.clone(), data_integrity);
if pre_block_height_processed.is_none() {
new_activated_wallet_registries.push(ActivatedWalletRegistry {
canister_seq_id: wallet_registry_seq_id.clone(),
canister_total_count: total_vaults_count,
canister_index: wallet_registry_vault_index,
hash_of_sharding: data_integrity,
});
}
},
);
(
wallet_registry_data_integrity,
new_activated_wallet_registries,
)
}
pub fn get_wallet_registry_seq_id_by_wallet_id(
&self,
wallet_id: &WalletRegisterId,
) -> WalletRegistrySeqId {
let wallet_id = wallet_id
.checked_sub(&Ex3Uint::from(self.start_wallet_id.clone()))
.unwrap();
(wallet_id.0 / self.max_wallets_per_registry_canister.clone())
.to_u64()
.unwrap()
}
}