1use std::{collections::BTreeMap, str::FromStr};
5
6use penumbra_sdk_asset::STAKING_TOKEN_ASSET_ID;
7use penumbra_sdk_compact_block::{CompactBlock, StatePayload};
8use penumbra_sdk_keys::FullViewingKey;
9use penumbra_sdk_num::Amount;
10use penumbra_sdk_shielded_pool::{Note, NotePayload};
11use penumbra_sdk_stake::{
12 rate::{BaseRateData, RateData},
13 DelegationToken,
14};
15use penumbra_sdk_tct::StateCommitment;
16
17#[derive(Debug, Clone)]
18pub struct FilteredGenesisBlock {
19 #[allow(dead_code)]
21 pub notes: BTreeMap<String, BTreeMap<StateCommitment, Note>>,
22 pub balances: BTreeMap<String, Amount>,
24}
25
26pub async fn scan_genesis_block(
31 CompactBlock {
32 height,
33 state_payloads,
34 ..
35 }: CompactBlock,
36 fvks: Vec<FullViewingKey>,
37) -> anyhow::Result<FilteredGenesisBlock> {
38 assert_eq!(height, 0);
39
40 let mut notes = BTreeMap::new();
41 let mut balances = BTreeMap::new();
42
43 let base_rate = BaseRateData {
45 epoch_index: 0,
46 base_reward_rate: 0u128.into(),
47 base_exchange_rate: 1_0000_0000u128.into(),
48 };
49
50 for fvk in fvks {
52 let trial_decrypt_note =
54 |note_payload: NotePayload| -> tokio::task::JoinHandle<Option<Note>> {
55 let fvk2 = fvk.clone();
56 tokio::spawn(async move { note_payload.trial_decrypt(&fvk2) })
57 };
58
59 let mut note_decryptions = Vec::new();
62
63 for payload in state_payloads.iter() {
65 if let StatePayload::Note { note, .. } = payload {
66 note_decryptions.push(trial_decrypt_note((**note).clone()));
67 }
68 }
69
70 let mut notes_for_this_fvk = BTreeMap::new();
71 for decryption in note_decryptions {
72 if let Some(note) = decryption
73 .await
74 .expect("able to join tokio note decryption handle")
75 {
76 notes_for_this_fvk.insert(note.commit(), note.clone());
77
78 let note_value = note.value();
80 if note_value.asset_id == *STAKING_TOKEN_ASSET_ID {
81 balances
82 .entry(fvk.to_string())
83 .and_modify(|existing_amount| *existing_amount += note.amount())
84 .or_insert(note.amount());
85 } else if let Ok(delegation_token) =
86 DelegationToken::from_str(¬e_value.asset_id.to_string())
87 {
88 let rate_data = RateData {
90 identity_key: delegation_token.validator(),
91 validator_reward_rate: 0u128.into(),
92 validator_exchange_rate: base_rate.base_exchange_rate,
93 };
94 let um_equivalent_balance = rate_data.unbonded_amount(note.amount());
95
96 balances
97 .entry(fvk.to_string())
98 .and_modify(|existing_amount| *existing_amount += um_equivalent_balance)
99 .or_insert(um_equivalent_balance);
100 } else {
101 tracing::warn!(
102 "ignoring note with unknown asset id: {}",
103 note_value.asset_id
104 );
105 }
106 }
107 }
108
109 notes.insert(fvk.to_string(), notes_for_this_fvk);
111 }
112
113 let result = FilteredGenesisBlock { notes, balances };
115
116 Ok(result)
117}